Aumentar la atención con memoria de decaimiento exponencial mejora la esparcidad KV consciente de consultas

Resumen

La inferencia eficiente es fundamental para los modelos de lenguaje de contexto largo, donde el cálculo de la atención y el acceso a la caché KV dominan el costo. El trabajo reciente RAT+ introduce un backbone de atención aumentada por recurrencia que permite una atención dilatada flexible en tiempo de inferencia. En este artículo, investigamos si esta memoria de decaimiento exponencial también puede mejorar los métodos existentes de inferencia dispersa consciente de la consulta. Utilizando métodos representativos como Quest, MoBA y SnapKV, demostramos que RAT+ mejora consistentemente la precisión en comparación con la atención estándar en diversos presupuestos de dispersión en ocho tareas de aguja en el pajar. Validamos estas mejoras tanto en los checkpoints publicados del artículo de RAT+ como en OLMo2-7B, al cual continuamos preentrenando con el módulo de memoria añadido durante 10 mil millones de tokens. Finalmente, proponemos dos hipótesis que explican por qué este módulo de memoria beneficia la inferencia dispersa consciente de la consulta y diseñamos experimentos dirigidos para respaldarlas.

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.