Дополнение внимания экспоненциально затухающей памятью улучшает разреженность KV, учитывающую запросы

Аннотация

Эффективный инференс имеет решающее значение для языковых моделей с длинным контекстом, где доминирующую роль в вычислительных затратах играют вычисление внимания и доступ к KV-кэшу. Недавняя работа RAT+ предлагает рекуррентно-усиленный каркас внимания, который обеспечивает гибкое дилатированное внимание на этапе инференса. В данной статье мы исследуем, может ли эта экспоненциально затухающая память также улучшить существующие методы разреженного инференса с учетом запроса. Используя репрезентативные методы, включая Quest, MoBA и SnapKV, мы показываем, что RAT+ последовательно повышает точность по сравнению со стандартным вниманием в условиях ограниченного бюджета разреженности на восьми задачах поиска иголки в стоге сена. Мы подтверждаем эти улучшения как на опубликованных контрольных точках из статьи RAT+, так и на модели OLMo2-7B, для которой мы продолжили предварительное обучение с добавленным модулем памяти на 10B токенов. Наконец, мы предлагаем две гипотезы, объясняющие, почему данный модуль памяти способствует разреженному инференсу с учетом запроса, и разрабатываем целенаправленные эксперименты для их подтверждения.

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.