Technische Samenvatting: QQWorld

Probleemstelling

Latente wereldmodellen maken efficiënte planning mogelijk door toekomstige toestanden te voorspellen binnen een compacte representatieruimte. Een cruciale ontwerpkeuze in recente architecturen, zoals LeWorldModel (LeWM), is het regulariseren van de marginale distributie van geleerde latents naar een isotrope Gaussische verdeling. Dit wordt doorgaans bereikt met behulp van de Epps–Pulley (EP) test als een differentieerbare straf om normaliteit af te dwingen.

De auteurs observeren echter dat LeWM, ondanks EP-regularisatie, vaak geleerde latente distributies met uitgesproken zware staarten (heavy tails) leert. Dit is problematisch omdat extreme latente waarden de geleerde dynamica naar slecht gerepresenteerde regio's kunnen duwen, wat fouten tijdens meerstaps-rollouts versterkt en de mismatch met de doel-Gaussische prior vergroot.

Het artikel identificeert een fundamentele beperking in de optimalisatiegeometrie van de EP-test:

1. Verdwijnende Gradiënten: De EP-statistiek is equivalent aan een gekwadrateerde Maximum Mean Discrepancy (MMD) met een Gaussische kernel. De auteurs demonstreren dat de corrigerende gradiënten geïnduceerd door dit objectief exponentieel snel afnemen (super-exponentieel) voor samples in de staart (outliers).
2. Ineffectieve Controle: Zodra een latente coördinaat zich buiten de interactieschaal van de Gaussische kernel beweegt, biedt de EP-regularisator een verwaarloosbare herstellende kracht. Bijgevolg worden afwijkingen met zware staarten onvoldoende gecontroleerd, waardoor ze kunnen voortbestaan en groeien, zelfs wanneer het model expliciet wordt gestraft voor niet-normaliteit.

Methodologie

Om het probleem van de verdwijnende gradiënt aan te pakken, stellen de auteurs QQWorld voor, die de EP-regularisator vervangt door een Quantile–Quantile (QQ) matching objectief.

Kernmechanisme: QQ-regularisatie

In plaats van momenten te matchen of kernel-gebaseerde discrepanties te gebruiken, stemt QQWorld geprojecteerde latente samples direct af met rang-gematchte Gaussische kwantielen.

• Objectief: Voor een batch geprojecteerde samples $X = \{x_n\}$ is het verlies gedefinieerd als de som van de gekwadrateerde verschillen tussen de orde statistieken van de samples ($\hat{x}_n$) en de overeenkomstige theoretische Gaussische kwantielen ($q_n = \Phi^{-1}(\frac{n-0.5}{N})$):
  $$L_{QQ}(X) = \sum_{n=1}^N (\hat{x}_n - q_n)^2$$
• Gradiëntgedrag: In tegenstelling tot de EP-gradiënt, die verdwijnt voor extreme waarden, is de gradiënt van de QQ-loss met betrekking tot een sample $x_n$ gelijk aan $2(x_n - q_{\rho(n)})$, waarbij $\rho(n)$ de rang is. De grootte van deze gradiënt groeit lineair met de afwijking van het doelkwantiel. Dit zorgt ervoor dat staart-samples sterke, symmetrische corrigerende krachten ontvangen, wat zware staarten effectief onderdrukt.
• Theoretische Garantie: De auteurs bewijzen dat het reduceren van de QQ-loss naar nul ook de EP-loss naar nul drijft (eenrichtingscontrole), maar het omgekeerde niet waar is. Een kleine EP-loss kan samengaan met willekeurig grote staartafwijkingen, terwijl QQ-matching een striktere distributionele uitlijning afdwingt.

Cross-Batch QQ

Om de geheugenbeperkingen die gepaard gaan met grote batchgroottes (die de nauwkeurigheid van de rangschikking verbeteren) aan te pakken, introduceren de auteurs Cross-Batch QQ.

• Strategie: Deze methode vergroot de effectieve ranking-pool door gedetacheerde features van de vorige $K$ iteraties op te nemen in een FIFO-wachtrij.
• Mechanisme: Ranks worden berekend over de gepoolde set van grootte $M = N(K+1)$, maar gradiënten worden alleen teruggepropageerd door de huidige batch van grootte $N$.
• Bias-Variance Trade-off: De auteurs karakteriseren een trade-off waarbij het vergroten van de wachtrijlengte $K$ de variantie van de rangschikkingsschatting vermindert (gunstig voor kleine batches), maar een "representatie-verouderingsbias" (representation-staleness bias) introduceert als de latente distributie aanzienlijk verschuift over de tijd.

Belangrijkste Bijdragen

1. Identificatie van EP-beperking: Het artikel analyseert en demonstreert formeel dat de corrigerende kracht van de EP-regularisator verdwijnt voor staart-samples, wat verklaart waarom LeWM faalt in het onderdrukken van zware staarten ondanks expliciete normaliteitsstraffen.
2. Voorstel van QQWorld: Een drop-in vervanging voor LeWM die een quantile–quantile matching objectief gebruikt. Het behoudt effectieve corrigerende gradiënten voor staartafwijkingen zonder nieuwe hyperparameters te introduceren.
3. Geheugenefficiënte Training: De ontwikkeling van Cross-Batch QQ, die een nauwkeurige kwantielschatting mogelijk maakt met grote effectieve batchgroottes, terwijl de GPU-geheugenoverhead laag blijft door de grootte van de ranking-pool te ontkoppelen van de backpropagation batchgrootte.
4. Theoretische Analyse: Karakterisering van de bias-variance trade-off in Cross-Batch QQ en het bewijs van de eenrichtingscontrole-relatie tussen de QQ- en EP-objectieven.

Experimentele Resultaten

De auteurs evalueerden QQWorld over vier controle-omgevingen: Two-Room, PushT, Reacher en OGBench-Cube.

• Planning Prestaties: QQWorld behaalde een gemiddeld planning succespercentage van 85,08%, waarmee het de LeWM-baseline (79,75%) met 5,33 procentpunten versloeg. Het overtrof consequent LeWM en andere recente varianten (Sub-JEPA, SD-JEPA, SMWM) in alle vier de omgevingen.
• Distributionele Uitlijning:
  • QQ RMSE: Verminderde de gemiddelde fout tussen empirische en Gaussische kwantielen van 0,157 (LeWM) naar 0,121 (QQWorld).
  • Staartkans (Tail Probability): Verminderde de kans dat samples de 90e percentiel van de chi-kwadraatverdeling overschrijden (een maat voor radiale staartzwaarte) van 0,315 naar 0,123, naderend bij de Gaussische referentie van 0,10.
  • Statistische Tests: QQWorld verminderde de Kolmogorov–Smirnov (KS) statistiek met 15,8% en de EP-statistiek met 31,4% vergeleken met LeWM, ook al werd LeWM direct getraind met de EP-objectief.
• Fysieke Toestand Probing: In de PushT-omgeving presteerde QQWorld gelijk aan of iets beter dan LeWM bij het decoderen van fysieke toestanden (agentlocatie, bloklocatie, blokhoek), wat aangeeft dat de sterkere regularisatie de codering van taakrelevante informatie niet heeft verslechterd.
• Cross-Batch Efficiëntie: Door een kleine batchgrootte ($N=32$) met een wachtrijlengte van 3 ($K+1=3$) te gebruiken, behaalde QQWorld een succespercentage van 83,50%, wat beter is dan LeWM getraind met een batchgrootte van 128 (79,75%), terwijl het GPU-geheugengebruik met 73% werd verminderd.

Betekenis en Claims

Het artikel stelt dat de effectiviteit van een statistische test voor het meten van distributionele discrepantie niet de effectiviteit ervan als trainingsdoel garandeert. Specifiek betogen de auteurs dat:

• Gradiëntkwaliteit Is Cruciaal: Naast het louter onderscheiden van distributies, moet een regularisatie-objectief informatieve en goed gedraaide gradiënten bieden gedurende het hele optimalisatieproces. De EP-test faalt op dit criterium voor staart-samples, terwijl het QQ-objectief daarin slaagt.
• Controle van de Staart Is Cruciaal: Effectieve onderdrukking van zware staarten in latente ruimtes is essentieel voor stabiele meerstaps-planning in wereldmodellen.
• Schaalbaarheid: De Cross-Batch QQ-strategie biedt een schaalbare mechanisme voor het trainen van grote wereldmodellen waarbij het vergroten van de fysieke batchgroottes prohibitief duur is qua geheugen.

Het werk concludeert dat QQWorld een eenvoudige, hyperparameter-vrije verbetering biedt voor bestaande latente wereldmodellen door direct de optimalisatiegeometrie van distributionele regularisatie aan te pakken.