PaperlayerView exact revision on arXiv

Source-linked AI summary

When Does LeJEPA Learn a World Model?

David Klindt, Yann LeCun, Randall Balestriero

arXiv:2605.26379v1stat.MLcs.LG

TL;DR

The paper addresses whether LeJEPA representations recover the world's latent structure rather than merely support narrow tasks. It analyzes alignment with Gaussian regularization in stationary additive-noise worlds and proves that Gaussian latents uniquely yield linear identifiability, with approximate recovery and planning consequences. The theory is validated across synthetic, high-dimensional, distributional, and robotic settings.

  • Problem

    Whether JEPA representations recover world latents remained open because prior collapse-prevention methods did not specify an embedding distribution.

  • Method

    The paper analyzes LeJEPA's alignment and Gaussian regularization under stationary, independent, additive-noise latent transitions using spectral and Sturm–Liouville arguments.

  • Results

    LeJEPA is linearly identifiable if and only if the latent variables are Gaussian, with equality characterized by h(z) = Qz and approximate recovery bounded by objective deviations.

  • Takeaways & Limitations

    Linear orthogonal identifiability makes the learned representation usable for optimal latent-space planning and transfers orthogonally invariant costs from the true world.

  • Takeaways & Limitations

    The guarantee is population-level, assumes matching encoder and latent dimensions, and does not address sample-size scaling or training dynamics.

Abstract

from arXiv · show

A representation that scrambles the true degrees of freedom of the world cannot support reliable planning or compositional generalization. We prove that LeJEPA (alignment plus Gaussian regularization) linearly recovers the world's latent variables from nonlinear observations, a property known as linear identifiability, in a broad class of worlds where latents evolve under stationary, additive-noise transitions. Our main result is that among all such worlds, the Gaussian is the unique latent distribution for which this guarantee holds. The forward direction rests on a spectral decomposition in which each degree of nonlinearity is strictly penalized by alignment, making the linear map the optimum; the converse rules out every non-Gaussian alternative. We further prove an approximate identifiability result where the guarantee degrades gracefully, and show that linear, orthogonal identifiability enables optimal latent-space planning. We validate the theory with experiments ranging from 2D examples to 1024-dimensional latents, including distributional ablations and pixel-based robotic control. Our theory turns an empirically successful recipe into a mathematical guarantee, providing the foundation for building World Models that provably recover the structure of the world.

1 Introduction

The paper asks when self-supervised representations become faithful World Models rather than merely useful task representations. It closes the JEPA identifiability gap by proving conditions under which LeJEPA recovers latent variables linearly.

  • Motivation: Self-supervised learning seeks useful representations from unlabeled observations, and JEPAs align embeddings of related views while preventing collapse.These representations have shown effectiveness across images, video, and latent-space planning.
  • Motivation: A representation can entangle unrelated factors and still perform well on narrow tasks, motivating linear identifiability as a guarantee of latent recovery up to simple symmetries.
  • Problem: Prior JEPA methods lacked identifiability results because their collapse-prevention mechanisms left the embedding distribution unspecified.LeJEPA changes this by explicitly regularizing embeddings toward an isotropic Gaussian with SIGReg.
  • Contributions: LeJEPA is proved linearly identifiable if and only if latent variables are Gaussian in a broad class of stationary, independent, additive-noise worlds.The theory also covers approximate objectives, optimal latent-space planning, and empirical validation from 2D examples to 1024-dimensional latents and robotic control.

2 Related Work

Related work situates LeJEPA at the intersection of representation learning, World Models, and identifiability. The paper distinguishes its result by proving latent recovery for JEPA objectives under a broad world-process model.

  • Representation Learning: JEPAs predict in representation space, while SSL methods prevent collapse through negatives, stop-gradient teachers, covariance regularization, or self-distillation.LeJEPA adds explicit Gaussianity regularization through SIGReg.
  • World Models: World Models span cognitive science, control theory, cybernetics, and neural latent-dynamics models that support reasoning or control.
  • Identifiability: Nonlinear ICA generally requires additional structure for identifiability, including temporal dependence, auxiliary variables, augmentations, interventions, or supervision.The paper instead uses a Hermite approach applicable to arbitrary measurable maps.
  • LeJEPA Theory: LeJEPA's theory uses attraction and Gaussian regularization to force the learned representation to be a rotation of the true latents.

3 The World and the Learner

The paper models observations as nonlinear transformations of latent variables and trains a representation to invert that transformation up to unavoidable symmetries. Its theory focuses on independent Gaussian latents with stationary, additive-noise transitions and Gaussianized embeddings.

  • World and Learner: The world generates observations x = g(z) from unobserved latents z, while the learned representation f produces h = f ◦ g that should recover z up to rotation.The condition h(z) = Qz is necessary for linear probes to work.
  • World Assumptions: Positive pairs are generated by a stationary process whose views share the same marginal distribution and use additive noise transitions.The framework also assumes independent latent variables.
  • Gaussian World: The paper specializes the framework to z ∼ N(0, I_n), the maximum-entropy distribution for fixed mean and covariance.Under stationarity and additive noise, this leads to Ornstein–Uhlenbeck transitions with correlation parameter ρ.
  • Learner: LeJEPA combines an invariance loss that aligns positive pairs with a distributional regularizer that prevents collapse by shaping h(z) toward a Gaussian.Whitening fixes the representation's expected squared norm, so minimizing distance is equivalent to maximizing positive-pair correlation.

4 Spectral Analysis of the World

The spectral analysis explains why alignment favors linear recovery in Gaussian worlds and why the Gaussian case is exceptional. Transition-operator eigenfunctions expose a strict penalty for nonlinear components and motivate the converse result.

  • Transition Operator: The transition process induces a linear operator whose eigenfunctions and eigenvalues determine how strongly functions of the latent state remain correlated across views.
  • Gaussian World: For Gaussian worlds, Hermite polynomials form the eigenbasis, and degree-d components have eigenvalue ρ^d.
  • Gaussian World: E[hi(z′) hi(z)] ≤ ρ, with equality only for a linear hi, so every nonlinear distortion strictly reduces positive-pair correlation.
  • General Case: For general distributions, Sturm–Liouville theory gives a monotonic first non-constant eigenfunction, while linear identifiability requires it to be affine.The paper uses this constraint to show that only the Gaussian satisfies the needed condition.

5 Theory

The theory characterizes when LeJEPA recovers world latents linearly: Gaussian latent variables yield orthogonal identifiability, Gaussianity is uniquely necessary, and approximate optimization causes controlled error. Under orthogonal identifiability and rotation-invariant costs, latent-space planning is optimal.

  • 5.1 Forward Direction: LeJEPA Learns the World Model: LeJEPA linearly identifies Gaussian latent variables exactly up to an orthogonal transformation, including their transition dynamics.At equality, h(z)=Qz and h(z′)|h(z) remains Gaussian with the corresponding transformed transition.
  • 5.2 Converse: Gaussian Uniqueness: The Gaussian distribution is uniquely characterized by the requirement that every whitened LeJEPA minimizer be linear.An affine eigenfunction forces the latent score function to be linear, whose normalizable solution is Gaussian.
  • 5.3 Approximate Identifiability: Approximate alignment and whitening yield graceful identifiability degradation, with recovery error typically dominated by the alignment gap.The bound separates a nonlinearity term D from an orthogonality-distortion term (ε + D)^2; in practice, recovery scales as δ/[2ρ(1−ρ)].
  • 5.4 Optimal Latent Planning: Orthogonal identifiability makes latent-space planning mathematically equivalent to planning in the true world when costs are rotation-invariant.The learned and true problems have identical optimal actions and value under the stated invariance condition.

6 Experiments

Experiments support the theory across nonlinear 2D observations, dimensions up to 1024, latent-distribution sweeps, approximate-identifiability bounds, and pixel-based control. Recovery is strongest for Gaussian latents, while planning performance tracks linear identifiability.

  • 6.1 Forward: Linear Identifiability from Gaussian Latent Variables: LeJEPA inverts four nonlinear 2D mixing functions up to rotation, recovering the isotropic Gaussian structure predicted by Theorem 1.The tested mixings include norm-dependent rotation, sinusoidal and parabolic shears, and a RealNVP coupling layer.
  • 6.2 Converse: Non-Gaussian Latent Variables Break Linear Identifiability: Linear recovery peaks sharply at α = 2, the Gaussian member of the generalized normal family, across all three objectives.SIGReg and InfoNCE retain a wider plateau than VICReg for heavy-tailed latents.
  • 6.2 Converse: Non-Gaussian Latent Variables Break Linear Identifiability: OU-trained pixel encoders achieve R2 = 0.95 at ρ = 0.99, while trajectory data violate several assumptions simultaneously.The trajectory condition has non-Gaussian marginals, anisotropic ρ, and joint-limit wrapping.
  • 6.1 Forward: Linear Identifiability from Gaussian Latent Variables: SIGReg and VICReg maintain R2 > 0.999 up to N=1024, whereas InfoNCE degrades at scale under fixed kernel width σ=1.All methods use the same RealNVP mixing and matched encoder, isolating optimization behavior across dimensions.
  • 6.3 Approximate Identifiability: The approximate-identifiability bound holds across grid search, 2D mixings, scaling, and the latent-distribution sweep.The measured bound uses whitening error ε, alignment gap δ, and D; training loss is reported as a practical proxy for identifiability.
  • 6.4 Optimal Latent Planning: Gaussian-encoder latent plans are statistically indistinguishable from the oracle, while control cost decreases with linear-identifiability R2.The experiments test goal-reaching with rotation-invariant costs and compare learned latent plans against oracle behavior.

7 Limitations

The paper’s guarantees depend on assumptions about latent distributions, dimensionality, sampling, and optimization that may not hold in practice. Experiments also show reduced identifiability under goal-directed RL trajectories.

  • Real-world latent Gaussianity is unknowable from observations alone, so the theory’s distributional premise cannot be empirically established without additional assumptions.
  • RL-policy trajectories show reduced identifiability because anisotropy and non-Gaussian transitions violate the Gaussian OU setting.
  • When the encoder dimension is smaller than the latent dimension, Gaussianity does not determine the selected subspace or prevent superposition; larger dimensions require collapse or redundancy.
  • The population-level global-optimum result does not establish how alignment gaps, covariance deviations, or bound violations scale with sample size and training dynamics.

8 Discussion

The discussion connects identifiability to planning and practical data collection while documenting the paper’s mathematical machinery and remaining World Model scope. It emphasizes that state identifiability does not yet solve action-conditioned dynamics.

  • Linear identifiability makes learned representations usable for latent-space planning, while action-conditioned transition learning remains a separate challenge.
  • The Reacher experiment suggests isotropic exploration preserves the theory’s regime, whereas goal-directed policies concentrate data in a low-entropy latent region.
  • The framework defines the composed map h = f ◦ g from latent variables to representations and characterizes exact recovery up to an orthogonal transformation.
  • The proof develops Hermite-polynomial spectral decomposition, with degree capturing the order of nonlinear structure.
  • Approximate identifiability is parameterized by alignment gap, covariance deviation, and a normalized alignment gap.

A.3 Proof

The proof expands Gaussian-preserving encoder components in the Hermite basis and analyzes how the Ornstein–Uhlenbeck transition attenuates each degree. Since higher degrees lose more correlation, equality forces an orthogonal linear map.

  • The transition preserves linear structure best because Eη[He_k(z′)] = ρ^k He_k(z), so every higher-degree component is more attenuated than degree one.
  • The Hermite basis diagonalizes the transition: degree-d components contribute correlation ρ^d, with cross-terms vanishing across degrees.
  • Nonlinearity strictly reduces correlation, making the alignment optimum require all encoder variance to lie in the degree-one Hermite component.
  • The Gaussianity constraint converts the resulting unit-norm linear rows into an orthogonal matrix Q, yielding h(z) = Qz.
  • At the optimum, the learned representation preserves the latent conditional transition up to a global rotation.

B Proof of Theorem 2 (Gaussian Uniqueness)

The Gaussian-uniqueness proof links the slowest transition eigenfunction to the latent density’s score function. Requiring that eigenfunction to be affine forces the score, and therefore the density, to be Gaussian.

  • Theorem 2 states that if every covariance-constrained minimizer is linear, then the latent distribution must be Gaussian.
  • The transition operator’s leading nonconstant eigenfunction determines whether the encoder recovers the latent variable linearly or through a nonlinear transformation.
  • Sturm–Liouville theory guarantees that the first nonconstant eigenfunction is monotonic, but linear identifiability requires it to be affine.
  • Because the differential operator encodes the score function, an affine first eigenfunction forces a linear score and uniquely characterizes a Gaussian density.

C Proof of Theorem 3 (Approximate Identifiability)

Theorem 3 shows that approximate alignment and whitening force an encoder close to an orthogonal linear map in a Gaussian world. The proof bounds nonlinear energy through the OU spectral gap, then controls the linear component via covariance whitening.

  • Theorem 3: Approximate identifiability guarantees an orthogonal Q with h(z) close to Qz under approximate alignment and whitening.The Gaussian setup assumes alignment error δ and whitening error ε; the theorem then establishes recovery up to an orthogonal transformation.
  • Decomposition: Hermite expansions split h into a degree-1 linear part Mz and a nonlinear residual ν(z) containing degrees at least 2.The residual is uncorrelated with z, and Hermite orthogonality separates the two components without cross terms.
  • Spectral control: Mehler’s formula makes Hermite degree d an OU eigenspace with eigenvalue ρ^d, so higher-degree components incur larger alignment penalties.The penalty weight 1 − ρ^d increases with degree because higher-order components decorrelate faster under the OU transition.
  • Spectral control: The alignment gap δ controls nonlinear energy through the spectral-gap denominator 2ρ(1 − ρ).The proof compares the exact linear optimum with the approximate alignment bound to derive the nonlinear-energy estimate.
  • Orthogonal recovery: Approximate whitening forces M near orthogonal, while Hermite orthogonality combines linear and nonlinear errors into a Pythagorean recovery bound.The covariance decomposition separates the linear covariance from the nonlinear Gram matrix before applying the polar decomposition of M.
  • Planning consequence: The planning theorem applies when stage and terminal costs are O(n)-invariant, yielding equal optimal values under orthogonal latent reparameterization.Goal-reaching and LQR are stated as examples, while transition-model identifiability remains outside the theorem’s scope.

E.2 Proof

The second proof converts the small-noise alignment objective into Dirichlet energy and shows that its minimizers must be orthogonal maps. It uses Jacobian rigidity, measure preservation, and the Mazur–Ulam theorem to establish linearity.

  • Theorem 5: Theorem 5 proves identifiability via Dirichlet energy for C1 diffeomorphisms preserving the standard Gaussian measure.This proof provides a geometric alternative to the Hermite-based spectral argument.
  • Dirichlet reduction: An infinitesimal additive-noise expansion reduces the alignment objective to the expected squared Frobenius norm of the Jacobian.The Taylor expansion and dominated convergence justify the small-noise limit under the stated regularity assumptions.
  • Energy lower bound: Measure preservation makes the expected log-determinant of the Jacobian zero, enabling an AM-GM and Jensen lower bound on the energy.The Gaussian density is substituted into the change-of-variables formula to establish the determinant constraint.
  • Equality case: The energy is minimized at n, and every orthogonal map h(z) = Qz attains this lower bound.Equality conditions force unit singular values of the Jacobian at every point.
  • Linearity: An everywhere orthogonal Jacobian makes h a surjective isometry, so Mazur–Ulam yields h(z) = Qz + b; measure preservation then forces b = 0.The inverse map supplies the reverse distance inequality needed to establish global isometry.
  • Generalization: The geometric proof extends beyond Gaussian latent densities when the density is preserved by h, concluding that minimizers are isometries compatible with that density.The Gaussian assumption is unnecessary for the later rigidity steps, while the noise covariance condition remains required.

F Prior Work: Connection to Slow Feature Analysis

The paper connects LeJEPA's simultaneous alignment-and-whitening objective to Slow Feature Analysis, while showing that isotropic transitions and Gaussian latents are central to linear identifiability. It combines spectral theory, empirical tests, and formal verification to characterize this guarantee and its boundaries.

  • Connection to SFA: LeJEPA's alignment objective is equivalent to maximizing temporal correlation under whitening, matching the optimization problem of Slow Feature Analysis.For unit-variance outputs, minimizing SFA's slowness measure is equivalent to maximizing temporal correlation.
  • Identifiability class: The simultaneous objective identifies representations only up to orthogonal rotation because whitening is invariant under orthogonal transformations.Optimizing all output dimensions jointly cannot select individual latent axes, establishing orthogonal rather than permutation identifiability.
  • Transition assumptions: Isotropic transitions are necessary for simultaneous identifiability: anisotropy can replace the fastest latent's first Hermite component with a slower latent's second Hermite component.The observed failure follows the predicted eigenvalue interleaving under distinct transition rates.
  • Gaussian uniqueness: For arbitrary latent distributions, the first Sturm–Liouville eigenfunction is monotonic, but only Gaussian latents make it affine and therefore linearly identifiable.Non-Gaussian distributions support identifiability only up to a shared monotonic nonlinearity when latent variables share the same transition structure.
  • Formal verification: The theory is formally verified in Lean 4, with all five stated results machine-checked between their axiomatized premises and conclusions.The formalization includes the Hermite proof, Gaussian uniqueness, Dirichlet energy, approximate identifiability, and planning equivalence results.
  • Empirical evidence: The spiral mixing experiment highlights the limits of second-order constraints because it preserves the Gaussian measure while remaining highly nonlinear.For this mixing, the contrast between whitening and SIGReg is especially pronounced because covariance information alone does not constrain recovery.
  • Empirical evidence: Experiments show that recovery depends strongly on alignment quality, while stronger Gaussian regularization can produce a collapsed representation despite near-perfect whitening.The best reported grid-search regime reaches R2 > 0.97 with orthogonality error approximately 0.15; high regularization instead yields R2 approximately 0.

H.10 Scaling Experiment

The scaling experiment evaluates LeJEPA across latent dimensions up to 1024 and contrasts SIGReg, VICReg, and InfoNCE. Batch-statistic methods preserve linear identifiability at scale, whereas InfoNCE degrades because its kernel underflows in high dimensions.

  • Scaling setup: Across N = 2 to 1024, the RealNVP mixing remains consistently nonlinear, with R2(x →z) ≈0.73–0.78 for all methods.The matched inverse-coupling encoder removes expressivity as a confounder.
  • SIGReg: R2 > 0.9995 for SIGReg at every dimension, while orthogonality error grows from ∼10−5 at N=2 to ∼2 × 10−2 at N=1024.The slow orthogonality degradation matches the dimension-dependent approximate bound.
  • VICReg: VICReg is numerically indistinguishable from SIGReg across dimensions in both linear identifiability and orthogonality error.Its whitening loss remains small and stable at ∼4 × 10−7.
  • InfoNCE: InfoNCE reaches ∼0.9999 at N=16 but collapses to ∼0.65–0.72 for N ≥64, with substantially larger orthogonality error.At N ≥32, the Gaussian kernel underflows for positive and negative pairs, saturating the loss at zero and killing the gradient.
  • Planning: Gaussian-encoder plans decode to near-straight, oracle-quality joint-space trajectories, whereas warped trajectory encoders produce curved paths and higher control cost.This geometric result connects linear identifiability to planning quality.
Loading 2605.26379v1…