Your Autoregressive Model Already Reveals the Causal Graph

AI summary

Authors: Hugo Math, Rainer Lienhart (Department of Machine Learning & Computer Vision, University of Augsburg, Germany) arXiv: 2602.01135v4 [cs.LG], 09 Jun 2026 Keywords: Machine Learning, ICML

Overview

Research area: Causal discovery in discrete event sequences, using pretrained autoregressive (AR) models as density estimators. The paper sits at the intersection of causal inference, information theory, and sequence modelling.

Technical level: Advanced. The paper relies on conditional mutual information, structural causal models, faithfulness assumptions, and information-theoretic error bounds, and assumes familiarity with concepts such as DAGs, back-door paths, and do-operators.

Scope (one sentence): The paper introduces TRACE, a framework that turns any pretrained autoregressive model into a causal discovery engine for a single observed sequence over a massive event vocabulary, with theoretical guarantees and evaluation on synthetic SCMs and real-world vehicle diagnostic logs.

What This Paper Is About

In many systems — a patient's disease trajectory, a manufacturing line's tests, a vehicle's diagnostic codes — one observes only a single stream of discrete events and wants to know which events causally influence which others. Classic causal discovery methods (PCMCI, DYNOTEARS, VARLiNGAM, TCDF, Hawkes processes) assume either many parallel streams or long multivariate series, and they scale poorly with variable count, so they become inaccurate or computationally intractable when the vocabulary is massive. TRACE's goal is to recover the causal graph within one observed sequence by reusing the learned probability distributions of an already-pretrained AR model, with no task-specific retraining.

Key Contributions

1. Prediction–Causality Duality. The authors prove that causal identifiability is achievable for any AR model that sufficiently approximates the data-generating process. They relax the standard Oracle CI-test assumption to an ϵ-regime and derive bounds showing the causal graph is recoverable up to a noise floor determined by the model's convergence. Consequently, minimizing the standard cross-entropy pretraining loss directly minimizes an upper bound on causal identification error.

2. Amortized Single-Sequence Discovery. TRACE is presented as the first framework designed to recover the summary causal graph from a single high-dimensional sequence, by amortizing the learning of system dynamics through a pretrained AR model. The paper describes this as an underexplored area.

3. Backbone Agnosticism. The framework strictly decouples density estimation (Phase 1) from causal discovery (Phase 2), so TRACE can leverage any state-of-the-art autoregressive backbone (Transformers, Mamba, RNNs).

4. Real-World Applicability. The method is validated on synthetic SCMs with challenging vocabulary and memory requirements, and applied to real-world vehicle diagnostic logs.

Main Findings

• Single-stream performance: On synthetic SCMs with |𝒳| = 1000, L = 64, ϵ = 0.05, τ = 3·10⁻⁵ and N = 128, TRACE (LLaMA) reaches F1 = 0.91 ± 0.01 with SHD 28.6 ± 2.8 and precision 0.89 ± 0.01, versus the strongest baseline (Neural Granger) at F1 = 0.69 ± 0.04, SHD 100.2 ± 14.6 and precision 0.71 ± 0.04. Across the paper's abstract-level results, this is reported as outperforming the strongest baseline by over 20 F1 points.

• Attention is a poor causal proxy: Attention (BERT) achieves F1 = 0.50 ± 0.01, precision 0.35 ± 0.01, SHD 321.0 ± 15 — consistent with the paper's claim that attention scores fail to distinguish correlation from causation.

• Attribution methods fall short: Saliency (Input x Gradient on LLaMA) reaches F1 = 0.67 ± 0.01, precision 0.51 ± 0.01, SHD 160.2 ± 6.55. Shapley Value Sampling (LLaMA) reaches F1 = 0.60 ± 0.01, precision 0.55 ± 0.01, SHD 148.0 ± 5.09. Naive baselines are near-zero (Random: F1 = 0.01 ± 0.00; Frequency: F1 = 0.09 ± 0.00), indicating the task is non-trivial.

• CMI is a more robust causality signal than probability fluctuation: The paper states that measuring the conditional mutual information is far more robust than monitoring the probability fluctuation of a single target token (Neural Granger) or relying on metrics not anchored in causal discovery (saliency, attention scores).

• Stability across vocabulary size: In the scalability analysis (ϵ̂ = 0.01, L = 64, N = 64, τ = 10⁻⁴), discovery performance is reported as stable across vocabulary size, with F1 remaining at approximately 0.81 even as the search space grows exponentially and SCM entropy increases. Predictability scores (Pred = H(P)/H_max) are used to confirm success in high-entropy regimes.

• Identifiability precedes full convergence: The paper states that causal structure becomes identifiable well before the autoregressive model fully converges.

• Error vanishes with model accuracy: Theorem 4.3 bounds the asymptotic error of the true CMI as lim sup_{N→∞} |I − Î_N| ≤ 2√(ϵ/2) ln(2) + 2(1 + √(ϵ/2)) h_b(√(ϵ/2)/(1 + √(ϵ/2))), where h_b is the binary entropy function. As ϵ → 0, the causal identification error vanishes.

• Scale of evaluation: Synthetic nonlinear SCMs use |𝒳| ∈ [200, 8000]; real-world vehicle diagnostic logs use |𝒳| = 29,100. The paper states TRACE is the first applicable method at this scale.

• Real-world vehicle logs: The abstract reports applying TRACE to vehicle diagnostic logs at |𝒳| = 29,100. The provided content is truncated before the per-metric results for this dataset, so specific numbers are not reported here.

Methodology in Plain English

The framework has two phases.

Phase 1 — learn the dynamics. An autoregressive model (the paper's experiments use a LLaMA architecture) is trained by next-token prediction on a corpus of event sequences. This is standard pretraining: the model outputs the probability of the next event given everything before it. The paper's key premise is that this ordinary training objective already forces the model to encode the conditional independence structure of the underlying data-generating process.

Phase 2 — extract the graph. The pretrained model is frozen and used as a density estimator. For a single observed sequence, TRACE asks, for each candidate pair of events at times t and t′, whether knowing the earlier event changes the model's predicted distribution of the later event, given the observed history. This is measured as conditional mutual information, which reduces to a comparison of two distributions: the baseline prediction and the prediction when the earlier event's context is fixed. If the difference exceeds a threshold, an edge exists.

To detect delayed effects, TRACE simulates interventions on the intermediate events (mediators) between a candidate cause and effect — a randomized do-operator that averages over counterfactual realizations sampled uniformly from the vocabulary — so spurious paths through mediators are blocked.

The computation is restructured as a single broadcasted tensor: each row fixes progressively more of the observed history while randomizing the future in a staircase pattern, and the model processes all rows in parallel on GPUs. Adjacent rows are then compared to produce the intervention-versus-baseline distribution difference. Histories are truncated to a context of c = max(0.1L, 20) to make parallel sampling feasible on GPUs.

The core tension is cost: testing every time step against every prior time step requires L(L+1)/2 CI-tests, which grows quadratically with sequence length and becomes difficult even on multiple GPUs for L > 100. A sparse variant bounds lagged influence to a memory window m ≪ L, reducing complexity from 𝒪(N·(L−c)·L·|𝒳|) to 𝒪(N·m·L·|𝒳|) — linear in the sequence length and in the vocabulary size.

The theoretical grounding includes assumptions of temporal precedence, causal sufficiency (no hidden confounders), and an ϵ-Oracle model in which the KL divergence between the true next-token distribution and the model's is at most ϵ. Identifiability uses ϵ-Strong Faithfulness — an active edge must have true CMI exceeding 2τ_ϵ, i.e. it must dominate the model's approximation error. The resulting graph, the Sample Time Causal Graph, is projected onto event types to yield the Sample Summary Causal Graph, which is permitted to contain cycles because it is an abstraction of the time-unrolled graph.

Evaluation on synthetic SCMs proceeds by atomic interventions: uniformly randomizing an event and measuring the average KL divergence over 10 counterfactuals between post-intervention and observational distributions, with a threshold of τ > 0.05 for declaring an edge, then comparing to ground truth via Precision, Recall, and Structural Hamming Distance.

Why This Matters

Impact on research. The paper reframes causal discovery as a byproduct of ordinary sequence prediction. If minimizing cross-entropy training loss minimizes an upper bound on causal identification error, then causal structure is not something that needs a separate, specialized training pipeline — it is already latent in any well-trained autoregressive model. This also reframes what scale buys: the paper argues the causal structure becomes identifiable well before the model fully converges, and that the resulting CI tests scale linearly with vocabulary size, putting regimes previously considered intractable within reach.

Real-world applications:

• Vehicle diagnostics: the paper's own real-world evaluation uses diagnostic trouble codes logged asynchronously over time, at |𝒳| = 29,100 event types, to infer which codes cause which.
• Healthcare: electronic health records encode temporal sequences of symptoms, test results, and treatments, where understanding event-to-event causation supports diagnosis and downstream decision-making.
• Manufacturing: a production line's tests form an event sequence where root-cause analysis depends on knowing which events influence others.
• Session-based user interactions: the paper positions its setting as the "session-based" regime common to NLP and system logs, where many short independent sequences share a massive vocabulary.

Industry relevance. Because TRACE is backbone-agnostic and requires no task-specific retraining, organizations that already have a pretrained sequence model for logs, telemetry, or user sessions can apply the method as an inference-time procedure. The sparse variant's linear scaling with sequence length and vocabulary is what makes deployment viable in domains where the event alphabet runs into the tens of thousands — the exact regime the paper's comparison table (Table 1) marks as out of reach for constraint-based, score-based, Granger, noise-based, Hawkes/TPP, and prior information-theoretic methods.

Future Directions

• Relaxing causal sufficiency. The randomized interventional do-operator only marginalizes out intermediate mechanisms if no hidden confounders exist (Assumption 3.5). Handling unobserved confounders in the single-stream setting remains open; the paper defers analysis of assumption violations to Appendix D.

• Beyond the truncated context. The context truncation to c = max(0.1L, 20) is acknowledged as a potential break of the Markovian assumption for long sequences. The paper reports empirical robustness and an ablation on unseen sequence lengths, but a principled treatment of long-range dependence is left open.

• Scaling to longer sequences. Because the full method requires L(L+1)/2 CI-tests, sequences with L > 100 become difficult even on multiple GPUs. The sparse memory-window variant addresses this, but how the window choice interacts with genuinely long-range causes is a natural next question.

• Extending beyond the evaluated backbones and domains. The paper emphasizes architectural modularity and evaluates with LLaMA-style autoregressive models plus a BERT attention baseline. Whether the same guarantees hold for other backbones named in the paper (Mamba, RNNs) and for other event-log domains beyond vehicle diagnostics is an open empirical question; the provided content is truncated before any explicit future-work section.

Target Audience

Researchers and practitioners working on causal discovery, causal inference, or sequence modelling who deal with discrete event data — especially those facing high-cardinality vocabularies and single-stream observations. It will be most useful to readers comfortable with conditional mutual information, structural causal models, and information-theoretic bounds. Applied engineers in diagnostics, healthcare analytics, manufacturing, and log analysis will find the architectural modularity and linear scaling the most immediately actionable parts, though they should expect to engage with the theoretical sections to interpret the threshold τ_ϵ and the ϵ-Oracle assumption correctly.

Authors’ abstract

Autoregressive models trained via next-token prediction implicitly learn the conditional independence structure of their data-generating process. We exploit this observation to perform scalable causal discovery from a single observed sequence of discrete events -- without any task-specific retraining. Such single-stream settings arise naturally in vehicle diagnostics, manufacturing systems, and patient trajectories, yet they remain largely unsolved: the absence of repeated samples, massive event vocabularies, and long-range temporal dependencies render existing methods either inaccurate or computationally intractable. We introduce TRACE, a framework that repurposes any pretrained autoregressive model as a density estimator for conditional mutual information, the fundamental primitive for conditional independence testing. By constructing parallelized CI tests on GPUs, TRACE recovers both the sample-level time causal graph and its summary projection, scaling linearly with the vocabulary size while naturally handling delayed causal effects. Crucially, we prove that minimizing the standard cross-entropy pretraining loss directly minimizes an upper bound on the causal identification error, establishing a duality between sequence prediction and causal discovery. On nonlinear SCMs (|X| = 8000) and real-world vehicle diagnostic logs (|X| = 29100), TRACE is the first applicable method at this scale, outperforming the strongest baseline by over 20 F1 points.