Research map / Label-efficient and robust learning / Self-supervised and contrastive
How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes gradient-flow training dynamics of deep linear JEPA (stop-grad self-distillation) versus MAE, showing JEPA prioritizes high-regression-coefficient (influential) features, with the bias growing with encoder depth.
Problem
Understanding the implicit bias that makes JEPA-style latent prediction learn more semantic, less noisy features than pixel-reconstruction MAE objectives.
Approach
Derives closed-form ODEs and critical-time analysis for deep linear encoders trained with JEPA (stop-gradient) versus MAE in a simultaneously diagonalizable setting, showing JEPA's learning order depends on regression coefficients rho, while MAE's depends on covariance lambda; the effect strengthens with depth.
Key results
- JEPA critical time ratio for features with equal covariance but different regression coefficients deviates from 1 by (2L-1)/(2L-2) * Delta_rho * eps^(1/L), so JEPA learns high-rho features first, with the effect growing with depth; MAE shows no such dependence on rho at leading order.
- Fixed points: MAE converges to rho^(L/(L+1)) and JEPA to rho^L; a 1-layer JEPA matches the infinite-depth MAE dynamics.
- In simulations with depth L=5 and reversed rho/lambda trends, JEPA reverses the feature learning order relative to MAE, confirmed also with SGD on deep linear MLPs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Prioritize features with higher regression coefficient ρ | JEPA · Gaussian data with fixed λ and varying ρ | feature learning order | Helped |
| Prioritize features by covariance λ rather than regression coefficient ρ | MAE · Gaussian data with fixed λ and varying ρ | feature learning order | No clear effect |
| Use JEPA rather than MAE | Deep linear encoders · Gaussian data, Distribution 2 (λ increases and ρ decreases) | feature learning order | Mixed |
| Increase encoder depth (ablation) | JEPA · Gaussian data with features sharing λ but differing in ρ | temporal spread between feature learning | Helped |
| Increase encoder depth (ablation) | MAE · Gaussian data with features sharing λ but differing in ρ | feature learning time dependence on ρ | No clear effect |
| Use the default Gaussian initialization and stochastic gradient descent instead of the special initialization and gradient flow setup | JEPA and MAE · Gaussian data | objective-dependent feature learning order | Helped |
| Increase encoder depth beyond one layer | JEPA · Deep linear model | suppression of low-ρ feature projections at convergence | Helped |
| Use JEPA rather than MAE | Deep linear encoder · Features with equal λ and different ρ | feature learning priority by ρ | Helped |
| Use MAE rather than JEPA | Shallow encoder · Deep linear model with L=1 | training dynamics similarity to infinite-depth MAE | No clear effect |
Surprising findings
- The JEPA-vs-MAE difference in learning order appears only when the encoder is deep; shallow JEPA behaves like infinite-depth MAE.
- JEPA can learn features in the opposite order to MAE, prioritizing high-regression-coefficient features even when their covariance is lower.
Limitations
- Results hold only for simultaneously diagonalizable covariances and a restricted balanced, orthogonal-scaled initialization.
- Analysis is for deep linear networks; nonlinear models not covered and generalization to practical settings is unverified.
Experimental setup
Datasets: synthetic Gaussian data linear temporal generative model
Metrics: feature projection norm ||W e_i|| over training critical time to learn feature
Baselines: MAE
Builds on: JEPA MAE BYOL I-JEPA V-JEPA
Closest papers by meaning
- Denoising with a Joint-Embedding Predictive Architecture (ICLR 2025 · similarity 0.89)
- Connecting Joint-Embedding Predictive Architecture with Contrastive Self-supervised Learning (NeurIPS 2024 · similarity 0.89)
- You Don’t Need Domain-Specific Data Augmentations When Scaling Self-Supervised Learning (NeurIPS 2024 · similarity 0.88)
- seq-JEPA: Autoregressive Predictive Learning of Invariant-Equivariant World Models (NeurIPS 2025 · similarity 0.87)
- Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures (NeurIPS 2026 · similarity 0.87)
- T-JEPA: Augmentation-Free Self-Supervised Learning for Tabular Data (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.