PaperlayerView exact revision on arXiv

Source-linked AI summary

How Transformers Learn Causal Structure with Gradient Descent

Eshaan Nichani, Alex Damian, Jason D. Lee

arXiv:2402.14735v2cs.LGcs.ITmath.ITstat.ML

TL;DR

The paper addresses the poorly understood question of how gradient descent teaches transformers to learn causal structure. It introduces an in-context task and proves that a simplified two-layer transformer encodes tree-structured causal graphs in its first attention layer, with induction heads as a Markov-chain special case.

  • Problem

    How gradient descent teaches transformers to learn causal structure remains poorly understood, despite the role of that structure in in-context learning.

    Source: §1 · pp. 1–2

  • Method

    The paper introduces a latent-causal-structure in-context task and analyzes gradient descent on a simplified two-layer transformer using token mutual information.

    Source: §1.1 · p. 2, §5.1.2 · p. 14

  • Results

    For tree graphs, gradient descent learns the task by encoding the graph in the first attention layer, and in-context Markov chains yield an induction head.

    Source: §1.1 · p. 2

  • Takeaways & Limitations

    The result connects induction heads to the special case of learning in-context Markov chains.

    Source: §1.1 · p. 2, §1 · p. 2

    For non-tree graphs, the paper reports empirical learning of a multi-head construction but leaves its optimization dynamics for future work because the heads must break symmetry.

    Source: §6 · p. 17

Abstract

from arXiv · show

The incredible success of transformers on sequence modeling tasks can be largely attributed to the self-attention mechanism, which allows information to be transferred between different parts of a sequence. Self-attention allows transformers to encode causal structure which makes them particularly suitable for sequence modeling. However, the process by which transformers learn such causal structure via gradient-based training algorithms remains poorly understood. To better understand this process, we introduce an in-context learning task that requires learning latent causal structure. We prove that gradient descent on a simplified two-layer transformer learns to solve this task by encoding the latent causal graph in the first attention layer. The key insight of our proof is that the gradient of the attention matrix encodes the mutual information between tokens. As a consequence of the data processing inequality, the largest entries of this gradient correspond to edges in the latent causal graph. As a special case, when the sequences are generated from in-context Markov chains, we prove that transformers learn an induction head (Olsson et al., 2022). We confirm our theoretical findings by showing that transformers trained on our in-context learning task are able to recover a wide variety of causal structures.

1 Introduction

The paper asks how gradient descent teaches transformers to encode causal structure, a process not explained by prior work on circuits or in-context learning. It introduces a latent-structure task, proves a two-layer transformer learns tree graphs through its first attention layer, and relates the result to prior work on transformer representations and induction heads.

1 Introduction

  • Prior in-context setups include induction heads [27] and function-class learning [17], where token pairings impose causal structure.In function-class prompts, each output is paired with its input, while the sampled function supplies a sequence-specific transition.Source: §1 · p. 2
  • The paper identifies how gradient-based training learns causal structure as an open question despite transformers’ ability to model it.Prior work studied circuits in trained attention layers, but not how those circuits arise during training from scratch.Source: §1 · pp. 1–2

1.1 Our Contributions

  • The paper introduces random sequences with causal structure, where each sequence is sampled from a different distribution respecting a fixed, hidden graph.The task makes the latent causal structure unknown to the transformer.Source: §1.1 · p. 2
  • For tree graphs, gradient descent solves the task by encoding the graph in the first attention layer and estimating transition distributions in context.For in-context Markov chains, this result implies that the transformer learns an induction head.Source: §1.1 · p. 2
  • The first attention-layer gradient computes χ2-mutual information between token pairs; data processing makes its largest entries correspond to graph edges.The first attention layer therefore converges to the graph’s adjacency matrix.Source: §1.1 · pp. 2–3
  • For non-tree graphs, the paper constructs a multi-head transformer that distributes the latent graph across heads and reports that gradient-trained transformers learn this construction.Analyzing optimization dynamics for multiple heads remains future work because the heads must break symmetry.Source: §1.1 · p. 3

1.2 Related Work

  • Prior work showed transformers can represent in-context learning algorithms [2, 4, 35], but did not establish that gradient descent learns those constructions.The paper distinguishes its focus on learning dynamics from representational-capability studies.Source: §1.2 · p. 3
  • Compared with Bietti et al. [5], this task covers more general causal structures and provides an end-to-end guarantee; Edelman et al. [13] analyze only two steps for two-state Markov chains.Bietti et al. [5] assume alphabet size S is much larger than context length T, whereas this paper assumes T ≫ S.Source: §1.2 · p. 4

2 Setup

The paper uses attention-only and disentangled transformer formulations to study sequences generated from a latent causal graph. Its task samples a sequence and predicts a target from a new token, with Markov-chain and in-context function-learning settings as examples.

2.1 Transformer Architecture

  • The disentangled transformer appends each attention layer’s output to the residual stream and is equivalent to a decoder-based attention-only transformer.Causal masking restricts attention to previous tokens, with each attention row normalized to sum to 1.Source: §2.1 · pp. 5–6

2.2 Problem Setup: Random Sequences with Causal Structure

  • Task 1 samples tokens according to a latent directed acyclic graph, then predicts a target drawn from the transition distribution conditioned on the final input token.Root positions are sampled from the chain’s stationary measure; non-roots are sampled conditioned on their parent token.Source: §2.2 · pp. 6–7

2.3 Examples

  • When each position’s parent is the preceding position, the task becomes in-context Markov-chain estimation, solvable by averaging tokens following prior occurrences of the final token.For the in-context function-learning graph, the model instead associates input-output pairs to learn the sampled function in-context.Source: §2.3 · p. 7

3 What does the Transformer Learn?

Experiments show that trained transformers’ first attention weights recover the latent graph’s adjacency pattern. The paper constructs a two-layer solution and studies a reduced model in which only its attention matrices are trained.

3.1 Experiments

  • For the in-context learning graph, the trained first-layer attention block converges to the graph adjacency matrix, while highlighted second-layer and output blocks converge to identity.The experiment uses a two-layer disentangled transformer with S = 10 and T = 20.Source: §3.1 · p. 7, Figure 2 · p. 8

3.2 Construction

  • The construction makes the first layer copy each token’s parent and the second average prior tokens whose parent token matches the current token, producing an empirical transition estimate.The output layer reads that average; the construction uses attention matrices scaled toward hard selection as β1 and β2 increase.Source: §3.2 · pp. 8–10

3.3 The Reduced Model

  • The reduced model fixes the architecture’s sparsity pattern and trains only A(1) and A(2), with A(1) restricted to lower-diagonal entries by causal masking.A small positive perturbation is added to predictions because otherwise an unseen target token would give infinite cross-entropy loss.Source: §3.3 · pp. 10–11

4 Main Results

The authors prove that stage-wise gradient descent approximately minimizes the task loss by encoding the latent graph in the first attention layer, provided the effective sequence length is sufficiently large. They also report good prediction on out-of-distribution transitions under a lower bound on transition probabilities.

4.1 Training Algorithm

  • Training proceeds in two stages: gradient descent first updates A(1), then updates A(2), with A(1) initialized to zero and A(2) to a small scaled identity.The algorithm uses separate learning rates and durations for the two stages.Source: §4.1 · p. 11
  • The analysis assumes conditions on the Markov-chain prior and that a non-vanishing fraction of graph nodes have parents.The effective sequence length roughly captures the number of independent samples and is T for the Markov-chain and in-context-learning examples.Source: §4.1 · pp. 11–12, §4.2 · p. 12

4.2 Main Theorem

  • When Teff ≥ poly(γ^-1, S), Algorithm 1 approximately minimizes loss by encoding the causal graph’s adjacency matrix in the first attention layer.The theorem states that suitable perturbation, learning rates, and stage durations exist.Source: §4.2 · p. 12
  • The trained model achieves good prediction on out-of-distribution transitions when the test transition probabilities are bounded below by γ/S.The theorem setup states a probability of at least 0.99 over the generated sequence.Source: §4.2 · pp. 12–13

5 Proof Sketch

The proof shows how gradient descent recovers a latent causal graph in the first attention layer, then increases the second-stage parameter β to reduce the loss. Its key link is that first-layer gradient entries measure token dependence, with the true parent attaining the largest score.

5.1.1 The Oracle Algorithm

  • For each non-root node, the oracle recovers its parent by maximizing conditional mutual information with earlier nodes; if all such values are zero, the node is a root.The resulting tree is the maximum-weight spanning forest when the topological ordering is known [10].Source: §5.1.1 · p. 14

5.1.2 The Gradient Descent Dynamics

  • The true-parent gradient entry equals χ2-mutual information, and data processing makes it larger than non-parent entries, so it grows fastest until its attention weight is approximately 1.The gradient dynamics therefore mimic the oracle algorithm.Source: §5.1.2 · pp. 14–15

5.1 Stage 1: Learning the Causal Graph

  • During stage 1, the first attention layer learns the latent causal graph because its gradient is aligned with the graph’s adjacency matrix.The largest gradient entry for a node corresponds to its parent.Source: §5.1 · p. 13, §5.1.2 · p. 15

5.2 Stage 2: Decreasing the Loss

  • In stage 2, the β-gradient is approximated by an f-mutual information that stays positive until β reaches Θ(log Teff).The proof then uses an approximation to the target conditional distribution to obtain the desired loss bound.Source: §5.2 · p. 16

6 Causal Graphs with Multiple Parents

The multiple-parent setting generalizes the causal-structure task to nodes with k parents, and the paper gives a transformer construction for estimating transitions from those parent sets. Trained transformers empirically converge to this construction, while analyzing the multi-head optimization dynamics remains open.

6 Causal Graphs with Multiple Parents

  • Task 2 allows each non-root node to have k parents, and includes in-context n-gram learning when each token depends on the preceding n − 1 tokens.The task estimates the shared transition in-context after learning the causal structure.Source: §6 · p. 16

B Multiple Parents Construction

  • The multiple-parent construction uses one first-layer head per parent to copy parent tokens, then compares parent tuples to attend to matching contexts and estimate the transition.The resulting empirical transition counts tokens with the target value and matching parent values.Source: §6 · pp. 16–17

6 Causal Graphs with Multiple Parents

  • Transformers trained on Task 2 empirically converge to a construction in which each first-layer head learns one parent set, but the gradient dynamics remain unanalyzed.The analysis challenge is breaking symmetry among heads that attend to different parents.Source: §6 · p. 17

A Disentangled Transformer Equivalence

  • A transformer can be represented by an equivalent disentangled transformer with the same depth and number of heads.The theorem also states a reverse representation with hidden dimension d(L).Source: A Disentangled Transformer Equivalence · p. 22

C Additional Experiments and Details

The additional experiments test whether first-layer attention recovers causal parent structure across graph types and transformer settings. The appendix then develops notation and population-gradient calculations used in its analysis.

C Additional Experiments and Details

  • With a decoder-based transformer, one head per layer, and intervening ReLU MLPs, average first-layer attention approximately matches the causal-graph adjacency matrix.The plotted attention is averaged over 1,024 sequences, using S=10, T=20, and d=30.Source: C Additional Experiments and Details · p. 26
  • Across in-context, Markov-chain, and random latent graphs, the first-layer attention matrices encode graph structure.Figure 6 repeats the comparison across these three graph types.Source: C Additional Experiments and Details · p. 26, Figure 6 · p. 27
  • In multiple-parent tasks, each attention head learns a single set of graph parents, including in 3-gram and 4-gram settings.The experiments vary the number and arrangement of parents across task rows.Source: C Additional Experiments and Details · p. 28
  • 0.837 (SD 0.054) was the mean avgattn across 20 random graphs, with parent-directed attention large across sequence positions.Figure 9 plots the mean parent-directed attention by sequence position, with shaded bars indicating one standard deviation.Source: C Additional Experiments and Details · p. 26, Figure 9 · p. 30

D.1 Proof of Lemma 1

  • The proof presents the first attention-layer output before deriving the second attention-layer output.These successive outputs are the forward-pass calculations used in the proof of Lemma 1.Source: §D.1 · p. 30

D.2 Notation

  • The proof introduces row-wise softmax notation for lower-triangular attention matrices and reparameterizes the second attention matrix.For each row, softmax is applied to the coordinates up to that row's position.Source: §D.2 · p. 31

D.3 Heuristic Derivation of Lemma 3

  • Under the stated assumptions on transition matrix π, long sequences mix so empirical token frequencies approach their stationary probabilities.This is the sequence-mixing step in the heuristic derivation.Source: §D.3 · p. 32

D.4 Gradient Computations

  • The appendix derives population-gradient expressions for the attention matrices.The passage introduces a lemma on population gradients and proceeds to the model gradients for A(1) and A(2).Source: §D.4 · p. 33

D.5 Gradient of A(1) (Stage 1)

This stage analyzes how gradient descent shapes the first attention matrix. It relates an idealized gradient to the latent graph, then studies how the attention weights evolve during training.

D.5 Gradient of A(1) (Stage 1)

  • The idealized gradient is maximized at the true parent p(i), and is zero for root nodes; the proof uses the data processing inequality.For a node with a parent, the idealized gradient at other candidate positions is lower than at p(i).Source: §D.5 · p. 34
  • The true gradient with respect to A(1) is approximated by the idealized gradient, aligning it with the graph adjacency matrix.The overall result states that A(1) converges to the adjacency matrix of G during stage 1.Source: §D.5 · pp. 33–34
  • For most root nodes, A(1) changes little during the analyzed training period.The proof separately bounds the movement of root-node attention rows.Source: §D.5 · p. 36
  • For nodes with parents, the parent-directed attention weight increases during training.The proof establishes that the parent entry remains at least as large as competing entries throughout training.Source: §D.5 · p. 37

D.6 Gradient of A(2) (Stage 2)

Stage 2 analyzes the second attention matrix after stage 1 has encoded the graph. Its dynamics preserve a symmetric form and drive a scalar parameter upward under the stated conditions.

D.6 Gradient of A(2) (Stage 2)

  • With the stage-1 first attention matrix fixed, A(2) retains a symmetric form with equal diagonal and off-diagonal entries.The parameter β describes the change from the initialized matrix during training.Source: §D.6 · pp. 38–39
  • When A(1) approximates the graph adjacency matrix, the stage-2 analysis shows that β grows large.The proof gives a time bound for β to reach the threshold β*.Source: §D.6 · pp. 39, 43
  • Convexity of h_s makes the idealized second-stage gradient nonnegative.The argument expresses the gradient as a linear combination of h_s-divergences.Source: §D.6 · p. 42

E Markov Chain Preliminaries

  • For Markov chains whose transition probabilities are bounded below by γ/S, the spectral quantity λ(π) is at most 1−γ/S.The appendix also states that the stationary measure has minimum mass at least γ/S under this condition.Source: E Markov Chain Preliminaries · p. 45

F Concentration

This section develops concentration bounds for sequences generated by the causal graph. It uses graph structure and Markov-chain mixing to control empirical frequencies and counts.

F Concentration

  • Effective sequence length Teff(λ) summarizes sequence length after accounting for the Markov-chain parameter λ.The section also defines graph distance as shortest-path length in the undirected version of the causal graph, with disconnected pairs at infinite distance.Source: F Concentration · p. 46
  • Under the stated lower bound on transition probabilities, the effective-length analysis gives Teff(λ)^−1 ≲ T^−1.The assumption is min π(s′ | s) ≥ γ/S.Source: F Concentration · p. 47
  • The empirical frequency estimate has mean-squared error bounded by a constant times 1/Teff(λ).The bound compares the empirical token frequency with its stationary probability.Source: F Concentration · p. 48
  • For positions with parents, the empirical pair count has expectation μ_π(s)π(s′|s), and the section bounds its variance.The passages state the expectation and introduce a corresponding concentration lemma.Source: F Concentration · p. 49

G.1 Strong Data Processing Inequality

The supplied proof excerpts give auxiliary inequalities for the analysis, including an application of the strong data processing inequality and bounds on softmax-Jacobian gradients.

G.1 Strong Data Processing Inequality

  • The strong data processing argument uses conditional independence of s_i and s_j given their least common ancestor, p(i, j).For i > j, the excerpt invokes Lemma 17 when d(i, p(i, j)) ≥ 1.Source: §G.1 · p. 51

G.3 Concentration

  • Lemma 24 is stated under a spectral-gap condition and a lower bound on the stationary probability of s′.The excerpt specifies 1 − λ(π) ≥ 1 − λ and μ_π(s′) ≥ γ/S.Source: §G.3 · p. 54

H.2 Auxiliary Dynamics Lemmas

  • For nonnegative u, the softmax Jacobian satisfies ||J(S(v))u||_1 ≤ 2S(v)^T u.The proof also bounds a softmax-gradient norm by 10(1 + β).Source: §H.2 · pp. 56–57

H.4 Proof of Theorem 2

The proof uses a union bound to obtain a probability-0.99 statement.

H.4 Proof of Theorem 2

  • A union bound yields a statement holding with probability 0.99.The supplied excerpt does not specify the event or quantity bounded.Source: §H.4 · p. 66

I Finite Sample Analysis

The appendix sketches how the population-loss analysis extends to finite samples, with both training stages succeeding when the prompt count is sufficiently large.

I Finite Sample Analysis

  • The finite-sample extension is presented as a sketch, while the main theory focuses on gradient descent on population loss.The finite-sample analysis uses an independent dataset for a single large Stage 2 gradient step because the update no longer preserves symmetry.Source: I Finite Sample Analysis · p. 66, §I.2 · p. 68

I.1 Stage 1

  • Stage 1 succeeds with high probability when N ≳ Cγ,S Teff log T.The argument concentrates empirical gradients and uses a union bound over i, j ∈ [T].Source: §I.1 · p. 68

I.2 Stage 2

  • N ≳ Teff log T suffices for both stages to succeed on the finite-sample loss.The Stage 2 output also achieves small population loss.Source: §I.2 · p. 70
Loading 2402.14735v2…