PaperlayerView exact revision on arXiv

Source-linked AI summary

Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures

José Lucas De Melo Costa, Seong Woo Ahn, Fabrice Popineau, Arpad Rimmel, Bich-Liên Doan

arXiv:2610.02344v1cs.LG

TL;DR

JEPAs can collapse even while minimizing their prediction objective, and existing remedies are largely empirical. This paper analyzes coupled JEPA dynamics through a drive–decay stability ratio and tests the resulting account across configurations and predictors. The framework links unstable modes to representation rank and supports a theory-guided predictor, while its criterion is proven in the linear regime and extended empirically to images.

  • Problem

    JEPA collapse remains a concern, while a complete characterization of the original objective under realistic update rules is lacking and many stabilizers rely on empirical heuristics.

    Source: §1 · pp. 1–2

  • Method

    The paper linearizes coupled JEPA dynamics around collapse and uses approximate spectral decoupling to define a per-mode drive-to-decay ratio, connecting it to rank and stabilization mechanisms.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    Across more than 800 configurations, the framework relates stability to effective rank and downstream performance, while ResidualPred improves effective rank and linear-probe accuracy on brittle tabular settings and I-JEPA image benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The drive–decay view interprets EMA, stop-gradient, and predictor design as mechanisms for shaping timescale separation and amplifying informative modes.

    Source: §8 · p. 10

    The criterion is proven in the linear regime; its image extension is an empirical prediction measured after predictor equilibration, and the quantitative initialization-budget relation is claimed only on tabular data.

    Source: §7 · p. 10, §6.2 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

JEPAs can collapse to trivial or low-rank representations even as they minimize their prediction objective, while existing stabilizers are largely empirical. This paper develops a dynamical theory of early-training stability and tests it across Tabular-JEPA configurations and transformer predictors.

1 Introduction

  • JEPAs predict target-view representations from context views, but their embeddings can collapse to trivial or low-rank solutions with little downstream information.Non-contrastive methods avoid explicit negative sampling [Grill et al., 2020, Bardes et al., 2022, Zhang et al., 2022].Source: §1 · p. 1
  • The paper addresses the lack of a complete characterization of the original JEPA objective under realistic update rules, beyond theories focused on explicit diversity regularizers.Many existing stabilizers are justified empirically, including asymmetric predictors, stop-gradient, and EMA targets.Source: §1 · p. 2
  • A drive–decay framework models stability as competition between view- and alignment-induced drive and decay induced by predictor capacity and data covariance.Figure 1 marks learning for µ = γ/σ > 1 and collapse for µ < 1.Source: Figure 1 · p. 1, §1 · p. 2
  • Across more than 800 configurations, the paper uses this framework to unify predictor scale, context masking, and EMA as controls on stability, then introduces identity-biased ResidualPred.The stated contributions include improvements in effective rank and linear-probe accuracy on brittle tabular settings and I-JEPA image benchmarks.Source: §1 · p. 2

2 Related Work

Prior work has proposed several practical ways to prevent JEPA collapse and has analyzed related non-contrastive dynamics. This paper positions identity-biased predictor initialization as a different application of an established idea, paired with a new JEPA stability framework.

2 Related Work

  • JEPA practice uses EMA, variance–covariance regularization, frozen teachers, or distributional constraints to prevent trivial representations.The cited work also connects regularizers to embedding-distribution control and density estimation [Huang, 2026, Balestriero et al., 2025].Source: §2 · p. 2
  • Analyses of non-contrastive SSL connect predictor and stop-gradient dynamics to representation diversity, while other work interprets EMA and stop-gradient as predictor orthonormalisation.The cited studies include Tian et al. [2021], Zhang et al. [2022], Halvagal et al. [2023], and Richemond et al. [2023].Source: §2 · pp. 2–3
  • Identity-biased initialization is established for stabilizing deep-network signal propagation, but this work studies its effect on JEPA training dynamics and collapse.The authors claim no architectural novelty for identity-biased attention and distinguish their per-mode stability ratio, supercriticality count, and rank budget from prior work.Source: §2 · p. 3

3 Drive vs. Decay: Theoretical Analysis of Coupled JEPA Dynamics

The analysis linearizes coupled encoder and EMA-target dynamics around collapse, then uses approximate spectral decoupling to characterize escape mode by mode. It relates unstable modes to an emerging rank budget and clarifies how the criterion applies when the predictor is trained jointly.

3.1 Linear Model: Coupled Dynamical System Formulation

  • The linear model represents context and target encoders and a predictor as weight matrices, applies masks to the views, and minimizes squared distance to a stop-gradient target.The target weights evolve by EMA while context and predictor weights follow negative-gradient dynamics.Source: §3.1 · p. 3

3.2 Jacobian Stability Analysis

  • Approximate shared eigenbases let the coupled dynamics decompose into scalar modes, and factorize each stability ratio into independent data-side and predictor-side terms.The data-side cross-view retention coefficient is governed by data structure and masking; the predictor-side alignment ratio is governed by predictor geometry.Source: §3.2 · pp. 4–5
  • Under approximate spectral decoupling, a mode escapes collapse when its drive-to-decay ratio µ_i exceeds 1; when µ_i is much less than 1, the trivial fixed point is stable.For µ_i > 1, the characteristic equation has a positive real eigenvalue; µ_i = γ_i/σ_i compares the driving and decay effects.Source: §3.2 · p. 4

3.3 Rank budget and joint training

  • The count of supercritical modes bounds the rank that can emerge after escape, with the effective-rank bound applying to the growing component rather than generic algebraic rank.The stated result assumes orthonormal eigenbases and real spectra; under its conditions the excess singular values decay exponentially.Source: §3.3 · p. 5
  • Near collapse, the predictor gradient vanishes, so joint training has the fixed-predictor Jacobian criterion plus a zero block at the collapsed set.Under stated conditions and a margin from µ = 1, the predictor moves only by O(δ^2) as the encoder state grows to norm δ, preserving the number of unstable eigenvalues locally.Source: §3.3 · p. 5
  • The predictor is quasi-static near collapse because of initialization scale, not learning-rate choice; its adaptation after escape is treated separately.The predictor’s state still conditions the encoder-side escape analysis.Source: §3.1 · p. 4, §3.3 · p. 5

4 Methodology

The experiments use tabular and image datasets to test the linear theory and its extension to standard JEPA models. The methodology specifies training protocols, model variants, and representation and downstream-accuracy evaluations.

4 Methodology

  • The evaluation covers three tabular datasets and five image datasets, with specified train-validation-test splits for tabular data and model-specific image training protocols.Image runs include CIFAR-10, CIFAR-100, STL-10, ImageNet-100, and ImageNet-1k.Source: §4 · p. 5
  • The models include T-JEPA, linear-encoder Tabular-JEPA variants for isolating stabilization mechanisms, and I-JEPA with specified ViT encoders for image experiments.Image downstream evaluation uses linear probes; ImageNet adds nearest-neighbour and attentive probes.Source: §4 · p. 6
  • The methodology evaluates embedding spectral diversity with effective rank and downstream accuracy with linear probes.The supplied metric description introduces effective rank as a measure of concentration in a low-dimensional subspace.Source: §4 · p. 6

5 Controlling the Training Dynamics: Example of Mechanisms

The section treats predictor scale, view masking, and EMA as controls over the stability ratio μ, and tests how its threshold and rank budget relate to learning and effective rank.

Predictor scaling

  • Predictor scaling makes μ inversely proportional to α: sufficiently small initialization variance can put the model in the supercritical regime, while large values can lead to collapse.Large initialization increases the quadratic self-regularization penalty faster than the linear driving force.Source: §5 · p. 6

EMA as a timescale

  • EMA decay rate β acts as a timescale parameter rather than deciding escape for a real drive spectrum: smaller β slows collapse when μ is low and throttles feature growth when μ is high.Source: §5 · p. 7

Empirical phase transitions

  • In high-context-ratio variants, increasing predictor scale drops normalized test accuracy from 0.872 to 0.586; the phase diagram relates performance to the μ = 1 boundary.For context ratios pm1 ≤ 0.4, empirical μ tracks the inverse-scale curve; at pm1 ≥ 0.5, curves shift downward and deviate from that law.Source: §5 · pp. 6–7, Figure 4 · p. 7

View correlation (masking strategy)

  • For T-JEPA, low context ratios (higher masking) yield median μ values of 0.05–0.11, effective rank about 4.6, and test accuracy about 0.464.Higher masking weakens context–target correlation and the cross-covariance contribution to the driving force, lowering μ.Source: §5 · pp. 6, 8

Rank-budget validation

  • Across 810 runs, the rank budget tracks effective rank on Helena (ρ = +0.91), exceeds 0.9 on ALOI, and remains positive but weaker on Jannis.Jannis has a near-degenerate top spectrum, which the paper identifies as distinguishing its weaker correlation.Source: §5 · p. 8, Figure 5 · p. 8

Other mechanisms as doses on μ

  • Five further anti-collapse mechanisms are derived as doses on μ, and their predicted escaping-mode counts are numerically verified in the linear model.The mechanisms include encoder weight decay, SIGReg and VICReg, mimetic initialization, and a frozen identity predictor.Source: §5 · p. 8, Table 3 · p. 8

6 ResidualPred: The Residual Transformer Predictor

ResidualPred biases the predictor toward the identity so early training is driven by view alignment rather than arbitrary predictor scale. It improves accuracy across the reported settings and raises representation rank in both brittle and non-brittle configurations, with image-domain gains also reported.

6.1 Theoretical Motivation

  • ResidualPred parameterizes the predictor as an identity plus a small residual, preventing arbitrary predictor-induced contraction at initialization.With the predictor near identity, decay is set by masked-view energy and drive by view alignment rather than predictor scale.Source: §6 · p. 8, §6.1 · p. 9
  • A diagonal attention-logit bias makes self-attention approximate the identity, so each token initially attends to itself rather than averaging spatial signals.The paper uses a large initialization bias, for example 10.0.Source: §6.1 · p. 9

6.2 ResidualPred Results

  • ResidualPred improves accuracy in every Table 4 setting and leads on CIFAR-100 throughout 300 epochs; its ImageNet-1k pilot advantage appears at both tested resolutions.Tripling the ImageNet-100 schedule to 90 epochs attenuates the margin.Source: §6.2 · p. 9, Figure 6 · p. 9
  • In brittle and non-brittle configurations, Standard T-JEPA collapses while ResidualPred increases rank, keeping rank above Standard’s throughout training on Helena.The comparison is shown in Figure 6(a) and Figure 13.Source: §6.2 · p. 9, Figure 6 · p. 9
  • ResidualPred beats SIGReg on the ImageNet-1k pilot, and the methods compose; it trails ResidualPred by 4.9 points with Mimetic initialisation.ResidualPred uses one unchanged hyperparameter and adds no loss term or per-step cost.Source: §6.2 · p. 10

7 Discussion

The discussion interprets JEPA stability as competition between alignment-driven growth and predictor-induced decay, and connects that account to existing collapse controls. Its scope is qualified by linear-theory assumptions and the empirical status of its image extension.

7 Discussion

  • The drive–decay framework interprets collapse as spectral competition between alignment-driven growth and predictor-induced decay, with low µ associated with low rank and degraded downstream performance.ResidualPred is presented as a structural alternative that initializes the predictor near identity so view alignment can dominate early training.Source: §7 · p. 10
  • The linear model recovers seven published collapse-prevention mechanisms as distinct adjustments to µ, including DirectPred, IsoLoss, Halvagal et al.’s rate equation, and predictor orthonormalisation.The predicted escaping count is reproduced by the linear model.Source: §7 · p. 10
  • The theory depends on linear dynamics and approximate spectral decoupling, while its image extension remains an empirical prediction and ImageNet-1k results are matched-budget pilots.The official full-resolution, full-schedule I-JEPA recipe is untested; effective rank on images is necessary for accuracy but not sufficient.Source: §7 · p. 10

8 Conclusion

The conclusion frames JEPA collapse as a dynamical trade-off between representation growth and decay, with the spectral ratio µ as the critical coefficient. The appendices derive the coupled flow and its stability analysis, including qualifications on when a scalar per-mode boundary applies.

8 Conclusion

  • The paper identifies µ = γ/σ as the critical coefficient governing the trade-off between representation growth and collapse in JEPAs.The conclusion presents this as a dynamical-systems perspective on JEPA stability.Source: §8 · p. 10

A Detailed Derivation of Dynamics and Stability

  • The appendix derives the coupled JEPA dynamics, vectorizes the matrix differential equations, and analyzes their spectrum to obtain the main-text stability conditions.The derivation begins from the expected squared prediction error between predictor output and target projection.Source: A Detailed Derivation of Dynamics and… · p. 14, §A.1 · p. 14

A.1 Derivation of the Coupled Gradient Dynamics

  • The context-encoder gradient follows from the prediction error while stop-gradient holds the target term constant; the target encoder follows an EMA update, expressed as a differential equation in continuous time.The appendix defines the error vector before taking the expected gradient.Source: §A.1 · p. 14

A.2 Vectorization and Jacobian Construction

  • The derivation stacks vectorized encoder weights into a state vector, constructs a block Jacobian for the linear ODE, and confirms the main-text equation.The EMA component is linear and acts element-wise.Source: §A.2 · pp. 14–15

A.3 Characteristic Equation and Eigenvalue Analysis

  • Under eigenmode decomposition, each mode’s stability is analyzed through the characteristic equation, with µ_i defined as the driving-to-decay ratio γ_i/σ_i.The appendix obtains the eigenvalues by solving the Jacobian’s characteristic equation.Source: §A.3 · p. 15

A.4 Reality and Positivity of the Driving Spectrum

  • The scalar criterion µ_i > 1 is exact only when the driving operator is symmetric and the drive and decay operators share orthonormal eigenbases; neither condition alone suffices.Without symmetry, γ_i may be complex, so µ_i > 1 is not a universal boundary.Source: §A.4 · pp. 15–16
  • A singular-value formulation makes the driving strength and ratio real and nonnegative, but its independently sorted values agree with the paired-mode form only when the orderings coincide.The appendix gives an example where the paired form identifies one supercritical mode while the independently sorted formulation identifies none.Source: §A.4 · p. 16

A.5 Stability of the full nonsymmetric block Jacobian

The appendix bounds stability of the full, potentially nonsymmetric linearised system, then shows how exact spectral decoupling recovers the per-mode criterion. It also tests the decoupling assumption using commutators, finding it especially tight on real CIFAR-10 patches.

A.5 Stability of the full nonsymmetric block Jacobian

  • The framework’s full-system stability analysis requires no symmetry or decoupling assumption, while its exact per-mode results are recovered when spectral decoupling holds.The general bounds become exact as Assumption 3.2 becomes exact.Source: §A.5 · pp. 16–18

Proofs

  • The quadratic pencil Q(λ) has the same eigenvalues, with algebraic multiplicities, as the full Jacobian J.The pencil identity also gives det J = β^n det(S − G).Source: §A.5 · p. 17, Proofs · p. 18

A.5 Stability of the full nonsymmetric block Jacobian

  • For arbitrary nonsymmetric drive, σmax(G) < λmin(S) certifies collapse, but failure of this sufficient bound does not establish instability.With symmetric drive, stability is equivalent to S − G ≻ 0; an unstable direction occurs when w⊤Gw > w⊤Sw.Source: §A.5 · p. 17

B Spectral-Decoupling Regime: Analytical and Empirical Evidence

  • Under the masking model, the context-ratio factor and data-side covariance factor separate, making the commutator identity operationally testable.The appendix evaluates this identity across spectra and reports measurements on real datasets.Source: B Spectral-Decoupling Regime:… · p. 19, §B.2 · p. 19

B.1 Commutator definition and what it controls

  • Approximate joint diagonalisation is measured by relative Frobenius commutators: they vanish under exact joint diagonalisation, and rank-control accuracy degrades smoothly with their size.The scalar μ-prediction is exact when both commutators vanish.Source: §B.1 · p. 19

B.3 Empirical commutators across spectra

  • At context ratio 0.85, CIFAR-10 raw-patch covariance has CΣ on the order of 10^-4, two to three orders of magnitude below rotated power-law surrogates.Across context ratios, the CIFAR-10 raw-patch curve also remains well below the synthetic benchmarks.Source: Table 6 · p. 20, §B.3 · p. 20, Figure 7 · p. 20

B.2 Closed-form factorisation of the data-side commutator

  • The data-side commutator factorises into a masking multiplier and a covariance-shape term; the multiplier p^2(1 − p) peaks at p = 2/3 and shrinks toward zero as p approaches 1.The covariance-shape term depends on how far ΣXX is from diagonal in the canonical basis and is invariant to rescaling ΣXX.Source: §B.2 · pp. 19–20

B.4 Tabular benchmarks: aloi, helena, jannis

The appendix checks spectral decoupling on tabular benchmarks and details how the rank-control corollary follows from independent mode dynamics. It bounds the rank of the unstable component, not the generic finite-time algebraic rank of the full predictor.

B.4 Tabular benchmarks: aloi, helena, jannis

  • On jannis, near-degenerate top eigenvalues coincide with weaker per-mode correlation (+0.619) than helena (+0.906) and aloi (+0.905), while signed and loose-count correlations remain positive.The authors present this three-dataset pattern as a measured, mechanism-consistent hypothesis, not a fitted law.Source: §B.4 · p. 21

C.1 Step 1 – Joint diagonalisation in the shared basis

  • In the shared eigenbasis, the transformed flow separates into independent 2 × 2 blocks, one for each predictor–data mode pair.The full 2dN-dimensional flow splits into dN planar flows.Source: §C.1 · p. 21

C.2 Step 2 – The 2 × 2 spectrum and saddle condition

  • Each mode’s characteristic polynomial uses μkℓ = (ckqℓ)/(akpℓ); subcritical modes have two eigenvalues with negative real parts, while a saddle block has an unstable eigenvector.The mode polynomial is λ² + (β + akpℓ)λ + βakpℓ(1 − μkℓ) = 0.Source: §C.2 · p. 22, §C.3 · p. 22

C.3 Step 3 – Mode-wise unstable manifolds; C.4 Step 4 – Reassembly into the original basis

  • The proof selects unstable directions within saddle blocks, then reassembles them in the original basis.The excerpt identifies the saddle-block unstable eigenvector as e+k.Source: §C.3 · p. 22, §C.4 · p. 22

C.5 Step 5 – The rank bound (unstable-component statement)

  • The rank bound applies to W+(t), the unstable component: generic initialisation can leave W(t) at full algebraic rank at every finite time.Experiments measure embedding effective rank, for which the appendix gives corresponding finite-time bounds.Source: §C.5 · p. 22

C.6 Finite-time restatement of the rank budget

This section restates the finite-time rank budget under exact spectral-decoupling assumptions, transfers it to embedding spectra, and compares signed and absolute-value formulations. It also examines how diagonal-attention bias and zero-initializing the MLP output affect predictor enhancement in token space, with zero-initialization treated as an ablated candidate.

C.6 Finite-time restatement of the rank budget

  • For all t ≥ 0, rank W+(t) ≤ rtight; after T(ε), rankε(W(t)) ≤ rtight, and ERank(W(t)) ≤ rtight e^φ(θ̄(t)) past T0.The exact-rank bound is attained when an axis projection is injective on U and the corresponding amplitudes are nonzero; the additive effective-rank form additionally requires φ(θ̄(t)) ≤ 1. These are bounds for the linear flow with a fixed predictor, not predictions of effective rank; the fastest escaping mode can eventually dominate.Source: §C.6 · p. 23, Proofs · p. 24

Proofs

  • The proof decouples the fixed-predictor flow into dN planar mode systems, with unstable modes spanning W+(t) and stable components decaying to control the thresholded-rank tail.The planar semigroup bound handles transient growth, and the unstable span yields the exact rank bound.Source: Proofs · pp. 23–24

C.8 Signed versus absolute-value formulation

  • On aloi, ρ(ERank, ˜rtight) = +0.905 versus ρ(ERank, rtight) = +0.672; helena gives +0.906 versus +0.908, while jannis gives +0.62 to +0.65.The sweep covered 810 cells, with 804 completed; near-zero eigenvalue sign changes can destabilize the signed count, so the paper reports the absolute-value form in the main text.Source: §C.8 · p. 25, Setup · p. 25, Figure 8 · p. 26

C.9.1 The covariance as a readout

  • The covariance C(t) = W(t)ΣW(t)⊤ does not close on its own, but the joint second moment Ξ = zz⊤ follows a closed linear flow and yields C as a readout.Specifically, Ξ̇ = JΞ + ΞJ⊤ and Ξ(t) = e^JtΞ(0)e^J⊤t.Source: §C.9.1 · p. 26

C.9.2 Growth at twice the rate

  • When a data-eigenbasis mode is supercritical, its embedding variance grows at twice the mode’s positive growth rate; otherwise it decays, making collapsed covariance unstable exactly when U is nonempty.This per-mode statement requires Σ to share the data eigenbasis; the subsequent transfer theorem does not.Source: §C.9.2 · p. 26

C.9.3 The spectral transfer, two-sided

  • For positive-definite Σ, the singular values of WΣ^1/2 are bounded above and below by those of W scaled by σmax(Σ)^1/2 and σmin(Σ)^1/2, respectively.Thus rank is preserved, and thresholded ranks transfer with corresponding rescaled thresholds; the condition-number factor affects the head-to-tail ratio, not rtight.Source: §C.9.3 · pp. 26–27

C.9.4 The budget, at the embedding level

  • Under the stated escape, non-degeneracy, positive-definite covariance, and generic-initialisation conditions, ERank of W(t)Σ^1/2 obeys the rank budget after an explicit time as its tail mass decays exponentially.For a fixed full-rank evaluation set, the embedding effective rank equals that of W(t)Σ̂^1/2, so the theorem applies without requiring Σ̂ to share the data eigenbasis.Source: §C.9.4 · pp. 27–28

D Choice of the Diagonal-Attention Bias βinit

  • ResidualPred uses distinct diagonal-attention bias designs: images learn a per-head β on context-token identity positions, while tabular models use fixed shared β = 10 with full identity bias.The image design leaves mask tokens to read from context; the tabular design uses a fixed value because small token counts already yield near-saturated diagonals and tokens are homogeneous.Source: D Choice of the Diagonal-Attention… · p. 28

D.1 Saturation criterion

  • With zero-mean off-diagonal logits, βinit ≈ log nctx + 5 guarantees Aii > 0.99; βinit = 10 yields values above 0.997 for the stated CIFAR and STL context sizes.The stated limits are nctx ≤54 on CIFAR-10/100 (ViT-S/4, N = 64) and nctx ≤122 on STL-10 (ViT-S/8, N = 144), with keep probability pm1 = 0.85.Source: §D.1 · p. 29

E Zero-Init MLP Output: An Ablated Design Choice

  • The appendix documents a 2 × 2 ablation of zero-initialized MLP output and diagonal-attention bias, and says the main paper retains only the diagonal bias.The supplied passages do not report the ablation’s empirical results.Source: E Zero-Init MLP Output: An Ablated… · p. 29

E.1 Motivation: a natural extension of the identity-bias principle

  • An exact identity-like predictor at initialization would require both transformer residual sub-blocks to behave as identity; zeroing the MLP output projection is a complementary candidate to diagonal-attention bias.The passages identify ReZero [Bachlechner et al., 2021], Fixup [Zhang et al., 2019], SkipInit [De and Smith, 2020], and IDInit [Pan et al., 2025] as methods using MLP output zero-initialization for deep residual signal propagation.Source: §E.1 · p. 29

E.2 The candidate formula

  • The candidate zero-initializes only Wfc2 and bfc2, leaving the attention output projection standard so context tokens can still inform mask tokens.Despite zero initialization, Wfc2 receives a non-zero gradient at the first optimization step, making this a delayed-engagement scheme rather than a frozen-parameter scheme.Source: §E.2 · pp. 29–30

E.3 What zero-init does to the predictor block at t = 0

  • In diffuse attention, zero-initialized MLP leaves most of the LN token-feature modes at their unenhanced baseline; diagonal bias supports predictor enhancement across all LN modes.The diffuse attention contribution is rank-one in token space, so the predictor amplifies the drive–decay ratio on at most N modes; the analysis therefore predicts a small accuracy effect from zero-init alone in that regime.Source: §E.3 · p. 30

E.4 Empirical results

The appendix derives the linear JEPA gradients and stationary-point condition, then gives an adiabatic stability threshold for non-trivial solutions. The threshold depends on the dominant eigenvalues of the data- and predictor-side operators.

E.4 Empirical results

  • The CIFAR-100 grid evaluates four configurations at 100- and 300-epoch budgets using ViT-Small/4, L = 64, and pm1 = 0.85.Source: §E.4 · p. 30, Table 8 · p. 31

F.1 Derivation of Linear JEPA Dynamics (Expected Loss)

  • The expected-loss derivation computes gradients with respect to encoder W and predictor Wp, simplifying terms under independent target masking and a fixed target encoder.The target encoder is treated as fixed within the expectation, as in an EMA update.Source: §F.1 · pp. 31–32

F.2 Stationary Point Equation Derivation

  • Non-trivial stationary points exist when the matrix condition is singular, equivalently when 1 is an eigenvalue of the scaled operator product.The derivation notes that the relevant eigenvalues may be complex, so the threshold uses the eigenvalue of largest magnitude.Source: §F.2 · p. 32

F.3 The adiabatic stability theorem

  • In the adiabatic limit, non-trivial stationary points exist only when the effective coupling crosses a critical value set by the largest-magnitude eigenvalues of AL and AR.The target weights are assumed to track the encoder proportionally; Figure 10 depicts this stability analysis.Source: §F.3 · p. 33, Figure 10 · p. 33

F.4.1 Bifurcation Phenomenon

The appendix links a spectral threshold to bifurcation in JEPA dynamics and the loss landscape. It also shows that the fixed-predictor local criterion remains exact under joint predictor training, with predictor drift bounded until encoder escape.

F.4.1 Bifurcation Phenomenon

  • For τ below τcrit, only the trivial stationary point exists; above it, non-trivial solutions emerge, marking a bifurcation in the loss landscape.Figure 11 plots encoder Frobenius norm and condition number against τ; its empirically estimated threshold is approximately 0.551.Source: §F.4.1 · p. 33, Figure 11 · p. 34

F.5 Loss Landscape Transformation

  • At τ = 0.200, the projected landscape has its global minimum at the origin; at τ = 0.600, the origin has higher loss and non-trivial minima appear away from it.These values illustrate a specific JEPA-like setup with τcrit ≈ 0.400.Source: §F.5 · pp. 34–35, Figure 12 · p. 35

F.6 The criterion under joint predictor training

  • Under joint training, the fixed-predictor criterion remains the exact local criterion, and the predictor’s drift is bounded until the encoder escapes.Source: §F.6 · p. 35

F.6.1 The collapsed set is a manifold of equilibria

  • The collapsed states form a d2-dimensional manifold of equilibria, and the linearisation there is block diagonal with d2 neutral directions tangent to that manifold.Source: §F.6.1 · p. 35

F.6.2 How far the predictor drifts before the encoder escapes

  • Under the stated initial decoupling and growth conditions, the encoder reaches threshold δ by time Tδ ≤ T̄, while the predictor’s drift remains bounded up to escape.The theorem’s bound does not depend on the initialisation scale ε.Source: §F.6.2 · pp. 36–37

F.6.3 Numerical check with the predictor trained concurrently

  • With concurrent predictor training, the frozen-predictor Jacobian’s leading-eigenvalue sign classified collapse versus escape in 104/104 runs at equal and 3× predictor rates, and 102/104 at 10×.Both disagreements occurred when |λinit| ≤ 0.004.Source: §F.6.3 · p. 38

G Training Details; G.1 Programming environment

G.2 Tabular-domain experiments

  • Tabular experiments use Jannis, Helena, and ALOI, spanning dataset scale, feature dimensionality, class cardinality, and classification complexity.Source: §G.2 · p. 38

G.3 Image-domain experiments

The image-domain experiments include CIFAR-10, CIFAR-100, and STL-10, while the appendix defines effective rank as a spectral diagnostic for collapse. Detailed results cover tabular ResidualPred runs and the image-experiment protocol.

G.3 Image-domain experiments

H.1 Quantifying Representational Collapse.

  • Effective rank is the exponential of the entropy of the normalized singular-value distribution, bounded between 1 and algebraic rank and invariant to global rescaling.It continuously penalizes concentration in a few singular modes, making it sensitive to soft collapse.Source: §H.1 · p. 39

I Full Results

  • The experiments use effective rank alongside downstream performance to assess representation dimensionality; detailed tabular results report collapse statistics, training loss, and test accuracy across configurations.Source: §H.1 · p. 40, Table 11 · p. 40

I.3 T-JEPA with ResidualPred – per-run details

  • On the subcritical stress grid, no Standard run escapes; ResidualPred escapes are discontinuous, and on Helena all six ResidualPred runs exceed all six Standard runs in probe accuracy.Probe accuracy and effective rank co-sign on every reported row; the recurring Jannis probe values are class priors.Source: §I.3 · p. 40

I.4 I-JEPA Detailed Results on CIFAR-10, CIFAR-100, and STL-10

  • The detailed CIFAR-10, CIFAR-100, and STL-10 runs use I-JEPA with ViT-Small encoders, context retention pm1 = 0.85, and a 100-epoch cosine schedule.The listed token counts are L = 64 for CIFAR-10/CIFAR-100 and L = 144 for STL-10.Source: §I.4 · p. 41

I.5 Image experiments beyond CIFAR: protocol

Across the image experiments, ResidualPred improves effective rank and probe performance over the standard predictor in matched settings. The appendix also tests predictor controls, loss-side comparisons, and seven mechanisms within the linear model.

I.5 Image experiments beyond CIFAR: protocol

  • The ImageNet experiments use a shared I-JEPA pipeline and differ only in predictor, except for loss-side arms; the ImageNet-1k nearest-neighbour and attentive readouts use a class-balanced bank.The ImageNet-1k pilot uses 30 epochs and five seeds per predictor, while stresses change resolution or training duration.Source: §I.5 · pp. 41–42

I.6 ImageNet-100 at two depths and three probes

  • At matched 100-epoch, depth-6 settings, ResidualPred raises ViT-B/16 linear-probe accuracy from 0.339 to 0.398 (+6.0 points) and ViT-S/16 accuracy by 4.4 points.ResidualPred effective rank is higher from epoch five onward in every matched seed; the ViT-B/16 backbone has four times the parameters.Source: §I.6 · p. 43

I.7 A thousand-class pilot and its two stresses

  • At depth 6 in the 128-pixel ImageNet-1k pilot, the parameter-free margin is +8.58 points versus +5.5 points at 100 classes; at 224 pixels it persists, while over 90 epochs it attenuates to +3.45 points.The 90-epoch comparison reports effective-rank seed means of 111 to 179 for the standard predictor and 200 to 251 for ResidualPred; matching ResidualPred’s 30-epoch accuracy costs up to three times the compute for the standard predictor.Source: §I.7 · p. 44

I.8 Controls: frozen and skip-only predictors

  • Freezing ResidualPred at initialisation reduces linear-probe accuracy from 0.427 to 0.099 and effective rank from 82.1 to 1.67; a skip-only predictor has zero encoder gradient.On CIFAR-100 BYOL runs, effective rank and accuracy rise with predictor-to-encoder learning-rate ratio across both seeds.Source: §I.8 · pp. 44–45

J Seven anti-collapse mechanisms as doses on one ratio: verification on the linear model

I.9 Comparison with loss-side mechanisms at the pilot setting

  • On the ImageNet-1k pilot setting (128 pixels, 30 epochs, depth 6), SIGReg beats the standard predictor by 3.98 points and ResidualPred beats SIGReg by 2.50 points, both with complete separation.Combining SIGReg and ResidualPred adds 18.2 to ResidualPred-alone effective rank; the tested VICReg arm destabilised in both seeds and was excluded.Source: §I.9 · pp. 45–46, Table 24 · p. 48

J.1 The seven doses

  • The seven mechanisms shift the ratio through distinct terms: weight decay raises decay, predictor scaling divides the ratio, and EMA changes escape speed rather than which modes escape.Loss-side regularisers, mimetic initialisation, and a frozen identity predictor also alter predictor-side factors or make escape depend on the data spectrum.Source: §J.1 · pp. 48–49

J.2 Four published results recovered as limiting cases

The appendix recovers published collapse analyses as limiting cases of the drive–decay criterion, then measures how ResidualPred changes the target-aligned drive after predictor equilibration. The image measurements also bound what can be inferred: the initialization-time estimator is null, and effective rank is necessary but not sufficient to track accuracy gains.

J.2 Four published results recovered as limiting cases

  • DirectPred [Tian et al., 2021] holds surviving modes at µk = 1; the Halvagal et al. [2023] rate equation is the criterion under predictor slaving.IsoLoss equalises escape rates and compresses the tabular sweep’s spread of escape epochs from 73 to 7.Source: §J.2 · p. 49
  • Orthonormalising the predictor [Richemond et al., 2023] yields an escape boundary that matched direct eigensolves in all 20,000 random draws.The boundary is g cos θℓ+βg^2 sin^2 θℓ/(β + ak)^2 > ak.Source: §J.2 · p. 49

K.1 The initialisation-time estimator is null, on the whole predictor block

  • The image estimator based on the predictor’s two linear maps cannot distinguish ResidualPred because the maps are initialized identically.For the whole-block Jacobian estimator, seed ranges overlap at every tested depth; at depth 6, they are 0.313–0.326 for ResidualPred and 0.312–0.315 for the standard predictor.Source: §K.1 · p. 49

K.2 The target-aligned drive after a few predictor-only steps

  • After predictor-only equilibration, ResidualPred’s target-aligned drive exceeds the standard predictor’s in all 9 matched initializations across depths 2, 6, and 12.The leading real eigenvalue is also higher in all 9 comparisons, and paired drive ratios exceed one in all 36 cells at depths 2 and 6.Source: §K.2 · pp. 49–50
  • Effective rank is claimed to track accuracy gains only as a necessary, depth-limited quantity, not as a sufficient condition or a mediator through µ.On CIFAR-100 at predictor depth 12, SIGReg raises effective rank from 34.5 to 61.1 without changing accuracy.Source: §K.2 · p. 50

NeurIPS Paper Checklist

  • The authors report discussing linearity, spectral-decoupling assumptions, dataset-dependent tightness, and the local nature of µ = 1 as limitations.The checklist identifies these as topics in the paper’s limitations paragraph.Source: NeurIPS Paper Checklist · p. 51

Related papers

Loading 2610.02344v1…