Drive vs. Decay：JEPA 联合嵌入预测架构的训练动力学分析

View PDF HTML (experimental)

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 ($\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 this https URL.

Comments:	Accepted at NeurIPS 2026 (Main Track). 57 pages, 10 pages of main text; appendices and the NeurIPS paper checklist included. Code: this https URL
Subjects:	Machine Learning (cs.LG)
Cite as:	arXiv:2610.02344 [cs.LG]
	(or arXiv:2610.02344v1 [cs.LG] for this version)
	https://doi.org/10.48550/arXiv.2610.02344 arXiv-issued DOI via DataCite (pending registration)

Submission history

From: José Costa [view email]
[v1] Thu, 1 Oct 2026 18:16:38 UTC (2,307 KB)