Résumé Technique : Mémoire Tensorielle : État Récurrent de Taille Fixe pour les Transformers à Long Horizon

1. Énoncé du Problème

Les Transformers traitent les images et les vidéos en aplatissant les dimensions spatiales et temporelles en longues séquences de tokens. Bien que les mécanismes d'attention et la mise en cache des clés et valeurs (KV) préservent les caractéristiques passées, ils souffrent de deux limitations principales à mesure que l'horizon de la séquence s'allonge :

1. Échelle de la Mémoire : L'inférence pratique repose sur le stockage des clés et valeurs passées, ce qui évolue linéairement avec le nombre de tokens. Cela devient prohibitif en termes de calcul et de mémoire pour les vidéos et documents longs.
2. Absence d'État Persistant : Les Transformers standards manquent d'un état explicite et persistant. Ils s'attendent aux caractéristiques passées mais ne maintiennent pas une représentation structurée pouvant être mise à jour et interrogée au fur et à mesure que l'entrée change. Cela entraîne une fragilité face aux occlusions, aux observations manquantes et à la nécessité de reconstruire à répétition le contexte global à partir d'une histoire sans cesse grandissante, plutôt que de maintenir un état persistant et compact.

Les approches existantes, telles que l'attention parcimonieuse, la compression de l'historique ou les contextes augmentés par la récupération, étendent souvent la longueur du contexte mais stockent l'information sous forme de collections non structurées de vecteurs de tokens ou d'emplacements. Elles ne préservent pas explicitement la structure spatiale ni ne maintiennent un « état du monde » compact adapté aux tâches basées sur l'état comme la compréhension vidéo.

2. Méthodologie : Mémoire Tensorielle

Les auteurs proposent la Mémoire Tensorielle, un module léger qui enrichit les blocs Transformer standards avec un tenseur mémoire récurrent 3D de taille fixe. Ce module découple la capacité de l'état de la longueur de l'entrée tout en préservant un biais inductif spatial.

Architecture Principale

La mémoire est définie comme une paire de volumes (état caché $h_t$ et état de cellule $c_t$) avec des dimensions $C \times D \times H \times W$, où les dimensions spatiales sont constantes quelle que soit la longueur de l'entrée. Le module fonctionne selon un cycle par tranche de tokens :

1. Découpage : Les tokens d'entrée sont regroupés en tranches. Une « requête de lecture » est dérivée du premier token de la tranche, et un « résumé d'écriture » est une projection apprise de l'ensemble de la tranche.
2. Adressage Continu : Un prédicteur de coordonnées partagé (utilisant des poids liés) génère des coordonnées 3D continues pour la lecture ($\mu_{read}$) et l'écriture ($\mu_{write}$) au sein de la grille de voxels.
3. Lecture (Échantillonnage Continu) : Le modèle lit dans l'état caché $h_t$ à $\mu_{read}$ en utilisant une interpolation trilinéaire. La caractéristique récupérée est projetée de nouveau vers la dimension du token et fusionnée dans le flux de tokens via une connexion résiduelle à porte. Un scalaire appris $\gamma$ contrôle la force de cette fusion, permettant au modèle de supprimer le chemin mémoire si celui-ci n'est pas bénéfique pour une tâche spécifique.
4. Écriture (Écriture Douce Différentiable) : Le modèle écrit de l'information en prédisant un vecteur de contenu et un paramètre de dispersion $\sigma$. Il dépose ce contenu sous forme d'un volume pondéré par une Gaussienne autour de l'emplacement 3D prédit $\mu_{write}$. Cette « écriture douce » évite l'assignation de voxels rigide, permettant aux gradients de façonner l'endroit où l'information est placée.
5. Mise à Jour Récurrente : Le volume d'écriture est concaténé à l'état caché actuel. Un opérateur 3D factorisé léger (convolutions 3D en profondeur suivies d'un mélange ponctuel) calcule des portes de style ConvLSTM ($i_t, f_t, o_t, g_t$). Ces portes mettent à jour les états de cellule et cachés, permettant à l'information de persister, de se propager localement et de rester disponible même lorsque les observations sont manquantes.

Choix de Conception Clés

• État de Taille Fixe : La taille du tenseur mémoire est constante, empêchant la croissance linéaire de la mémoire associée à la mise en cache KV.
• Structure Spatiale : La grille de voxels 3D préserve explicitement les relations spatiales, contrairement aux emplacements de tokens non structurés.
• Écritures Douces Différentiables : L'utilisation de volumes pondérés par une Gaussienne permet des mises à jour continues et lisses plutôt que des assignations discrètes et rigides.
• Fusion à Porte : La porte apprise $\gamma$ assure que le module peut être attaché à des modèles pré-entraînés sans déstabiliser l'entraînement ; si la mémoire est inutile, la porte supprime le chemin.

3. Contributions Clés

• Mécanisme de Mémoire Novel : Introduction d'un état de voxels 3D récurrent de taille fixe auquel les tokens peuvent écrire et depuis lequel ils peuvent lire via des coordonnées continues, agissant comme un état spatial explicite.
• Intégration : Le module s'intègre proprement dans les pipelines d'entraînement Transformer standards et peut être attaché ou retiré des blocs existants sans modifications architecturales.
• Suite de Diagnostic : Développement d'une suite de diagnostic « jouet » contrôlée conçue pour isoler les scénarios où un état persistant est bénéfique (par exemple, liaison de coordonnées, occlusion, construction de carte) par rapport aux scénarios où il ne l'est pas.

4. Résultats Expérimentaux

Les auteurs ont évalué la Mémoire Tensorielle sur des benchmarks standards couvrant les domaines du langage, de l'image et de la vidéo, ainsi que sur la suite de diagnostic.

• Modélisation du Langage (WikiText-2) : La Mémoire Tensorielle a obtenu les gains absolus les plus importants, réduisant la perplexité de test de 130,30 (Transformer de base) à 100,07. Elle a surpassé les références incluant Transformer-XL, l'attention locale et l'attention linéaire.
• Reconstruction de Patchs d'Image (CUB-200-2011) : Dans une tâche de reconstruction de patchs masqués, la variante Mémoire Tensorielle (Spatial ViT) a montré des améliorations constantes du PSNR et du SSIM sur les ratios de masquage (25 %, 50 %, 75 %) par rapport au ViT de base, au ViT local et au ViT avec registres.
• Reconnaissance d'Actions Vidéo (UCF-101) : Le modèle a surpassé à la fois une référence Transformer spatiotemporelle conjointe et une forte référence « token registre » sur la précision Top-1 (89,21 % contre 82,12 %) et Top-5. Les auteurs notent que cela s'est fait au prix d'un temps d'entraînement significativement plus élevé en temps réel dû à la récurrence de voxels par étape.
• Diagnostics Jouet :
  • Liaison de Coordonnées : La Mémoire Tensorielle a maintenu une haute précision (75–91 %) à mesure que le nombre d'écritures augmentait, tandis que les références basées sur l'attention s'effondraient vers une performance proche du hasard (7–14 %).
  • Contrôle Sans Nuisance : Sur des tâches ne nécessitant aucune mémoire, la porte apprise $\sigma(\gamma)$ est restée proche de son initialisation (supprimant le chemin mémoire), et tous les modèles ont atteint 100 % de précision.
  • Construction de Carte : La Mémoire Tensorielle a nettement surpassé les références dans des tâches de navigation sur grille à long horizon où l'information doit être retenue dans le temps.

5. Importance et Revendications

L'article revendique que la Mémoire Tensorielle offre une manière pratique de gérer le raisonnement à long horizon en fournissant aux Transformers un état explicite, compact et structuré.

• Découplage de la Capacité et de la Longueur : En utilisant un tenseur de taille fixe, le modèle évite l'échelle linéaire de la mémoire des caches KV, le rendant théoriquement plus évolutif pour les séquences longues.
• Robustesse aux Occlusions : L'état persistant permet au modèle de maintenir un « état de croyance » sur la scène, le rendant plus robuste aux observations manquantes ou aux occlusions par rapport aux modèles qui reposent uniquement sur la reconstruction du contexte à partir d'un flux de tokens.
• Polyvalence : Bien que motivée par la vidéo à long horizon, le mécanisme s'avère bénéfique pour le langage (accumulation de contexte global) et les images (agrégation spatiale structurée).
• Fondation pour le Travail Futur : Les auteurs considèrent l'implémentation actuelle comme une fondation. Ils suggèrent que les directions futures incluent des mémoires de dimensions supérieures (par exemple, 4D pour une mémoire spatiale consciente du temps), des dynamiques plus riches (par exemple, warping de flux optique) et des mises à jour parcimonieuses pour améliorer l'efficacité.

Les auteurs restent modestes quant aux limitations actuelles, reconnaissant que les mises à jour de voxels denses constituent un goulot d'étranglement pour le débit et que l'état de taille fixe peut saturer si une tâche nécessite le stockage d'un grand nombre de détails distincts. Cependant, ils affirment que l'ajout d'un petit état structuré est une voie viable vers des systèmes de perception et de raisonnement à long horizon plus performants.