ControlJEPA: Principled Trajectory Regularization via Lyapunov Tube Loss

Next-token prediction leaves hidden-state trajectories largely unconstrained, motivating recent geometric regularization methods based on angular penalties. However, these approaches control only the direction of motion and provide no explicit bound on trajectory deviation. We introduce the Controlled Joint Embedding Predictive Architecture (ControlJEPA), a trajectory regularization framework grounded in discrete-time Input-to-State Stability theory. Our novel Lyapunov Tube Loss enforces a contraction condition on a normalized measure of transversal deviation, yielding a provable guarantee that confines hidden states within a tube of explicit width determined by two interpretable hyperparameters. ControlJEPA integrates seamlessly with standard training and requires no architectural changes. Across six benchmarks and seven model families, it consistently outperforms standard fine-tuning and prior trajectory regularizers, improving accuracy and data efficiency while preserving output diversity. Notably, the induced geometric structure persists at inference time, indicating that the method reshapes the representation space rather than memorizing training trajectories.