Technische Samenvatting: Wanneer leert LeJEPA een Wereldmodel?

Probleemstelling

Het artikel adresseert een fundamentele kloof in zelftoezicht leren (SSL): terwijl Joint-Embedding Predictive Architectures (JEPAs) empirisch succes hebben bewezen bij het leren van representaties voor visie, video en planning, is er geen theoretische garantie dat deze representaties een "Wereldmodel" vormen. Specifiek blijft het onduidelijk onder welke voorwaarden een geleerde representatie de onderliggende latente structuur van de wereld (bijvoorbeeld positie, snelheid, kleur) trouw herstelt in plaats van deze factoren te verweven.

Een representatie die de ware vrijheidsgraden verdraait (bijvoorbeeld het mengen van de positie van een object met zijn kleur) kan goed presteren op smalle taken, maar faalt wanneer de wereld verandert of wanneer deze wordt gebruikt voor compositionele generalisatie. De kernvraag is: Wanneer herstelt een geleerde representatie lineair de latente variabelen van de wereld? De auteurs definiëren deze eigenschap als lineaire identificeerbaarheid, een noodzakelijke voorwaarde voor betrouwbare lineaire probing en optimale planning in de latente ruimte.

Methodologie en Theoretisch Kader

Het Wereldmodel

De auteurs formaliseren de "Wereld" als een systeem met latente variabelen $z \in \mathbb{R}^n$ die niet worden waargenomen. De waargenomen data $x$ wordt gegenereerd door een onbekende, potentieel sterk niet-lineaire mengfunctie $g$, zodat $x = g(z)$. Het leerdoel omvat positieve paren $(z, z')$ gegenereerd door een stationair proces met onafhankelijke, additieve ruis-overgangen.

Het artikel richt zich op een specifieke klasse werelden waar:

1. Onafhankelijkheid: Latente variabelen en hun overgangen zijn onafhankelijk over dimensies.
2. Stationariteit: De marginale verdeling van de latente variabelen blijft constant tussen views.
3. Additieve Ruis: Overgangen volgen $z' = m(z) + \eta$, waarbij $\eta$ onafhankelijke ruis is.

Cruciaal specialiseren de auteurs deze klasse tot de Gaussische Wereld, waarbij $z \sim \mathcal{N}(0, I_n)$. Onder deze aannames corresponderen de overgangsdynamica met een Ornstein-Uhlenbeck (OU) proces: $z' = \rho z + \sqrt{1-\rho^2}\eta$.

De Leraar: LeJEPA

De leerarchitectuur is LeJEPA (LeCun's JEPA), die een encoder $f$ traint om embeddings $h(z) = f(g(z))$ te produceren door een verliesfunctie te minimaliseren die bestaat uit twee termen:

1. Alignementverlies: Minimaliseert de afstand tussen embeddings van positieve paren, $E[\|h(z') - h(z)\|^2]$. Dit moedigt de representatie aan om langzame, voorspelbare kenmerken vast te leggen.
2. SIGReg (Sketched Isotropic Gaussian Regularization): Regulariseert expliciet de embeddingsverdeling tot een isotrope Gaussische verdeling, $h(z) \sim \mathcal{N}(0, I_n)$, om ineenstorting van de representatie te voorkomen.

Het optimalisatieprobleem is:
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{onder voorwaarde} \quad h(z) \sim \mathcal{N}(0, I_n)$$

Theoretische Aanpak

De auteurs maken gebruik van spectrale analyse van de overgangoperator $T$, gedefinieerd als $(Th)(z) = E[h(z')|z]$.

• Hermite-ontwikkeling: Voor Gaussische latenten zijn de eigenfuncties van de overgangoperator de Hermite-polynomen. De eigenwaarden nemen geometrisch af als $\rho^d$, waarbij $d$ de graad van het polynoom is.
• Spectrale Straf: Het alignementverlies straft hogere-graads (niet-lineaire) componenten zwaarder dan lineaire componenten, omdat deze sneller decorreleren onder de overgang.
• Sturm-Liouville-theorie: Voor algemene (niet-Gaussische) verdelingen worden de eigenfuncties gekarakteriseerd door de Sturm-Liouville-theorie. De eerste niet-constante eigenfunctie is altijd monotoon, maar lineariteit is alleen gegarandeerd als de verdeling Gaussisch is.

Belangrijkste Bijdragen

1. Voorwaartse Richting: Lineaire Identificeerbaarheid voor Gaussische Werelden

Stelling 1 bewijst dat als de latente variabelen Gaussisch zijn en de overgang een OU-proces is, elke minimizer van het LeJEPA-doel (onderhevig aan Gaussische beperkingen) een lineaire orthogonale transformatie van de ware latenten moet zijn: $h(z) = Qz$ voor een zekere orthogonale matrix $Q$.

• Mechanisme: Het bewijs rust op een spectrale decompositie waarbij elke niet-lineariteit de correlatie tussen positieve paren strikt verlaagt. Omdat het alignementverlies ernaar streeft deze correlatie te maximaliseren, dwingt het globale optimum de representatie om puur lineair te zijn (Hermite-polynomen van graad 1).
• Resultaat: De geleerde representatie herstelt de latente variabelen tot op een globale rotatie, en de overgangsdynamica worden exact hersteld in de latente ruimte.

2. Omgekeerde Richting: Uniekheid van de Gaussische

Stelling 2 stelt vast dat de Gaussische verdeling de enige latente verdeling is waarvoor lineaire identificeerbaarheid geldt in deze klasse van werelden.

• Mechanisme: Met behulp van de Sturm-Liouville-theorie tonen de auteurs aan dat het eisen dat de eerste eigenfunctie van de overgangoperator affien (lineair) is, de scorefunctie van de latente verdeling lineair dwingt. Deze voorwaarde karakteriseert de Gaussische verdeling uniek.
• Implicatie: Als de latente variabelen niet-Gaussisch zijn, kan LeJEPA geen lineaire identificeerbaarheid garanderen; de geleerde representatie zal noodzakelijkerwijs een niet-lineaire transformatie van de latenten zijn.

3. Benaderende Identificeerbaarheid

Stelling 3 biedt een kwantitatieve bovengrens voor de herstelerror wanneer de doelen slechts bij benadering worden voldaan (d.w.z. imperfect alignement en imperfect whitening).

• De error schaalt vlot met de alignementkloof $\delta$ en de whitening-error $\epsilon$.
• De bovengrens benadrukt dat de alignementkloof de dominante factor is in de herstelerror, terwijl benaderend whitening relatief "gratis" te bereiken is.

4. Optimale Latente Planning

Stelling 4 verbindt lineaire identificeerbaarheid met control theory. Het bewijst dat als de encoder lineair identificeerbaar is ($h(z) = Qz$), elk optimaal controleprobleem met kosten die invariant zijn onder orthogonale transformaties (bijvoorbeeld Euclidische afstand tot een doel) optimaal kan worden opgelost in de geleerde latente ruimte. De optimale acties en waardenfuncties in de latente ruimte映射en direct naar de ware wereld, zonder dat aanpassing van de kostenfunctie nodig is.

Experimentele Resultaten

De auteurs valideren hun theorie in verschillende experimentele settings:

1. 2D Simulaties: Met behulp van diverse niet-lineaire mengfuncties (spiraal, schuif, RealNVP) herstelt LeJEPA de latente variabelen tot op rotatie wanneer de latenten Gaussisch zijn, wat Stelling 1 bevestigt.
2. Schalen naar Hoge Dimensies: Experimenten tot $N=1024$ dimensies tonen aan dat regularizers op basis van batch-statistieken (SIGReg, VICReg) een bijna perfecte lineaire identificeerbaarheid behouden ($R^2 > 0.999$), terwijl contrastieve methoden (InfoNCE) bij schaal degradatie door mismatch in kernelbreedte.
3. Ablaties op Verdeling: Het doorlopen van de vorm van de latente verdeling (van zware staarten tot uniform) toont aan dat lineair herstel scherp piekt bij de Gaussische verdeling ($\alpha=2$) en degradeert voor niet-Gaussische verdelingen, wat Stelling 2 valideert.
4. Pixel-gebaseerde Robotische Controle: Met behulp van de DeepMind Control Suite Reacher vergelijken de auteurs training op synthetische Gaussische overgangen (OU) versus echte RL-trajectoria.
   • OU-voorwaarde: Bereikt hoge lineaire identificeerbaarheid ($R^2 \approx 0.95$).
   • Trajectoria-voorwaarde: Real-world beleidslijnen induceren niet-Gaussische, anisotrope marginaal, wat resulteert in aanzienlijk lagere identificeerbaarheid ($R^2 < 0.5$).
5. Validatie van Planning: In de Reacher-taak decodeert rechte lijn-interpolatie in de geleerde latente ruimte van de Gaussisch-getrainde encoder naar bijna rechte paden in de gewrichtenruimte (optimale planning). Daarentegen produceert de op trajectoria getrainde encoder gebogen, suboptimale paden, wat Stelling 4 bevestigt.

Betekenis en Claims

Het artikel claimt het eerste identificeerbaarheidsresultaat voor JEPAs te bieden, waardoor een empirisch succesvol recept wordt omgezet in een wiskundige garantie.

• Fundamentele Garantie: Het werk stelt vast dat LeJEPA onder specifieke, goed gedefinieerde voorwaarden (Gaussische latenten, stationaire additieve ruis) de structuur van de wereld (latente variabelen en dynamica) bewijsbaar herstelt.
• Rol van Gaussichtheid: Het artikel keert het klassieke ICA-verhaal om. In lineaire ICA is Gaussichtheid de verdeling waar bronseparatie faalt. In deze niet-lineaire, zelftoezicht setting is Gaussichtheid de enige verdeling die lineaire identificeerbaarheid mogelijk maakt.
• Praktische Implicaties: De theorie verklaart waarom exploratiestrategieën die isotrope willekeurige wandelingen benaderen (data in het Gaussische regime houden) cruciaal zijn voor het leren van nuttige wereldmodellen, terwijl doelgerichte beleidslijnen (die data in laag-entropie gebieden laten instorten) de identificeerbaarheid kunnen verbreken.
• Fundering voor Planning: Door te bewijzen dat lineaire identificeerbaarheid optimale planning met rotatie-invariante kosten mogelijk maakt, biedt het artikel een theoretische rechtvaardiging voor het gebruik van geleerde latente ruimtes voor controle zonder dat terugmapping naar de ware fysieke coördinaten nodig is.

De auteurs merken beperkingen op: de theorie gaat ervan uit dat de encoderdimensie overeenkomt met de latentendimensie en rust op optimaliteit op populatieniveau. Echter, ze tonen aan dat de garanties vlot degradatie bij eindige steekproeven en optimalisatiefouten, en ze verifiëren formeel alle hoofdstellingen met de Lean 4 bewijstool.