Table of Contents
Fetching ...

Transformer Is Inherently a Causal Learner

Xinyue Wang, Stephen Wang, Biwei Huang

TL;DR

This work demonstrates that decoder-only transformers trained for autoregressive forecasting inherently learn time-delayed causal structure in their representations. By introducing the score-gradient energy $H_{j,i}^{\ell}$ and aggregating gradient attributions via Layer-wise Relevance Propagation (LRP), the authors establish causal identifiability under standard assumptions and provide a practical graph-extraction pipeline that surpasses state-of-the-art causal-discovery methods across nonlinear, long-range, and non-stationary dynamics. The approach enables a data-efficient, scalable perspective on causal discovery, and situates foundation models as interpretable causal learners when viewed through the lens of gradient-based causality. Importantly, the framework supports extensions with domain indicators and latent-variable post-processing, promoting robustness and transfer across diverse environments. Overall, the paper reframes causal discovery as a by-product of scalable representation learning, with significant implications for interpretable, environment-aware foundation-model development.

Abstract

We reveal that transformers trained in an autoregressive manner naturally encode time-delayed causal structures in their learned representations. When predicting future values in multivariate time series, the gradient sensitivities of transformer outputs with respect to past inputs directly recover the underlying causal graph, without any explicit causal objectives or structural constraints. We prove this connection theoretically under standard identifiability conditions and develop a practical extraction method using aggregated gradient attributions. On challenging cases such as nonlinear dynamics, long-term dependencies, and non-stationary systems, this approach greatly surpasses the performance of state-of-the-art discovery algorithms, especially as data heterogeneity increases, exhibiting scaling potential where causal accuracy improves with data volume and heterogeneity, a property traditional methods lack. This unifying view lays the groundwork for a future paradigm where causal discovery operates through the lens of foundation models, and foundation models gain interpretability and enhancement through the lens of causality.

Transformer Is Inherently a Causal Learner

TL;DR

This work demonstrates that decoder-only transformers trained for autoregressive forecasting inherently learn time-delayed causal structure in their representations. By introducing the score-gradient energy and aggregating gradient attributions via Layer-wise Relevance Propagation (LRP), the authors establish causal identifiability under standard assumptions and provide a practical graph-extraction pipeline that surpasses state-of-the-art causal-discovery methods across nonlinear, long-range, and non-stationary dynamics. The approach enables a data-efficient, scalable perspective on causal discovery, and situates foundation models as interpretable causal learners when viewed through the lens of gradient-based causality. Importantly, the framework supports extensions with domain indicators and latent-variable post-processing, promoting robustness and transfer across diverse environments. Overall, the paper reframes causal discovery as a by-product of scalable representation learning, with significant implications for interpretable, environment-aware foundation-model development.

Abstract

We reveal that transformers trained in an autoregressive manner naturally encode time-delayed causal structures in their learned representations. When predicting future values in multivariate time series, the gradient sensitivities of transformer outputs with respect to past inputs directly recover the underlying causal graph, without any explicit causal objectives or structural constraints. We prove this connection theoretically under standard identifiability conditions and develop a practical extraction method using aggregated gradient attributions. On challenging cases such as nonlinear dynamics, long-term dependencies, and non-stationary systems, this approach greatly surpasses the performance of state-of-the-art discovery algorithms, especially as data heterogeneity increases, exhibiting scaling potential where causal accuracy improves with data volume and heterogeneity, a property traditional methods lack. This unifying view lays the groundwork for a future paradigm where causal discovery operates through the lens of foundation models, and foundation models gain interpretability and enhancement through the lens of causality.
Paper Structure (67 sections, 7 theorems, 20 equations, 25 figures, 14 tables)

This paper contains 67 sections, 7 theorems, 20 equations, 25 figures, 14 tables.

Key Result

Theorem 1

Under A1--A4 and regularity conditions, the lagged causal graph $\mathcal{G}^*$ is uniquely identifiable via the score gradient energy: edge $j\,\stackrel{\ell}{\longrightarrow}\, i$ exists iff $H_{j,i}^{\ell} := \mathbb{E}[(\partial_{x_{j,t-\ell}} \log p(X_{i,t} \mid X_{<t}))^2] > 0$.

Figures (25)

  • Figure 1: Data generation and transformer-based causal discovery.Left: A decoder-only transformer trained for next-step prediction. Tokens are lagged observations from $t\!-\!L$ to $t\!-\!1$; the model predicts $X_t$ from $X_{t-1:t-L}$. Right: A lagged data-generating process with $N\!=\!3$ and window $L\!=\!3$. Each $X_{i,t}$ depends on selected past values $X_{j,t-\ell}$ per the true graph $\mathcal{G}^*$. The trained transformer learns the process, and relevance attributions help recover the causal structure.
  • Figure 2: F1 score analysis across regimes.(A) Mean F1 across all experiments (averages exclude timeout cases). (B) High-dimensional input: F1 averaged across scales and seeds vs. the number of nodes. (C) Long-range dependencies: F1 averaged across scales and seeds vs. maximum lag. (D) Nonlinearity: F1 averaged across scales and seeds vs. different types of functional forms. (E) Non-stationarity: F1 averaged across scales and seeds vs. the number of domains. We run each method with three seeds. Missing results indicate method timeouts due to computational limits. DOT stands for Decoder-only Transformer. PL and M stand for piecewise linear and monotonic functions.
  • Figure 3: Nonlinear dependencies. F1 scores averaged across seeds vs. sample size in different nonlinear settings.
  • Figure 4: Non-stationary dependencies. F1 scores averaged across seeds vs. sample size in different non-stationary settings.
  • Figure 5: Robustness to latent variables and noise.Left: F1 scores on scenarios including different amounts of latent variables. Right: F1 scores on different kinds of noise (equal variance and non-equal variance).
  • ...and 20 more figures

Theorems & Definitions (16)

  • Theorem 1
  • Definition 1: Lagged Causal Structure Identifiability
  • Lemma 1: Zero weak partial implies no dependence
  • proof
  • Lemma 2: Conditional exogeneity implies Causal Markov
  • proof
  • Lemma 3: Zero score partial and conditional independence
  • proof
  • Theorem 2: Score-based characterization of lagged parents
  • proof
  • ...and 6 more