L'Attention Complète Contre-Attaque : Transfert de l'Attention Complète en Attention Éparse en Cent Étapes d'Entraînement

Résumé

L'inférence en contexte long dans les grands modèles de langage est limitée par le coût quadratique de l'attention complète. Les alternatives efficaces existantes reposent souvent soit sur un entraînement clairsemé natif, soit sur une éviction heuristique de tokens, créant un compromis indésirable entre efficacité, coût d'entraînement et précision. Dans ce travail, nous montrons que les LLMs à attention complète sont déjà intrinsèquement clairsemés et peuvent être transformés en modèles hautement clairsemés avec une adaptation minimale. Notre approche repose sur trois observations : (1) seule une petite sous-partie des têtes d'attention nécessite réellement un traitement complet en contexte long ; (2) la récupération à longue portée est principalement régie par un sous-espace de faible dimension, permettant de récupérer efficacement les tokens pertinents à l'aide d'un indexeur à 16 dimensions ; et (3) le budget de tokens utiles est fortement dépendant de la requête, ce qui rend la sélection dynamique top-p plus adaptée que la sparsification fixe top-k. Sur la base de ces observations, nous proposons RTPurbo, qui conserve l'intégralité du cache KV uniquement pour les têtes de récupération et introduit un indexeur de tokens léger pour l'attention clairsemée. En exploitant la clairsemé intrinsèque du modèle, RTPurbo parvient à une sparsification en seulement quelques centaines d'étapes d'entraînement. Les expériences sur des benchmarks de contexte long et des tâches de raisonnement montrent que RTPurbo préserve une précision quasi sans perte tout en offrant des gains d'efficacité substantiels, notamment une accélération du préremplissage jusqu'à 9,36 fois pour un contexte de 1M et une accélération du décodage d'environ 2,01 fois. Ces résultats suggèrent qu'une inférence clairsemée performante peut être obtenue à partir d'un entraînement standard à attention complète, sans nécessiter un pré-entraînement clairsemé natif coûteux.

English

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-p selection more suitable than fixed top-k sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36times prefill speedup at 1M context and about a 2.01times decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.