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

TL;DR

This paper introduces TRACE, a framework for causal representation learning that models continuous mechanism transitions as convex combinations of atomic mechanisms, enabling accurate recovery of evolving causal dynamics in real-world systems.

Contribution

The paper formalizes continuous mechanism transitions in causal learning, proves joint identifiability of variables and trajectories, and proposes TRACE, a Mixture-of-Experts model for trajectory recovery.

Findings

TRACE achieves up to 0.99 correlation in recovering mechanism trajectories.

Outperforms discrete-switching baselines on synthetic and real data.

Generalizes to unseen intermediate mechanism states.

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.