Résumé technique : L'apprentissage par renforcement sensible à l'architecture rend l'attention à fenêtre glissante compétitive pour le raisonnement mathématique

Problématique

L'avancement rapide des modèles de langage (LLM) de raisonnement et d'agents a créé une demande critique pour l'inférence à contexte long. Cependant, les mécanismes d'auto-attention (SA) standards augmentent de manière quadratique ($O(N^2)$) avec la longueur du contexte, créant un goulot d'étranglement fondamental pour l'efficacité des modèles basés sur les transformers. Bien que les alternatives à complexité linéaire, comme l'attention à fenêtre glissante (SWA), offrent une voie prometteuse pour le passage à l'échelle, elles peinent généralement à égaler les performances de la SA sur des tâches exigeant une modélisation de dépendances à longue portée, telles que le raisonnement mathématique (entrée courte, génération longue).

Les approches existantes pour combler cet écart font face à deux défis majeurs :

1. Acquisition du modèle de base : Les checkpoints à haute performance de complexité linéaire sont rares, et les pré-entraîner à partir de zéro est prohibitif en termes de coûts.
2. Écart de performance : La simple conversion d'un modèle pré-entraîné en SA vers la SWA suivie d'un ajustement fin supervisé (SFT) entraîne une baisse significative des performances. Les auteurs émettent l'hypothèse que cela est dû à un décalage entre les données et l'architecture : les jeux de données SFT standards sont préparés pour des modèles SA et contiennent des dépendances à longue portée que la SWA ne peut pas modéliser efficacement, menant à une imitation sous-optimale.

Méthodologie : SWARR

Le papier propose SWARR (Sliding-Window Attention with Reinforced Adaptation for Math Reasoning), une recette empirique en deux étapes conçue pour adapter les modèles SWA au raisonnement mathématique tout en préservant l'efficacité.

Étape 1 : Conversion efficace avec SFT

• Initialisation : Un modèle SA pré-entraîné (ex: Qwen2.5-Math-1.5B-Base) est converti en SWA en remplaçant l'attention globale par une fenêtre glissante locale tout en réutilisant les projections QKV pré-entraînées.
• Ajustement fin supervisé (SFT) : Le modèle subit un SFT sur des données de raisonnement mathématique.
• Limitation : Les auteurs notent que le SFT est « agnostique à l'architecture ». Puisque les trajectoires cibles dans le jeu de données ont été générées par des modèles SA, elles peuvent contenir des dépendances s'étendant au-delà de la fenêtre SWA, forçant le modèle SWA à imiter des schémas qu'il ne peut pas naturellement supporter.

Étape 2 : Adaptation de la politique avec RL

• Apprentissage par renforcement (RL) : Le modèle SWA converti par SFT est ensuite adapté via le RL (spécifiquement CISPO avec des récompenses vérifiables).
• Sensibilité à l'architecture : Contrairement au SFT, qui s'entraîne sur des trajectoires externes fixes, le RL optimise les trajectoires générées par le modèle lui-même sous ses propres contraintes architecturales. Cette optimisation « on-policy » permet au modèle de déplacer sa distribution de génération vers des schémas plus faciles à produire avec une attention locale (c'est-à-dire des dépendances plus locales), atténuant ainsi le décalage entre les données et l'architecture.

Métriques de localité

Pour valider l'hypothèse selon laquelle le RL adapte le modèle pour le rendre plus « local », les auteurs introduisent une métrique d'écart d'information basée sur la probabilité ($g_t(w)$) et un score de localité au niveau de la séquence ($L(w)$). Ces métriques quantifient la quantité d'information prédictive pour le jeton suivant capturée dans une fenêtre locale par rapport au contexte complet.

Principales contributions

1. Recette SWARR : Un pipeline pratique en deux étapes (conversion SA-vers-SWA + adaptation par RL) qui permet aux modèles SWA à fenêtre large de rivaliser avec la SA en raisonnement mathématique sans pré-entraînement coûteux.
2. Réévaluation empirique de la SWA : Le papier démontre que si le SFT seul laisse un écart substantiel entre la SWA et la SA, l'adaptation par RL réduit considérablement cet écart, contestant la conclusion selon laquelle la SWA est inadaptée au raisonnement à génération longue.
3. Analyse du décalage données-architecture : À travers des expériences de SFT croisé, les auteurs montrent que les données générées par SA nuisent davantage aux performances de la SWA que les données générées par SWA (appariées à l'architecture). Cela soutient l'hypothèse selon laquelle le RL améliore les performances en alignant la distribution d'entraînement avec la contrainte SWA.
4. Localité de la pensée : L'étude fournit des preuves que le raisonnement mathématique efficace peut émerger de dépendances de jetons plus locales, s'alignant sur le phénomène de « localité de la pensée ».

Résultats expérimentaux

Les expériences ont été menées sur le modèle Qwen2.5-Math-1.5B-Base sur des benchmarks incluant Math500, AMC, OlympiadBench, AIME24 et AIME25.

• Performance post-SFT : Après conversion et SFT, les modèles SWA accusaient un retard par rapport à la SA. L'écart s'accentuait à mesure que la taille de la fenêtre diminuait (ex: SWA2k-SFT était 17,8 points en dessous de SA-SFT).
• Performance post-RL :
  • Entraînement à étapes égales : SWA8k-RL-900 a atteint une performance comparable à SA-RL-900, avec un écart de précision moyen de seulement -0,4 %.
  • Entraînement à temps égal : En raison de la complexité linéaire de la SWA, l'entraînement est plus rapide. Entraînés pour un temps de calcul similaire (~500 heures GPU), SWA4k-RL-1400 a égalé la performance de SA-RL-900 (+0,1 % d'écart), et SWA8k-RL-1200 a même légèrement surpassé celle-ci (+0,7 %, bien que dans la marge d'incertitude statistique).
  • Petites fenêtres : SWA2k est resté contraint, bien que l'écart se soit nettement réduit (passant de -17,8 à -2,5 points sous entraînement à temps égal).
• Efficacité : Les modèles SWA ont démontré des gains d'efficacité substantiels. Par exemple, SWA4k a obtenu une accélération d'environ 6,2× par rapport à la SA à une longueur de contexte de 32k et a supporté des tailles de lots (batch sizes) nettement plus grandes sous des contraintes de mémoire fixes.
• Validation par SFT croisé : Lorsqu'ils ont été entraînés sur des données générées par des modèles SWA (appariées à l'architecture), l'écart de performance entre la SWA et la SA a presque disparu, confirmant que la chute initiale était largement due au décalage avec les données orientées SA.

Signification et affirmations

La contribution centrale du papier est une découverte empirique : l'apprentissage par renforcement change fondamentalement l'évaluation de la viabilité de la SWA à fenêtre large pour le raisonnement mathématique.

• Changement de conclusion : Sans RL, la conversion de la SA en SWA suivie d'un SFT suggère que la SWA est inférieure pour les tâches de génération longue. Avec le RL, la SWA à fenêtre large devient compétitive avec la SA.
• Mécanisme : L'amélioration est attribuée à la capacité du RL à adapter la trajectoire de génération du modèle pour qu'elle soit plus « locale », s'ajustant mieux aux contraintes de l'architecture SWA qu'un jeu de données SFT statique ne pourrait le faire.
• Portée et limites : Les auteurs précisent explicitement que leurs conclusions sont spécifiques au raisonnement mathématique (une tâche de type entrée courte, génération longue). Ils ne prétendent pas que SWARR est une recette universelle pour toutes les tâches, notant que différents domaines peuvent nécessiter des stratégies d'optimisation différentes. De plus, en tant qu'étude académique, les résultats sont basés sur des modèles de 1,5B et 4B ; le passage à l'échelle vers des modèles commerciaux massifs (ex: 200B+) reste un travail futur.

En résumé, SWARR offre une voie pratique et efficace pour déployer des modèles d'attention à complexité linéaire pour des tâches de raisonnement complexe en exploitant un RL sensible à l'architecture pour surmonter les limitations du fine-tuning standard.