LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

Dit paper introduceert LeWorldModel, het eerste stabiele end-to-end Joint-Embedding Predictive Architecture (JEPA) dat direct van ruwe pixels leert met slechts twee verliesfuncties, waardoor het aanzienlijk sneller en efficiënter is dan bestaande wereldmodellen terwijl het toch robuuste fysische structuren leert voor controle en anomaliedetectie.

De kern

Stel je voor dat je een kind bent dat leert hoe de wereld werkt. Je ziet een bal die rolt, je duwt hem, en hij stopt ergens anders. Je brein bouwt een onzichtbaar model op van hoe dingen bewegen, zonder dat iemand je de wiskundige formules uitlegt. Je "voelt" gewoon dat als je harder duwt, de bal verder gaat.

Dit is precies wat kunstmatige intelligentie (AI) probeert te doen, maar dan met camera's in plaats van ogen. Het nieuwe onderzoek van LeWorldModel (of LeWM) is een grote stap in die richting. Hier is de uitleg, vertaald naar alledaags taal met een paar leuke vergelijkingen.

Het Probleem: De "Slapende" AI

Vroeger waren AI-modellen die de wereld probeerden te begrijpen, erg kwetsbaar. Ze hadden een vervelende gewoonte: ze werden "lui".
Stel je een student voor die voor een examen moet leren. In plaats van de stof echt te begrijpen, leert hij één antwoord uit zijn hoofd en geeft dat antwoord op elke vraag. Dat is makkelijk, maar het werkt niet als de vraag verandert.

In de AI-wereld noemen we dit instorting (collapse). Het model leert dat het makkelijkst is om elke afbeelding om te zetten in exact hetzelfde puntje in zijn hoofd. Dan kan het wel voorspellen wat er gebeurt, maar het begrijpt niets. Om dit te voorkomen, moesten onderzoekers tot nu toe heel veel ingewikkelde regels en "trucs" gebruiken (zoals een tweede brein dat nooit verandert, of complexe straffen). Dit maakte het trainen van deze modellen erg moeilijk, duur en onstabiel.

De Oplossing: LeWorldModel (LeWM)

De auteurs van dit papier hebben een nieuwe manier bedacht om deze AI te trainen. Ze noemen het LeWorldModel. Het is als het geven van een slimme, simpele opdracht aan een leerling, in plaats van hem met duizenden regels te overladen.

Hoe werkt het? (De Twee Regels)
In plaats van een ingewikkeld spel met veel regels, gebruikt LeWM slechts twee simpele regels:

1. De Voorspeller: "Als ik nu dit zie en ik doe dit, wat zie ik dan een seconde later?"
   • Vergelijking: Stel je voor dat je een film kijkt, maar dan met de audio uit. Je probeert te raden wat er in de volgende scène gebeurt op basis van wat je nu ziet en wat je net hebt gedaan.
2. De "Niet-Vermoeidheids"-Regel: "Zorg dat je antwoorden niet allemaal hetzelfde zijn."
   • Vergelijking: Stel je een klaslokaal voor. Als elke leerling precies hetzelfde antwoord geeft, is er geen discussie. LeWM dwingt de AI om zijn antwoorden (zijn "latente ruimte") te verspreiden, alsof het een regenboog van verschillende kleuren is in plaats van één grijze massa. Dit zorgt ervoor dat de AI echt verschillende dingen onderscheidt.

Het mooie is: dit werkt end-to-end. Dat betekent dat de AI leert direct vanuit de ruwe beelden van de camera, zonder dat iemand eerst handmatig moet uitleggen wat een "wiel" of een "muur" is. Het leert alles zelf, net als een kind.

Waarom is dit zo speciaal?

• Het is goedkoper en sneller: Vroeger hadden zulke modellen enorme rekenkracht nodig (zoals een hele fabriek aan computers). LeWM kan trainen op één enkele grafische kaart (een GPU) in een paar uur. Het is alsof je van een zware vrachtwagen naar een snelle elektrische scooter overstapt.
• Het is sneller in plannen: Als de AI een taak moet uitvoeren (bijvoorbeeld een blokje duwen naar een doel), moet het eerst nadenken: "Als ik hier duw, wat gebeurt er dan?" LeWM doet dit 48 keer sneller dan de beste concurrenten. Het is alsof het AI-brein in een seconde een heel leven aan scenario's heeft doorgespeeld.
• Het begrijpt de natuurwetten: De onderzoekers hebben getest of de AI echt begrijpt hoe de wereld werkt. Ze deden een "verrassingstest".
  • Het experiment: Ze lieten de AI een video zien waarin een blokje plotseling verdween en ergens anders weer verscheen (alsof het teleporteerde).
  • Het resultaat: De AI werd "verbaasd". Zijn voorspelling klopte niet meer met de werkelijkheid. Dit betekent dat de AI een innerlijk gevoel heeft voor fysica: het weet dat dingen niet zomaar teleporteren.

De Vergelijking: De Bouwvakker vs. De Architect

• Oude methoden (zoals DINO-WM): Dit is als een bouwvakker die een huis bouwt, maar alleen mag werken met de blauwdrukken van een ander die hij niet zelf heeft getekend. Hij kan het huis wel bouwen, maar als er iets onverwachts gebeurt, weet hij niet hoe hij moet reageren omdat hij de basisprincipes niet zelf heeft ontdekt.
• LeWM: Dit is een architect die zelf het huis ontwerpt, van de grond af. Hij begrijpt waarom de muren staan waar ze staan. Als hij een nieuwe kamer moet toevoegen, kan hij dat snel en slim doen, omdat hij de logica van het hele gebouw begrijpt.

Conclusie

LeWorldModel is een doorbraak omdat het laat zien dat je geen ingewikkelde trucs nodig hebt om een AI een wereldmodel te laten bouwen. Met twee simpele regels kan een AI leren hoe de wereld werkt, plannen maken en zelfs "verrassing" voelen als de natuurwetten worden overtreden.

Het is een stap dichter bij AI die niet alleen data onthoudt, maar echt begrijpt hoe de wereld in elkaar zit, zodat het in de toekomst zelfstandig nieuwe vaardigheden kan leren, net als wij mensen.

Technische samenvatting

Probleemstelling

Het ontwikkelen van agents die vaardigheden kunnen leren vanuit ruwe sensorische input (pixels) zonder handmatig ontworpen toestandsrepresentaties is een centraal doel in kunstmatige intelligentie. World Models (WM) zijn een krachtige methode hiervoor, waarbij agents de gevolgen van acties voorspellen om te plannen in een "verbeeldingsruimte". Een populaire benadering hiervoor is de Joint Embedding Predictive Architecture (JEPA), die observaties codeert naar een compacte, latente ruimte en de dynamiek modelleert door toekomstige latente toestanden te voorspellen.

Echter, bestaande JEPA-methoden kampen met ernstige stabiliteitsproblemen:

1. Representatie-instorting (Collapse): Het model leert vaak een triviale oplossing waarbij alle inputs worden gemapt naar dezelfde representatie om de voorspellingsfout te minimaliseren.
2. Complexiteit: Om dit te voorkomen, vertrouwen bestaande methoden op complexe heuristieken, zoals stop-gradient, exponentiële bewegende gemiddelden (EMA), voorgeïmplementeerde encoders (foundation models), of complexe multi-term verliesfuncties met vele hyperparameters.
3. Resource-intensiteit: Veel methoden vereisen grote rekenkracht of pre-trained modellen, wat de toegankelijkheid voor onderzoek beperkt.

Methodologie: LeWorldModel (LeWM)

De auteurs introduceren LeWorldModel (LeWM), de eerste JEPA die stabiel end-to-end kan worden getraind vanaf ruwe pixels, zonder heuristieken, en met een extreem eenvoudige doelstelling.

1. Architectuur:

• Encoder: Een Vision Transformer (ViT) die een frame-observatie ($o_t$) mapt naar een compacte, lage-dimensionale latente representatie ($z_t$).
• Predictor: Een transformer die de omgevingdynamica modelleert door de volgende latente toestand ($\hat{z}_{t+1}$) te voorspellen op basis van de huidige toestand ($z_t$) en de uitgevoerde actie ($a_t$).

2. Trainingsdoelstelling (Loss Function):
LeWM gebruikt slechts twee verliestermen, wat een drastische vereenvoudiging is ten opzichte van concurrenten:

• Voorspellingsverlies ($L_{pred}$): De Mean Squared Error (MSE) tussen de voorspelde latente toestand en de werkelijke volgende toestand.
  $$L_{pred} = \|\hat{z}_{t+1} - z_{t+1}\|_2^2$$
• Regularisatieverlies (SIGReg): Om representatie-instorting te voorkomen, wordt een Sketched-Isotropic-Gaussian Regularizer toegepast. Deze term dwingt de verdeling van de latente embeddings af naar een isotrope Gaussische verdeling.
  • In plaats van normaliteit direct in hoge dimensies te testen, projecteert SIGReg de embeddings op $M$ willekeurige eenheidsrichtingen.
  • Op deze 1D-projecties wordt de Epps-Pulley teststatistiek toegepast om de normaliteit te optimaliseren.
  • Volgens de Cramér-Wold stelling impliceert het matchen van alle 1D-marginaalverdelingen het matchen van de volledige gezamenlijke verdeling.

3. Training en Planning:

• End-to-End: Alle parameters (encoder en predictor) worden gezamenlijk geoptimaliseerd zonder stop-gradient of EMA.
• Hyperparameters: Het systeem heeft slechts één effectieve hyperparameter om te tunen: het gewicht ($\lambda$) van de SIGReg-term. De andere parameter (aantal projecties) heeft een verwaarloosbaar effect.
• Latente Planning: Tijdens inferentie wordt Model Predictive Control (MPC) gebruikt. De agent optimaliseert een reeks acties in de latente ruimte om een doeltoestand te bereiken, gebruikmakend van de Cross-Entropy Method (CEM) als solver.

Belangrijkste Bijdragen

1. Stabiliteit en Eenvoud: LeWM is de eerste JEPA die stabiel end-to-end leert vanaf pixels met slechts twee loss-termen, zonder heuristieken zoals stop-gradient of pre-trained encoders.
2. Efficiëntie: Het model (15M parameters) kan op één GPU in een paar uur worden getraind. Het is tot 48x sneller in planning dan foundation-model-gebaseerde wereldmodellen (zoals DINO-WM) omdat het veel minder tokens nodig heeft voor encoding.
3. Task-Agnostisch en Reward-Free: Het leert generieke wereldmodellen zonder toegang tot beloningssignalen (reward-free) of specifieke taakdefinities tijdens training.
4. Fysiek Begrip: De auteurs tonen aan dat de latente ruimte betekenisvolle fysieke structuren encodeert, wat wordt bewezen door het detecteren van "fysiek onmogelijke" gebeurtenissen (violation-of-expectation).

Resultaten

De auteurs evalueren LeWM op diverse 2D en 3D taken (navigatie, manipulatie, locomotie) zoals PushT, OGBench-Cube, Two-Room en Reacher.

• Prestaties: LeWM presteert beter dan bestaande end-to-end methoden (zoals PLDM) en is concurrerend met foundation-model-benaderingen (zoals DINO-WM), zelfs wanneer DINO-WM extra proprioceptieve input heeft.
  • Op de PushT-taak behaalt LeWM een 18% hogere success rate dan PLDM.
  • Het overtreft DINO-WM op PushT, ondanks dat DINO-WM gebruikmaakt van een voorgeïmplementeerde encoder.
• Snelheid: LeWM voltooit planning in minder dan 1 seconde (48x sneller dan DINO-WM) bij behoud van vergelijkbare prestaties.
• Stabiliteit: De training vertoont een gladde, monotoon dalende verliescurve, in tegenstelling tot de ruis en instabiliteit van multi-term verliesfuncties bij concurrenten.
• Fysiek Inzicht:
  • Probing: Lineaire en niet-lineaire probes kunnen fysieke grootheden (positie, hoek) nauwkeurig voorspellen uit de latente ruimte.
  • Verwachtenschending (VoE): Het model herkent betrouwbaar fysiek onwaarschijnlijke gebeurtenissen (bijv. teleportatie van objecten) door een sterke piek in "surprise" (voorspellingsfout), terwijl visuele veranderingen (kleur) minder impact hebben.

Significantie

LeWorldModel vertegenwoordigt een belangrijke doorbraak in het veld van wereldmodellen en zelftoezicht (self-supervised learning):

• Verwijdering van Barrières: Het maakt het mogelijk om wereldmodellen te trainen op één GPU zonder afhankelijkheid van grote foundation models of complexe hyperparameter-tuning.
• Principiële Aanpak: Door het gebruik van een wiskundig onderbouwde regularisatie (SIGReg) in plaats van heuristieken, biedt het een meer robuuste en interpreteerbare training.
• Toekomstperspectief: Het bewijst dat een eenvoudige, end-to-end architectie voldoende is om complexe dynamica en fysiek inzicht te leren, wat de weg vrijmaakt voor schaalbare, task-agnostische agents die direct vanuit pixels kunnen plannen.

Kortom, LeWM lost het probleem van instabiele training bij end-to-end wereldmodellen op door een elegante, wiskundig gefundeerde regularisatie toe te passen, wat leidt tot snellere, stabielere en effectievere agents.