ControlJEPA: Principled Trajectory Regularization via Lyapunov Tube Loss
Authors: Thai-Khanh Nguyen, Nhi Y. Pham, Uyen N. B. Vo, Thieu N. Vo, Cuong Pham, Tan M. Nguyen
Organizations: Faculty of Information Technology 1, Posts & Telecommunications Institute of Technology · Center for AI Research, VinUniversity · Faculty of Artificial Intelligence, Posts & Telecommunications Institute of Technology · Department of Computer Science, University of Bath · Department of Mathematics, National University of Singapore
Next-token prediction leaves hidden-state trajectories largely unconstrained, while existing geometric regularizers primarily control the direction rather than the magnitude of trajectory deviations. We introduce ControlJEPA, a trajectory regularization framework based on a discrete-time Lyapunov contraction condition from practical Input-to-State Stability theory. Its Lyapunov Tube Loss penalizes violations of this condition on a normalized, scale-invariant measure of transversal deviation. When the contraction condition holds at every step, the hidden-state trajectory remains within an explicit tube around the reference secant. In the partial-ISS regime, we bound mean transversal energy, temporal variance, and average tube width using the Lyapunov loss and contraction parameters, with horizon-independent variance bounds under uniformly bounded violations. ControlJEPA requires no architectural modifications. Across six benchmarks and six model families, it consistently improves accuracy and data efficiency over standard fine-tuning and prior trajectory regularizers. The resulting smoother transversal profile also persists empirically during inference, when the regularizer is inactive. The code is publicly available at: https://github.com/t-khanusa/ControlJEPA
Figures & tables
Figure 1: Lyapunov Tube. The displacement dt=ht−htstart is decomposed into a longitudinal component along the reference secant vgeo and a transversal error et⊥vgeo . The Lyapunov Tube Loss penalizes violations of Vt+1≤γVt+τ , confining hidden-state trajectories within the tube. (Definition 3.2 ).
Figure 2: (a) Unembedding norm and transversal energy across λLyap . (b) Impact of contraction rate γ on geometric components and downstream accuracy.
Figure 3: Accuracy and generalization performance of LNTP , LLM-JEPA, STP, and ControlJEPA across (a) Datasets, (b) Model families, and (c) Model sizes.
Figure 4: (a) Data Efficiency and Signal-to-Noise ratio. (b) Training Dynamics and Hyperparameter Robustness.
Figure 5: (a) Inference-time trajectory of squared transversal error ∥et∥22 across generated tokens in NL-RX-SYNTH dataset. (b) Without normalization, the representation scale collapses, shrinking both ∥et∥ and ∥vgeo∥ decrease over time. (Right) Normalizing by local displacement introduces severe heteroscedastic instability and volatile loss dynamics.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Step-by-step evolution of the unnormalized transversal error ∥et∥22 (top panels) and the dimensionless Lyapunov candidate Vt (bottom panels) across multiple generated sequences. Operating in an open-loop regime without any active loss constraints, ControlJEPA (green) uniquely maintains a tight, stable semantic tube. In contrast, regular fine-tuning (orange) exhibits severe geometric drift, and STP (blue) suffers from high trajectory volatility.
Figure 7: The accuracy landscape across diverse datasets demonstrates a broad stability region. Tasks with rigid, rule-based structures (e.g., SYNTH) peak at tighter tube configurations ( γ∈[0.84,0.9] ), while semantically diverse open-domain tasks (e.g., NQ-Open) require a more relaxed envelope ( γ≥0.9 with larger τ ) to accommodate complex manifold variations.