A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures

Probleemstelling

Intelligente systemen hebben interne modellen van hun omgeving nodig voor planning en actie. Hoewel recente doorbraken in videogenatie en interactieve wereldsimulatoren indrukwekkend zijn, lijden ze onder fundamentele beperkingen:

1. Rekenkracht: Generatieve modellen moeten alle pixels reconstrueren, inclusief irrelevante details, wat enorme rekenbronnen vereist.
2. Toegankelijkheid: Bestaande implementaties van Joint-Embedding Predictive Architectures (JEPAs) zijn vaak complex, gericht op grote schaal, en moeilijk te navigeren voor onderzoekers die de theoretische principes willen begrijpen of snel willen experimenteren.
3. Kennislacune: Er ontbreekt een toegankelijke, modulaire bibliotheek die de brug slaat tussen de theorie van JEPAs en praktische toepassingen, van statische afbeeldingen tot actie-geconditioneerde wereldmodellen.

Methodologie: EB-JEPA Framework

Het paper introduceert EB-JEPA, een open-source bibliotheek die gebaseerd is op Energy-Based Models (EBMs) en Joint-Embedding Predictive Architectures (JEPAs). In plaats van pixels te reconstrueren, voorspellen JEPAs in een geleerde representatieruimte (latent space).

Kernprincipes:

• Energie-functie: Het model leert een energie-functie $E(x, y)$ te minimaliseren waarbij lage energie overeenkomt met compatibiliteit tussen input en output. Het doel is om de energie van correcte input-output paren lager te maken dan die van incorrecte paren.
• Voorspellingsverlies: De trainingsoptimalisatie bestaat uit een voorspellingsverlies ($L_{pred}$) in de representatieruimte, gecombineerd met een regularisatieterm ($R$) om "collapse" (het leren van triviale, constante representaties) te voorkomen:
  $$L = L_{pred}(g_\phi(z, u), z') + \lambda R(z)$$
  Waarbij $z$ de representatie is, $u$ conditionele informatie (zoals acties), en $g_\phi$ de predictor.

Drie Progressieve Implementaties:
De bibliotheek biedt drie voorbeelden van toenemende complexiteit:

1. Image-JEPA (Zelftoezicht): Leert view-invariantie door twee verschillende weergaven van dezelfde afbeelding te vergelijken. De energie meet de gelijkenis tussen de representaties van deze weergaven.
2. Video-JEPA (Temporele Voorspelling): Voorspelt de volgende frame-representatie op basis van een context van eerdere frames, zonder toegang tot de toekomstige frames tijdens de voorspelling. Dit leert temporele dynamiek.
3. Action-Conditioned Video-JEPA (Wereldmodellen voor Planning): Voegt acties toe aan de voorspelling. Het model leert hoe een actie ($a_t$) de volgende staat ($z_{t+1}$) beïnvloedt. Dit maakt het mogelijk om doelgerichte planning uit te voeren door een reeks acties te vinden die de energie (voorspellingsfout) naar een doeltoestand minimaliseert.

Regularisatie en Collapse Preventie:
Een kritiek onderdeel is het voorkomen van representatie-collapse. EB-JEPA implementeert twee families van regularisatie:

• VICReg: Gebruikt variantie- en covariantie-termen om te zorgen dat features voldoende spreiding hebben en niet gecorreleerd zijn.
• SIGReg: Een theoretisch onderbouwde alternatief dat de embeddingsverdeling forceert om een isotrope Gaussische verdeling te benaderen (via de Epps-Pulley test). Dit is stabieler en vereist minder hyperparameters.
• Aanvullende Regularisatie voor Wereldmodellen: Voor actie-geconditioneerde modellen worden extra termen toegevoegd:
  • Temporele Similariteit: Zorgt voor gladde trajecten in de representatieruimte.
  • Inverse Dynamics Model (IDM): Voorspelt acties op basis van opeenvolgende representaties. Dit is cruciaal om te voorkomen dat het model leert op spurious correlaties in willekeurige omgevingen.

Planning:
Voor planning wordt Model Predictive Path Integral (MPPI) gebruikt. Het optimaliseert een reeks acties door de cumulatieve energie over een geïmagineerde rollout naar een doeltoestand te minimaliseren. Het gebruik van cumulatieve kosten (over alle tijdstappen) werkt beter dan alleen het kijken naar de eindtoestand.

Belangrijkste Bijdragen

1. Toegankelijke Implementaties: Drie volledig werkende voorbeelden (afbeelding, video, planning) die binnen enkele uren op één GPU kunnen worden getraind.
2. Modulaire Architectuur: Herbruikbare componenten (encoders, predictors, regularizers, planners) die eenvoudig kunnen worden gecombineerd voor nieuwe toepassingen.
3. Uitgebreide Evaluatie: Systematische experimenten en ablatiestudies die het belang van elke component aantonen, inclusief praktische richtlijnen voor hyperparameter-selectie.
4. Educatieve Hulpbron: Duidelijke documentatie en code-structuur ontworpen om onderzoekers te helpen de principes van energie-gebaseerd zelftoezicht te begrijpen.

Resultaten

De auteurs evalueren EB-JEPA op drie benchmarks:

• CIFAR-10 (Afbeeldingen):
  • Lineaire probing resulteert in 91% nauwkeurigheid met SIGReg (en 90.12% met VICReg).
  • Het gebruik van een geleerde projector verbetert de prestaties met ongeveer 3 punten.
  • SIGReg toont zich stabieler over verschillende hyperparameter-configuraties dan VICReg.
• Moving MNIST (Video):
  • Multi-step autoregressive rollouts (k-stappen) verbeteren de voorspellingskwaliteit aanzienlijk ten opzichte van single-step voorspellingen.
  • Het model behoudt coherente voorspellingen van cijferbewegingen over lange horizons.
• Two Rooms Navigatie (Actie-geconditioneerd):
  • Het model bereikt een 97% succesrate bij het plannen van paden in een omgeving met willekeurige muren.
  • Ablatie-studies tonen kritieke inzichten:
    • Zonder de IDM-term stort het systeem in tot een succesrate van 1% (door spurious correlaties).
    • Variantie- en covariantie-termen dragen elk ongeveer 50% bij aan de verbetering.
    • Cumulatieve kosten voor planning werken 8% beter dan alleen kosten voor de eindtoestand.

Betekenis en Toekomstperspectief

EB-JEPA is een belangrijke stap in het democratiseren van onderzoek naar wereldmodellen en JEPAs.

• Snel Prototyping: Het stelt onderzoekers in staat om fundamentele vragen over regularisatie en architectuur te beantwoorden zonder de complexiteit van grote, gedistribueerde clusters.
• Theoretische Inzichten: Het bevestigt dat regularisatie (vooral IDM en variance/covariance) essentieel is voor het voorkomen van collapse in wereldmodellen.
• Schaalbaarheid: Hoewel het nu gericht is op kleine schaal, vormt het een basis voor toekomstige werk op grotere schaal, zoals hiërarchische wereldmodellen en het leren van kostfuncties.

Het paper concludeert dat EB-JEPA zowel een praktisch toolkit is voor onderzoekers als een educatieve bron om de overgang van generatieve pixel-modellen naar efficiëntere, representatie-gebaseerde wereldmodellen te begrijpen.