💬 NLP

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

Cet article introduit une méthode de réglage fin pour les modèles transformer qui permet une inférence de contexte long efficace avec une attention éparse sur du matériel modeste en permettant aux modèles de co-adapter avec diverses politiques de cache KV, surpassant souvent les approches d'attention exacte.

Auteurs originaux : Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Publié 2026-08-21

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

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

Les systèmes d'intelligence artificielle modernes qui génèrent du texte semblable à celui de l'humain reposent sur un mécanisme qui agit comme une mémoire à court terme, leur permettant de se souvenir de ce qui a été dit plus tôt dans une conversation ou un document long. Cette mémoire est stockée dans un tampon numérique qui s'agrandit à chaque nouveau mot que le système traite. Pour que ces systèmes fonctionnent bien sur de longues séquences de texte, cette mémoire doit être vaste, mais le matériel informatique requis pour la contenir est coûteux et limité. Lorsque le tampon de mémoire est plein, le système doit décider quelles anciennes informations il doit écarter pour faire de la place aux nouvelles. S'il jette les mauvaises informations, le système perd sa capacité à raisonner ou à répondre avec précision. Cela crée un arbitrage difficile : garder une mémoire petite permet d'économiser de l'argent et permet au système de fonctionner sur du matériel standard, mais cela risque de perdre le contexte nécessaire pour être intelligent.

Les chercheurs ont longtemps tenté de résoudre ce problème en apprenant au système à être sélectif sur ce qu'il conserve, un processus connu sous le nom d'attention parcimonieuse (sparse attention). Cependant, une nouvelle étude révèle une faille critique dans la manière dont ces systèmes ont été entraînés jusqu'à présent. La plupart des méthodes existantes entraînent l'IA en utilisant une mémoire parfaite et illimitée, puis tentent de la forcer à fonctionner avec une mémoire limitée plus tard. Les chercheurs ont découvert que cette approche échoue car l'IA n'a jamais appris à fonctionner sous les contraintes spécifiques auxquelles elle serait réellement confrontée. En apprenant au modèle à oublier intentionnellement et à s'adapter à une taille de mémoire fixe dès le début, l'équipe a démontré que le système pouvait obtenir des performances nettement supérieures à celles des modèles entraînés avec des ressources illimitées, même en fonctionnant sur une seule puce informatique modérément puissante.

L'équipe, dirigée par des scientifiques d'Amazon Web Services et de l'Université d'Amsterdam, a développé une nouvelle façon d'affiner ces grands modèles de langage. Au lieu d'utiliser des supercalculateurs massifs pour simuler une mémoire parfaite, ils ont appris aux modèles à co-adapter avec une politique de gestion de la mémoire spécifique. Imaginez un bibliothécaire qui est formé pour organiser des livres dans une bibliothèque dotée d'étagères infinies, pour ensuite être informé qu'il devra travailler dans une toute petite pièce avec une seule étagère. Le bibliothécaire formé dans la grande bibliothèque aurait probablement du mal à prioriser ce qu'il faut garder dans la petite pièce. En revanche, la méthode proposée dans cet article entraîne le bibliothécaire directement dans la petite pièce, lui apprenant exactement quels livres garder et lesquels écarter en fonction des règles de cet espace spécifique. Cela permet au modèle d'apprendre le rythme de ses propres limites, plutôt que d'essayer de désapprendre les habitudes liées à l'avoir trop d'espace.

Les chercheurs ont testé cette approche sur un modèle de quatre milliards de paramètres, une taille substantielle mais gérable. Ils ont mené leurs expériences sur une seule carte graphique de 40 gigaoctets de mémoire, une configuration abordable pour de nombreuses organisations comparée aux clusters de dizaines de cartes requis par les méthodes précédentes. Ils ont comparé cette nouvelle technique d'entraînement à la méthode standard, qui utilise une technique appelée parallélisme de séquence pour répartir la charge de la mémoire sur plusieurs dispositifs coûteux. Les résultats ont montré que les modèles entraînés avec la nouvelle méthode surpassaient souvent les modèles standards, particulièrement lorsque la tâche exigeait que le système génère des réponses spécifiques et concises plutôt que des textes longs et décousus. Dans plusieurs tests impliquant des questions complexes et de l'extraction de données, la méthode standard produisait des sorties qui étaient bien trop longues et remplies de chiffres aléatoires et sans aucun sens, tandis que la nouvelle méthode apprenait à s'arrêter au bon moment et à fournir la valeur unique correcte.

Une part clé de ce succès a été l'amélioration de l'« oracle heavy-hitter », une stratégie populaire pour décider quelles informations conserver. Cette stratégie fonctionne en suivant les éléments d'information auxquels le modèle prête le plus d'attention au fil du temps. Les chercheurs ont constaté que la version originale de cette stratégie était lente et inefficace. Ils ont réécrit le code sous-jacent pour qu'il soit beaucoup plus rapide, permettant au système de calculer ces scores d'importance sans ralentir l'ensemble du processus. Cette optimisation a permis au système de prendre des décisions intelligentes sur ce qu'il fallait oublier en temps réel, sans avoir besoin de la puissance de calcul massive qui accompagne habituellement de telles tâches. L'équipe a également publié une nouvelle bibliothèque logicielle en open-source pour rendre ces techniques accessibles aux autres, visant à abaisser la barrière pour quiconque souhaite construire des systèmes d'IA à contexte long sans avoir besoin d'une fortune en matériel.

L'étude souligne que la manière dont un modèle est entraîné est tout aussi importante que le matériel sur lequel il fonctionne. Lorsque les chercheurs ont forcé le modèle à s'entraîner avec les mêmes contraintes de mémoire qu'il rencontrerait lors de l'utilisation, il a appris à naviguer efficacement dans ces contraintes. Dans un test spécifique impliquant des données JSON, la méthode standard a totalement échoué, incapable de trouver les points de données corrects, tandis que la nouvelle méthode a réussi à les identifier environ la moitié du temps. Cela suggère que la capacité à gérer de longs contextes n'est pas seulement une question d'avoir plus de mémoire, mais d'apprendre au système comment gérer la mémoire dont il dispose. Les conclusions indiquent que pour de nombreuses applications, la voie la plus efficace ne consiste pas à construire des ordinateurs plus gros, mais à apprendre au logiciel à être plus efficace avec les ressources qu'il possède déjà.

Résumé technique : Apprendre à oublier : Fine-tuning pour l'attention parcellaire à contexte long

1. Énoncé du problème

Les modèles de langage de grande taille (LLM) modernes nécessitent la capacité de traiter des contextes extrêmement longs pour des tâches telles que les conversations multi-tours, l'utilisation d'outils et le raisonnement par chaîne de pensée (chain-of-thought). Bien que les implémentations standards de transformeurs passent à une complexité quadratique en calcul et linéaire en mémoire avec la largeur du contexte, des méthodes d'attention parcellaire ont été proposées pour atteindre une mise à l'échelle temporelle linéaire et une mise à l'échelle mémoire constante. Ces méthodes utilisent un cache Clé-Valeur (KV) de taille fixe où l'information est évincée selon des politiques spécifiques (par exemple, l'oracle Heavy-Hitter, H2O) une fois que le cache est plein.

Cependant, un écart important existe entre l'inférence et l'entraînement dans ce domaine :

Contraintes d'entraînement : Le fine-tuning de modèles avec des politiques d'attention parcellaire est notoirement difficile sur des budgets matériels modérés. L'autodifférentiation standard (autograd) nécessite de stocker l'intégralité du graphe de calcul, y compris les buffers du cache KV pour chaque segment de la séquence. Pour les contextes longs, cette exigence de mémoire dépasse la mémoire GPU disponible, nécessitant souvent des configurations multi-GPU coûteuses avec parallélisme de séquence ou de contexte (SP/CP).
Écart de co-adaptation : Les modèles entraînés avec une attention exacte (utilisant SP/CP) sont souvent nettement moins performants lorsqu'ils sont déployés avec une attention parcellaire lors de l'inférence. Cela est dû au fait que le modèle n'est jamais exposé aux contraintes de mémoire spécifiques et à la logique d'éviction auxquelles il sera confronté lors de l'inférence.
Inefficacité d'implémentation : Les principales politiques d'attention parcellaire comme H2O sont sous-utilisées dans les bibliothèques d'inférence de pointe (par exemple, vLLM) car les implémentations actuelles sont lentes, manquant de kernels optimisés pour les exigences spécifiques de l'attention parcellaire (par exemple, le retour des poids d'attention sommés).

2. Méthodologie

Le papier propose une nouvelle méthode de fine-tuning qui permet aux modèles de transformeurs de co-adapter avec des politiques de cache KV arbitraires sur du matériel modéré (démontré sur un seul Nvidia A100 40GB).

2.1 Checkpointing d'activation imbriqué avec déchargement CPU (Offloading)

Pour remédier au goulot d'étranglement de la mémoire lors de l'entraînement sur de longues séquences ( $N \gg N_C$ , où $N_C$ est la taille du cache), les auteurs emploient une stratégie de checkpointing d'activation imbriquée :

Boucle externe (Couches) : La passe arrière itère sur les couches du modèle.
Boucle interne (Cellules) : La séquence est divisée en "cellules" (groupes de segments). Pour chaque cellule, une passe avant est exécutée pour stocker les buffers de cache KV nécessaires dans la mémoire CPU.
Mécanisme de Replay : Au lieu de différencier la politique de cache KV (qui est souvent non-différentiable ou coûteuse), la passe avant enregistre toutes les décisions d'éviction. Pendant la passe arrière, ces décisions sont rejouées à l'aide de "caches de replay" qui imitent la logique d'éviction de la passe avant sans nécessiter de calcul de gradient à travers la politique elle-même.

2.2 Encodage Delta pour les buffers de cache KV

L'optimisation de mémoire la plus critique concerne la manière dont les buffers de cache KV sont stockés dans le graphe autograd.

Récurrence Linéaire : Les auteurs observent que le buffer de cache KV pour un nouveau segment ($keys'$) est une transformation linéaire du buffer précédent ($keys$), ne différant que par les valeurs des nouveaux tokens et les indices spécifiques des emplacements écrasés par la politique d'éviction.
Stockage Delta : Au lieu de stocker les tenseurs $keys$ et $values$ complets pour chaque segment dans le graphe de calcul, la méthode stocke uniquement le delta (la différence) et les indices des emplacements écrasés.
Hooks de tenseurs sauvegardés Autograd : En utilisant les autograd_saved_tensors_hooks de PyTorch, les auteurs implémentent un mécanisme de "packing". Pendant la passe avant, les tenseurs complets sont remplacés par des références à des tuples (index, delta). Pendant la passe arrière, les tenseurs complets sont reconstruits ("unpacked") à partir des deltas et des buffers subséquents. Cela réduit l'empreinte mémoire du graphe autograd d'un facteur $k$ (le nombre de segments par cellule), la ramenant à des niveaux comparables à l'inférence.

2.3 Améliorations de la politique H2O

Le papier introduit des améliorations spécifiques à la politique Heavy-Hitter Oracle (H2O) :

Indépendance de Batch : L'implémentation originale de H2O agrégeait les scores à travers les dimensions de batch. La nouvelle méthode permet de prendre des décisions d'éviction de manière indépendante pour chaque position de batch.
Normalisation : Un score H2O normalisé est introduit pour éviter le biais envers les tokens qui sont restés plus longtemps dans le cache.
Optimisation de Kernel : Pour permettre un entraînement et une inférence efficaces, les auteurs fournissent du code Triton qui retourne les poids d'attention sommés aux côtés d'un kernel FlashInfer Scaled Dot Product Attention (SDPA). Cela permet à H2O de fonctionner efficacement sans recourir à des implémentations SDPA naïves et lentes.

3. Contributions clés

Nouvelle méthode de Fine-Tuning : Une méthode pour effectuer le fine-tuning de modèles transformeurs avec une attention parcellaire et des politiques de cache KV arbitraires sur des ressources comparables à l'inférence d'attention parcellaire (par exemple, un seul A100). Elle combine le checkpointing d'activation imbriqué, le déchargement CPU et l'encodage delta via des hooks autograd.
Implémentation et Optimisation de H2O : Améliorations méthodologiques de la politique H2O (indépendance de batch, normalisation) et une implémentation haute performance utilisant des kernels SDPA dédiés qui retournent les poids d'attention sommés, marquant une étape significative vers des latences compétitives avec les bibliothèques d'inférence de pointe.
Évaluation Complète : Expériences sur la suite de benchmarks Helmet (fenêtres de contexte de 64k et 128k) démontrant que, bien que les performances varient selon le jeu de données, la méthode proposée permet aux modèles de co-adapter avec des politiques de cache KV, surpassant souvent les modèles entraînés avec le parallélisme de séquence lors d'une évaluation sous inférence d'attention parcellaire.
Bibliothèque KeysAndValues : La publication d'une bibliothèque open-source (awslabs/keys_values) fournissant un code performant et facile à utiliser pour l'inférence et le fine-tuning à contexte long avec une attention parcellaire.

4. Résultats Expérimentaux

Les auteurs ont évalué un modèle Qwen3-4B-Instruct sur le benchmark Helmet pour des longueurs de contexte de 64k et 128k.

Comparaison : Les modèles ont été entraînés en utilisant :
- SP : Parallélisme de Séquence avec attention exacte (ligne de base SotA).
- US : La méthode proposée avec la politique de cache KV spécifique en place.
- NO : Le checkpoint de base sans fine-tuning.
Constatations :
- Résultats mixtes sur les jeux de données de la Table 1 : Pour les jeux de données de la Table 1, les résultats étaient mitigés et peu concluants. La ligne de base SP (attention exacte) a obtenu les meilleurs résultats pour nq et hotpot qa. Pour trivia qa, le checkpoint de base sans fine-tuning (no) a été le plus performant, indiquant que le fine-tuning n'a pas aidé dans ce cas précis. La méthode proposée (us) a surpassé SP sur pop qa.
- Forte performance sur les jeux de données de la Table 2 : Pour les jeux de données de la Table 2, us a largement surpassé à la fois sp et no.
- Modes de défaillance de SP : Les modèles entraînés avec SP ont présenté un mode de défaillance où les sorties étaient excessivement longues et contenaient des absurdités aléatoires lorsqu'ils étaient soumis à l'éviction d'attention parcellaire, probablement parce qu'ils n'avaient jamais été entraînés à gérer le mécanisme d'oubli. En revanche, us a appris à s'arrêter correctement et produit généralement un nombre unique ou la cible correcte.
- Variance de la politique : Les résultats variaient considérablement selon le jeu de données. Bien que les variantes de H2O aient bien performé sur plusieurs tâches, la politique lastrec (lr) a surpassé H2O sur des jeux de données tels que nq et hotpot qa. De plus, pour nq et hotpot qa, la ligne de base SP (attention exacte) a souvent obtenu les meilleurs résultats, tandis que pour pop qa et d'autres jeux de données de la Table 2, la méthode US avec des politiques d'attention parcellaire était supérieure.
- Efficacité : L'implémentation de H2O avec le nouveau support de kernel représente une amélioration majeure de l'efficacité, faisant un grand pas vers des latences compétitives avec les bibliothèques SotA, bien que le papier reconnaisse que la latence de l'attention parcellaire reste actuellement nettement supérieure aux méthodes utilisant le parallélisme de séquence ou de contexte.

5. Signification et Revendications

Le papier affirme que l'attention parcellaire est sous-utilisée en pratique principalement en raison de problèmes de latence et de la difficulté de fine-tuner les modèles pour qu'ils s'adaptent à des politiques d'éviction spécifiques.

Faisabilité : Les auteurs démontrent qu'un fine-tuning de haute qualité pour l'attention parcellaire à contexte long est réalisable sur du matériel modéré, éliminant le besoin de clusters d'entraînement distribués massifs uniquement pour adapter un modèle à une politique de cache.
Nécessité de Co-adaptation : Les résultats suggèrent fortement que les modèles doivent être entraînés sous les mêmes contraintes (politiques d'éviction) qu'ils rencontreront lors de l'inférence pour éviter une dégradation des performances.
Direction Future : Le papier soutient qu'en simplifiant la structure du cache KV et en clarifiant les exigences des kernels SDPA (par exemple, le retour des poids sommés), la communauté peut mieux intégrer l'attention parcellaire dans les bibliothèques d'inférence SotA.
Limites : Les auteurs reconnaissent que bien que leur méthode améliore la latence, la nature séquentielle des décisions d'attention parcellaire peut encore entraîner des latences plus élevées que RingAttention ou le Parallélisme de Séquence pour des budgets matériels très importants. Ils prévoient d'aborder cela via la fusion de kernels et le déchargement asynchrone CPU dans des travaux futurs.

En résumé, le papier fournit une voie pratique pour "apprendre à oublier", permettant aux modèles de se spécialiser pour l'inférence à contexte long sur du matériel accessible. Bien que les performances par rapport au parallélisme de séquence varient selon le jeu de données, la méthode surpasse souvent l'entraînement par attention exacte dans les scénarios nécessitant une adaptation à des politiques d'attention parcellaire spécifiques.