When Does LeJEPA Learn a World Model? — part 8/9
Cleaned-text code points [109302, 124914). Bibliography omitted here; source.md includes it.
Figure 6 shows the results. Three regimes are visible. First, when Gaussianity regularization is too strong ( λ = 0.5 \lambda=0.5 ), the encoder collapses to a non-informative Gaussian representation with R 2 ≈ 0 R^{2}\approx 0 , regardless of ρ \rho : the regularizer overwhelms the alignment signal. Second, at low λ \lambda ( ≤ 10 − 2 \leq 10^{-2} ), identifiability improves monotonically with ρ \rho , consistent with the theoretical prediction that the correlation bound tightens as ρ → 1 \rho\to 1 . The best performance ( R 2 > 0.97 R^{2}>0.97 , orthogonality error ≈ 0.15 \approx 0.15 ) is achieved at λ ∈ { 10 − 3 , 5 × 10 − 3 } \lambda\in\{10^{-3},5\times 10^{-3}\} with ρ ∈ { 0.9 , 0.95 } \rho\in\{0.9,0.95\} . Third, at intermediate λ \lambda ( 5 × 10 − 2 5\times 10^{-2} to 10 − 1 10^{-1} ), there is a non-monotonic interaction: performance peaks at moderate ρ \rho and degrades at ρ = 0.99 \rho=0.99 . This occurs because very high correlation makes the invariance loss trivially small, shifting the loss balance toward Gaussianity and pushing the system into the collapse regime.
### H.7 Latent Variable Distribution Sweep Across Mixings
The main-text gennorm figure (Fig. 4 b) shows linear recovery on the spiral mixing across latent distribution shapes α ∈ { 2 − 3 , … , 2 5 } \alpha\in\{2^{-3},\ldots,2^{5}\} in the generalized normal family ( α = 2 \alpha=2 Gaussian, α = 1 \alpha=1 Laplace, α → ∞ \alpha\to\infty uniform). We extend this sweep to all four 2D mixings used in our paper (Figs. 7 , 8 ).
Two patterns are robust across mixings. First, both methods peak at α = 2 \alpha=2 , confirming Thm. 5.2 empirically: linear identifiability is maximal precisely at the Gaussian latents. Second, SIGReg’s plateau is consistently wider than whitening’s. Whitening collapses sharply for α < 2 \alpha<2 (heavier-than-Gaussian tails), whereas SIGReg retains high recovery across a broad range of latent distribution shapes. This confirms that SIGReg’s Gaussianization of h h is doing more than enforcing covariance: it absorbs latent-side non-Gaussianity that whitening cannot see.
The contrast between methods is most pronounced for the spiral mixing, the only one of the four that is approximately measure-preserving and therefore unconstrained by second-order statistics alone. For the banana, sinusoid, and NVP mixings, the mixing structure itself partially constrains the recovery, narrowing the gap between methods. The orthogonality error (Fig. 8 ) tells the same story from the geometric side: whitening achieves a sharp minimum at α = 2 \alpha=2 where the second-order constraint exactly matches the latents, while SIGReg’s basin is broader.
### H.8 Approximate Bound Decomposition
Figure 9 decomposes the recovery error into its two sources: the whitening error ε \varepsilon and the alignment gap δ \delta . The two quantities have distinct roles. Large ε \varepsilon alone does not predict poor recovery: the λ = 5 × 10 − 1 \lambda=5\times 10^{-1} runs (crosses) achieve near-zero ε \varepsilon because SIGReg dominates, yet their recovery error is maximal because the alignment signal is absent ( δ \delta is large). Conversely, the alignment gap δ \delta is a much stronger predictor of actual recovery error, consistent with the interpretation that the D D term in the bound captures the nonlinear energy that alignment fails to suppress. This asymmetry confirms the practical takeaway from Thm. 5.3 : approximate whitening is easy to achieve in practice, so the binding constraint on identifiability is the quality of the alignment objective.
### H.9 Loss Predictivity
Figure 10 pools all trained encoders across our four experiments to examine whether the training losses are predictive of linear identifiability. We restrict the view to converged runs ( R 2 > 0.9 R^{2}>0.9 ) to focus on the regime where the theory applies. The alignment loss is the strongest single predictor of R 2 R^{2} , consistent with its direct connection to the alignment gap δ \delta in Thm. 5.3 .
### H.10 Scaling Experiment
We test whether identifiability holds at scale by sweeping the latent dimension N ∈ { 2 , 4 , 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 } N\in\{2,4,8,16,32,64,128,256,512,1024\} with 5 seeds each. The mixing function is a RealNVP-style coupling layer architecture (4 layers with orthogonal weight matrices) [ 81 ] , and the encoder is a matched inverse-coupling-layer architecture with learnable weights. This architectural matching removes encoder expressivity as a confounder: any failure of identifiability is due to the optimization landscape, not the encoder’s function class. For N ≤ 32 N\leq 32 , we train K = 3 K=3 encoders per seed and select the one with lowest final loss; for N > 32 N>32 , all runs converge to equivalent solutions ( K = 1 K=1 ).
Tabs. 5 – 7 report the per-method results. The mixing difficulty is consistently nonlinear across dimensions ( R 2 ( x → z ) ≈ 0.73 R^{2}(x\to z)\approx 0.73 – 0.78 0.78 ) for all methods, since the RealNVP mixing is shared. The three methods differ sharply in their failure modes:
SIGReg (Tab. 5 ). Linear identifiability is essentially perfect at all dimensions ( R 2 > 0.9995 R^{2}>0.9995 ). Alignment and SIGReg losses are stable across N N . The only quantity that drifts is the orthogonality error ‖ Q ^ ⊤ Q ^ − I ‖ F / N \|\hat{Q}^{\top}\hat{Q}-I\|_{F}/\sqrt{N} , which grows from ∼ 10 − 5 \sim 10^{-5} at N = 2 N{=}2 to ∼ 2 × 10 − 2 \sim 2\times 10^{-2} at N = 1024 N{=}1024 . This is the slow degradation predicted by Thm. 5.3 : at fixed alignment quality, the bound on the linear-part distortion grows with the dimension being whitened.
VICReg (Tab. 6 ). Numerically indistinguishable from SIGReg across the entire sweep, on both linear identifiability and orthogonality error. The whitening loss is small ( ∼ 4 × 10 − 7 \sim 4\times 10^{-7} ) and stable. Confirms that for Gaussian latents at the population optimum, second-moment whitening is sufficient (the Gaussianity assumption on h ( z ) h(z) in Thm. 5.1 can be relaxed to whitening, see remark after Thm. 5.3 ).
InfoNCE (Tab. 7 ). Qualitatively different. R 2 R^{2} is non-monotone in N N : high at N = 16 N{=}16 ( ∼ 0.9999 \sim 0.9999 ), then collapsing to ∼ 0.65 \sim 0.65 – 0.72 0.72 for N ≥ 64 N\geq 64 . Orthogonality error is one to three orders of magnitude larger than for the batch-statistic methods. The InfoNCE loss column reveals the mechanism: at N ≥ 32 N\geq 32 , the loss saturates at zero, indicating that the Gaussian kernel exp ( − ∥ ⋅ ∥ 2 / σ 2 ) \exp(-\|\cdot\|^{2}/\sigma^{2}) underflows for both positive and negative pairs once ‖ z ‖ 2 ≈ N \|z\|^{2}\approx N at σ = 1 \sigma{=}1 , killing the gradient. The non-monotone R 2 R^{2} is the fingerprint of this kernel-width mismatch, not a property of the InfoNCE objective itself; per-dimension kernel tuning would presumably restore performance, at the cost of a hyperparameter sweep absent from the batch-statistic methods.
These three patterns make the main-text story precise: the practical gap between methods is not in their theoretical identifiability guarantees (Thm. 5.1 applies to all), but in how each method’s optimization behaves at scale.
### H.11 Reacher Experiment
##### Environment and rendering.
We use the DeepMind Control Suite Reacher (“hard” variant) [ 80 ] with two revolute joints: a shoulder (unbounded) and a wrist (limited to [ − 2.79 , 2.79 ] [-2.79,2.79] rad). The latent state is z = ( θ 0 , θ 1 ) ∈ ℝ 2 z=(\theta_{0},\theta_{1})\in\mathbb{R}^{2} (Fig. 11 ). Images are rendered at 64 × 64 64\times 64 pixels via MuJoCo with EGL headless rendering. The target position is fixed at ( 0.1 , 0.1 ) (0.1,0.1) across all experiments to eliminate it as a confound. For each experimental condition, we pre-render 100,000 image pairs and 10,000 evaluation images (i.i.d. Gaussian samples, shared across all conditions). Images are stored as uint8 and normalized per-channel at training time.
OU condition. For each ρ ∈ { 0.3 , 0.5 , 0.7 , 0.8 , 0.9 , 0.95 , 0.99 } \rho\in\{0.3,0.5,0.7,0.8,0.9,0.95,0.99\} , we sample z ∼ 𝒩 ( 0 , I 2 ) z\sim\mathcal{N}(0,I_{2}) , generate z ′ z^{\prime} via ( 1 ), and render. The theory’s assumptions (Gaussian marginals, isotropic OU transitions, additive noise) are exactly satisfied. Wrapping beyond ± π \pm\pi occurs for < 0.2 % <0.2\% of samples.
Trajectory condition. We use 10,000 episodes (201 steps each) from the LeWorldModel [ 9 ] dataset, collected by a trained SAC policy. For each temporal stride δ ∈ { 1 , 2 , 4 , 8 , 16 , 32 , 64 } \delta\in\{1,2,4,8,16,32,64\} , we subsample 10 pairs per episode (100,000 total), extract the joint angles ( θ 0 t , θ 1 t ) (\theta_{0}^{t},\theta_{1}^{t}) and ( θ 0 t + δ , θ 1 t + δ ) (\theta_{0}^{t+\delta},\theta_{1}^{t+\delta}) , and re-render with the fixed target. The marginal distributions (Figure 12 ) violate the Gaussian assumption: the shoulder is broad and heavy-tailed ( 13.6 % 13.6\% beyond ± π \pm\pi ), while the wrist is nearly uniform over [ − π , π ] [-\pi,\pi] . Autocorrelations are anisotropic (e.g., at δ = 8 \delta=8 : ρ 0 = 0.992 \rho_{0}=0.992 , ρ 1 = 0.982 \rho_{1}=0.982 ).
##### Encoder architecture.
We use a CNN with BatchNorm: Conv ( 3 → 32 , 4 × 4 , s = 2 ) → BN → GELU → Conv ( 32 → 64 ) → BN → GELU \displaystyle\texttt{Conv}(3\to 32,4\times 4,s{=}2)\to\texttt{BN}\to\texttt{GELU}\to\texttt{Conv}(32\to 64)\to\texttt{BN}\to\texttt{GELU} → Conv ( 64 → 128 ) → BN → GELU → Conv ( 128 → 256 ) → BN → GELU \displaystyle\to\texttt{Conv}(64\to 128)\to\texttt{BN}\to\texttt{GELU}\to\texttt{Conv}(128\to 256)\to\texttt{BN}\to\texttt{GELU} → AvgPool ( 4 ) → Flatten → Linear ( 256,256 ) → BN → GELU → Linear ( 256 , 2 ) . \displaystyle\to\texttt{AvgPool}(4)\to\texttt{Flatten}\to\texttt{Linear}(256,256)\to\texttt{BN}\to\texttt{GELU}\to\texttt{Linear}(256,2). This architecture has ∼ 1.1 {\sim}1.1 M parameters. BatchNorm was critical for stable training: without it, ∼ 36 % {\sim}36\% of runs collapsed to zero-variance outputs. Gradient clipping ( ‖ ∇ ‖ max = 1.0 \|\nabla\|_{\max}=1.0 ) is applied after each backward pass.
##### Training.
We train with AdamW (lr = 3 × 10 − 3 =3\times 10^{-3} , weight decay 10 − 4 10^{-4} ) and cosine annealing for 100 epochs, batch size 256, with SIGReg using 256 random slices. The loss is ℒ = λ ℒ SIG + ( 1 − λ ) ℒ inv \mathcal{L}=\lambda\mathcal{L}_{\mathrm{SIG}}+(1-\lambda)\mathcal{L}_{\mathrm{inv}} . We sweep λ ∈ { 10 − 3 , 5 × 10 − 3 , 10 − 2 , 5 × 10 − 2 } \lambda\in\{10^{-3},5\times 10^{-3},10^{-2},5\times 10^{-2}\} with 3 seeds per setting, reporting the best λ \lambda per condition (mean ± \pm std over seeds). The final model (after 100 epochs) is always used; no early stopping or checkpoint selection is applied.
##### Evaluation.
We fit a linear regression from embeddings to joint angles on 10,000 training images and score on 10,000 held-out evaluation images (proper train/test split). We report bidirectional R 2 R^{2} , per-dimension R 2 R^{2} , and an R 2 R^{2} diagnostic for sin / cos \sin/\cos targets ( sin θ 0 , cos θ 0 , sin θ 1 , cos θ 1 ) (\sin\theta_{0},\cos\theta_{0},\sin\theta_{1},\cos\theta_{1}) . The sin / cos \sin/\cos diagnostic tests whether the encoder learns a trigonometric representation of the cyclic angles; negative values indicate the 2D embedding encodes raw angles (linear in θ \theta ) rather than their trigonometric functions.
##### Latent distributions.
Figure 12 shows the stationary joint-angle distribution and, for each stride δ \delta , the 2D transition-difference distribution together with the per-dimension autocorrelation scatter. The stationary marginal (leftmost panel) reveals a striking asymmetry: the shoulder ( z 0 z_{0} ) is broad and mildly platykurtic, while the wrist ( z 1 z_{1} ) is nearly bimodal, reflecting a policy-induced preference for two joint-limit configurations. The transition-difference panels (top row) grow from tight, narrow clouds at small δ \delta to fully-developed isotropic clouds at δ = 64 \delta=64 . The autocorrelation scatter (bottom row) confirms that both joints have ρ ≈ 1 \rho\approx 1 at δ = 1 \delta=1 (the transition is nearly trivial) and decorrelate monotonically, with the wrist decorrelating faster. The per-dimension R 2 R^{2} values annotated on each panel already hint at the core phenomenon examined in the next paragraph: both small and large δ \delta are bad, and identifiability peaks in an intermediate regime.
##### Gaussianity and autocorrelation jointly determine identifiability.
Thm. 5.1 requires two conditions on the data-generating process: Gaussian latent marginals and a non-trivial autocorrelation ρ ∈ ( 0 , 1 ) \rho\in(0,1) . Thm. 5.3 relaxes these to approximate whitening and a non-vanishing spectral gap. The trajectory condition satisfies neither exactly, and the temporal stride δ \delta controls how close each condition comes to being met. To disentangle these two effects, we measure the (zscored) SIGReg value of each transition-difference distribution and of their 2D joint, and plot it against the Pearson autocorrelation, colored by identifiability R 2 R^{2} (Figure 13 ). We compare to a finite-sample floor obtained by evaluating SIGReg on iid standard Gaussian samples of matched size; this floor is approximately 1.2 1.2 , so any value near 1 1 – 2 2 is indistinguishable from Gaussian at this sample size.
Three regimes are visible. At δ = 1 , 2 \delta=1,2 (top right of each panel), ρ \rho is essentially 1 1 and SIGReg is two or three orders of magnitude above the floor: the transition is too trivial to carry an identifiability signal and the differences inherit the non-Gaussian structure of the stationary marginals, because at this stride the differences z t + δ − z t z_{t+\delta}-z_{t} are tiny perturbations that cannot Gaussianize via any averaging. At δ = 64 \delta=64 (far left), ρ \rho has dropped to 0.86 0.86 – 0.92 0.92 and the shoulder differences are near the Gaussian floor, but the wrist differences climb again because the increments at this stride start to resolve the bimodal stationary distribution (they approach a symmetrized sum of two draws from a bimodal law). At δ = 8 , 16 \delta=8,16 , both constraints are satisfied simultaneously: ρ ∈ [ 0.96 , 0.99 ] \rho\in[0.96,0.99] gives a meaningful spectral gap 2 ρ ( 1 − ρ ) 2\rho(1-\rho) , and SIGReg drops to within a few times the floor for both marginals and the joint. This is precisely the regime where R 2 R^{2} is maximal, as shown by the yellow points clustered in the bottom-middle region of each panel. The best OU ρ \rho value (vertical dashed line, ρ = 0.95 \rho=0.95 ) sits near these high- R 2 R^{2} points, confirming that the trajectory condition most closely approaches the theoretical assumptions at intermediate stride.
The joint panel (rightmost) tells a consistent story at slightly higher absolute values: the 2D SIGReg of the transition-difference distribution is always larger than either marginal because it additionally penalizes the cross-dimension dependence that neither marginal can see. The R 2 R^{2} ordering is the same as in the per-dimension panels.
##### Additional results on OU and Trajectory conditions.
Fig. 14 directly contrasts the two: at matched ρ \rho , Gaussian (OU) latents achieve substantially higher R 2 R^{2} than SAC-policy trajectories (right panel); an empirical instance of Thm. 5.2 on raw pixels. Within OU alone (left), R 2 R^{2} rises monotonically in ρ \rho , and stronger λ \lambda helps only at low ρ \rho where the alignment signal is weak.