TRACE: Trajectory Recovery for Continuous Mechanism Evolution in Causal Representation Learning
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 test-time recovery of mechanism trajectories, including 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.
Lay Summary
Many systems in the real world change their behavior gradually rather than flipping a switch. For example, when a car makes a turn, the relationships between steering, speed, and direction shift smoothly from straight-line driving to cornering. Similarly, a person's gait transitions continuously from walking to running. Current AI methods for understanding these hidden cause-and-effect relationships assume that changes happen suddenly between a fixed set of modes, missing the smooth transitions in between. We developed TRACE, a framework that models these gradual changes as smooth paths between a small number of basic behavior patterns. Think of it like color mixing: just as any color can be created by blending red, green, and blue, any intermediate behavior can be represented as a blend of a few fundamental patterns. Our method first learns these fundamental patterns from labeled examples, then figures out the blending recipe for new, unlabeled data where patterns shift over time. We proved mathematically that this blending recipe can be recovered reliably, and showed it works on both simulated systems and real data from vehicle turns and human walking-to-running transitions. This could help engineers and scientists detect early warning signs before a system's behavior fundamentally changes.