Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TRACE: Trajectory Recovery for Continuous Mechanism Evolution in Causal Representation Learning

Shicheng Fan, Kun Zhang, Lu Cheng

TL;DR

TRACE addresses the challenge of continuous mechanism evolution in causal representation learning by modeling intermediate mechanism states as convex combinations of a finite set of atomic mechanisms and proving identifiability of both the latent variables and the continuous mixing trajectory. It introduces a two-stage Mixture-of-Experts framework: Stage 1 learns a shared encoder and domain-specific experts to recover disentangled latents from pure-domain data, and Stage 2 recovers the mixing trajectory at test time via least-squares projection onto a basis formed from domain-conditioned expectations, augmented with temporal smoothing for stability. Theoretical guarantees include latent identifiability up to permutation and component-wise transformations, plus finite-sample and minimax-optimal error bounds for trajectory recovery that improve with trajectory smoothness. Empirically, TRACE achieves up to 0.99 correlation in mixing-trajectory recovery and up to 0.96 latent identifiability on synthetic data, while outperforming discrete-switch baselines on semi-synthetic CartPole and real-world vehicle and gait datasets, illustrating its practical potential for continuous mechanism-aware analysis.

Abstract

Temporal causal representation learning methods assume that causal mechanisms switch instantaneously between discrete domains, yet real-world systems often exhibit continuous mechanism transitions. For example, a vehicle's dynamics evolve gradually through a turning maneuver, and human gait shifts smoothly from walking to running. We formalize this setting by modeling transitional mechanisms as convex combinations of finitely many atomic mechanisms, governed by time-varying mixing coefficients. Our theoretical contributions establish that both the latent causal variables and the continuous mixing trajectory are jointly identifiable. We further propose TRACE, a Mixture-of-Experts framework where each expert learns one atomic mechanism during training, enabling recovery of mechanism trajectories at test time. This formulation generalizes to intermediate mechanism states never observed during training. Experiments on synthetic and real-world data demonstrate that TRACE recovers mixing trajectories with up to 0.99 correlation, substantially outperforming discrete-switching baselines.

TRACE: Trajectory Recovery for Continuous Mechanism Evolution in Causal Representation Learning

TL;DR

TRACE addresses the challenge of continuous mechanism evolution in causal representation learning by modeling intermediate mechanism states as convex combinations of a finite set of atomic mechanisms and proving identifiability of both the latent variables and the continuous mixing trajectory. It introduces a two-stage Mixture-of-Experts framework: Stage 1 learns a shared encoder and domain-specific experts to recover disentangled latents from pure-domain data, and Stage 2 recovers the mixing trajectory at test time via least-squares projection onto a basis formed from domain-conditioned expectations, augmented with temporal smoothing for stability. Theoretical guarantees include latent identifiability up to permutation and component-wise transformations, plus finite-sample and minimax-optimal error bounds for trajectory recovery that improve with trajectory smoothness. Empirically, TRACE achieves up to 0.99 correlation in mixing-trajectory recovery and up to 0.96 latent identifiability on synthetic data, while outperforming discrete-switch baselines on semi-synthetic CartPole and real-world vehicle and gait datasets, illustrating its practical potential for continuous mechanism-aware analysis.

Abstract

Temporal causal representation learning methods assume that causal mechanisms switch instantaneously between discrete domains, yet real-world systems often exhibit continuous mechanism transitions. For example, a vehicle's dynamics evolve gradually through a turning maneuver, and human gait shifts smoothly from walking to running. We formalize this setting by modeling transitional mechanisms as convex combinations of finitely many atomic mechanisms, governed by time-varying mixing coefficients. Our theoretical contributions establish that both the latent causal variables and the continuous mixing trajectory are jointly identifiable. We further propose TRACE, a Mixture-of-Experts framework where each expert learns one atomic mechanism during training, enabling recovery of mechanism trajectories at test time. This formulation generalizes to intermediate mechanism states never observed during training. Experiments on synthetic and real-world data demonstrate that TRACE recovers mixing trajectories with up to 0.99 correlation, substantially outperforming discrete-switching baselines.
Paper Structure (113 sections, 6 theorems, 66 equations, 9 figures, 12 tables, 1 algorithm)

This paper contains 113 sections, 6 theorems, 66 equations, 9 figures, 12 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Formulation
  4. Generative Model
  5. Mechanism Parameterization
  6. Problem Statement
  7. Identifiability Theory
  8. Identifiability of Latent Causal Processes
  9. Recovery of Mixing Coefficients
  10. Method
  11. Stage 1: MoE-based Representation Learning
  12. Stage 2: Mechanism Trajectory Inference
  13. Experiments
  14. Experimental Setup
  15. Synthetic Experiments
  16. ...and 98 more sections

Key Result

Theorem 4.1

Suppose $\mathbf{x}_t = g(\mathbf{z}_t)$ where $g$ is invertible, and the conditional distribution $p(z_{k,t} \mid \mathbf{z}_{t-1})$ varies across $K$ domains. Under sufficient variability conditions (Appendix app:thm1), any learned representation $\hat{\mathbf{z}}_t$ satisfying conditional indepen

Figures (9)

  • Figure 1: Continuous mechanism transitions as trajectories through a simplex. Atomic mechanisms $\{C_0, C_1, C_2\}$ are vertices; any interior point $s_t$ is a convex combination inducing a mixed causal graph (bottom). Red and dashed curves show distinct trajectories between the same endpoints, which discrete-switching models cannot distinguish.
  • Figure 2: Architecture of TRACE. Stage 1 (Training): Shared encoder and one-hot gating route to domain-specific experts. Stage 2 (Inference): Mixing coefficients $\boldsymbol{\alpha}$ recovered via least-squares projection (Algorithm \ref{['alg:inference']}).
  • Figure 3: Mechanism trajectory recovery with $K_{\text{active}}=3$ (domains 0, 2, 4). Calibrated predictions accurately track sequential domain activations including the intermediate peak (Corr $= 0.99$).
  • Figure 4: Empirical validation of Theorem \ref{['thm:recovery']}. Left: Recovery errors vs. theoretical bounds (all points below $y=x$). Middle:$\sigma_{\min}$ under varying noise and perturbation. Right: MAE vs. SNR$_{\mathrm{eff}}$.
  • Figure 5: Evaluation on modified CartPole. Top: Frames across domains showing distinct dynamics. Bottom left: Correlation matrix confirms latent identifiability (MCC $= 0.970$). Bottom right: Recovered mixing coefficients (Corr. $0.95$) capture the sequential activation pattern.
  • ...and 4 more figures

Theorems & Definitions (22)

  • Theorem 4.1: Identifiability of Latent Variables
  • Theorem 4.2: Pointwise Recovery
  • Theorem 4.3: Smooth Trajectory Recovery
  • Remark 4.4: Scale Calibration
  • Remark 4.5: Verifiability
  • Remark 4.6: Geometric Bottleneck
  • Definition 1.1: Component-wise Invertible Transformation
  • Definition 1.2: Identifiable Latent Causal Processes
  • proof
  • Lemma 1.3: Parameter Independence Implies Sufficient Variability
  • ...and 12 more