ControlJEPA: Principled Trajectory Regularization via Lyapunov Tube Loss

AI summary

Overview

Research area: Machine learning — specifically representation learning and trajectory regularization for large language models, drawing on discrete-time Input-to-State Stability (ISS) theory from control theory.

Technical level: Advanced. The paper combines a control-theoretic stability framework (Lyapunov functions, ISS contraction, steady-state bounds) with transformer fine-tuning experiments, and its core arguments depend on the proofs of Proposition 4.1 and Theorem 4.2.

Scope: The paper introduces ControlJEPA, a training-time regularizer that constrains the magnitude of hidden-state deviation from a reference direction in LLM representation space, with theory bounding mean deviation, temporal variance, and tube width, validated across six benchmarks and six model families.

What This Paper Is About

Next-token prediction optimizes each token position independently and places no constraint on the shape of the hidden-state trajectory that the model traces through representation space. Prior fixes such as Semantic Tube Prediction (STP) and Temporal Straightening penalize angular deviation between trajectory segments, so a trajectory can drift arbitrarily far from the intended path in magnitude without being penalized. ControlJEPA instead borrows a discrete-time Lyapunov contraction condition from ISS theory and penalizes violations of it on a scale-invariant measure of perpendicular ("transversal") deviation, giving an explicit, interpretable bound on how wide the tube around the reference path can be.

Key Contributions

1. A stability-based framework for trajectory regularization. The authors recast hidden-state regularization as a discrete-time ISS problem, giving explicit and interpretable bounds on trajectory deviation in terms of hyperparameters, rather than heuristic geometric penalties.
2. The Lyapunov Tube Loss. A practical, architecture-agnostic objective that penalizes violations of the contraction condition V_{t+1} ≤ γV_t + τ on a normalized transversal error, and is zero when the condition already holds. It is added to the standard NTP loss with a weight λ, requiring no architectural change.
3. Theory for the "partial-ISS" regime. Theorem 4.2 bounds the time-averaged Lyapunov value (mean confinement, Eq. 27), the temporal variance of the Lyapunov sequence (Eq. 28), and the average tube width (Eq. 29), with a horizon-independent variance bound whenever the maximal violation is uniformly bounded in T (remainder O(1/T)). Proposition 4.1 gives a spectral capacity bound linking vanishing transversal deviation to limited logit expressivity under a bounded unembedding norm.
4. Empirical validation across six benchmarks and six model families. ControlJEPA is compared against standard fine-tuning, LLM-JEPA, and STP, with gains in accuracy and data efficiency, a stable hyperparameter region, and a smoother transversal profile that persists at inference when the regularizer is inactive.

Main Findings

• Accuracy gains on reasoning tasks: On GSM8K and NQ-Open, ControlJEPA improves accuracy over standard fine-tuning by approximately 10 percentage points on each benchmark, and over Semantic Tube Prediction by approximately 6 and 4 percentage points respectively. The paper reports that gains are most pronounced on these complex reasoning tasks, where STP "suffers from unconstrained error magnitudes."
• Consistency across benchmarks, architectures, and scales: The reported advantage holds across NL-RX-SYNTH, NL-RX-TURK, GSM8K, Spider, NQ-Open, and HellaSwag, across Llama-3.2 (1B/3B/8B), Gemma-2-2b-it, OpenELM-1.1B, OLMo-2-1B, Qwen3-1.7B, and DeepSeek-R1-Distill-Qwen-1.5B, and scales smoothly across Llama-3 from 1B to 8B.
• Improved data efficiency: In moderate data regimes ControlJEPA outperforms standard fine-tuning and STP, and it matches full-dataset NTP performance using only a fraction of the data (the specific fraction is not reported in the truncated text). The authors attribute this to a higher median signal-to-noise ratio, defined as the ratio of aligned displacement energy to orthogonal residual energy.
• Optimization stability and a broad hyperparameter region: The NTP loss curves overlap whether trained alone or jointly with ControlJEPA. The Lyapunov loss stabilizes near 0.69, matching the softplus minimum ln(1+e^0) ≈ 0.693 once the ISS condition is satisfied. Accuracy remains consistently high for λ in [0.01, 0.08].
• Inference-time tube preservation: On unseen NL-RX-SYNTH trajectories with the penalty deactivated, standard fine-tuning shows substantial drift, STP reduces mean error but remains highly oscillatory, and ControlJEPA achieves the lowest and most stable squared transversal error.
• Strict contraction is harmful: With γ = 0.8 (approximating STP's strict geodesic constraint), |e_t| is driven toward zero and downstream accuracy falls to 28.7%. With γ = 0.95, the system reaches a joint equilibrium with positive slack and |v_geo| grows substantially, yielding 86.78% accuracy.
• Normalization is essential: Penalizing the unnormalized transversal energy causes representation scale collapse (d_t → 0, v_geo → 0) and reduces accuracy to 6.3%. Normalizing by local displacement instead introduces heteroscedastic instability, because ||d_t||² grows over time.

Methodology in Plain English

The method works on the hidden state at the final transformer layer, before the unembedding projection to vocabulary logits.

1. Pick two anchor states inside a contiguous semantic segment (the paper follows STP in relying on dataset-provided segmentation). The straight line between them is the "reference secant" v_geo, the direction a well-behaved trajectory should roughly follow.
2. Split each intermediate state's displacement from the start anchor into a component along that line (longitudinal progress) and a perpendicular leftover (the transversal error e_t).
3. Turn the error into a dimensionless number: V_t = ||e_t||² / ||v_geo||². Dividing by the chord length makes V_t scale-invariant, so the model can reduce relative deviation either by shrinking the perpendicular error or by spreading the anchors further apart — avoiding the collapse that plain error minimization causes.
4. Penalize only violations: The stability condition says V should shrink by factor γ each step plus an allowed slack τ. A ReLU or softplus activation is applied to V_{t+1} − γV_t − τ, so the loss is exactly (or nearly) zero when the condition already holds. Only valid interior transitions of each segment are penalized.
5. Add it to the language modeling loss with weight λ: L = L_NTP + λ·L_Lyap. No architecture changes; the only additions are one projection and a scalar comparison.

The theory explains why the constraint is applied softly: driving the loss to zero would force transversal error toward zero, which Proposition 4.1 shows collapses the effective rank of the logits toward a rank-≤2 structure. In practice the loss stays well above zero (the "partial-ISS regime"), and what the regularizer buys is a stable temporal profile — uniform rather than oscillating distance from the reference line — not a suppressed level.

Experimental protocol: Baselines (standard fine-tuning, LLM-JEPA, STP) use their published, dataset-tuned optimal epochs and learning rates; ControlJEPA fixes those shared hyperparameters and tunes only λ, γ, and τ, which the authors describe as conservatively favoring the baselines. All methods share a codebase, splits, decoding, and validation rules, with 5 random seeds (82, 23, 37, 84, 4), run on two NVIDIA H100 80GB GPUs.

Why This Matters

Impact on research: The paper reframes trajectory regularization as a stability problem with explicit bounds, replacing a heuristic cosine penalty with a condition whose hyperparameters map onto interpretable geometric quantities (tube radius, contraction rate, slack). It also supplies a negative result — that too-tight geodesic constraints destroy representational capacity — which explains why prior angle-based methods must be tuned carefully rather than pushed to their limit.

Real-world applications:

• Reasoning-heavy assistants: The largest reported gains are on GSM8K (multi-step math) and NQ-Open (open-domain QA), suggesting use in systems that must stay on-topic over long generated chains.
• Low-data domain adaptation: Better data efficiency matters for specialized domains (legal, clinical, industrial) where fine-tuning data is scarce.
• Structured generation: Spider (text-to-SQL) and the NL-RX benchmarks (natural-language-to-regex) benefit from keeping intermediate representations anchored to a semantic plan.
• Long-context generation: Because the tube bound is horizon-independent under uniformly bounded violations, the approach targets drift accumulation over long outputs.

Industry relevance: The method requires no architectural modification, adds only a projection and scalar comparison to standard fine-tuning, and works across Llama, Gemma, OpenELM, OLMo, Qwen, and DeepSeek-R1-Distill checkpoints up to 8B parameters. It is a drop-in auxiliary loss rather than a new model design, which lowers adoption cost. The wide stable λ range of [0.01, 0.08] also reduces the tuning burden typically associated with regularizers.

Future Directions

1. Out-of-distribution guarantees: The bounds hold only for trajectories on the training distribution. The observed inference-time persistence of the smooth transversal profile is empirical; a theoretical account of why training reshapes the manifold globally remains open.
2. Removing the segmentation dependency: Semantic anchors currently rely on dataset-provided segmentation, which limits applicability to data without such annotation.
3. Scaling beyond 8B parameters: Experiments are limited to fine-tuning models up to 8B parameters; whether the partial-ISS trade-off behaves the same at larger scale is untested.
4. Extending to other modalities and analyzing the loss interaction: The authors list extending the framework to vision and multimodal models, and studying the interaction between the NTP loss and the Lyapunov loss. (The truncated text cuts off a further direction about incorporating Lyapunov constraints.)

Target Audience

Researchers and graduate students in representation learning and LLM training who are comfortable with dynamical-systems or control-theoretic formalism; practitioners who fine-tune open-weight models and want a principled, low-overhead alternative to cosine-based trajectory regularizers; and control theorists interested in applications of discrete-time ISS to deep network training dynamics. Readers seeking a purely empirical "plug-in recipe" will find the theory-heavy Sections 3–4 demanding, though the loss itself is implementable in a few lines.

Authors’ abstract

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 \emph{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: \url{https://github.com/t-khanusa/ControlJEPA}