Technische Samenvatting: JEDI: Gezamenlijke Embedding Diffusie Wereldmodel voor Online Model-Based Versterkingsleren

Probleemstelling

Online Model-Based Versterkingsleren (MBRL) vertrouwt op geleerde wereldmodellen om toekomstige uitkomsten te voorspellen en de sample-efficiëntie te verbeteren. Hoewel diffusiemodellen zijn opgekomen als krachtige ruggegronden voor wereldmodellering in visueel complexe domeinen, staan huidige benaderingen voor een aanzienlijke spanning:

1. Pixel-Diffusie: Modellen zoals DIAMOND werken direct in pixelruimte. Hoewel ze effectief zijn, zijn ze rekenkundig duur wat betreft geheugen (VRAM), bemonsteringstijd en trainingsduur.
2. Latente Diffusie: Het verplaatsen van diffusie naar latente ruimte verbetert de efficiëntie, maar vertrouwt vaak op apart getrainde latenten (bijvoorbeeld via reconstructiedoelen) in plaats van end-to-end wereldmodel-doelen. Deze benaderingen, zoals Horizon Imagination (HI), hebben nog niet de prestaties van op pixels gebaseerde diffusiemodellen weten te evenaren en missen het end-to-end voorspellende representatieleren dat de vooruitgang in modern MBRL heeft gedreven.

De centrale vraag die wordt behandeld, is of een diffusiewereldmodel zijn latente representatie end-to-end kan leren uit het doelfunctie voor ruisreductie met behulp van een Joint Embedding Predictive Architecture (JEPA)-kader, waardoor de efficiëntie van latente diffusie wordt gecombineerd met de prestaties van end-to-end voorspellend leren.

Methodologie: JEDI

De auteurs introduceren JEDI (Joint Embedding Diffusion), het eerste online end-to-end latente diffusiewereldmodel gebouwd op een JEPA-kader.

Kernarchitectuur

JEDI leert gelijktijdig een latente ruimte-encoder ($E_\phi$) en een dynamisch model. Het systeem bestaat uit:

• Encoder: Comprimeert observaties $x_t$ tot schone latenten $z^0_t$.
• Latente Diffusiedynamiek: Voorspelt de volgende latente toestand via conditionele ruisreductie. In plaats van een voorspelling in één stap, gebruikt het een omgekeerde diffusieketen die is geconditioneerd op de huidige latente $z^0_t$ en eerdere acties om de schone toekomstige latente $z^0_{t+1}$ te voorspellen.
• Beloning- en Beëindigingskoppen: Voorspellen beloningen en beëindigingssignalen direct vanuit de latente toestand.

Trainingsmechanisme

In tegenstelling tot traditionele latente diffusiemodellen die vertrouwen op reconstructieverliezen (bijvoorbeeld VAE's) om de encoder te trainen, traint JEDI de encoder end-to-end met behulp van het diffusie-ruisreductieverlies.

• Gezamenlijke Embedding-doelstelling: De encoder wordt bijgewerkt via gradiënten van de conditionele ruisreductiedoelstelling. De toekomstige latente doelwit wordt behandeld als een "label" afgeleid uit de omgeving, en het model leert om te denoisen van een ruisversie van deze toekomstige latente terug naar de schone toestand, geconditioneerd op de huidige toestand.
• Stop-Gradiënt: Om representatie-inzakking te voorkomen en training te stabiliseren, wordt een stop-gradiënt toegepast op de toekomstige latente doelwit tijdens encoder-bijwerkingen, in lijn met JEPA- en TDMPC2-principes.
• Stabilisatietechnieken: Het model maakt gebruik van differentieerbare tanh-klemming op latente activaties en een willekeurig schakelmechanisme tussen de output van de ruisreductor en de output van de encoder voor conditionering in de volgende stap om consistentie en supervisie in evenwicht te brengen.

Theoretische Motivatie

Het artikel biedt een theoretische rechtvaardiging die JEDI koppelt aan het Informatieknelpunt-principe.

• JEPA als Knelpunt: De standaard JEPA-doelstelling blijkt een variational informatieknelpunt te induceren, wat representaties belonen die informatie behouden die nuttig is voor het voorspellen van de toekomst, terwijl irrelevante details worden gecomprimeerd.
• Diffusie-JEPA Decompositie: De auteurs tonen aan dat conditionele latente diffusie een vergelijkbare variational informatieknelpunt-decompositie toestaat. Het conditionele ruisreductieverlies fungeert als het voorspellende element binnen deze knelpuntstructuur. Dit suggereert dat diffusieruisreductie een principieel, end-to-end voorspellend leersignaal biedt voor de latente encoder, in plaats van slechts te dienen als een generatieve vervanger.

Belangrijkste Bijdragen

1. Eerste End-to-End Latente Diffusiewereldmodel: JEDI is het eerste online MBRL-model dat zijn latente encoder end-to-end traint met behulp van een diffusie-ruisreductiedoelstelling binnen een JEPA-kader, waardoor de noodzaak voor apart voortrainen of reconstructieverliezen wordt geëlimineerd.
2. Theoretische Connectie: Het artikel legt een wiskundig verband tussen conditionele diffusiedoelstellingen en het variational informatieknelpunt, en verklaart waarom ruisreductie effectief voorspellend representatieleren kan ondersteunen.
3. Efficiëntiewinsten: Door te verschuiven van pixelruimte naar een compacte geleerde latente ruimte ($16 \times 8 \times 8$) en end-to-end te trainen, reduceert JEDI de rekenkundige overhead aanzienlijk in vergelijking met op pixels gebaseerde diffusie.

Experimentele Resultaten

JEDI werd geëvalueerd op de Atari100k-benchmark (26 taken, 5 zaden) en de Craftium-omgeving (3D embodied decision-making).

Prestaties

• Atari100k: JEDI bereikt concurrerende prestaties, waarbij het sterke baselines (DIAMOND, DreamerV3, STORM) op Interquartile Mean (IQM) en Optimality Gap-metingen evenaart of overtreft. Het bereikt State-of-the-Art (SOTA) resultaten in 7 spellen en tweede plaats in het aantal supermenselijke taken.
• Craftium: JEDI presteert beter dan Horizon Imagination (HI) op alle vier geteste taken, wat aantoont dat end-to-end voorspellende latenten sterke prestaties in 3D-omgevingen kunnen ondersteunen zonder reconstructie-gebaseerde latente training.
• Stochasticiteit: In experimenten met stochastische frame-skipping overtreft JEDI DreamerV3 consistent, wat suggereert dat diffusiemodellen robuuster kunnen zijn tegen hoge aleatorische onzekerheid.

Efficiëntie

Ten opzichte van de op pixels gebaseerde baseline DIAMOND toont JEDI aanzienlijke efficiëntieverbeteringen:

• VRAM-gebruik: Gebruikt 43% minder VRAM.
• Bemonsteringssnelheid: Bereikt >3× snellere wereldmodel-bemonstering.
• Trainingsnelheid: Traint >2,5× sneller (38 uur versus 98 uur op A100 voor Atari100k).
• Modelgrootte: Behoudt een vergelijkbaar aantal parameters (13,5M) als DIAMOND, maar overtreft HI (97M parameters) aanzienlijk in efficiëntie terwijl het concurrerende prestaties behoudt.

Verschuiving in Prestatieprofiel

Een opvallende bevinding is dat JEDI een ander prestatieprofiel op taakniveau vertoont in vergelijking met DIAMOND. Waar DIAMOND uitblinkt in regimes met hoge HNS (Human Normalized Score), laat JEDI duidelijke verbeteringen zien in regimes met lage HNS (moeilijke spellen). Dit wordt vastgelegd door de Har-HNS-metriek, waarbij JEDI DIAMOND overtreft, wat aangeeft dat het minder ernstig faalt op de moeilijkste taken. De auteurs schrijven dit toe aan de end-to-end latente interface die de dimensionaliteit van de invoer voor de actor-critic verlaagt, waardoor betere optimalisatie mogelijk is in complexe spellen met een grote actieruimte.

Betekenis en Beweringen

Het artikel beweert dat JEDI aantoont dat voorspellende latente diffusie een levensvatbaar en efficiënt alternatief is voor zowel pixelruimtediffusie als apart getrainde latente modellen voor online MBRL.

• Voorbij Rekenkracht: De resultaten suggereren dat end-to-end voorspellende latenten het gedrag van het wereldmodel zelf veranderen, niet alleen de rekenkosten. De verschuiving in prestatieprofielen geeft aan dat de geleerde representatie fundamenteel verandert hoe de agent interacteert met complexe omgevingen.
• Principiële Aanpak: Door de methode te verankeren in de theorie van het informatieknelpunt, positioneert het artikel voorspellende latente diffusie als een principiële richting voor het schalen van online MBRL.
• Beperkingen: De auteurs merken bescheiden op dat het bewijs momenteel beperkt is tot Atari100k en Craftium, en dat de methode nog niet is getest op domeinen zoals DMControl of Procgen. Ze erkennen ook dat de prestaties afhankelijk zijn van specifieke ontwerpkeuzes (bijvoorbeeld klemming, willekeurig schakelen) en dat de theoretische afleiding representatie-eigenschappen motiveert in plaats van optimalisatie van het volledige variational doel te garanderen.