L'augmentation de l'attention par une mémoire à décroissance exponentielle améliore la sparsité KV consciente de la requête

Résumé

L'inférence efficace est cruciale pour les modèles de langage à long contexte, où le calcul de l'attention et l'accès au cache KV dominent le coût. Des travaux récents, RAT+, introduisent une architecture d'attention augmentée par récurrence qui permet une attention dilatée flexible au moment de l'inférence. Dans cet article, nous examinons si cette mémoire à décroissance exponentielle peut également améliorer les méthodes d'inférence parcimonieuse existantes qui tiennent compte des requêtes. En utilisant des méthodes représentatives telles que Quest, MoBA et SnapKV, nous montrons que RAT+ améliore constamment la précision par rapport à l'attention standard sur différents budgets d'inférence parcimonieuse pour huit tâches de type « aiguille dans une botte de foin ». Nous validons ces gains à la fois sur les points de contrôle publiés dans l'article RAT+ et sur OLMo2-7B, que nous continuons à pré-entraîner avec le module de mémoire ajouté pour 10 milliards de tokens. Enfin, nous proposons deux hypothèses expliquant pourquoi ce module de mémoire bénéficie à l'inférence parcimonieuse tenant compte des requêtes et nous concevons des expériences ciblées pour les étayer.

English

Efficient inference is critical for long-context language models, where attention computation and KV-cache access dominate the cost. Recent work RAT+, introduces a recurrence-augmented attention backbone that enables flexible dilated attention at inference time. In this paper, we investigate whether this exponentially decaying memory can also improve existing query-aware sparse inference methods. Using representative methods including Quest, MoBA, and SnapKV, we show that RAT+ consistently improves accuracy over standard attention across sparse budgets on eight needle-in-a-haystack tasks. We validate these gains both on the released checkpoints from the RAT+ paper and on OLMo2-7B, which we continue pretraining with the added memory module for 10B tokens. Finally, we propose two hypotheses explaining why this memory module benefits query-aware sparse inference and design targeted experiments to support them.