Résumé Technique : L'Attention à Fenêtre Courte Permet la Mémoire à Long Terme

Énoncé du Problème

Les modèles de langage modernes (LLM) basés sur l'attention softmax reposent sur un cache Key-Value (KV) qui croît linéairement avec la longueur de la séquence, entraînant des coûts de calcul et de mémoire prohibitifs pour les tâches à contexte long. À l'inverse, les réseaux de neurones récurrents (RNN) linéaires et les modèles d'espace d'état (SSM) maintiennent des coûts de mémoire et de calcul constants, mais souffrent souvent de capacités de rappel inférieures à celles des Transformers.

Les architectures hybrides récentes tentent de combiner l'attention à fenêtre glissante locale (SWA) avec une attention globale ou des couches de RNN linéaires pour équilibrer efficacité et performance. Cependant, l'impact de la longueur de la fenêtre glissante sur l'interaction entre la mémoire locale et la mémoire à long terme reste sous-étudié. Plus précisément, il est unclear comment la taille de la fenêtre influence la capacité du modèle à apprendre des dépendances à long terme en présence de RNN linéaires, et si des fenêtres plus grandes améliorent intrinsèquement la performance ou entravent la capacité du modèle à extrapoler vers des contextes plus longs.

Méthodologie

Les auteurs étudient des architectures hybrides qui entrelacent des couches d'attention à fenêtre glissante (SWA) avec des couches de RNN linéaires, utilisant spécifiquement xLSTM (une variante de mLSTM avec mémoire matricielle) comme composant de mémoire à long terme. L'architecture proposée, nommée SWAX, alterne entre des couches SWA et des couches xLSTM dans un rapport de 1:1.

L'étude se concentre sur deux contributions méthodologiques principales :

1. Analyse de la Taille de Fenêtre : Les auteurs entraînent des modèles SWAX avec des tailles de fenêtre glissante fixes allant de 128 à 2048 tokens. Ils évaluent ces modèles sur des benchmarks de raisonnement à contexte court (par exemple, ARC, Hellaswag, MBPP) et des tâches de récupération à contexte long (par exemple, RULER Needle-In-A-Haystack).
2. Entraînement Stochastique à Fenêtre : Pour répondre au compromis entre la précision à contexte court (favorisée par des fenêtres plus grandes) et l'extrapolation à contexte long (favorisée par des fenêtres plus petites), les auteurs proposent une stratégie d'entraînement où la taille de la fenêtre glissante est échantillonnée de manière stochastique pour chaque lot. Dans leurs expériences, la taille de la fenêtre alterne entre une taille courte (128) et une taille longue (2048) avec une probabilité spécifique (0,5 pour les modèles de 1,4B, 0,75 pour les modèles de 7B). De plus, ils emploient un calendrier d'annealing où l'échantillonnage stochastique est supprimé pour les 10 % finaux de l'entraînement afin de stabiliser les performances à contexte court.

Les expériences sont menées sur des modèles de 1,4B et 7B paramètres, entraînés sur 150 milliards de tokens avec une longueur de séquence fixe de 16k pendant le pré-entraînement.

Contributions Clés

1. Résultat Contre-Intuitif sur la Taille de Fenêtre : L'article démontre que pour les architectures hybrides combinant SWA et RNN linéaires, des fenêtres glissantes plus grandes peuvent dégrader les performances à contexte long. Contrairement à l'idée reçue selon laquelle des fenêtres plus grandes fournissent plus de contexte, les auteurs constatent que de grandes fenêtres (par exemple, 2048) permettent au modèle de trop compter sur l'attention softmax locale précise, entraînant un sous-entraînement des couches de RNN linéaires sur les dépendances à long terme.
2. Mécanisme du Bénéfice de la Fenêtre Courte : Des fenêtres plus courtes (par exemple, 128) forcent le modèle à s'appuyer sur les couches de RNN linéaires pour la récupération d'informations tombant en dehors de la fenêtre locale. Ce « signal de supervision » contraint les couches linéaires à se spécialiser dans la modélisation des dépendances à long terme, améliorant ainsi la capacité du modèle à extrapoler vers des longueurs de séquence dépassant largement la fenêtre d'entraînement.
3. Stratégie d'Entraînement Stochastique : Les auteurs introduisent une procédure d'entraînement qui fait varier stochastiquement la taille de la fenêtre. Cette approche force le modèle à exploiter à la fois la précision à court terme de la fenêtre et la mémoire à long terme du RNN, agissant efficacement comme une forme de dropout sur le mécanisme d'attention.

Résultats

• Performances à Contexte Long : Les modèles entraînés avec une grande fenêtre fixe (2048) présentent une chute catastrophique des performances sur les tâches de récupération à contexte long (par exemple, RULER NIAH) lorsqu'ils sont testés sur des séquences plus longues que la fenêtre. En revanche, les modèles entraînés avec des fenêtres plus courtes (128) maintiennent une haute précision même à des longueurs de séquence de 131k tokens. La méthode d'entraînement stochastique atteint des performances comparables ou supérieures à celles des modèles utilisant uniquement des fenêtres courtes sur ces tâches à contexte long.
• Performances à Contexte Court : Bien que des fenêtres courtes (128) seules sous-performent parfois sur les benchmarks de raisonnement à contexte court par rapport à des fenêtres plus grandes, la méthode d'entraînement stochastique récupère cette performance. Aux échelles de 1,4B et 7B, l'entraînement stochastique produit des performances à contexte court comparables, voire dans certains cas supérieures, à celles des modèles entraînés avec une fenêtre fixe de 2048.
• Généralisation : Les résultats tiennent pour différents backbones d'attention linéaire (xLSTM et Gated DeltaNet) et différentes configurations hybrides (alternance SWA avec des couches d'attention complète). La stratégie d'entraînement stochastique surpasse systématiquement les baselines à fenêtre fixe sur les benchmarks à contexte court et long.

Signification et Revendications

L'article affirme que le choix de la taille de la fenêtre glissante dans les architectures hybrides n'est pas simplement un compromis entre coût computationnel et performance, mais un facteur critique dans la manière dont le modèle apprend à distribuer les responsabilités de mémoire.

Les auteurs affirment que :

• Les fenêtres courtes sont essentielles pour entraîner la mémoire à long terme : En restreignant le champ d'attention local, le modèle est contraint d'utiliser ses couches de RNN linéaires pour les dépendances à longue portée, conduisant à une meilleure extrapolation de longueur.
• L'entraînement stochastique résout le compromis : En changeant dynamiquement la taille de la fenêtre pendant l'entraînement, le modèle apprend à utiliser à la fois l'attention locale haute fidélité pour les contextes courts et la mémoire linéaire globale pour les contextes longs.
• Une sur-reliance sur l'attention locale est préjudiciable : Si un modèle est entraîné avec une grande fenêtre, il échoue à apprendre les mécanismes nécessaires pour récupérer des informations au-delà de cette fenêtre, conduisant à une mauvaise généralisation sur des séquences plus longues.

Ce travail suggère que pour que les architectures hybrides atteignent des performances à contexte long de l'état de l'art sans coûts de calcul illimités, la procédure d'entraînement doit activement empêcher le modèle de trop compter sur les mécanismes d'attention locale, même si ces mécanismes sont plus précis pour les tâches à courte portée.