When Does LeJEPA Learn a World Model? — part 7/9

Cleaned-text code points [93753, 109302). Bibliography omitted here; source.md includes it.

The Hermite expansion and temporal correlation bounds used in our proof are closely related to classical results in probability theory. Specifically, the correlation bound (equation ( 4 )) is a consequence of the classical Hirschfeld–Gebelein–Rényi maximal correlation for jointly Gaussian variables [ 92 , 93 , 94 ] . Our contribution is not the bound itself but its application to identifiability via the measure-preservation constraint. As discussed above, Mehler’s formula, the key technical tool in our proof, is equivalent to the spectral expansion of the transition kernel associated with the Sturm–Liouville (SL) operator of Sprekeler et al. [72] in the Gaussian case. The SL route is more general (handles arbitrary latent variable distributions) but does not directly yield the quantitative bound (Thm. 5.3 ) that Mehler’s formula provides.

##### Practical differences: learned features versus fixed kernels.

While the theoretical optimization problems coincide, the practical implementations differ substantially. Sprekeler et al. [72] use fixed polynomial kernel expansions (degree 7 for two latent variables) followed by a linear eigenvalue problem, a two-stage procedure in which the feature space is chosen a priori and never adapted. In contrast, the LeJEPA/SIGReg pipeline learns the encoder end-to-end via gradient descent on a neural network, jointly optimizing the feature representation and the alignment plus regularization objectives. This learned-feature approach avoids the numerical instabilities of high-degree polynomial expansions noted by Sprekeler et al. [72] , scales to high-dimensional data, and is the standard paradigm in modern self-supervised learning. Our identifiability result thus provides theoretical grounding for an architecture and training procedure that is already in widespread use.

##### Connections to diffusion maps and spectral geometry.

Sprekeler [95] showed that SFA is closely related to Laplacian eigenmaps and diffusion maps, with the operator 𝒟 \mathcal{D} playing the role of a Laplace–Beltrami operator weighted by the temporal structure of the data. The eigenfunctions of 𝒟 \mathcal{D} are the same objects studied in spectral geometry and manifold learning (see also [ 73 ] ). Singer and Coifman [96] exploited a similar connection, using data-driven diffusion maps with a carefully chosen local metric to achieve nonlinear latent variable separation. This spectral-geometric perspective suggests that the identifiability phenomenon we study, encoders converging to specific linear transformations of the latents, may be an instance of a broader principle: that temporal or geometric structure on the data manifold, combined with appropriate regularization, constrains learned representations to align with the intrinsic coordinates of the generative process.

##### Summary of the relationship.

Table 3 summarizes the key similarities and differences. The mathematical core, Hermite eigenfunctions under Gaussian measure with temporal correlations yielding identifiability, is shared. Our contributions are (a) the simultaneous formulation with its clean orthogonal identifiability class, (b) the converse result identifying the Gaussian as the unique latent distribution yielding linear identifiability (Thm. 5.2 ), (c) the quantitative approximate identifiability bound under whitening alone (Thm. 5.3 ), (d) the observation that isotropic transitions are necessary for the simultaneous approach, and (e) the planning equivalence (Thm. 5.4 ) connecting linear identifiability to optimal latent-space planning.

## Appendix G Lean Verification

We formally verify all five theoretical results, namely the Hermite proof of Thm. 5.1 , the Gaussian uniqueness result (Thm. 5.2 ), the Dirichlet energy proof (Appendix E ), the approximate identifiability bound (Thm. 5.3 ), and the planning equivalence theorem (Thm. 5.4 ), in the Lean 4 theorem prover using the Mathlib mathematical library. The formalization compiles with zero sorry obligations: every logical step from axiomatized premises to stated conclusions is machine-checked.

##### Scope and methodology.

Lean 4 verification requires every inference step to be justified by a previously established lemma or axiom. When a standard mathematical result exists in Mathlib (e.g., Finset.sum_lt_sum for strict finite-sum monotonicity), we invoke it directly. When a result is standard but not yet available in Mathlib, either because the mathematical objects have not been formalized (Hermite polynomials) or because connecting available results to our specific statement forms would require substantial software engineering (Mazur–Ulam, AM-GM with uniform weights), we introduce it as a Lean axiom with a reference to its mathematical source. The complete reasoning chains between these axioms are fully verified. Table 4 gives the complete inventory. The formalization is organized into five files corresponding to the five results.

### G.1 What Is Axiomatized and Why

The axiomatized results fall into five categories:

##### 1. Hermite polynomial infrastructure.

Mathlib does not yet define Hermite polynomials or their properties under the Gaussian measure. The contraction property ( 21 ), Mehler’s formula ( 24 ), and the correspondence between degree-1 concentration and linearity are axiomatized for this reason. These are classical results in probability theory [ 79 , 94 ] ; their eventual formalization in Mathlib would allow these axioms to be discharged.

##### 2. Sturm–Liouville spectral theory.

The Gaussian uniqueness proof (Thm. 5.2 ) axiomatizes two results: that an affine score function uniquely determines a Gaussian density (standard exponential family characterization), and that a Gaussian density yields Hermite eigenfunctions with the identity as the first non-constant eigenfunction. The core algebraic step between these, namely extracting the affine score from the Sturm–Liouville eigenfunction equation via field arithmetic (including the sign argument that the score slope − λ 1 / K -\lambda_{1}/K is strictly negative), is fully verified.

##### 3. Standard analysis results requiring API software engineering.

The AM-GM inequality with uniform weights, Jensen’s inequality for exp \exp , and the Mazur–Ulam theorem are all available in Mathlib in closely related forms, but connecting them to our specific statement signatures requires nontrivial type-class and measure-theory software engineering. We axiomatize the exact statements we need and reference their Mathlib counterparts in the source code.

##### 4. Measure-theoretic and matrix-analytic facts.

The polar decomposition bound ( ‖ M − Q ‖ F ≤ ε + W nl \|M-Q\|_{F}\leq\varepsilon+W_{\mathrm{nl}} ) and the Pythagorean decomposition ( 𝔼 ⁡ [ ‖ h ⁡ ( z ) − Q ​ z ‖ 2 ] = ‖ M − Q ‖ F 2 + W nl \mathbb{E}[\|h(z)-Qz\|^{2}]=\|M-Q\|_{F}^{2}+W_{\mathrm{nl}} ) combine matrix analysis with integration under the Gaussian measure. These are stated as axioms with their mathematical content documented in the source.

##### 5. Trajectory pushforward for expected costs.

The planning equivalence theorem (Thm. 5.4 ) axiomatizes the pushforward relation between the learned-latent and true-latent expected per-step costs: for any test function c c , the expected value of c ⁡ ( z ^ t , a t ) c(\hat{z}_{t},a_{t}) under the pushforward dynamics starting from Q ​ z 0 Qz_{0} equals the expected value of c ⁡ ( Q ​ z t , a t ) c(Qz_{t},a_{t}) under the original dynamics starting from z 0 z_{0} . This is the measure-theoretic content of “the joint law of ( z ^ 0 , … , z ^ T ) (\hat{z}_{0},\ldots,\hat{z}_{T}) under the pushforward equals the joint law of ( Q ​ z 0 , … , Q ​ z T ) (Qz_{0},\ldots,Qz_{T}) under the original”; stating it at the level of expected values rather than joint distributions sidesteps the need to formalize stochastic dynamics in Lean while preserving the algebraic content of the theorem’s proof.

##### What is verified.

All reasoning chains between the axiomatized premises are fully machine-checked: the correlation bound and its strict equality characterization (the mathematical core of the Hermite proof), the algebraic chain from the SL eigenfunction equation to the affine score function and the biconditional assembly (the Gaussian uniqueness proof), the geometric chain from orthogonal Jacobian to linear isometry (the core of the Dirichlet proof), the spectral gap argument controlling nonlinear energy, the monotonicity and algebraic assembly of the approximate bound, and the recovery of the exact theorem as a special case. The inverse function theorem is incorporated as a structural assumption on the diffeomorphism rather than invoked as a separate lemma.

### G.2 Build Information

The project compiles against Lean 4 v4.28.0 with Mathlib v4.28.0 (8,032 build targets, zero errors, zero sorry obligations) and is available at https://github.com/klindtlab/lejepa-identifiability .

## Appendix H Experimental Details and Additional Results

### H.1 Nonlinear Mixing Functions

The four mixing functions g : ℝ 2 → ℝ 2 g:\mathbb{R}^{2}\to\mathbb{R}^{2} used in Figs. 1 , 3 are: 1. Spiral: g ⁡ ( z ) = R ⁡ ( π ​ ‖ z ‖ 2 ) ​ z g(z)=R(\pi\|z\|_{2})\,z , where R ⁡ ( θ ) = ( cos ⁡ θ − sin ⁡ θ sin ⁡ θ cos ⁡ θ ) R(\theta)=\bigl(\begin{smallmatrix}\cos\theta&-\sin\theta\\ \sin\theta&\cos\theta\end{smallmatrix}\bigr) . This norm-dependent rotation is a measure-preserving diffeomorphism of 𝒩 ⁡ ( 0 , I 2 ) \mathcal{N}(0,I_{2}) : at each point, g g acts as an orthogonal rotation ( | det J g | = 1 |\det J_{g}|=1 ), so it preserves the Gaussian density despite being highly nonlinear.

2. Sinusoidal shear: g ⁡ ( z 1 , z 2 ) = ( z 1 + sin ⁡ ( 1.5 ​ z 2 ) , z 2 ) g(z_{1},z_{2})=(z_{1}+\sin(1.5\,z_{2}),\;z_{2}) .

3. Parabolic shear: g ⁡ ( z 1 , z 2 ) = ( z 1 , z 2 + z 1 2 ) g(z_{1},z_{2})=(z_{1},\;z_{2}+z_{1}^{2}) .

4. RealNVP coupling layer: a single affine coupling block [ 81 ] of the form g ⁡ ( z 1 , z 2 ) = ( z 1 , z 2 ⊙ exp ⁡ ( s ⁡ ( z 1 ) ) + t ⁡ ( z 1 ) ) g(z_{1},z_{2})=(z_{1},\;z_{2}\odot\exp(s(z_{1}))+t(z_{1})) , where s , t s,t are small MLPs with random initialization. This is the same family used as the matched encoder in the scaling experiment (App. H.10 ); here we use it in 2D as an additional nonlinear diffeomorphism.

The first (spiral) map applies a rotation whose angle π ​ ‖ z ‖ 2 \pi\|z\|_{2} depends on the distance from the origin. Points at different radii are rotated by different angles, producing a characteristic spiral structure in observation space (Figure 3 , center). The map is a measure-preserving diffeomorphism of 𝒩 ⁡ ( 0 , I 2 ) \mathcal{N}(0,I_{2}) : at each point, g g acts as an orthogonal rotation ( | det J g | = 1 |\det J_{g}|=1 everywhere), so it preserves the Gaussian density. This is an interesting case because it entails SIGReg is already fulfilled in input space. However, g g is highly nonlinear; it entangles the radial and angular components of the latent space.

### H.2 Network Architecture and Training

The encoder f f is a 4-layer MLP with hidden dimension 256 and GELU activations. The architecture is: ℝ 2 → Linear ​ ( 2 , 256 ) → GELU → Linear ​ ( 256 , 256 ) → GELU → Linear ​ ( 256 , 256 ) → GELU → Linear ​ ( 256 , 2 ) ℝ 2 . \mathbb{R}^{2}\xrightarrow{\text{Linear}(2,256)}\xrightarrow{\text{GELU}}\xrightarrow{\text{Linear}(256,256)}\xrightarrow{\text{GELU}}\xrightarrow{\text{Linear}(256,256)}\xrightarrow{\text{GELU}}\xrightarrow{\text{Linear}(256,2)}\mathbb{R}^{2}. (74)

### H.3 Training Procedure

The loss function combines two terms with a balancing coefficient λ ∈ [ 0 , 1 ] \lambda\in[0,1] : ℒ = λ ​ ℒ SIG + ( 1 − λ ) ​ ℒ inv . \mathcal{L}=\lambda\,\mathcal{L}_{\mathrm{SIG}}+(1-\lambda)\,\mathcal{L}_{\mathrm{inv}}. (75) 1. Invariance loss: ℒ inv = 1 B ​ ∑ i = 1 B ‖ f ⁡ ( x i ) − f ⁡ ( x i ′ ) ‖ 2 \mathcal{L}_{\text{inv}}=\frac{1}{B}\sum_{i=1}^{B}\|f(x_{i})-f(x_{i}^{\prime})\|^{2} , where ( x i , x i ′ ) (x_{i},x_{i}^{\prime}) are positive pairs generated by sampling z i ∼ 𝒩 ⁡ ( 0 , I 2 ) z_{i}\sim\mathcal{N}(0,I_{2}) , generating z i ′ = ρ ​ z i + 1 − ρ 2 ​ η i z_{i}^{\prime}=\rho z_{i}+\sqrt{1-\rho^{2}}\eta_{i} via the Ornstein–Uhlenbeck channel, and mixing: x i = g ⁡ ( z i ) x_{i}=g(z_{i}) , x i ′ = g ⁡ ( z i ′ ) x_{i}^{\prime}=g(z_{i}^{\prime}) .

2. SIGReg regularization: The SIGReg regularizer of Balestriero and LeCun [2] encourages the empirical distribution of embeddings to match the standard Gaussian by penalizing deviations of the empirical sliced characteristic function from the Gaussian target.

### H.4 Hyperparameters

All experiments share the same training infrastructure and learning rate schedule: constant learning rate for the first half of training, followed by cosine decay to zero. Data is generated online (fresh samples each step, infinite data regime).

Hyperparameter 2D / Gennorm / Grid Scaling Optimizer AdamW AdamW Learning rate 3 × 10 − 3 3\times 10^{-3} 3 × 10 − 3 3\times 10^{-3} LR schedule warmup (10k) + cosine (10k) warmup (10k) + cosine (10k) Weight decay 0 0 Batch size 256 256 Training steps 20,000 20,000 Data regime online (infinite) online (infinite) Eval points 10,000 (fixed) 10,000 (fixed) Views per sample 2 2 ρ \rho 0.95 (2D/Gennorm), swept (Grid) 0.95 λ \lambda 10 − 3 10^{-3} (2D/Gennorm), swept (Grid) 10 − 6 10^{-6} Encoder MLP (spiral/banana/sinusoid) matched inverse-NVP [ 81 ] matched NVP (nvp mixing) Hidden dimension 256 (MLP) — MLP layers 4 — NVP coupling layers 8 (2D/Gennorm) 4 Activation GELU (MLP) tanh (NVP)

### H.5 Evaluation Metrics

After training, we evaluate the composed map h = f ∘ g h=f\circ g on a fixed evaluation set of 10,000 points. All metrics are computed on GPU via PyTorch for efficiency at large N N . We report:

1. Linear R 2 R^{2} (bidirectional): We fit linear maps z ^ = A ​ z + b \hat{z}=Az+b and h ^ = B ​ h + c \hat{h}=Bh+c via ordinary least squares and report R 2 ​ ( z → h ) R^{2}(z\to h) and R 2 ​ ( h → z ) R^{2}(h\to z) . Values close to 1 indicate that h h is well-approximated by a linear function of z z .

2. Orthogonality error: For the fitted linear map Q ^ \hat{Q} , we report ‖ Q ^ ⊤ ​ Q ^ − I n ‖ F / n \|\hat{Q}^{\top}\hat{Q}-I_{n}\|_{F}/\sqrt{n} . A value near 0 indicates that Q ^ \hat{Q} is orthogonal, consistent with Thm. 5.1 .

3. Approximate bound quantities: The covariance deviation ε = ‖ Cov ⁡ ( h ⁡ ( z ) ) − I n ‖ F \varepsilon=\|\mathrm{Cov}(h(z))-I_{n}\|_{F} , the alignment gap δ = ℒ ⁡ ( h ) − 2 ​ ( 1 − ρ ) ​ tr ​ ( Cov ⁡ ( h ⁡ ( z ) ) ) \delta=\mathcal{L}(h)-2(1-\rho)\,\mathrm{tr}(\mathrm{Cov}(h(z))) (clamped to ≥ 0 \geq 0 ), and the Orthogonal recovery error min Q ∈ O ⁡ ( n ) ⁡ 𝔼 ⁡ [ ‖ h ⁡ ( z ) − Q ​ z ‖ 2 ] \min_{Q\in O(n)}\mathbb{E}[\|h(z)-Qz\|^{2}] solved via SVD.

### H.6 Grid Search

We perform a 2D grid search over the regularization weight λ \lambda and the OU correlation ρ \rho : • λ ∈ { 10 − 6 , 10 − 5 , 10 − 4 , 10 − 3 , 5 × 10 − 3 , 10 − 2 , 5 × 10 − 2 , 10 − 1 , 5 × 10 − 1 } \lambda\in\{10^{-6},\;10^{-5},\;10^{-4},\;10^{-3},\;5\times 10^{-3},\;10^{-2},\;5\times 10^{-2},\;10^{-1},\;5\times 10^{-1}\}

• ρ ∈ { 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\}

For each of the 9 × 7 = 63 9\times 7=63 settings, we train 3 independent seeds, for a total of 189 runs. The extended λ \lambda range (down to 10 − 6 10^{-6} ) covers the regime where SIGReg has negligible weight, showing that Gaussianity regularization below λ = 10 − 4 \lambda=10^{-4} is insufficient for identifiability regardless of ρ \rho .