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

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

In the univariate case, using the law of iterated expectations: 𝔼 ⁡ [ h ^ k ​ ( z ′ ) ​ h ^ j ​ ( z ) ] = 𝔼 z ​ [ 𝔼 η ​ [ h ^ k ​ ( z ′ ) ] ​ h ^ j ​ ( z ) ] . \mathbb{E}\!\left[\hat{h}_{k}(z^{\prime})\,\hat{h}_{j}(z)\right]=\mathbb{E}_{z}\!\left[\mathbb{E}_{\eta}\!\left[\hat{h}_{k}(z^{\prime})\right]\,\hat{h}_{j}(z)\right]. (22) By the contraction property ( 21 ), 𝔼 η ​ [ h ^ k ​ ( z ′ ) ] = ρ k ​ h ^ k ​ ( z ) \mathbb{E}_{\eta}[\hat{h}_{k}(z^{\prime})]=\rho^{k}\,\hat{h}_{k}(z) . Substituting: 𝔼 ⁡ [ h ^ k ​ ( z ′ ) ​ h ^ j ​ ( z ) ] = ρ k ​ 𝔼 z ​ [ h ^ k ​ ( z ) ​ h ^ j ​ ( z ) ] = ρ k ​ δ j ​ k , \mathbb{E}\!\left[\hat{h}_{k}(z^{\prime})\,\hat{h}_{j}(z)\right]=\rho^{k}\,\mathbb{E}_{z}\!\left[\hat{h}_{k}(z)\,\hat{h}_{j}(z)\right]=\rho^{k}\,\delta_{jk}, (23) where the last step uses orthonormality of the h ^ k \hat{h}_{k} under γ \gamma .

In the multivariate case, since z z has independent components and the noise in the Ornstein–Uhlenbeck channel acts independently on each coordinate, the multivariate Hermite products factorize: 𝔼 ⁡ [ H α ​ ( z ′ ) ​ H β ​ ( z ) ] = ∏ j = 1 n 𝔼 ⁡ [ h ^ α j ​ ( z j ′ ) ​ h ^ β j ​ ( z j ) ] = ∏ j = 1 n ρ α j ​ δ α j ​ β j = ρ | α | ​ δ α ​ β . \mathbb{E}\!\left[H_{\alpha}(z^{\prime})\,H_{\beta}(z)\right]=\prod_{j=1}^{n}\mathbb{E}\!\left[\hat{h}_{\alpha_{j}}(z_{j}^{\prime})\,\hat{h}_{\beta_{j}}(z_{j})\right]=\prod_{j=1}^{n}\rho^{\alpha_{j}}\,\delta_{\alpha_{j}\beta_{j}}=\rho^{|\alpha|}\,\delta_{\alpha\beta}. (24)

##### Step 3: Nonlinearity strictly reduces correlation.

Now we can compute the total correlation of any function through the transition. Since the Hermite basis diagonalizes the problem, the answer is simply a weighted sum over degrees, and because ρ d \rho^{d} decreases with d d , putting variance on higher degrees always costs correlation.

Expanding h i h_{i} in the Hermite basis ( 13 ) and applying Mehler’s formula ( 24 ): 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] = ∑ | α | ≥ 1 ∑ | β | ≥ 1 c α ( i ) ​ c β ( i ) ​ 𝔼 ​ [ H α ​ ( z ′ ) ​ H β ​ ( z ) ] = ∑ | α | ≥ 1 ( c α ( i ) ) 2 ​ ρ | α | . \mathbb{E}\!\left[h_{i}(z^{\prime})\,h_{i}(z)\right]=\sum_{|\alpha|\geq 1}\sum_{|\beta|\geq 1}c_{\alpha}^{(i)}c_{\beta}^{(i)}\,\mathbb{E}\!\left[H_{\alpha}(z^{\prime})\,H_{\beta}(z)\right]=\sum_{|\alpha|\geq 1}(c_{\alpha}^{(i)})^{2}\,\rho^{|\alpha|}. (25) Note that all cross-terms vanish ( δ α ​ β = 0 \delta_{\alpha\beta}=0 for α ≠ β \alpha\neq\beta ), and each surviving term contributes its squared coefficient times ρ | α | \rho^{|\alpha|} .

Grouping by degree d = | α | d=|\alpha| with w d = ∑ | α | = d ( c α ( i ) ) 2 w_{d}=\sum_{|\alpha|=d}(c_{\alpha}^{(i)})^{2} : 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] = w 1 ⋅ ρ + w 2 ⋅ ρ 2 + w 3 ⋅ ρ 3 + ⋯ \mathbb{E}\!\left[h_{i}(z^{\prime})\,h_{i}(z)\right]=w_{1}\cdot\rho+w_{2}\cdot\rho^{2}+w_{3}\cdot\rho^{3}+\cdots (26) with w d ≥ 0 w_{d}\geq 0 and w 1 + w 2 + w 3 + ⋯ = 1 w_{1}+w_{2}+w_{3}+\cdots=1 . This is a weighted average of ρ , ρ 2 , ρ 3 , … \rho,\rho^{2},\rho^{3},\ldots with weights summing to 1. Since 0 < ρ < 1 0<\rho<1 , each ρ d ≤ ρ \rho^{d}\leq\rho with strict inequality for d ≥ 2 d\geq 2 : w 1 ⋅ ρ + w 2 ⋅ ρ 2 + w 3 ⋅ ρ 3 + ⋯ ≤ w 1 ⋅ ρ + w 2 ⋅ ρ + w 3 ⋅ ρ + ⋯ = ρ ⋅ ( w 1 + w 2 + ⋯ ) = ρ . w_{1}\cdot\rho+w_{2}\cdot\rho^{2}+w_{3}\cdot\rho^{3}+\cdots\leq w_{1}\cdot\rho+w_{2}\cdot\rho+w_{3}\cdot\rho+\cdots=\rho\cdot(w_{1}+w_{2}+\cdots)=\rho. (27)

Equality requires w d ⋅ ρ d = w d ⋅ ρ w_{d}\cdot\rho^{d}=w_{d}\cdot\rho for every d d , which means w d = 0 w_{d}=0 for all d ≥ 2 d\geq 2 (since ρ d < ρ \rho^{d}<\rho for d ≥ 2 d\geq 2 ). So all the variance is at degree 1: w 1 = 1 w_{1}=1 . The degree-1 Hermite basis consists of the coordinate functions { z 1 , … , z n } \{z_{1},\ldots,z_{n}\} , so: h i ​ ( z ) = ∑ j = 1 n c e j ( i ) ​ z j = ⟨ q i , z ⟩ h_{i}(z)=\sum_{j=1}^{n}c_{e_{j}}^{(i)}\,z_{j}=\langle q_{i},z\rangle (28) for some vector q i q_{i} with ‖ q i ‖ 2 = w 1 = 1 \|q_{i}\|^{2}=w_{1}=1 . That is, h i h_{i} must be a linear function.

No nonlinear map can match the correlation of a linear map. This closes both gaps from the information-theoretic argument. That argument failed because it assumed the joint distribution of ( h ⁡ ( z ) , h ⁡ ( z ′ ) ) (h(z),h(z^{\prime})) is Gaussian, which need not hold for nonlinear h h . The Hermite analysis avoids this entirely: Mehler’s formula computes 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] \mathbb{E}[h_{i}(z^{\prime})h_{i}(z)] exactly for any h i ∈ L 2 ​ ( γ n ) h_{i}\in L^{2}(\gamma_{n}) , using only the joint Gaussianity of ( z , z ′ ) (z,z^{\prime}) , not of ( h i ​ ( z ) , h i ​ ( z ′ ) ) (h_{i}(z),h_{i}(z^{\prime})) . The prior argument also left open the possibility that a nonlinear map could match the correlation of a linear one. The strict inequality above rules this out: any nonlinearity ( w d > 0 w_{d}>0 for some d ≥ 2 d\geq 2 ) strictly reduces correlation.

##### Step 4: Assembly.

Summing the correlation bound over all n n components: ∑ i = 1 n 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] ≤ ρ ​ n , \sum_{i=1}^{n}\mathbb{E}[h_{i}(z^{\prime})h_{i}(z)]\leq\rho n, (29) which by ( 3 ) gives ℒ ​ ( h ) ≥ 2 ​ ( 1 − ρ ) ​ n \mathcal{L}(h)\geq 2(1-\rho)n .

At equality, every h i h_{i} is linear: h i ​ ( z ) = ⟨ q i , z ⟩ h_{i}(z)=\langle q_{i},z\rangle with ‖ q i ‖ = 1 \|q_{i}\|=1 . So h ⁡ ( z ) = Q ​ z h(z)=Qz for some matrix Q Q with unit-norm rows. The Gaussianity constraint requires Cov ⁡ ( h ⁡ ( z ) ) = Q ​ Q ⊤ = I n \mathrm{Cov}(h(z))=QQ^{\top}=I_{n} , so Q Q is orthogonal. Conversely, any orthogonal Q Q achieves ∑ i 𝔼 ⁡ [ ⟨ Q i , z ′ ⟩ ​ ⟨ Q i , z ⟩ ] = ρ ​ n \sum_{i}\mathbb{E}[\langle Q_{i},z^{\prime}\rangle\langle Q_{i},z\rangle]=\rho n , giving equality.

##### Step 5: Recovery of the transition.

At the optimum, h ⁡ ( z ′ ) = Q ​ z ′ = ρ ​ Q ​ z + 1 − ρ 2 ​ Q ​ η = ρ ​ h ​ ( z ) + 1 − ρ 2 ​ Q ​ η h(z^{\prime})=Qz^{\prime}=\rho\,Qz+\sqrt{1-\rho^{2}}\,Q\eta=\rho\,h(z)+\sqrt{1-\rho^{2}}\,Q\eta . Since Q Q is orthogonal and η ∼ 𝒩 ⁡ ( 0 , I n ) \eta\sim\mathcal{N}(0,I_{n}) , we have Q ​ η ∼ 𝒩 ⁡ ( 0 , I n ) Q\eta\sim\mathcal{N}(0,I_{n}) independently of h ⁡ ( z ) h(z) . Therefore h ⁡ ( z ′ ) | h ⁡ ( z ) ∼ 𝒩 ⁡ ( ρ ​ h ​ ( z ) , ( 1 − ρ 2 ) ​ I n ) h(z^{\prime})\mid h(z)\sim\mathcal{N}(\rho\,h(z),(1-\rho^{2})I_{n}) , which is the same conditional as z ′ | z z^{\prime}\mid z . The encoder recovers both the latent variables and the transition dynamics, up to a global rotation. ∎

## Appendix B Proof of Theorem 5.2 (Gaussian Uniqueness)

Since the latent variables are independent, the transition operator separates, and it suffices to prove the result for a single component z i z_{i} whose stationary distribution has full support on ℝ \mathbb{R} .

### B.1 Background

The forward result (Thm. 5.1 ) shows that the optimal encoder extracts the first few eigenfunctions of the transition operator T T (Section 4 ): these are the functions of the latent variables with the highest temporal autocorrelation. If the first (slowest) non-constant eigenfunction is linear, the encoder recovers the latent variable itself; if it is nonlinear, the encoder recovers a nonlinear transformation of the latent variable, and linear identifiability fails. The question is therefore: for which latent variable distributions is the first eigenfunction linear?

The proof uses the Sturm–Liouville characterization of these eigenfunctions [ 72 ] . Assumption (iii) ( η i \eta_{i} independent of z i z_{i} ) implies the noise variance does not depend on the state; in continuous time, this is the constant-diffusion regime in which Sturm–Liouville machinery applies. The transition operator T T has an infinitesimal generator 𝒟 \mathcal{D} that shares the same eigenfunctions; 𝒟 \mathcal{D} is a second-order differential operator whose structure makes the eigenvalue problem tractable (see Appendix F for the full connection to Slow Feature Analysis). For a scalar latent variable evolving under constant diffusion K > 0 K>0 with stationary density p ⁡ ( z i ) > 0 p(z_{i})>0 , the generator, written in self-adjoint form, acts on test functions φ \varphi as 𝒟 ​ φ = 1 p ⁡ ( z i ) ​ d d ​ z i ​ [ K ​ p ​ ( z i ) ​ φ ′ ​ ( z i ) ] . \mathcal{D}\varphi=\frac{1}{p(z_{i})}\frac{d}{dz_{i}}\Big[K\,p(z_{i})\,\varphi^{\prime}(z_{i})\Big]. (30) The eigenfunctions φ k \varphi_{k} are properties of the latent process, not of the encoder h h ; the encoder inherits its structure from them.

When the score function diverges ( ( log ⁡ p ) ′ ​ ( z i ) → ∓ ∞ (\log p)^{\prime}(z_{i})\to\mp\infty as z i → ± ∞ z_{i}\to\pm\infty ), the operator has a purely discrete spectrum 0 < λ 1 < λ 2 < ⋯ 0<\lambda_{1}<\lambda_{2}<\cdots with eigenfunctions φ 1 , φ 2 , … \varphi_{1},\varphi_{2},\ldots forming a complete orthogonal basis of L 2 ​ ( ℝ , p ​ d ​ z ) L^{2}(\mathbb{R},p\,dz) . The eigenvalues are ordered by increasing “speed”: φ 1 \varphi_{1} is the slowest non-constant feature, φ 2 \varphi_{2} the next slowest, and so on. By classical Sturm–Liouville oscillation theory, φ k \varphi_{k} has exactly k − 1 k-1 interior zeros. In particular, φ 1 \varphi_{1} is always monotonic, so one always gets identifiability up to a monotonic transformation. Linear identifiability requires φ 1 \varphi_{1} to be affine, and since an affine function is monotonic (zero interior zeros), only φ 1 \varphi_{1} can be affine: no higher eigenfunction can be. (For heavy-tailed densities such as Laplace, 𝒟 \mathcal{D} develops continuous spectrum and the discrete picture does not apply; this does not affect our result, since assuming an affine eigenfunction exists forces p p to be Gaussian, which has discrete spectrum.) For the full connection between the Sturm–Liouville framework and Slow Feature Analysis, see Appendix F .

### B.2 Proof

The key observation is that the operator 𝒟 \mathcal{D} encodes the score function ( log ⁡ p ) ′ (\log p)^{\prime} of the latent variable distribution. Demanding that the first eigenfunction is affine forces the score to be linear in z i z_{i} , and a linear score uniquely characterizes the Gaussian.

Suppose φ ⁡ ( z i ) = a ​ z i + b \varphi(z_{i})=az_{i}+b is an eigenfunction of 𝒟 \mathcal{D} with eigenvalue λ > 0 \lambda>0 . Since φ ′ = a \varphi^{\prime}=a is constant, the operator simplifies: 𝒟 ⁡ [ a ​ z i + b ] = a ​ K ​ ( log ⁡ p ) ′ ​ ( z i ) . \mathcal{D}[az_{i}+b]=a\,K(\log p)^{\prime}(z_{i}). (31) The eigenfunction equation 𝒟 ​ φ = − λ ​ φ \mathcal{D}\varphi=-\lambda\varphi requires a ​ K ​ ( log ⁡ p ) ′ ​ ( z i ) = − λ ⁡ ( a ​ z i + b ) a\,K(\log p)^{\prime}(z_{i})=-\lambda(az_{i}+b) , giving K ​ ( log ⁡ p ) ′ ​ ( z i ) = − λ ​ z i − λ ​ b / a . K(\log p)^{\prime}(z_{i})=-\lambda z_{i}-\lambda b/a. (32) Integrating both sides: log ⁡ p ⁡ ( z i ) = − λ 2 ​ K ​ ( z i + b a ) 2 + C , \log p(z_{i})=-\frac{\lambda}{2K}\left(z_{i}+\frac{b}{a}\right)^{2}+C, (33) which is a Gaussian density with mean − b / a -b/a . No other density satisfies this equation. Under the zero-mean convention b = 0 b=0 , and a a is fixed by unit-variance normalization. ∎

## Appendix C Proof of Theorem 5.3 (Approximate Identifiability)

We prove the stability bound in three steps: (1) the alignment gap controls nonlinear energy, (2) the whitening constraint forces the linear part near orthogonality, and (3) orthogonality of Hermite degrees gives a clean Pythagorean decomposition.

##### Setup and notation.

Since z ∼ 𝒩 ⁡ ( 0 , I n ) z\sim\mathcal{N}(0,I_{n}) , each zero-mean component h i h_{i} of the encoder admits a Hermite expansion (Appendix A.2 ): h i ​ ( z ) = ∑ | α | ≥ 1 c α ( i ) ​ H α ​ ( z ) , h_{i}(z)\;=\;\sum_{|\alpha|\geq 1}c_{\alpha}^{(i)}\,H_{\alpha}(z), (35) where the constant term vanishes by the zero-mean assumption. We split this into a linear part and a nonlinear residual.

The linear part is M ​ z Mz , where M ∈ ℝ n × n M\in\mathbb{R}^{n\times n} is the matrix of correlations between encoder outputs and latent variables: M i ​ j = 𝔼 ⁡ [ h i ​ ( z ) ​ z j ] . M_{ij}\;=\;\mathbb{E}\bigl[h_{i}(z)\,z_{j}\bigr]. (36) This extracts exactly the degree-1 Hermite coefficients, since all higher-degree terms are orthogonal to the coordinate functions z j z_{j} under the Gaussian measure.

The nonlinear residual ν ⁡ ( z ) = h ⁡ ( z ) − M ​ z \nu(z)=h(z)-Mz contains everything at Hermite degree ≥ 2 \geq 2 . By construction, 𝔼 ⁡ [ ν ⁡ ( z ) ​ z ⊤ ] = 0 \mathbb{E}[\nu(z)\,z^{\top}]=0 : the residual is uncorrelated with the latent variables. We define the total nonlinear energy as W nl = 𝔼 ⁡ [ ‖ ν ⁡ ( z ) ‖ 2 ] . W_{\mathrm{nl}}\;=\;\mathbb{E}\bigl[\|\nu(z)\|^{2}\bigr]. (37)

Since the linear and nonlinear parts live on orthogonal Hermite subspaces, the total variance decomposes without cross terms: tr ⁡ ( Cov ⁡ ( h ⁡ ( z ) ) ) = ‖ M ‖ F 2 + W nl . \mathrm{tr}\bigl(\mathrm{Cov}(h(z))\bigr)\;=\;\|M\|_{F}^{2}\;+\;W_{\mathrm{nl}}. (38) This Pythagorean decomposition is the key identity underlying Step 3 below.

### C.1 Step 1: Alignment gap controls nonlinear energy

Recall that the alignment loss decomposes as ℒ ⁡ ( h ) = 2 ​ n − 2 ​ ∑ i = 1 n 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] , \mathcal{L}(h)\;=\;2n\;-\;2\sum_{i=1}^{n}\mathbb{E}\!\left[h_{i}(z^{\prime})\,h_{i}(z)\right], (39) where ( z , z ′ ) (z,z^{\prime}) is a pair from the OU process with autocorrelation ρ \rho . Mehler’s formula (equation ( 24 ) in Appendix A ) states that Hermite polynomials diagonalize the OU conditional expectation: 𝔼 ⁡ [ H α ​ ( z ′ ) | z ] = ρ | α | ​ H α ​ ( z ) , \mathbb{E}\bigl[H_{\alpha}(z^{\prime})\,\big|\,z\bigr]\;=\;\rho^{|\alpha|}\,H_{\alpha}(z), (40) so each Hermite degree | α | |\alpha| is an eigenspace with eigenvalue ρ | α | \rho^{|\alpha|} . Substituting the Hermite expansion of h h into the cross-correlation 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] \mathbb{E}[h_{i}(z^{\prime})\,h_{i}(z)] and applying this eigenstructure yields an exact decomposition over degrees: ℒ ⁡ ( h ) = 2 ​ ∑ i = 1 n ∑ | α | ≥ 1 ( 1 − ρ | α | ) ​ ( c α ( i ) ) 2 . \mathcal{L}(h)\;=\;2\sum_{i=1}^{n}\sum_{|\alpha|\geq 1}(1-\rho^{|\alpha|})\,(c_{\alpha}^{(i)})^{2}. (41) The weight 1 − ρ | α | 1-\rho^{|\alpha|} increases with degree: higher-order components are penalized more heavily because they decorrelate faster under the OU transition.

We now split the sum by degree. At degree 1 the weight is exactly 1 − ρ 1-\rho . At degree | α | ≥ 2 |\alpha|\geq 2 , the bound ρ | α | ≤ ρ 2 \rho^{|\alpha|}\leq\rho^{2} (valid for 0 < ρ < 1 0<\rho<1 ) gives 1 − ρ | α | ≥ 1 − ρ 2 = ( 1 − ρ ) ​ ( 1 + ρ ) 1-\rho^{|\alpha|}\geq 1-\rho^{2}=(1-\rho)(1+\rho) , so ℒ ⁡ ( h ) ≥ 2 ​ ( 1 − ρ ) ​ ‖ M ‖ F 2 + 2 ​ ( 1 − ρ ) ​ ( 1 + ρ ) ​ W nl . \mathcal{L}(h)\;\geq\;2(1-\rho)\,\|M\|_{F}^{2}\;+\;2(1-\rho)(1+\rho)\,W_{\mathrm{nl}}. (42) Substituting the Pythagorean decomposition ‖ M ‖ F 2 = tr ⁡ ( Cov ⁡ ( h ⁡ ( z ) ) ) − W nl \|M\|_{F}^{2}=\mathrm{tr}(\mathrm{Cov}(h(z)))-W_{\mathrm{nl}} from ( 38 ) and collecting terms: ℒ ⁡ ( h ) \displaystyle\mathcal{L}(h) ≥ 2 ​ ( 1 − ρ ) ​ [ tr ⁡ ( Cov ⁡ ( h ⁡ ( z ) ) ) − W nl ] + 2 ​ ( 1 − ρ ) ​ ( 1 + ρ ) ​ W nl \displaystyle\;\geq\;2(1-\rho)\bigl[\mathrm{tr}(\mathrm{Cov}(h(z)))-W_{\mathrm{nl}}\bigr]+2(1-\rho)(1+\rho)\,W_{\mathrm{nl}} = 2 ​ ( 1 − ρ ) ​ tr ​ ( Cov ⁡ ( h ⁡ ( z ) ) ) + 2 ​ ( 1 − ρ ) ​ [ ( 1 + ρ ) − 1 ] ​ W nl \displaystyle\;=\;2(1-\rho)\,\mathrm{tr}(\mathrm{Cov}(h(z)))+2(1-\rho)\bigl[(1+\rho)-1\bigr]\,W_{\mathrm{nl}} = 2 ​ ( 1 − ρ ) ​ tr ​ ( Cov ⁡ ( h ⁡ ( z ) ) ) + 2 ​ ρ ​ ( 1 − ρ ) ​ W nl . \displaystyle\;=\;2(1-\rho)\,\mathrm{tr}(\mathrm{Cov}(h(z)))+2\rho(1-\rho)\,W_{\mathrm{nl}}. (43)