REVIEW 5 cited by
How Transformers Learn Causal Structure with Gradient Descent
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
read the original 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.
Forward citations
Cited by 5 Pith papers
-
Metastable Dynamics of Chain-of-Thought Reasoning: Provable Benefits of Search, RL and Distillation
In a metastable Markov model of reasoning, search and RL improve hitting times by boosting rare hard edges, and distilling clusters into a meta-chain preserves escape dynamics.
-
Transformers as Multi-task Learners: Decoupling Features in Hidden Markov Models
A hand-constructed O(ln L + T)-layer Transformer is shown to approximate low-rank hidden Markov models in-context, with lower layers extracting local history features and upper layers performing regression-based decoding.
-
The Features at Convergence Theorem: a first-principles alternative to the Neural Feature Ansatz for how networks learn representations
FACT is a first-order stationarity identity for weight matrices that matches or beats the Neural Feature Ansatz as a description of learned features at convergence.
-
Spectral Journey: How Transformers Predict the Shortest Path
Two-layer transformers learn shortest paths on small graphs by building embeddings that correlate with spectral decomposition of the line graph, yielding an approximate spectral path-finding algorithm.
-
Transformers and Their Roles as Time Series Foundation Models
Transformers can implement autoregressive least-squares regression in-context on time series, and pretraining on weakly dependent data gives test error decaying as 1 divided by the square root of the number of pretrai...
Discussion (0). Continue with ORCID to comment.