SAS entraîne le ranking sparse d'attention avec la LM loss

Le papier SAS attaque un défaut connu des méthodes d'attention sparse entraînables, thème récurrent dans nos flux open-source ces derniers mois. Les selectors existants distillent les poids d'attention denses du modèle original, mais ce classement n'est pas aligné avec l'impact réel des tokens sur la prédiction sous budget d'attention fixe, et gaspille donc le budget sur des unités moins utiles.

La proposition de SAS : injecter les scores continus du selector directement dans les logits d'attention pendant l'entraînement, pour que la LM loss propage ses gradients à travers la sélection via une rétro-propagation standard. Les auteurs identifient trois choix techniques nécessaires : "placing the gate inside the attention softmax in log form, using normalized softmax gates to calibrate historical context against the always-retained current block, and preserving continuous selector scores".

Zhiwei Li, Lei Zhu, Hao Gu et six co-auteurs implémentent le mécanisme via un "memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation". Sur du raisonnement, du long-contexte et des tâches agentiques, "SAS consistently outperforms trainable sparse attention baselines across attention budgets, with especially large gains under tight budgets".

L'abstract ne publie ni chiffres de latence par budget ni comparaison de qualité chiffrée face à Full Attention. Le seul graphique référencé couvre le speedup de Sparse Attn contre Full Attn sur des séquences allant jusqu'à 512k mots.