TRACE: Scalable Amortized Causal Discovery from Single Sequences via Autoregressive Density Estimation

TL;DR

The paper tackles causal discovery from a single high-dimensional sequence of discrete events, a setting where repeated samples are unavailable and long-range dependencies complicate inference. It introduces TRACE, which repurposes pretrained autoregressive density estimators to estimate conditional mutual information and perform parallel, instance-level causal discovery, yielding a summary causal graph over event types with linear complexity in vocabulary size. The authors establish identifiability under an epsilon-strong faithfulness framework when the autoregressive model approximates the true distribution (an $\epsilon$-oracle), and they demonstrate robustness to imperfect models, long horizons, and high-dimensional vocabularies. Empirically, TRACE outperforms baselines on synthetic data and scales to real-world vehicle diagnostics with $|\mathcal{X}|\approx 29{,}000$, enabling scalable, GPU-accelerated causal discovery for industrial and healthcare sequences.

Abstract

We study causal discovery from a single observed sequence of discrete events generated by a stochastic process, as encountered in vehicle logs, manufacturing systems, or patient trajectories. This regime is particularly challenging due to the absence of repeated samples, high dimensionality, and long-range temporal dependencies of the single observation during inference. We introduce TRACE, a scalable framework that repurposes autoregressive models as pretrained density estimators for conditional mutual information estimation. TRACE infers the summary causal graph between event types in a sequence, scaling linearly with the event vocabulary and supporting delayed causal effects, while being fully parallel on GPUs. We establish its theoretical identifiability under imperfect autoregressive models. Experiments demonstrate robust performance across different baselines and varying vocabulary sizes including an application to root-cause analysis in vehicle diagnostics with over 29,100 event types.

Figures (11)

• Figure 1: Methodological Shift.(A) Traditional Causal Discovery in Sequences (e.g., PCMCI, Hawkes, Granger) relies on iterative solvers (CI-tests) over long multivariate time series ($T \to \infty$). (B) Our TRACE Approach processes a single sequence (e.g., event logs, user interactions, patient trajectories) through a pretrained autoregressive (AR) model as density estimator to compute the Conditional Mutual Information (CMI) in parallel, enabling scalable causal discovery over massive vocabularies ($|\mathcal{X}| > 1000$).
• Figure 2: TRACE Methodology.Phase 1 (Training): An autoregressive (AR) model (e.g., LM, RNN) is pretrained on a corpus of event sequences via next-token prediction to learn the process dynamics ($P_\theta$). Phase 2 (Inference): A single sequence $s$ is passed through the frozen model. We then estimate conditional mutual information (Parallelized CMI module) to prune non-causal edges and form the Instance Time Causal Graph$\mathcal{G}_{t,s}$. Finally, this graph is projected onto the event types to recover the Summary Causal Graph$\mathcal{G}_s$.
• Figure 3: Overview of TRACE Parallel CI-tests. We construct a single broadcasted tensor $\mathbf{X}_{do}$ where each row $j$ incrementally fixes the history $x_{\le j}$ while randomizing the future (staircase pattern). The model processes this tensor in parallel to produce raw probabilities $\mathbf{P}_{raw}$ (grey). We then compute the Causal Mutual Information by comparing adjacent rows: the distribution at row $j-1$ serves as the baseline ($\mathbf{P}_{base}$, blue) for the intervention at row $j$ ($\mathbf{P}_{do}$, red).
• Figure 4: Scalability to High-Dimensional Event Spaces. Evaluation of structural identifiability across exponentially growing vocabulary sizes. Top: Evolution of discovery metrics. TRACE exhibits performance invariance, maintaining F1 $\approx 0.81$ even as the combinatorial search space explodes. Bottom: Visual examples of recovered summary graphs $\mathcal{G}_s$ at scale. Predictability scores (Pred = $H(P)/H_{max}$) confirm that TRACE succeeds even in high-entropy regimes. ($\hat{\epsilon}=0.01, L=64, N=64, \tau=10^{-4}$).
• Figure 5: Robustness and Scalability Analysis ($|\mathcal{X}|=1000, N=128, \tau=10^{-4}, L=64$). Evolution of causal discovery performance (F1, Precision, Recall, SHD). (a) Robustness to Generative Error: Performance as a function of the model's oracle score $\epsilon$. TRACE exhibits a phase transition, recovering structure even for imperfect models ($\epsilon < 0.1$) and maintaining high Precision even as fidelity degrades. (b) Scalability to Length: Performance and GPU memory usage vs. sequence length $L$. The Sparse variant demonstrates linear memory scaling ($O(mL)$), enabling inference on sequences far exceeding the training length ($L=64$), whereas the Full variant scales quadratically. (c) Long-Range Dependencies: Robustness to increasing delayed-effects $m$. TRACE maintains F1 $>0.8$ even as dependencies span one third of the sequence ($m=20$), confirming the method's ability to capture distant causal mechanisms.

Theorems & Definitions (18)

• Definition 3.1: Instance Time Causal Graph
• Definition 3.2: Instance Summary Causal Graph
• Remark 3.3: Cyclicity in Summary Graphs
• Remark 4.1
• Proposition 4.2: Convergence to the $\epsilon$-Proxy
• Theorem 4.3: Total Error Bound in the $\epsilon$-Regime
• Definition 4.4: $\epsilon$-Strong Faithfulness
• Lemma 4.5: Identifiability of the Instance Time Causal Graph
• Definition 4.6: Randomized Interventional Do-Operator
• Remark 4.7