Résumé Technique : QQWorld

Énoncé du Problème

Les modèles de monde latents permettent une planification efficace en prédisant les états futurs au sein d'un espace de représentation compact. Un choix de conception critique dans les architectures récentes, telles que LeWorldModel (LeWM), consiste à régulariser la distribution marginale des latents appris vers une gaussienne isotrope. Cela est généralement réalisé en utilisant le test d'Epps–Pulley (EP) comme une pénalité différentiable pour imposer la normalité.

Cependant, les auteurs observent que malgré la régularisation EP, LeWM apprend souvent des distributions latentes présentant des queues épaisses prononcées. Cela est problématique car les valeurs extrêmes des latents peuvent pousser les dynamiques apprises vers des régions mal représentées, amplifiant les erreurs lors des déroulements multi-étapes et augmentant le décalage avec la priorité gaussienne cible.

Le papier identifie une limitation fondamentale dans la géométrie d'optimisation du test EP :

1. Gradients Disparates : La statistique EP est équivalente à une MMD (Maximum Mean Discrepancy) au carré avec un noyau gaussien. Les auteurs démontrent que les gradients correctifs induits par cet objectif décroissent rapidement (super-exponentiellement) pour les échantillons dans la queue (outliers).
2. Contrôle Inefficace : Une fois qu'une coordonnée latente dépasse l'échelle d'interaction du noyau gaussien, le régularisateur EP fournit une force de rappel négligeable. Par conséquent, les écarts à queue épaisse sont insuffisamment contrôlés, ce qui permet à ces derniers de persister et de croître même lorsque le modèle est explicitement pénalisé pour sa non-normalité.

Méthodologie

Pour remédier au problème de disparition des gradients, les auteurs proposent QQWorld, qui remplace le régularisateur EP par un objectif d'appariement Quantile–Quantile (QQ).

Mécanisme Central : Régularisation QQ

Au lieu d'apparier les moments ou d'utiliser des divergences basées sur des noyaux, QQWorld aligne directement les échantillons latents projetés avec les quantiles gausiens appariés par rang.

• Objectif : Pour un lot d'échantillons projetés $X = \{x_n\}$, la perte est définie comme la somme des différences au carré entre les statistiques d'ordre des échantillons ($\hat{x}_n$) et les quantiles gausiens théoriques correspondants ($q_n = \Phi^{-1}(\frac{n-0.5}{N})$) :
  $$L_{QQ}(X) = \sum_{n=1}^N (\hat{x}_n - q_n)^2$$
• Comportement du Gradient : Contrairement au gradient EP, qui disparaît pour les valeurs extrêmes, le gradient de la perte QQ par rapport à un échantillon $x_n$ est $2(x_n - q_{\rho(n)})$, où $\rho(n)$ est le rang. La magnitude de ce gradient croît linéairement avec l'écart par rapport au quantile cible. Cela garantit que les échantillons de queue reçoivent des forces correctives fortes et symétriques, supprimant efficacement les queues épaisses.
• Garantie Théorique : Les auteurs prouvent que pousser la perte QQ vers zéro conduit également la perte EP vers zéro (contrôle unidirectionnel), mais la réciproque n'est pas vraie. Une faible perte EP peut coexister avec des écarts de queue arbitrairement grands, tandis que l'appariement QQ impose un alignement distributionnel plus strict.

QQ Cross-Batch (Inter-lots)

Pour traiter les contraintes de mémoire associées aux tailles de lots importantes (qui améliorent la précision de l'estimation des rangs), les auteurs introduisent le Cross-Batch QQ.

• Stratégie : Cette méthode élargit le pool de classement effectif en inclissant des caractéristiques détachées des $K$ itérations précédentes dans une file FIFO.
• Mécanisme : Les rangs sont calculés sur le pool agrégé de taille $M = N(K+1)$, mais les gradients ne sont rétropropagés qu'à travers le lot actuel de taille $N$.
• Compromis Biais-Variance : Les auteurs caractérisent un compromis où l'augmentation de la longueur de la file $K$ réduit la variance de l'estimation du rang (bénéfique pour les petits lots) mais introduit un biais de "lassitude de représentation" si la distribution latente change significativement au fil du temps.

Contributions Clés

1. Identification de la Limite de l'EP : Le papier analyse et démontre formellement que la force correctrice du régularisateur EP disparaît pour les échantillons de queue, expliquant pourquoi LeWM échoue à supprimer les queues épaisses malgré les pénalités explicites de normalité.
2. Proposition de QQWorld : Un remplacement direct pour LeWM qui utilise un objectif d'appariement quantile-quantile. Il maintient des gradients correctifs efficaces pour les écarts de queue sans introduire de nouveaux hyperparamètres.
3. Entraînement Efficace en Mémoire : Le développement du Cross-Batch QQ, qui permet une estimation précise des quantiles en utilisant de grandes tailles de lots effectives tout en maintenant une faible surcharge de mémoire GPU en découplant la taille du pool de classement de la taille du lot de rétropropagation.
4. Analyse Théorique : Caractérisation du compromis biais-variance dans le Cross-Batch QQ et preuve de la relation de contrôle unidirectionnelle entre les objectifs QQ et EP.

Résultats Expérimentaux

Les auteurs ont évalué QQWorld à travers quatre environnements de contrôle : Two-Room, PushT, Reacher, et OGBench-Cube.

• Performance de Planification : QQWorld a atteint un taux de succès de planification moyen de 85,08 %, surpassant la ligne de base LeWM (79,75 %) de 5,33 points de pourcentage. Il a systématiquement surpassé LeWM et d'autres variantes récentes (Sub-JEPA, SD-JEPA, SMWM) dans les quatre environnements.
• Alignement Distributionnel :
  • RMSE QQ : A réduit l'erreur moyenne entre les quantiles empiriques et gausiens de 0,157 (LeWM) à 0,121 (QQWorld).
  • Probabilité de Queue : A réduit la probabilité que les échantillons dépassent le 90e percentile de la distribution chi-deux (une mesure de l'épaisseur radiale de la queue) de 0,315 à 0,123, se rapprochant de la référence gaussienne de 0,10.
  • Tests Statistiques : QQWorld a réduit la statistique Kolmogorov–Smirnov (KS) de 15,8 % et la statistique EP de 31,4 % par rapport à LeWM, bien que LeWM soit entraîné directement avec l'objectif EP.
• Sondage de l'État Physique : Sur l'environnement PushT, QQWorld égale ou surpasse légèrement LeWM pour le décodage des états physiques (position de l'agent, position du bloc, angle du bloc), indiquant que la régularisation plus forte n'a pas dégradé l'encodage des informations pertinentes pour la tâche.
• Efficacité Cross-Batch : En utilisant un petit lot ($N=32$) avec une longueur de file de 3 ($K+1=3$), QQWorld a atteint un taux de succès de 83,50 %, surpassant LeWM entraîné avec un lot de 128 (79,75 %), tout en réduisant l'utilisation de la mémoire GPU de 73 %.

Signification et Revendications

Le papier affirme que l'efficacité d'un test statistique pour mesurer la divergence de distribution ne garantit pas son efficacité en tant qu'objectif d'entraînement. Plus précisément, les auteurs soutiennent que :

• La Qualité du Gradient Importante : Au-delà de la simple distinction entre les distributions, un objectif de régularisation doit fournir des gradients informatifs et bien comportés tout au long du processus d'optimisation. Le test EP échoue à ce critère pour les échantillons de queue, tandis que l'objectif QQ réussit.
• Le Contrôle de la Queue est Critique : La suppression efficace des queues épaisses dans les espaces latents est essentielle pour une planification multi-étapes stable dans les modèles de monde.
• Scalabilité : La stratégie Cross-Batch QQ offre un mécanisme évolutif pour l'entraînement de grands modèles de monde où l'augmentation des lots physiques est prohibitive en termes de mémoire.

Le travail conclut que QQWorld fournit une amélioration simple et sans hyperparamètre par rapport aux modèles de monde latents existants en traitant directement la géométrie d'optimisation de la régularisation distributionnelle.