Beyond Gaussian Worlds: Latent Geometry Matters for JEPAs

2026-09-18

AI Summary

本文探讨了非欧几里得潜在几何对JEPAs的影响，通过匹配适当的分布来解决表示崩溃问题，并证明了在某些条件下可以实现线性恢复。

Abstract

Recent Joint-Embedding Predictive Architectures (JEPAs) prevent representation collapse by constraining learned representations to follow a prescribed target distribution, such as an isotropic Gaussian or the uniform distribution on a hypersphere. Klindt et al. (2026) showed that, under their Euclidean assumptions, matching a Gaussian target can recover Gaussian latent variables up to a linear transformation, and that the Gaussian is the unique distribution with this guarantee. We extend their analysis to latent variables supported on embedded Riemannian manifolds and derive conditions on the latent geometry and positive-pair dynamics under which alignment and exact distribution matching guarantee linear recovery. In particular, when the latent variables are uniformly distributed on a sphere and the representations are matched to the same spherical distribution, every optimal representation recovers the latent state up to an orthogonal transformation. This shows that Gaussian uniqueness is not a universal property of distribution-matched JEPAs: non-Euclidean latent geometries can admit other linearly recoverable distributions. We further derive an approximate-recovery bound that is strictly tighter for the spherical world than for the Gaussian world. Experiments on Gaussian, spherical, and toroidal latent spaces show that geometrically compatible targets yield better linear recovery when optimization succeeds, whereas mismatched targets distort the latent structure. This advantage persists in high-dimensional Clifford-torus worlds.

Problem

• JEPAs
• latent geometry
• Riemannian manifolds
• linear recovery
• distribution matching

Innovation

• Riemannian manifolds
• latent geometry
• linear recovery
• spherical distribution
• distribution-matched JEPAs