JEPA · 2026-10-05
Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures
Jos\'e Lucas De Melo Costa, Seong Woo Ahn, Fabrice Popineau, Arpad Rimmel, Bich-Li\^en Doan
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors develop a drive–decay stability theory for Joint-Embedding Predictive Architectures (JEPAs), analyzing training dynamics near the trivial fixed point via a per-mode stability ratio μ_i=γ_i/σ_i. They predict a phase boundary (collapse when μ<1, learning when μ>1) and confirm it across 800+ Tabular-JEPA configurations. Theory suggests stabilizers; they propose ResidualPred, an identity-biased transformer predictor, improving effective rank and linear-probe accuracy on tabular and I-JEPA benchmarks.
Why it matters
JEPAs are prone to representational collapse, and prior fixes are largely empirical. The paper provides an early-training dynamical picture (drive vs decay) that unifies multiple collapse-prevention heuristics (predictor scaling, masking ratio, EMA) through how they shift μ. It also motivates an architectural change (ResidualPred) and reports empirical evidence on effective rank and downstream linear-probe accuracy.
Method
- Linearize coupled JEPA gradient-flow + EMA dynamics around the trivial fixed point and study Jacobian stability per eigenmode.
- Under approximate spectral decoupling, derive a per-mode stability ratio μ_i=γ_i/σ_i and a mode-wise escape criterion μ_i>1; count unstable modes to bound effective rank.
- Introduce ResidualPred: a transformer predictor with attention biased toward the identity at initialization, then evaluate effective rank and linear-probe accuracy on tabular and image JEPA settings.
Limitation
The phase boundary and rank-budget results rely on the theory assumptions, including approximate spectral decoupling (Assumption 3.2) and the specified conditions used for corollary 3.3 (e.g., exact form of Assumption 3.2 with orthonormal eigenbases and real spectra).
Abstract (from arXiv)
Joint-Embedding Predictive Architectures (JEPAs) are prone to representation collapse, typically mitigated through empirical heuristics. We develop an early-training stability theory that unifies these heuristics. Linearising the coupled JEPA gradient flow around the trivial fixed point reveals two competing effects: a driving force ($\gamma$) and a decay effect ($\sigma$). Under approximate spectral decoupling, a per-mode stability ratio $\mu_i = \gamma_i / \sigma_i$ factorises into independent data-side and predictor-side terms and the count of unstable modes tracks the rank of representations that can emerge. The framework predicts a phase boundary, which we confirm empirically across more than 800 Tabular-JEPA configurations. It also unifies predictor scaling, masking ratio, and EMA as distinct mechanisms for shifting $\mu$. Guided by this analysis, we introduce ResidualPred, a transformer predictor whose attention is biased toward the identity at initialisation; it improves both effective rank and downstream accuracy on tabular benchmarks and in I-JEPA pretraining on CIFAR-10, CIFAR-100, STL-10, and ImageNet. Our framework connects empirical collapse-avoidance heuristics to an explicit dynamical picture, yielding theory-driven stabilizers. Code is available at https://github.com/jose-melo/drive-vs-decay.
Related papers
- Does Joint-Embedding Predictive Architecture Pretraining Help Time Series Forecasting?
- Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State
- Predictive Closure: A Set-Level Characterization of JEPA's Feature Selection
- JEPA Learns What the Mask Leaves Unrecoverable
- Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-01 · Indexed by Pub-AI: 2026-10-05 UTC
Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures
Preprint posted: 2026-10-01 · Original record ↗
- arxiv
- 2610.02344
Explicit source versions (1)
- v1 · 2026-10-01
- Preprint available
Source date: 2026-10-01 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
Joint-Embedding Predictive Architectures (JEPAs) are prone to representation collapse, typically mitigated through empirical heuristics. We develop an early-training stability theory that unifies these heuristics. Linearising the coupled JEPA gradient flow around the trivial fixed point reveals two competing effects: a driving force ($γ$) and a decay effect ($σ$). Under approximate spectral decoupling, a per-mode stability ratio $μ_i = γ_i / σ_i$ factorises into independent data-side and predictor-side terms and the count of unstable modes tracks the rank of representations that can emerge. The framework predicts a phase boundary, which we confirm empirically across more than 800 Tabular-JEPA configurations. It also unifies predictor scaling, masking ratio, and EMA as distinct mechanisms for shifting $μ$. Guided by this analysis, we introduce ResidualPred, a transformer predictor whose attention is biased toward the identity at initialisation; it improves both effective rank and downstream accuracy on tabular benchmarks and in I-JEPA pretraining on CIFAR-10, CIFAR-100, STL-10, and ImageNet. Our framework connects empirical collapse-avoidance heuristics to an explicit dynamical picture, yielding theory-driven stabilizers. Code is available at https://github.com/jose-melo/drive-vs-decay.