🤖 machine learning

Composition of Memory Experts for Diffusion World Models

Ce papier introduit un modèle du monde basé sur la diffusion qui surmonte le compromis entre fidélité et mémoire des architectures existantes en composant des experts spécialisés de mémoire à court terme, épisodique et spatiale via une formulation de produit d'experts contrastive, permettant ainsi d'atteindre une prédiction d'avenir évolutive et de haute fidélité sans coûts de calcul quadratiques.

Auteurs originaux : Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

Publié 2026-05-20

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayiez de raconter une histoire sur un long voyage que vous avez effectué à travers un labyrinthe. Vous voulez que l'histoire soit parfaite : les tournants doivent avoir du sens, les pièces que vous avez visitées il y a des heures doivent apparaître exactement de la même manière lorsque vous y retournez, et les détails du paysage doivent être nets.

Le problème avec les « conteurs » d'IA actuels (appelés Modèles du Monde) est qu'ils souffrent d'un problème de mémoire.

L'IA « Court Terme » est comme une personne ayant une excellente mémoire pour les dernières phrases qu'elle a entendues, mais qui oublie ce qui s'est passé il y a une heure. Elle peut décrire parfaitement l'étape immédiate suivante, mais elle perd le fil de tout le voyage.
L'IA « Long Terme » est comme un historien qui se souvient de tout le trajet mais dont les détails de l'étape immédiate suivante deviennent flous. Elle connaît la forme générale du labyrinthe, mais pourrait confondre les couleurs des murs de la pièce où vous vous trouvez en ce moment même.

Tenter de construire une seule IA géante qui se souvient de tout parfaitement, c'est comme essayer de transporter une bibliothèque dans son sac à dos tout en courant un marathon ; c'est trop lourd, trop lent, et éventuellement, vous trébuchez sur les livres.

La Solution : Une Équipe de Spécialistes

Les auteurs de cet article proposent une nouvelle façon de construire ces conteurs d'IA. Au lieu d'un seul cerveau géant, ils utilisent une équipe d'experts spécialisés qui travaillent ensemble pour raconter l'histoire. Ils appellent cela la Composition d'Experts de la Mémoire (CoME).

Pensez-y comme à une équipe de production pour un film :

L'Expert « Court Terme » (Le Scénariste) : Cet expert se concentre sur la scène immédiate. Il regarde les dernières images vidéo et dit : « D'accord, le personnage vient de tourner à gauche, donc l'image suivante doit montrer un mur sur la droite. » Il est excellent pour les détails fins et le mouvement immédiat, mais ne peut pas se souvenir de ce qui s'est passé il y a 100 images.
L'Expert « Long Terme » (L'Archiviste) : Cet expert possède une immense bibliothèque de tout le trajet. Il ne regarde pas chaque image en temps réel (ce qui serait trop lent). Au lieu de cela, il « étudie » l'histoire du voyage et met à jour ses propres notes internes (ses « poids ») pour se souvenir de l'ensemble. Si vous demandez : « À quoi ressemblait cette pièce il y a 5 minutes ? », il peut s'en souvenir parfaitement car il a mémorisé l'épisode.
L'Expert « Spatial » (Le Cartographe) : Cet expert garde une trace de où se trouvent les choses. Si l'IA se déplace en boucle (comme un couloir qui semble identique partout), les autres experts pourraient se confondre. Le Cartographe dit : « Attendez, nous sommes à l'entrée Nord, pas à l'entrée Sud », assurant ainsi que l'histoire reste au bon endroit.

Comment Ils Travaillent Ensemble : Le « Produit d'Experts »

Habituellement, si vous demandez conseil à trois personnes, vous pourriez simplement prendre la moyenne. Mais en IA, la moyenne peut parfois créer un résultat trouble et confus.

Les auteurs utilisent une astuce ingénieuse appelée le Produit d'Experts. Imaginez que les trois experts tiennent tous des pancartes avec leurs prédictions.

Si le Scénariste, l'Archiviste et le Cartographe sont tous d'accord sur à quoi devrait ressembler l'image suivante, l'IA est très confiante et dessine cette image.
S'ils ne sont pas d'accord, l'IA ne devine pas au hasard ; elle cherche le « terrain d'entente ».

Cependant, il y a un piège. Parfois, les experts sont d'accord sur les mauvaises choses simplement parce qu'ils se font écho les uns aux autres. Pour corriger cela, les auteurs ont inventé une méthode « Contrastive ».

L'Analogie : Imaginez que les experts chantent une chanson. Parfois, ils fredonnent tous la même fausse note parce qu'ils s'écoutent trop les uns les autres. La méthode « Contrastive » agit comme un chef d'orchestre qui dit : « Arrêtez de fredonner le bruit de fond ! Concentrez-vous uniquement sur les notes uniques qui diffèrent de vos habitudes habituelles. » Cela garantit que l'IA ne reste pas coincée dans une boucle de répétition d'erreurs et maintient l'histoire fraîche et précise.

Les Résultats

L'article a testé cette approche d'équipe sur :

Des Labyrinthes Virtuels : Où un agent doit naviguer dans un labyrinthe 3D et se souvenir de l'endroit où il a été.
Des Vidéos du Monde Réel : Comme marcher dans une maison ou conduire un robot.

Ils ont constaté que cette approche d'équipe était bien meilleure que d'essayer d'utiliser une seule IA géante.

Meilleure Mémoire : L'IA pouvait se souvenir de détails provenant de centaines d'images auparavant sans se confondre.
Plus Rapide : Elle n'avait pas besoin de porter une « bibliothèque » massive de données dans sa tête à chaque étape ; elle consultait simplement ses spécialistes.
Plus Cohérente : Lorsque l'IA rentrait dans une pièce qu'elle avait visitée plus tôt, la pièce avait exactement la même apparence, et non un flou confus.

En Résumé

L'article soutient que, au lieu d'essayer de construire un seul super-cerveau capable de tout faire, nous devrions construire un comité de cerveaux spécialisés. En laissant un expert court terme gérer les détails, un expert long terme gérer l'histoire, et un expert spatial gérer l'emplacement, puis en les faisant voter sur l'avenir en utilisant une règle « contrastive » spéciale, nous pouvons créer une IA qui se souvient parfaitement du passé tout en prédisant avec précision l'avenir, le tout sans avoir besoin d'un superordinateur pour le faire.

Résumé Technique : Composition d'Experts Mémoire pour les Modèles de Monde Diffusifs

Énoncé du Problème

Les modèles de monde (WM) sont essentiels pour l'apprentissage par renforcement (RL) et la planification, car ils prédisent des états futurs plausibles cohérents avec les observations passées. Cependant, les architectures existantes font face à un compromis fondamental entre la fidélité de la mémoire et l'efficacité computationnelle. Les modèles basés sur les transformateurs préservent les détails locaux mais souffrent de coûts d'attention quadratiques, limitant la longueur du contexte. À l'inverse, les réseaux récurrents et les modèles d'espace d'état (SSM) s'étendent efficacement mais compressent l'histoire dans des états cachés, jetant inévitablement des détails fins au fil du temps. Augmenter simplement la longueur du contexte dans une architecture unique conduit à un entraînement instable et à des coûts d'inférence prohibitifs. De plus, la composition naïve de plusieurs modèles aboutit souvent à un « effondrement de mode », où la distribution combinée exagère les régions communes et perd en diversité.

Méthodologie

Les auteurs proposent la Composition d'Experts Mémoire (CoME), un cadre basé sur la diffusion qui découple la cohérence futur-passé d'une architecture monolithique unique. Au lieu de cela, CoME exploite un ensemble d'experts spécialisés, chacun gérant un rôle mémoire spécifique, et les fusionne de manière probabiliste au moment de l'inférence.

1. Produit d'Experts Contrastifs (PoCE)

Pour pallier les limites du Produit d'Experts (PoE) standard, qui peut supprimer la diversité et amplifier des modes spurius, les auteurs introduisent une formulation de Produit d'Experts Contrastifs.

Mécanisme : Pour un ensemble d'experts hétérogènes $p_i(x)$ , la composition est définie comme $p(x) \propto \prod \tilde{p}_i(x)$ , où $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$ .
Référence Contrastive : $\bar{p}_i(x)$ représente la distribution de référence inconditionnelle de l'expert. Le coefficient $\alpha_i > 1$ agit comme un facteur de contraste.
Justification Théorique : Contrairement au tempérage naïf (qui réduit la variance du noyau et déforme la géométrie locale), ce mélange contrastif répondère les composantes du mélange tout en préservant la forme des noyaux. Il supprime efficacement les modes spurius (où un expert attribue une probabilité élevée sans preuve forte) tout en maintenant des modes secondaires valides, à condition que l'expert conditionnel attribue une probabilité plus élevée aux modes valides que la référence.

2. Experts Mémoire Spécialisés

Le cadre instancie trois experts complémentaires :

Mémoire à Court Terme (STM) : Un modèle de diffusion conditionné sur une fenêtre de contexte récente ( $10\text{--}100$ images). Il capture les dynamiques locales fines et les dépendances temporelles immédiates en utilisant une architecture avec une longueur de contexte étendue (par exemple, une attention à fenêtre glissante).
Mémoire à Long Terme (LTM) : Un modèle de diffusion distinct qui stocke l'histoire épisodique ( $100\text{--}1000$ images) directement dans ses poids. Cela est réalisé par un finetuning au moment du test d'adaptateurs LoRA légers sur le contexte à long terme. Cette approche amortit le coût de la récupération, permettant une réutilisation constante dans le temps des expériences passées sans alimenter l'ensemble de l'histoire dans le modèle à chaque étape. Elle évite l'oubli catastrophique en gelant le socle (backbone) et en mettant à jour uniquement les adaptateurs.
Mémoire Spatiale à Long Terme (SLTM) : Une extension de la LTM qui intègre des priors spatiaux (par exemple, poses de caméra, cartes de points) extraits de l'histoire. Cet expert impose une cohérence géométrique et spatiale, désambiguïsant les lieux dans des environnements répétitifs ou visuellement similaires.

3. Génération Unifiée

La distribution finale est un produit de l'a priori préentraîné et des trois experts contrastifs (STM, LTM, SLTM). Lors de l'échantillonnage, les fonctions de score de ces experts sont combinées de manière additive (dans l'espace logarithmique) pour guider le processus de débruitage, garantissant que le futur généré respecte les dynamiques récentes, l'histoire à long terme et la cohérence spatiale.

Contributions Clés

Architecture Mémoire Découplée : L'article propose de résoudre le compromis mémoire en distribuant la mémoire sur des experts spécialisés plutôt que sur un seul socle, formulé comme un problème de Produit d'Experts.
Produit d'Experts Contrastifs (PoCE) : Une stratégie novatrice pour fusionner des experts hétérogènes. Les auteurs fournissent des preuves théoriques et empiriques que cette méthode prévient l'effondrement de mode et la surconfiance en factorisant les modes redondants, contrairement au PoE standard.
Mémoire Diffusive Externe : L'introduction d'un module LTM utilisant le finetuning au moment du test d'adaptateurs LoRA pour stocker des connaissances épisodiques dans les poids. Cela est étendu à la SLTM pour intégrer une structure spatiale, améliorant la précision dans les tâches de navigation.

Résultats Expérimentaux

La méthode a été évaluée sur des benchmarks synthétiques (Memory Maze, DMLab, Minecraft) et du monde réel (RealEstate10K, RECON).

Qualité de Reconstruction : Sur le jeu de données Memory Maze, CoME a atteint des performances de pointe en LPIPS (0,097), SSIM (0,892) et PSNR (23,07), surpassant significativement les références incluant les transformateurs à attention complète, les SSM et les modèles à fenêtre glissante.
Évolutivité : Contrairement aux modèles à attention complète qui nécessitent $60\times$ plus de calcul que la référence STM, CoME atteint une cohérence supérieure avec une fraction du coût d'entraînement et d'inférence.
Longueur de Contexte : Le modèle n'a montré aucune saturation dans la cohérence des performances jusqu'à 480 images de contexte, démontrant l'efficacité du composant LTM.
Planification et Navigation : Dans les tâches de navigation conditionnées par un objectif (benchmark RECON), CoME a réduit l'Erreur de Trajectoire Absolue (ATE) à 0,96, surpassant les modèles de monde de navigation spécialisés (NWM) et d'autres références.
Capacité de Rappel : Dans une expérience de trajectoire inversée sur RealEstate10K, CoME a réussi à rappeler les images initiales après des déroulements vers l'avant, tandis que les références sans mémoire ont généré des scènes plausibles mais incohérentes.
Études d'Ablation : La suppression de la formulation contrastive a provoqué un effondrement des performances, confirmant la nécessité du mécanisme PoCE. L'augmentation du rang LoRA et de la longueur de contexte pour la LTM a constamment amélioré les résultats sans surapprentissage.

Importance et Revendications

L'article revendique l'établissement d'un paradigme novateur pour construire des modèles de monde diffusifs augmentés par la mémoire. En distribuant la mémoire sur des experts spécialisés et en utilisant une stratégie de composition contrastive, CoME équilibre l'efficacité avec une cohérence temporelle de haute fidélité. Les auteurs soutiennent que cette approche permet aux modèles de monde de s'étendre à de longs contextes sans encourir de coûts quadratiques ou subir la perte de détails inhérente à la compression récurrente. La méthode est présentée comme un cadre général applicable à divers modèles de monde basés sur la diffusion, améliorant leur fiabilité pour la planification et la prise de décision dans des contextes d'apprentissage par renforcement. L'ouvrage souligne que la mémoire doit être une propriété de système distribué plutôt qu'une contrainte d'une architecture unique.