How Transformers Learn Causal Structure with Gradient Descent

Eshaan NichaniAlex DamianJason D. Lee

article2024ICML132 citations

Proves that gradient descent enables two-layer transformers to learn latent causal graphs and form induction heads by showing that the attention matrix gradients naturally track token-level mutual information.

Listen

Modern sequence modeling relies heavily on transformer architectures, which exhibit an exceptional capacity for in-context learning—the ability to adapt to new tasks and infer relationships purely from prompt context without updating internal model parameters. While empirical studies show that trained models form specialized internal wiring such as induction heads to capture sequence dependencies, foundational understanding has lagged regarding why and how standard training algorithms actually discover these underlying causal mechanisms from scratch.

To bridge this gap, the article investigates the training dynamics of autoregressive transformers trained with gradient descent on synthetic sequence tasks governed by hidden causal structures. The primary objective is to prove mathematically and demonstrate empirically that standard gradient descent naturally forces a two-layer transformer to uncover latent causal graphs and execute accurate in-context predictions.

The authors analyze an in-context learning task where sequences are generated from latent, directed acyclic graphs over discrete alphabets, encompassing structures such as Markov chains and function classes. Using a simplified two-layer attention-only model, the study evaluates two-stage gradient descent on the population cross-entropy loss, tracing the evolution of attention weights across training iterations. These theoretical derivations are paired with empirical simulations across varied causal tree graphs and multi-parent structures using standard neural network optimization.

The article establishes several key findings. First, gradient descent provably drives the first attention layer to converge directly to the adjacency matrix of the hidden causal graph, successfully attending to true parent tokens with probability approaching one. Second, the mechanism driving this alignment is informational: the mathematical gradient of the attention matrix naturally computes the mutual information between token pairs, which, due to information theory constraints, strictly peaks at direct causal parent connections. Third, the second attention layer learns to aggregate matching historical contexts, forming an induction circuit that achieves near-optimal statistical predictions with vanishing loss errors scaling inversely with effective sequence length. Fourth, the resulting model generalizes effectively to out-of-distribution transition distributions, achieving bounded estimation error with high probability. Finally, the authors show that in tasks where tokens depend on multiple causal parents, multi-head architectures distribute individual parent associations across distinct heads.

These findings provide rigorous theoretical justification for why attention mechanisms excel at structured sequence modeling. Rather than acting as black boxes, transformer layers optimized via gradient methods systematically mirror classical tree-recovery algorithms by extracting maximum mutual information. This mechanistic insight reduces architectural uncertainty and confirms that standard optimization reliably guides attention layers toward robust, interpretable causal representations without requiring explicit graph supervision.

Organizations developing sequence models should leverage multi-layer and multi-head attention structures tailored to the dependency complexity of their domain, ensuring context length and head capacity match the underlying data-generating graphs. For complex dependency modeling, practitioners should ensure adequate prompt length relative to vocabulary size to allow mixing and accurate in-context estimation. Future engineering and research efforts should explore the optimization dynamics of multi-head symmetry breaking under multiple parent dependencies and extend these dynamic guarantees to deeper architectures with non-linear feedforward layers.

The analysis relies on specific theoretical boundary conditions, including tree-structured latent dependencies, two-layer attention-focused architectures, and stage-wise gradient updates on well-behaved stationary measures. Confidence in the core finding—that gradient descent naturally aligns attention weights with maximum mutual information—is very high for the studied regimes, though readers should exercise appropriate caution when extrapolating specific convergence rates directly to massive, highly non-linear production models without further empirical validation.

Princeton University attention mechanism encoder-decoder markov chains causal structure gradient descent attention-only transformer

eshnich/transformers-learn-causal-structure Nichani et al (2024).pdf

Paper: Attention Is All You Need, Ashish Vaswani et al. (2017). Introduces the foundational Transformer architecture and multi-head self-attention mechanism whose optimization dynamics and causal structure learning are theoretically analyzed in the source.
Paper: Quantifying Attention Flow in Transformers, Samira Abnar et al. (2020). Establishes how information propagates across Transformer layers using directed acyclic graph representations of attention, providing key context for understanding layer-wise causal graph encoding.
Paper: What Does BERT Look at? An Analysis of BERT’s Attention, Kevin Clark et al. (2019). Analyzes how individual attention heads specialize into functional sub-circuits, directly relating to the induction heads and attention structures analyzed in the source.

Paper: One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context, Skanda Athreya et al. (2026). Extends the theoretical study of gradient descent dynamics on shallow attention models from learning latent causal graphs to learning in-context nearest-neighbor classification rules.
Paper: Next-Latent Prediction Transformers Learn Compact World Models, Jayden Teoh et al. (2025). Builds upon how Transformers learn internal causal graphs and representations by introducing auxiliary latent prediction objectives to reinforce compact world model representations.
Paper: Learn from your own latents and not from tokens: A sample-complexity theory, Daniel J. Korchinski et al. (2026). Provides a complementary sample-complexity theory for how neural networks learn latent hierarchical structures during gradient-based training rather than simple token prediction.
Paper: Test-time regression: a unifying framework for designing sequence models with associative memory, Ke Alexander Wang et al. (2025). Generalizes the mechanisms behind in-context associative recall and induction phenomena into a formal test-time regression framework across sequence model architectures.

Cover for How Transformers Learn Causal Structure with Gradient Descent

Abstract

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.

Topics

attention mechanism encoder-decoder markov chains mutual informations training algorithms latent factor latent variable

Keywords

causal structure gradient descent attention-only transformer optimization dynamics causal self-attention causal graph induction heads linear attention decoder-based transformer Markov chains

1. Introduction
1.1. Our Contributions
1.2. Related Work
2. Setup
2.1. Transformer Architecture
2.2. Problem Setup: Random Sequences with Causal Structure
2.3. Examples
3. What does the Transformer Learn?
3.1. Experiments
3.2. Construction
3.3. The Reduced Model
4. Main Results
4.1. Training Algorithm
4.2. Main Theorem
5. Proof Sketch
5.1. Stage 1: Learning the Causal Graph
5.1.1. THE ORACLE ALGORITHM
5.1.2. THE GRADIENT DESCENT DYNAMICS
5.2. Stage 2: Decreasing the Loss
6. Causal Graphs with Multiple Parents
Impact Statement
Acknowledgements
References
A. Disentangled Transformer Equivalence
B. Multiple Parents Construction
C. Additional Experiments and Details
D. Analyzing the Dynamics
D.1. Proof of Lemma 3.2
D.2. Notation
D.3. Heuristic Derivation of Lemma 5.3
D.4. Gradient Computations
D.5. Gradient of A (1) (Stage 1)
D.6. Gradient of A (2) (Stage 2)
D.7. Proof of Theorem 4.4
E. Markov Chain Preliminaries
F. Concentration
G. Lemmas for Stage 1
G.1. Strong Data Processing Inequality
G.2. Auxiliary Dynamics Lemmas
G.3. Concentration
H. Lemmas for Stage 2
H.1. Idealized Gradient
H.2. Auxiliary Dynamics Lemmas
H.3. Concentration
H.4. Proof of Theorem 4.5
I. Finite Sample Analysis
I.1. Stage 1
I.2. Stage 2

Knowls

Knowl 1 — Convergence and Causal Structure Recovery Guarantee for Two-Layer Transformers
theoretical result
Let $G = ([T], E)$ be a directed acyclic graph representing a latent causal structure where each node $i \in [T]$ has at most one parent $p(i) \in [i-1]$ , and $R = \{i \in [T] : p(i) = \emptyset\}$ denotes the set of root nodes. Let $T_{\text{eff}} := T / \max_i L_i$ be the effective sequence length, where $L_i$ is the number of leaves in tree $T_i$ of the disjoint tree decomposition $G = \bigcup_{i=1}^k T_i$ .

Consider the reduced two-layer autoregressive transformer $f_\theta(s_{1:T}) = X^\top S(S(A^{(1)}) X A^{(2)} x_T) \in \mathbb{R}^S$ , where $X = [e_{s_1}, \dots, e_{s_T}]^\top \in \mathbb{R}^{T \times S}$ is the sequence token embedding matrix, $A^{(1)} \in \mathbb{R}^{T \times T}$ is a lower-triangular attention matrix, $A^{(2)} \in \mathbb{R}^{S \times S}$ is a token-token attention matrix, and $S(\cdot)$ denotes the row-wise softmax operator. The model is trained to minimize the perturbed population cross-entropy loss
$L(\theta) = -\mathbb{E}_{\pi, s_{1:T}} \left[ \sum_{s' \in [S]} \pi(s' \mid s_T) \log (f_\theta(s_{1:T})_{s'} + \epsilon) \right]$
with perturbation $\epsilon = T_{\text{eff}}^{-1/2}$ , where the sequence $s_{1:T}$ is generated from a Markov transition matrix $\pi$ sampled from prior $P_\pi$ .

Assume the prior $P_\pi$ satisfies regularity conditions with parameter $\gamma > 0$ , the root fraction satisfies $|R|/T \le 1 - \gamma$ , and $T_{\text{eff}} \ge \text{poly}(\gamma^{-1}, S)$ . When training $f_\theta$ via two-stage gradient descent initialized at $A^{(1)}(0) = 0_{T \times T}$ and $A^{(2)}(0) = \beta_0 I_S$ with $\beta_0 \le c_{\gamma, S} T_{\text{eff}}^{-3/2}$ , there exist step sizes $\eta_1, \eta_2$ and step counts $\tau_1, \tau_2$ such that the resulting parameter $\hat{\theta} = (\hat{A}^{(1)}, \hat{A}^{(2)})$ satisfies:
1. Graph Recovery in First Layer: For all non-root nodes $i \notin R$ , the first attention layer places almost all attention weight on the true parent $p(i)$ :
$S(\hat{A}^{(1)})_{i, p(i)} \ge 1 - O\left(\frac{1}{T}\right).$
1. Population Loss Convergence: The population loss reaches the minimal possible loss $L^* := -\mathbb{E}_\pi [\frac{1}{S} \sum_{s, s'} \pi(s' \mid s) \log \pi(s' \mid s)]$ up to a polynomial error in effective context length:
$L(\hat{\theta}) - L^* \le C_{\gamma, S} \left( \frac{\log^2 T}{T_{\text{eff}}} \right)^{c\gamma / 24}$
for a universal constant $c > 0$ .
Knowl 2 — Random Sequences with Latent Causal Structure
definition
Let $[S] = \{1, \dots, S\}$ be a finite alphabet, and let $G = ([T], E)$ be a directed acyclic graph over positions $[T] = \{1, \dots, T\}$ representing a latent causal structure, where directed edges $(j \to i) \in E$ only exist if $j < i$ . For each position $i \in [T]$ , $p(i) := \{j : (j \to i) \in E\}$ denotes the set of parent nodes, and $R := \{i \in [T] : p(i) = \emptyset\}$ denotes the set of root nodes, with $|p(i)| \le 1$ for all $i \in [T]$ .

Let $P_\pi$ be a prior distribution over irreducible and aperiodic Markov transition matrices $\pi$ on $[S]$ , and let $\mu_\pi$ denote the unique stationary distribution of $\pi$ . A random sequence with causal structure $G$ and target token $y$ is generated according to the following procedure:
1. Sample a transition matrix $\pi \sim P_\pi$ .
2. For $i = 1, \dots, T - 1$ $i = 1, \dots, T - 1$ :
  - If $p(i) = \emptyset$ , sample $s_i \sim \mu_\pi$ .
  - If $p(i) = \{j\}$ , sample $s_i \sim \pi(\cdot \mid s_j)$ .
3. Draw the context query token $s_T \sim \text{Unif}([S])$ and the target token $y = s_{T+1} \sim \pi(\cdot \mid s_T)$ . (Hence $T \in R$ is always a root node).
4. Return prompt sequence $x = s_{1:T} = (s_1, \dots, s_T)$ and target label $y = s_{T+1}$ .
Special cases include:
- In-context Markov chain estimation (induction heads): $p(i) = i - 1$ for all $i > 1$ .
- In-context function learning: $p(2k - 1) = \emptyset$ and $p(2k) = 2k - 1$ for $k \ge 1$ , where transitions encode input-output pairs $(x_k, f(x_k))$ .
Knowl 3 — Attention Layer Gradients Compute Conditional Chi-Squared Mutual Information
theoretical result
For a two-layer transformer with first-layer lower-triangular attention matrix $A^{(1)} \in \mathbb{R}^{T \times T}$ and initialization scale $A^{(2)} = \beta_0 I_S$ , the population gradient with respect to row $i$ of $A^{(1)}$ is given by
$\nabla_{A^{(1)}_i} L(\theta) = -\frac{\beta_0}{S T} J(S(A^{(1)}_i)) \left( g_i + O(T_{\text{eff}}^{-1/2}) \right)$
where $J(v) = \text{diag}(v) - v v^\top$ is the Jacobian matrix of the softmax operator $S(\cdot)$ , and the entries of the idealized gradient vector $g_i \in \mathbb{R}^i$ are
$g_{i, j} := \mathbb{E}_\pi \left[ \sum_{s, s' \in [S]} \frac{\pi(s' \mid s)}{\mu_\pi(s')} \mathbb{P}[s_j = s, s_i = s'] \right] - 1.$
For any non-root position $i \notin R$ with true parent $p(i)$ :
1. The gradient entry at the true parent equals the conditional $\chi^2$ -mutual information between $s_i$ and $s_{p(i)}$ given $\pi$ :
$g_{i, p(i)} = I_{\chi^2}(s_i ; s_{p(i)} \mid \pi) := \mathbb{E}_\pi \left[ \sum_{s, s'} \frac{(\mathbb{P}[s_i = s', s_{p(i)} = s \mid \pi])^2}{\mu_\pi(s') \mu_\pi(s)} - 1 \right].$
1. For any non-parent position $j \ne p(i)$ , the sequence conditioned on $\pi$ forms a Markov chain $s_j \to s_{p(i)} \to s_i$ . By the Data Processing Inequality for $\chi^2$ -divergence, the mutual information is strictly strictly contracted:
$g_{i, p(i)} - g_{i, j} \ge \frac{1 - \alpha(\pi)}{2} \|B_\pi\|_F^2 \ge \frac{\gamma^3}{2 S}$
where $\alpha(\pi) \le 1 - \gamma$ is the Markov contraction coefficient and $\|B_\pi\|_F^2 \ge \gamma^2/S$ is the non-degeneracy of the chain.
1. For root positions $i \in R$ , $s_i$ is independent of all preceding tokens $s_j$ given $\pi$ , leading to $g_{i, j} = 0$ , which preserves an approximately uniform attention distribution $S(A^{(1)}_i)_j \approx 1/i$ for all $j \le i$ .
Knowl 4 — Out-of-Distribution Generalization of Learned Causal Transformers
theoretical result
Let $\hat{\theta} = (\hat{A}^{(1)}, \hat{A}^{(2)})$ be the two-layer transformer parameters obtained from running two-stage gradient descent (Algorithm 1) on sequences generated with prior distribution $P_\pi$ .

Let $\tilde{\pi} \in \mathbb{R}^{S \times S}$ be an arbitrary target Markov transition matrix that may lie completely outside the support of the training prior $P_\pi$ , satisfying only the entrywise lower bound
$\min_{s, s' \in [S]} \tilde{\pi}(s' \mid s) \ge \frac{\gamma}{S}.$
If a sequence $s_{1:T}$ is generated according to the latent causal DAG $G$ using the out-of-distribution transition matrix $\tilde{\pi}$ , then with probability at least $0.99$ over the draw of $s_{1:T}$ , the prediction $f_{\hat{\theta}}(s_{1:T}) \in \mathbb{R}^S$ uniformly approximates the true transition distribution $\tilde{\pi}(\cdot \mid s_T)$ :
$\sup_{s' \in [S]} \left| f_{\hat{\theta}}(s_{1:T})_{s'} - \tilde{\pi}(s' \mid s_T) \right| \le C_{\gamma, S} \frac{\log T}{T_{\text{eff}}^{c\gamma}}$
where $T_{\text{eff}}$ is the effective sequence length of $G$ , and $c > 0, C_{\gamma, S} > 0$ are constants independent of $T$ .
Knowl 5 — Two-Stage Gradient Descent for Causal Transformer Training
algorithm
The training procedure uses stage-wise gradient descent on the population cross-entropy loss $L(\theta)$ using a reduced two-layer transformer model $f_\theta$ , parameterized by lower-triangular position attention matrix $A^{(1)} \in \mathbb{R}^{T \times T}$ and token attention matrix $A^{(2)} \in \mathbb{R}^{S \times S}$ .

In Stage 1, $A^{(1)}$ is trained with learning rate $\eta_1$ for $\tau_1$ steps while $A^{(2)}$ is frozen at small initialization $\beta_0 I_S$ , driving $A^{(1)}$ to encode the causal adjacency matrix. In Stage 2, $A^{(2)}$ is trained with learning rate $\eta_2$ for $\tau_2$ steps while $A^{(1)}$ is held fixed at $A^{(1)}(\tau_1)$ , growing the magnitude of $A^{(2)}$ along $I_S - \frac{1}{S} \mathbf{1}_S \mathbf{1}_S^\top$ to minimize predictive cross-entropy.
Input: initialization scale $\beta_0$ , learning rates $\eta_1, \eta_2$ , step counts $\tau_1, \tau_2$
Initialize $A^{(1)}(0) = 0_{T \times T}$ , $A^{(2)}(0) = \beta_0 \cdot I_{S \times S}$
for $t = 1, \dots, \tau_1$ do
$A^{(1)}(t) \leftarrow A^{(1)}(t-1) - \eta_1 \nabla_{A^{(1)}} L(A^{(1)}(t-1), A^{(2)}(0))$
end for
for $t = \tau_1, \dots, \tau_1 + \tau_2 - 1$ do
$A^{(2)}(t+1) \leftarrow A^{(2)}(t) - \eta_2 \nabla_{A^{(2)}} L(A^{(1)}(\tau_1), A^{(2)}(t))$
end for
$\hat{\theta} \leftarrow (A^{(1)}(\tau_1), A^{(2)}(\tau_1 + \tau_2))$
Output: $\hat{\theta}$
Knowl 6 — Reduced Two-Layer Disentangled Transformer Model
model/method
Let $[S]$ be a finite alphabet of size $S$ , and let $s_{1:T} = (s_1, \dots, s_T) \in [S]^T$ be an input sequence. Let $X = [e_{s_1}, \dots, e_{s_T}]^\top \in \mathbb{R}^{T \times S}$ denote the one-hot token embedding matrix of the sequence.

The reduced two-layer transformer model $f_\theta : [S]^T \to \mathbb{R}^S$ parameterized by $\theta = (A^{(1)}, A^{(2)})$ with lower-triangular matrix $A^{(1)} \in \mathbb{R}^{T \times T}$ and symmetric matrix $A^{(2)} \in \mathbb{R}^{S \times S}$ is defined as:
$f_\theta(s_{1:T}) = X^\top S\left( S(\text{MASK}(A^{(1)})) X A^{(2)\top} e_{s_T} \right)$
where $\text{MASK}(M)_{i,j} = M_{i,j}$ for $i \ge j$ and $-\infty$ for $i < j$ , and $S(\cdot)$ is the softmax function applied row-wise.

To prevent infinite loss when tokens do not appear in context, predictions are perturbed by $\epsilon = T_{\text{eff}}^{-1/2}$ . The perturbed population cross-entropy loss under prior $P_\pi$ is
$L(\theta) = -\mathbb{E}_{\pi, s_{1:T}} \left[ \sum_{s' \in [S]} \pi(s' \mid s_T) \log (f_\theta(s_{1:T})_{s'} + \epsilon) \right].$
Knowl 7 — Equivalence of Decoder-Based and Disentangled Transformers
theoretical result
A standard decoder-based attention-only transformer $\text{TF}_\theta$ of depth $L$ with embedding dimension $d$ , token embeddings $E \in \mathbb{R}^{d \times S}$ , positional embeddings $P \in \mathbb{R}^{d \times T}$ , head parameters $(Q_i^{(\ell)}, K_i^{(\ell)}, V_i^{(\ell)}) \in (\mathbb{R}^{d \times d})^3$ for $i \in [m_\ell]$ , and output projection $W_O \in \mathbb{R}^{d_{\text{out}} \times d}$ computes residual updates $h^{(\ell)} = h^{(\ell-1)} + \sum_{i=1}^{m_\ell} \text{attn}(h^{(\ell-1)}; Q_i^{(\ell)} K_i^{(\ell)\top}) V_i^{(\ell)\top}$ .

A disentangled transformer $\widetilde{\text{TF}}_{\tilde{\theta}}$ initializes $h^{(0)} = \tilde{X} = [e_{s_t}, e_t]_{t=1}^T \in \mathbb{R}^{T \times d_0}$ with $d_0 = S + T$ , and at each layer $\ell \in [L]$ concatenates the attention outputs along the feature dimension to produce $h^{(\ell)} = [h^{(\ell-1)}, \text{attn}(h^{(\ell-1)}; \tilde{A}_1^{(\ell)}), \dots, \text{attn}(h^{(\ell-1)}; \tilde{A}_{m_\ell}^{(\ell)})] \in \mathbb{R}^{T \times d_\ell}$ with $d_\ell = (1 + m_\ell)d_{\ell-1}$ , parameterized by single attention matrices $\tilde{A}_i^{(\ell)} \in \mathbb{R}^{d_{\ell-1} \times d_{\ell-1}}$ and output projection $\tilde{W}_O \in \mathbb{R}^{d_{\text{out}} \times d_L}$ .

For any standard transformer $\text{TF}_\theta$ , there exists a disentangled transformer $\widetilde{\text{TF}}_{\tilde{\theta}}$ of the same depth and number of heads such that $\text{TF}_\theta(s_{1:T}) = \widetilde{\text{TF}}_{\tilde{\theta}}(s_{1:T})$ for all $s_{1:T} \in [S]^T$ . Conversely, any disentangled transformer $\widetilde{\text{TF}}_{\tilde{\theta}}$ can be exactly represented by a standard transformer $\text{TF}_\theta$ with hidden dimension $d = d_L$ .
Knowl 8 — Exact Construction for In-Context Causal Transition Estimation
model/method
For a latent causal DAG $G = ([T], E)$ with parent function $p(i)$ for $i \notin R$ and root nodes $R = \{i : p(i) = \emptyset\}$ , there exists a two-layer disentangled transformer whose output approximates the empirical transition estimator
$\hat{\pi}_{s_{1:T}}(s' \mid s_T) := \frac{|\{(j \to i) \in E : (s_j, s_i) = (s_T, s')\}|}{|\{(j \to i) \in E : s_j = s_T\}|}.$
The construction sets attention matrices $\tilde{A}^{(1)}$ and $\tilde{A}^{(2)}$ as:
$\tilde{A}^{(1)} = \begin{bmatrix} 0_{S \times S} & 0_{S \times T} \\ 0_{T \times S} & A^{(1)} \end{bmatrix}, \quad \tilde{A}^{(2)} = \begin{bmatrix} 0_{S \times S} & 0_{S \times T} & A^{(2)} & 0_{S \times T} \\ 0_{T \times S} & 0_{T \times T} & 0_{T \times S} & 0_{T \times T} \\ 0_{S \times S} & 0_{S \times T} & 0_{S \times S} & 0_{S \times T} \\ 0_{T \times S} & 0_{T \times T} & 0_{T \times S} & 0_{T \times T} \end{bmatrix}$
where $A^{(1)}_{i, j} = \beta_1 \mathbf{1}(j = p(i))$ for $i \notin R$ (and $0$ for $i \in R$ ), and $A^{(2)} = \beta_2 I_S$ , with $\beta_1, \beta_2 \to \infty$ .

In the forward pass:
1. The first attention layer copies the parent token embedding $\tilde{x}_{p(i)}$ into position $i$ for non-root positions $i \notin R$ , and averages all prior tokens $\frac{1}{i}\sum_{j \le i} \tilde{x}_j$ for root positions $i \in R$ .
2. The second attention layer compares the final token $s_T$ against the copied parent tokens in the residual stream, attending equally to all tokens $i$ where $s_{p(i)} = s_T$ .
3. The output layer $\tilde{W}_O = [0_{S \times d}, 0_{S \times d}, I_S, 0_{S \times T}, 0_{S \times d}]$ extracts the averaged token vector, yielding $f_\theta(s_{1:T})_{s'} = \hat{\pi}_{s_{1:T}}(s' \mid s_T)$ .
Knowl 9 — In-Context Estimation for Multi-Parent Causal Graphs and n-Grams
model/method
Let $G = ([T+1], E)$ be a directed acyclic graph where each non-root node $i \in [T+1] \setminus R$ has exactly $k$ parent nodes $p(i) = \{p(i)_1, \dots, p(i)_k\} \subset [i-1]$ with $p(i)_1 < \dots < p(i)_k$ . Transitions are governed by $k$ -parent transition tensors $\pi(\cdot \mid a_1, \dots, a_k)$ sampled from prior $P_\pi^k$ . For $n$ -gram models, $k = n - 1$ and $p(i) = \{i - n + 1, \dots, i - 1\}$ .

A two-layer transformer with $k$ heads in the first layer computes the empirical multi-parent transition estimate
$\hat{\pi}_{s_{1:T}}(s' \mid s_{p(T+1)_1}, \dots, s_{p(T+1)_k}) = \frac{|\{i < T : s_i = s', s_{p(i)_1} = s_{p(T+1)_1}, \dots, s_{p(i)_k} = s_{p(T+1)_k}\}|}{|\{j < T : s_{p(j)_1} = s_{p(T+1)_1}, \dots, s_{p(j)_k} = s_{p(T+1)_k}\}|}.$
In this construction:
- The $\ell$ -th attention head in layer 1 uses scaled position-position weights $(A_\ell^{(1)})_{i, j} = \beta_1 \mathbf{1}(j = p(i)_\ell)$ for $i < T$ and $(A_\ell^{(1)})_{T, j} = \beta_1 \mathbf{1}(j = p(T+1)_\ell)$ , copying the $\ell$ -th parent of $i$ into the residual stream of $i$ , and the $\ell$ -th parent of $T+1$ into the residual stream of $T$ .
- The second attention layer uses block diagonal weights $A^{(2)} = \text{diag}(0_{d \times d}, \beta_2 I_S, \dots, \beta_2 I_S)$ to compute inner products $\beta_2 \sum_{\ell=1}^k \mathbf{1}(s_{p(i)_\ell} = s_{p(T+1)_\ell})$ , selecting tokens whose parent tuples match $(s_{p(T+1)_1}, \dots, s_{p(T+1)_k})$ .
Knowl 10 — Assumptions on Transition Prior and Graph Mixing
assumption
The theoretical convergence guarantees for causal structure recovery and in-context learning require the following structural conditions on the prior $P_\pi$ over Markov transition matrices $\pi$ on $[S]$ and on the causal graph $G = ([T], E)$ :
1. Transition Lower Bound: Almost surely over $\pi \sim P_\pi$ , $\min_{s, s' \in [S]} \pi(s' \mid s) > \gamma / S$ for a constant $\gamma > 0$ , ensuring stationary measures satisfy $\min_s \mu_\pi(s) \ge \gamma / S$ .
2. Non-Degeneracy: The Markov chain does not mix in a single step, satisfying $\sum_s \|\pi(\cdot \mid s) - \mu_\pi(\cdot)\|_2^2 \ge \gamma^2 / S$ .
3. Permutation Symmetry and Uniform Mean: For any permutation $\sigma$ on $[S]$ , $\sigma^{-1} \pi \sigma \stackrel{d}{=} \pi$ , and $\mathbb{E}_\pi[\pi] = \frac{1}{S} \mathbf{1}_S \mathbf{1}_S^\top$ .
4. Non-Vanishing Non-Root Fraction: The fraction of root nodes $r := |R|/T$ satisfies $r \le 1 - \gamma$ .
5. Effective Sequence Length: For graph decomposition $G = \bigcup_{i=1}^k T_i$ into disjoint trees where tree $T_i$ has $L_i$ leaves, the effective sequence length $T_{ ext{eff}} := T / \max_i L_i$ satisfies $T_{ ext{eff}} \ge \text{poly}(\gamma^{-1}, S)$ .

Coverage note — Omitted detailed finite-sample concentration proofs (Section I) and specific empirical heatmap figures, as their primary takeaways and sample complexity scaling ($N \gtrsim T_{\text{eff}} \log T$) are directly reflected in the theoretical and architectural knowls.

References

1.Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.00297, 2023.
2.Aky"urek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=0g0X4H8yN4I.
3.Aky"urek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Arhitectures and algorithms. arXiv preprint arXiv:2401.12973, 2024.
4.Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637, 2023.
5.Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L. Birth of a transformer: A memory viewpoint. arXiv preprint arXiv:2306.00802, 2023.
6.Boix-Adsera, E., Littwin, E., Abbe, E., Bengio, S., and Susskind, J. Transformers learn through gradual rank increase. arXiv preprint arXiv:2306.07042, 2023.
7.Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.
8.Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877–1901, 2020.
9.Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084–15097, 2021.
10.Chow, C. and Liu, C. Approximating discrete probability distributions with dependence trees. IEEE transactions on Information Theory, 14(3):462–467, 1968.
11.Cohen, J. E., Iwasa, Y., Rautu, G., Beth Ruskai, M., Seneta, E., and Zbaganu, G. Relative entropy under mappings by stochastic matrices. Linear Algebra and its Applications, 179:211–235, 1993. ISSN 0024-3795. doi: https://doi.org/10.1016/0024-3795(93)90331-H. URL https://www.sciencedirect.com/science/article/pii/002437959390331H.
12.Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
13.Edelman, B. L., Edelman, E., Goel, S., Malach, E., and Tsilivis, N. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint arXiv:2402.11004, 2024.
14.Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https://transformer-circuits.pub/2021/framework/index.html.
15.Friedman, D., Wettig, A., and Chen, D. Learning transformer programs. arXiv preprint arXiv:2306.01128, 2023.
16.Fu, D., Chen, T.-Q., Jia, R., and Sharan, V. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17086, 2023.
17.Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583–30598, 2022.
18.Giannou, A., Rajput, S., Sohn, J.-y., Lee, K., Lee, J. D., and Papailiopoulos, D. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196, 2023.
19.Huang, Y., Cheng, Y., and Liang, Y. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
20.Jelassi, S., Sander, M., and Li, Y. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822–37836, 2022.
21.Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Žídek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583–589, 2021.
22.Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization, 2017.
23.Li, Y., Li, Y., and Risteski, A. How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning, pp. 19689–19729. PMLR, 2023.
24.Lu, H., Mao, Y., and Nayak, A. On the dynamics of training attention models. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=1OCTOShAmqB.
25.Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023.
26.Nguyen, T. and Grover, A. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In International Conference on Machine Learning, pp. 16569–16594. PMLR, 2022.
27.Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads. Transformer Circuits Thread, 2022. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
28.Reddy, G. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. arXiv preprint arXiv:2312.03002, 2023.
29.Snell, C., Zhong, R., Klein, D., and Steinhardt, J. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601, 2021.
30.Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. Transformers as support vector machines. arXiv preprint arXiv:2308.16898, 2023a.
31.Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. Max-margin token selection in attention mechanism. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b.
32.Tian, Y., Wang, Y., Chen, B., and Du, S. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305.16380, 2023a.
33.Tian, Y., Wang, Y., Zhang, Z., Chen, B., and Du, S. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535, 2023b.
34.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
35.Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151–35174. PMLR, 2023.
36.Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
37.Zhang, R., Frei, S., and Bartlett, P. L. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023.

Citation

MLA

Nichani, E., et al. “How Transformers Learn Causal Structure with Gradient Descent”. arXiv, 2024, https://doi.org/10.48550/arxiv.2402.14735.

APA

Nichani, E., Damian, A., & Lee, J. D. (2024). How Transformers Learn Causal Structure with Gradient Descent. arXiv. https://doi.org/10.48550/arxiv.2402.14735

Chicago

Nichani, E., A. Damian, and J. D. Lee. 2024. “How Transformers Learn Causal Structure with Gradient Descent”. Preprint, ArXiv. https://doi.org/10.48550/arxiv.2402.14735.

Harvard

Nichani, E., Damian, A. and Lee, J.D. (2024) “How Transformers Learn Causal Structure with Gradient Descent”. arXiv. Available at: https://doi.org/10.48550/arxiv.2402.14735.

Vancouver

1. Nichani E, Damian A, Lee JD (2024) How Transformers Learn Causal Structure with Gradient Descent. https://doi.org/10.48550/arxiv.2402.14735

BibTeX

@misc{https://doi.org/10.48550/arxiv.2402.14735,
  doi = {10.48550/ARXIV.2402.14735},
  url = {https://arxiv.org/abs/2402.14735},
  author = {Nichani, Eshaan and Damian, Alex and Lee, Jason D.},
  keywords = {Machine Learning (cs.LG), Information Theory (cs.IT), Machine Learning (stat.ML), FOS: Computer and information sciences, FOS: Computer and information sciences},
  title = {How Transformers Learn Causal Structure with Gradient Descent},
  publisher = {arXiv},
  year = {2024},
  copyright = {arXiv.org perpetual, non-exclusive license}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

How Transformers Learn Causal Structure with Gradient Descent

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Convergence and Causal Structure Recovery Guarantee for Two-Layer Transformers

Knowl 2 — Random Sequences with Latent Causal Structure

Knowl 3 — Attention Layer Gradients Compute Conditional Chi-Squared Mutual Information

Knowl 4 — Out-of-Distribution Generalization of Learned Causal Transformers

Knowl 5 — Two-Stage Gradient Descent for Causal Transformer Training

Knowl 6 — Reduced Two-Layer Disentangled Transformer Model

Knowl 7 — Equivalence of Decoder-Based and Disentangled Transformers

Knowl 8 — Exact Construction for In-Context Causal Transition Estimation

Knowl 9 — In-Context Estimation for Multi-Parent Causal Graphs and n-Grams

Knowl 10 — Assumptions on Transition Prior and Graph Mixing

References

Citation

BibTeX

Source Code

Access the Paper