Die Anreicherung der Aufmerksamkeit mit exponentiell abklingendem Gedächtnis verbessert die abfragebewusste KV-Sparsity

Zusammenfassung

Effiziente Inferenz ist entscheidend für Sprachmodelle mit langen Kontexten, bei denen die Aufmerksamkeitsberechnung und der KV-Cache-Zugriff die Kosten dominieren. Die aktuelle Arbeit RAT+ führt ein rekurrenzverstärktes Aufmerksamkeits-Backbone ein, das zur Inferenzzeit flexible dilatierte Aufmerksamkeit ermöglicht. In diesem Beitrag untersuchen wir, ob dieser exponentiell abklingende Speicher auch bestehende abfragebewusste sparse Inferenzmethoden verbessern kann. Anhand repräsentativer Verfahren wie Quest, MoBA und SnapKV zeigen wir, dass RAT+ die Genauigkeit im Vergleich zur Standardaufmerksamkeit bei unterschiedlichen Sparsity-Budgets in acht Needle-in-a-Haystack-Aufgaben konsistent verbessert. Wir validieren diese Verbesserungen sowohl an den veröffentlichten Checkpoints aus der RAT+-Arbeit als auch an OLMo2-7B, für das wir das Vortraining mit dem zusätzlichen Speichermodul für 10 Milliarden Tokens fortsetzen. Abschließend schlagen wir zwei Hypothesen vor, die erklären, warum dieses Speichermodul abfragebewusste sparse Inferenz begünstigt, und entwickeln gezielte Experimente, um diese zu untermauern.

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.