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.
Joint Embedding Predictive Architectures (JEPAs) are a promising paradigm for learning task-agnostic latent world models without visual reconstruction. However, standard JEPA training exhibits a strong inductive bias towards slow features, causing feature suppression and the collapse of latent representation. While inverse dynamics provides temporal anti-collapse, it relies on action labels and offers little incentive to embed general, unlabeled dynamics. We introduce Difference Image and Single image embedding Regularization (DISReg), a novel regularizer that builds on an inverse-dynamics-style module that predicts temporal difference image embeddings without any pixel reconstruction loss, encouraging balanced static and dynamic feature learning. DISReg consists of a static term that shapes the distribution of the image embedding and encourages slow features, and a dynamic term, which, unlike direct regularization on the embedding, imposes no constraint on the image embedding's shape or distribution and instead only incentivizes that dynamic features be present. By integrating this regularizer into a standard JEPA, we establish our new architecture, MotionJEPA. Latent probing demonstrates that MotionJEPA produces more complete representations than other methods, and our trajectory analysis shows it maintains geometrically simple latent embeddings with low curvature. We further show that MotionJEPA improves downstream planning success under static-background distractors across four environments.
Markus Karmann, Shile Li, Christian Internò +10
University of Oxford · vivo Tech Research GmbH · Bielefeld University +5
Large Language Models (LLMs) obey consistent scaling laws -- empirical power-law fits that predict how loss decreases with compute, data, and parameters. While predictive, these laws are descriptive rather than prescriptive: they characterize typical training, not optimal training. Surprisingly few works have successfully challenged the data-efficiency bounds implied by these laws -- which is our primary focus. To that end, we introduce the Geodesic Hypothesis, positing that token sequences trace geodesics on a smooth semantic manifold and are therefore locally linear. Building on this principle, we propose a novel Semantic Tube Prediction (STP) task, a JEPA-style regularizer that confines hidden-state trajectories to a tubular neighborhood of the geodesic. STP generalizes JEPA to language without requiring explicit multi-view augmentations. We show this constraint improves signal-to-noise ratio, and consequently preserves diversity by preventing trajectory collisions during inference. Empirically, STP allows LLMs to match baseline accuracy with 16× less training data on the NL-RX-SYNTH dataset, directly violating the data term of Chinchilla-style scaling laws and demonstrating that principled geometric priors can surpass brute-force scaling. Code is available at https://github.com/galilai-group/llm-jepa#stp.
Fine-tuning a vision-language-action model (VLA-JEPA) on a single GPU should be simple: load a pretrained checkpoint, run training, deploy. There is a hidden danger. Run the same fine-tuning code thirteen times -- same data, same architecture, different random seed -- and twelve runs produce a robot succeeding 91--94% of the time, while one run silently degrades to 65.2%: a 29 pp gap with no error message, no warning, and no way to predict which seed will fail. We call this the seed lottery. We trace the cause to output collapse: the action predictor quietly learns to produce nearly identical outputs regardless of what the robot sees. Existing weight-level methods (L2, EWC) are structurally blind to this collapse -- they penalize weight changes, but collapse occurs in directions weights can move freely without affecting outputs, a gap we formalize via the Jacobian null-space. Across 7 methods x up to 13 seeds x 3 LIBERO benchmarks, three output-level regularizers -- VICReg (n=12 seeds), Dropout (n=4), and a halved learning rate (n=5) -- each eliminate every catastrophic seed (0/21 combined collapses vs. 1/13 Baseline; F(12,11)=28.7, p<0.001), while weight-level methods (L2, EWC) preserve the lottery. The simplest fix is changing one number in your optimizer config.
Jeffrin Sam, Dzmitry Tsetserukou
Skolkovo Institute of Science and Technology (Skoltech), Moscow, Russia