Aprimoramento da Atenção com Memória de Decaimento Exponencial Melhora a Esparsidade KV Consciente da Consulta

Resumo

A inferência eficiente é fundamental para modelos de linguagem de contexto longo, onde o cálculo de atenção e o acesso ao cache KV dominam os custos. O trabalho recente RAT+ introduz um backbone de atenção aumentada por recorrência que possibilita atenção dilatada flexível durante a inferência. Neste artigo, investigamos se essa memória com decaimento exponencial também pode melhorar os métodos existentes de inferência esparsa conscientes da consulta. Utilizando métodos representativos como Quest, MoBA e SnapKV, mostramos que o RAT+ melhora consistentemente a acurácia em relação à atenção padrão em diversos orçamentos esparsos em oito tarefas de agulha no palheiro. Validamos esses ganhos tanto nos checkpoints disponibilizados no artigo do RAT+ quanto no OLMo2-7B, para o qual continuamos o pré-treinamento com o módulo de memória adicional por 10 bilhões de tokens. Por fim, propomos duas hipóteses que explicam por que esse módulo de memória beneficia a inferência esparsa consciente da consulta e projetamos experimentos direcionados para corroborá-las.

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.