Joint embedding model stability improved by balancing drive and decay forces
Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures
Machine Learning
Summary
Training certain machine learning models called Joint-Embedding Predictive Architectures (JEPAs) can fail because the model’s representations collapse and stop learning useful features. The authors explain why this happens by identifying two opposing effects during training: a force that drives learning and another that causes decay. They show mathematically how these effects interact and predict when stable, useful representations will appear. Using this insight, they create an improved model design that leads to better learning and performance on image and tabular data.
What this means in practice
- •For machine learning engineers: Build more stable self-supervised models for vision and tabular data by implementing training dynamics that prevent representation collapse.
- •For data platform teams: Improve unsupervised feature extraction in data pipelines by applying the ResidualPred predictor to increase representation quality.
Authors
José Lucas De Melo Costa, Seong Woo Ahn, Fabrice Popineau, Arpad Rimmel, Bich-Liên Doan
Abstract
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.