Technische Samenvatting: UWM-JEPA

Probleemstelling

Joint Embedding Predictive Architectures (JEPAs) voorspellen toekomstige toestanden in de latente ruimte in plaats van ruwe waarnemingen te reconstrueren, waardoor representatieleren effectief wordt gescheiden van generatieve decoding. Echter, in gedeeltelijk waarneembare omgevingen mist een standaard vectoriële latente representatie de interne structuur die nodig is om "geloof" over verborgen voortzettingen te dragen. Wanneer de huidige waarnemingsgeschiedenis de verborgen toestand niet uniek bepaalt, is de toekomst geen enkel punt in de representatieruimte, maar een verdeling over compatibele verborgen toestanden.

Standaard vectorvoorspellers die worden getraind tegen Exponential Moving Average (EMA)-doelen kunnen de verliezen minimaliseren, maar de verliesfunctie alleen specificeert niet welke onzekerheden behouden moeten blijven, welke gemiddeld moeten worden weggenomen, of hoe contrafactuele acties de geloofstoestand moeten verschuiven. Bijgevolg slagen vector-latente JEPAs er vaak niet in om actiesensitieve dynamiek te behouden of taakrelevante informatie te bewaren tijdens "blinde rollout" (het voorstellen van toekomstige toestanden zonder nieuwe waarnemingen).

Methodologie

De auteurs introduceren UWM-JEPA (Unitary World Model JEPA), een modificatie van het JEPA-raamwerk waarbij de latente variabele wordt behandeld als een expliciete geloofstoestand in plaats van een deterministisch punt.

1. Geloofstoestand-Latente

In plaats van een vector, beeldt de online encoder $E_\theta$ de waargenomen context af op een dichtheidsmatrix $\rho_t$ in een bipartiete latente ruimte $\mathcal{H}_T = \mathcal{H}_S \otimes \mathcal{H}_E$ (Systeem $\otimes$ Omgeving).

• Dimensies: De systeemfactor ($d_{sys}=8$) wordt gelezen door het verlies en de probes; de omgevingsfactor ($d_{env}=2$) fungeert als register om verbonden modi en systeem-omgeving correlaties te dragen.
• Beperkingen: De dichtheidsmatrix is beperkt tot Hermitisch, positief semi-definiet en met een eenheids-spur. Deze worden afgedwongen door constructie: de initiële toestand is maximaal gemengd, waarnemingsupdates passen positieve meetoperatoren toe, en voorspellingstappen gebruiken unitaire conjugatie.

2. Unitair Voorspeller

De voorspeller brengt de latente toestand vooruit via unitaire conjugatie:
$$\hat{\rho}_{t+k|t} = U_k \rho_t (U_k)^\dagger$$
waarbij $U = \exp(-iH\Delta t)$ en $H$ een Hermitische generator is.

• Actie-Conditionering: De generator is lineair in de actie $a$: $H(a) = H_0 + a H_1$. Dit maakt het mogelijk dat de unitaire baan wordt gestuurd door acties.
• Trainingsdoel: Het model gebruikt het standaard JEPA-raamwerk: een EMA-doel-encoder $E_\xi$ met stop-gradient, en een verlies in de latente ruimte dat wordt geëvalueerd na partiële trace naar de systeemfactor en een projectiehoofd.
• Contrafactuele Doelen: Om ervoor te zorgen dat de voorspeller leert de actie-kanaal te gebruiken, wordt het trainingsdoel gegenereerd door een simulator-rollout onder een contrafactuele actiesequentie $\tilde{a}$, in plaats van de door de leraar geforceerde waargenomen baan.

3. Theoretische Garantie

De kern theoretische claim is niet-vergeten van de gezamenlijke toestand. Omdat unitaire conjugatie een isometrie is, behoudt deze het spectrum (eigenwaarden), de zuiverheid en de von Neumann-entropie van de gezamenlijke systeem-omgeving dichtheidsmatrix exact tijdens blinde rollout.

• Beperking: Deze garantie geldt voor de volledige gezamenlijke toestand. De gereduceerde systeemtoestand (verkregen via partiële trace) is niet unitair, wat betekent dat informatie kan verschuiven naar systeem-omgeving correlaties en ontoegankelijk kan worden voor lineaire probes, zelfs als de gezamenlijke toestand niet-dissipatief blijft.

Belangrijkste Resultaten

1. Actiesensitiviteit Vereist Contrafactuele Doelen

• Falen door Leraar-Force: Wanneer getraind met door de leraar geforceerde doelen (het overeenkomen met de waargenomen toekomst), wordt het actie-kanaal inert. De verhouding van de actie-afhankelijke term tot de basis-dynamica term stort in tot $\|H_1\|/\|H_0\| \approx 0.03$. De voorspeller leert het doel te matchen ongeacht de actie, omdat het doel al het effect van de actie bevat.
• Succes met Contrafactuele Doelen: Het vervangen van het doel door een simulator-rollout onder bemonsterde contrafactuele acties herstelt de actiesensitiviteit, met $\|H_1\|/\|H_0\| = 1.00 \pm 0.15$.
• Controles met Perturbaties: UWM-JEPA getraind met contrafactuele doelen (UWM-JEPA-CF) toont monotoon afnemende voorspellingsnauwkeurigheid naarmate acties worden verstoord (omgekeerd, geschud, genegateerd), terwijl een op parameters afgestemde LSTM-JEPA-CF onder alle omstandigheden instort tot meerderheidsklasse-nauwkeurigheid (0.53).

2. Blinde Rollout en Korte-Horizon Behoud

• Indicator voor Verborgen Snelheid: In een taak die 5-staps voorwaartse simulatie vereist om het teken van verborgen snelheid te voorspellen (met het doelwaarneming gemaskeerd), bereikt UWM-JEPA-CF 0.77 nauwkeurigheid. De overeenkomstige LSTM-JEPA-CF faalt en voorspelt de meerderheidsklasse voor alle voorbeelden.
• Informatiebehoud ($\Delta R^2$): Het meten van het verlies aan informatie ten opzichte van een leraar-probe tijdens blinde rollout:
  • Bij korte horizons ($k=1, 3$) verliest UWM-JEPA minder dan 10 punten $R^2$ (respectievelijk 0.023 en 0.069).
  • De LSTM-JEPA-baseline verliest aanzienlijk meer (0.413 en 0.678).
  • Bij de getrainde horizon ($k=5$) heraligneert de LSTM, maar behoudt UWM superieure prestaties op tussenliggende stappen.
  • Bij lange horizons ($k>10$) degraderen beide modellen, wat aangeeft dat behoud op lange horizon een open uitdaging blijft.

3. Contextrepresentatie is Niet de Drijvende Kracht

• Encoder-Equivalentie: Een terughoudende lineaire probe op de bevroren encoder-representaties toont aan dat UWM-JEPA en LSTM-JEPA identiek presteren ($R^2 = 0.901 \pm 0.032$ versus $0.894 \pm 0.021$, $p=0.70$).
• Conclusie: Het prestatieverschil is niet te wijten aan dat UWM-JEPA een superieure context-encoder heeft. De scheiding ontstaat volledig uit het vermogen van de voorspeller om de geloofstoestandstructuur tijdens rollout te behouden.

Betekenis en Claims

Het artikel claimt UWM-JEPA als het eerste JEPA-familie wereldmodel waarbij de latente variabele expliciet wordt behandeld als een geloof over verborgen voortzettingen (een dichtheidsmatrix) en de voorspeller is beperkt tot het evolueren van dit geloof zonder het spectrum van de gezamenlijke toestand te dissiperen.

• Bescheiden Omvang: De auteurs stellen expliciet dat dit geen algemeen doel planner is, geen universeel sterkere encoder, of een concurrerende oplossing voor planning op lange horizon. De eigenschap "niet-vergeten" is exact alleen voor de gezamenlijke toestand, niet noodzakelijk voor de gereduceerde toestand die wordt gebruikt voor downstream taken.
• Kernbijdrage: Het werk toont aan dat voor JEPA wereldmodellen om te fantaseren onder gedeeltelijke waarneembaarheid, de latente geometrie en de dynamiek van de voorspeller belangrijker zijn dan alleen het bevroren vermogen tot contextcodering.
• Praktische Implicatie: De studie benadrukt dat trainen tegen contrafactuele doelen essentieel is voor het leren van actiesensitieve voorspellers, een bevinding die verder reikt dan de specifieke unitaire parameterisering.

De resultaten zijn gevalideerd op gestructureerde gedeeltelijk waarneembare oscillator- en CartPole-achtige omgevingen. De auteurs erkennen dat de huidige aanpak toegang tot een simulator vereist voor contrafactuele training en zich nog niet uitstrekt tot ongestructureerde visuele voorspelling of verborgen Markov-ketens.