🤖 machine learning

AdaSplash-2: Faster Differentiable Sparse Attention

Ce papier présente AdaSplash-2, une nouvelle implémentation efficace de l'attention $\alpha$ -entmax qui, grâce à une initialisation basée sur des histogrammes et une exécution GPU optimisée, réduit considérablement les coûts computationnels pour rivaliser avec FlashAttention-2 tout en offrant des performances supérieures sur les tâches à contexte long.

Auteurs originaux : Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

Publié 2026-04-17

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

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 êtes le chef d'un orchestre géant (un modèle d'intelligence artificielle) qui doit écouter des milliers de musiciens (les mots d'un texte) pour jouer une symphonie parfaite.

Le problème, c'est que dans les orchestres actuels (les modèles d'IA classiques), le chef doit écouter chaque musicien en même temps pour chaque note qu'il joue. Si l'orchestre a 100 musiciens, c'est gérable. Mais si vous avez 100 000 musiciens (un contexte long, comme un livre entier), le chef passe son temps à écouter tout le monde, même ceux qui ne jouent rien d'important. C'est lent, épuisant et ça coûte cher en énergie.

Voici l'histoire de ADASPLASH-2, la nouvelle méthode qui va changer la donne.

1. Le Problème : L'Orchestre qui écoute tout le monde

Dans les modèles actuels, le chef utilise une règle appelée "Softmax". C'est comme si le chef disait : "Je vais donner un peu d'attention à tout le monde, même au musicien qui dort dans le coin."
Résultat : Le chef est submergé. Il ne peut pas lire un livre entier en une seule fois sans s'effondrer.

2. La Solution Intelligente : L'Orchestre qui sait ignorer le bruit

Les chercheurs ont proposé une nouvelle règle appelée α-entmax. C'est comme un chef d'orchestre plus malin qui dit : "Je vais écouter seulement les musiciens qui jouent vraiment fort. Ceux qui chuchotent ? Je les ignore complètement."
C'est génial car cela rend le processus beaucoup plus rapide (on ne traite que les notes importantes). Mais il y a un hic : pour savoir qui ignorer, le chef doit faire un calcul mathématique très compliqué et lent à chaque fois. C'est comme essayer de trier des milliers de cartes à la main avant de pouvoir jouer.

3. L'Innovation ADASPLASH-2 : Le "Tri-Flash"

C'est ici qu'intervient ADASPLASH-2. Les auteurs ont inventé une astuce incroyable pour rendre ce tri ultra-rapide.

Imaginez que le chef a une boîte à outils magique (la mémoire rapide de l'ordinateur, appelée SRAM). Au lieu de trier chaque note une par une, il fait ceci :

Le Histogramme (Le Tri-Express) : Au lieu de regarder chaque musicien individuellement, le chef lance un filet rapide qui capture les sons par "paquets". Il remplit une boîte avec des compartiments (un histogramme) : "Combien de musiciens jouent fort ?", "Combien jouent moyen ?", "Combien chuchotent ?".
- Analogie : C'est comme si vous deviez compter les voitures sur une autoroute. Au lieu de les compter une par une, vous lancez un filet qui les attrape par groupes de 10, 20 ou 50. Vous avez une idée très précise du trafic en une seconde.
L'Estimation Rapide : Grâce à cette boîte à compartiments, le chef peut deviner presque instantanément qui doit être ignoré. Il n'a plus besoin de faire des calculs longs et pénibles. Il a une "devinette" très précise qui lui permet de sauter directement à la bonne réponse.
Le Nettoyage (La Mise à Jour) : Une fois qu'il a cette idée, il fait juste un ou deux petits ajustements rapides pour être sûr à 100% de ne rien rater.

4. Pourquoi c'est révolutionnaire ?

Vitesse : Grâce à cette astuce, le chef peut lire des livres entiers (des contextes longs) presque aussi vite qu'il lit une phrase courte.
Économie d'énergie : En ignorant les musiciens qui ne jouent rien (les zéros), l'ordinateur ne gaspille plus d'énergie à les écouter.
Résultat : Sur les tâches courtes, c'est aussi bon que les anciennes méthodes. Sur les tâches longues (lire un roman, analyser une vidéo), c'est beaucoup mieux et beaucoup plus rapide.

En résumé

ADASPLASH-2, c'est comme donner à votre IA des lunettes de vision nocturne et un radar. Au lieu de regarder toute la forêt et de se perdre dans les arbres, elle voit immédiatement où sont les animaux importants et ignore les feuilles mortes.

C'est une avancée majeure pour permettre aux intelligences artificielles de lire, comprendre et raisonner sur des documents très longs sans se fatiguer, tout en restant aussi précises que jamais.

1. Problématique

Les mécanismes d'attention dans les Transformers souffrent d'une complexité quadratique ( $O(n^2)$ ) en temps et en mémoire, ce qui constitue un goulot d'étranglement majeur pour l'entraînement sur de longues séquences. Bien que FlashAttention-2 ait résolu ce problème pour l'attention softmax (dense) en optimisant l'E/S et la fusion des opérations, l'attention softmax reste intrinsèquement dense. Cela entraîne une dispersion de la masse d'attention sur des tokens non pertinents et une perte de distinction des représentations dans les contextes longs.

Une alternative prometteuse est l'α-entmax, une transformation différentiable qui produit une attention creuse (sparse) dépendante de l'entrée (c'est-à-dire qu'elle attribue une probabilité exactement nulle à certains tokens). Cependant, le calcul du normalisateur $\tau$ requis par l'α-entmax nécessite des méthodes itératives de recherche de racines coûteuses. Les implémentations précédentes (comme ADASPLASH) nécessitaient plusieurs passages sur les données pour converger vers $\tau$ , ce qui annulait les gains de performance potentiels par rapport aux noyaux softmax optimisés, en particulier dans les régimes de sparsité modérée.

2. Méthodologie : ADASPLASH-2

L'objectif principal de ce travail est de réduire le coût computationnel du calcul de $\tau$ pour l'α-entmax, tout en exploitant efficacement la sparsité dynamique sur GPU. La méthode repose sur trois piliers techniques :

A. Initialisation par Histogramme On-Chip

Au lieu de calculer $\tau$ directement sur l'ensemble des scores d'attention, ADASPLASH-2 construit un histogramme compact des scores directement dans la mémoire SRAM on-chip (rapide) pendant le flux des données (streaming).

Fonctionnement : Les scores sont normalisés et discrétisés dans des "bins" (intervalles). Un histogramme $H$ compte le nombre de scores dans chaque bin.
Avantage : Ce histogramme fournit une borne inférieure prouvée pour le vrai normalisateur $\tau^*$ . Cela permet d'obtenir une estimation initiale de très haute qualité ( $\tau_h$ ) sans avoir à stocker ou parcourir les scores denses complets.

B. Solveur Hybride à Convergence Rapide

Grâce à l'initialisation précise fournie par l'histogramme, le solveur de racines (qui affine $\tau_h$ vers $\tau^*$ ) converge extrêmement vite.

Convergence : En pratique, une seule itération (ou deux au maximum) suffit pour atteindre la solution exacte, contrairement aux méthodes précédentes qui nécessitaient de multiples passages.
Solveur Sécurisé : Les auteurs proposent un solveur hybride adaptatif (Halley, Newton ou sécante) selon la valeur de $\alpha$ , garantissant la stabilité numérique tout en profitant de la convergence quadratique lorsque possible.

C. Implémentation GPU Économe en Sparsité

L'implémentation (en Triton) utilise une stratégie de masquage léger pour exploiter la sparsité :

Masquage par Bit-Packing : Les blocs d'attention non nuls sont encodés dans un masque binaire compact (un int32 pour 32 blocs).
Parcours Efficace : Lors de la phase de sortie (forward) et de la rétropropagation (backward), le noyau utilise des instructions GPU natives (find-next-set, popcount) pour sauter directement les blocs nuls, évitant ainsi les lectures inutiles en mémoire HBM (High Bandwidth Memory).

3. Contributions Clés

Normalisation via Histogramme On-Chip : Une méthode pour calculer des bornes rigoureuses sur le normalisateur $\tau$ en utilisant uniquement la mémoire SRAM, éliminant la nécessité de former des intermédiaires denses.
Convergence en Une Passe : La combinaison de l'initialisation par histogramme et du solveur hybride permet de converger vers la solution exacte en 1 à 2 itérations, réduisant drastiquement le coût de la phase de normalisation.
Exploitation Dynamique de la Sparsité : Une implémentation Triton optimisée qui saute les blocs nuls avec un surcoût négligeable, rendant l'approche compétitive même à des niveaux de sparsité modérés.
Résultats Empiriques Solides : Démonstration que l'approche est non seulement plus rapide que les versions précédentes d'α-entmax, mais aussi plus rapide que FlashAttention-2 dans des régimes de sparsité modérée à élevée.

4. Résultats Expérimentaux

Performance et Efficacité

Vitesse d'Entraînement : Sur des régimes de sparsité de blocs modérée à élevée (ex. > 60%, fréquent en contextes longs), ADASPLASH-2 égale ou dépasse FlashAttention-2 (implémenté en CUDA et Triton).
Gain de Vitesse : Dans les cas de forte sparsité, l'accélération peut atteindre 2x par rapport aux variantes de FlashAttention-2.
Analyse de Convergence : L'initialisation par histogramme place le solveur à des ordres de grandeur plus proches de la racine vraie dès la première itération, transformant le processus de raffinement en une correction quasi immédiate.

Performance des Modèles (Benchmarks)

Les auteurs ont entraîné des modèles de langage (LLaMA-3, 350M et 1B paramètres) avec ADASPLASH-2 ( $\alpha=1.5$ ).

Contextes Courts : Les modèles α-entmax atteignent des performances comparables, voire supérieures, aux baselines softmax sur des tâches standard (OLMES, ARC, etc.).
Contextes Longs (RULER & HELMET) :
- Les modèles α-entmax surpassent significativement les baselines softmax sur le benchmark RULER (jusqu'à 32K tokens).
- L'amélioration est particulièrement marquée sur les tâches de suivi de variables (Variable Tracking) et d'extraction de mots fréquents, où la capacité à ignorer les tokens non pertinents est cruciale.
- L'utilisation combinée de NAPE (No Positional Encoding pour certaines têtes) et d'α-entmax s'avère être une synergie puissante pour la généralisation en contexte long.

5. Signification et Impact

Ce travail résout le principal frein à l'adoption de l'attention α-entmax : son coût computationnel élevé. En rendant l'attention creuse différentiable aussi efficace, voire plus, que l'attention dense optimisée (FlashAttention) dans des scénarios réalistes, ADASPLASH-2 ouvre la voie à :

L'entraînement efficace de modèles sur des contextes extrêmement longs (au-delà de 100k tokens).
Une meilleure généralisation et une réduction de l'effondrement des représentations (representational collapse) observés avec le softmax.
Une alternative viable aux méthodes de sparsité heuristiques (comme le top-k), offrant une sparsité exacte, différentiable et adaptative sans instabilité d'entraînement.

En résumé, ADASPLASH-2 transforme l'α-entmax d'une curiosité théorique en une composante pratique et performante pour la prochaine génération de modèles de langage à long contexte.