Запоминание контекста для эффективной генерации длинного контекста

Аннотация

Современные приложения больших языковых моделей (LLM) все чаще полагаются на длинные обусловливающие префиксы для управления поведением модели на этапе инференса. Хотя инференс с дополнением префикса эффективен, он сопряжен с двумя структурными ограничениями: i) влияние префикса ослабевает по мере генерации, и ii) вычисление внимания к префиксу масштабируется линейно с его длиной. Существующие подходы либо оставляют префикс в механизме внимания, сжимая его, либо интернализируют его в параметры модели через градиентное обучение. Первый вариант по-прежнему требует обработки префикса при инференсе, а второй требует интенсивного обучения и плохо подходит для обновления префикса. Для решения этих проблем мы предлагаем память состояний внимания — бестренировочный подход, который экстернализирует префикс в легковесную память на основе поиска, содержащую предвычисленные состояния внимания между токенами префикса и запроса. На ManyICLBench с LLaMA-3.1-8B наш метод улучшает точность по сравнению с обучением в контексте при бюджете памяти от 1K до 8K, одновременно снижая задержку внимания в 1,36 раза при 8K, а на эталонном тесте NBA превосходит производительность RAG с полным вниманием, используя лишь 20% его занимаемой памяти.

English

Modern large language model (LLM) applications increasingly rely on long conditioning prefixes to control model behavior at inference time. While prefix-augmented inference is effective, it incurs two structural limitations: i) the prefix's influence fades as generation proceeds, and ii) attention computation over the prefix scales linearly with its length. Existing approaches either keep the prefix in attention while compressing it, or internalize it into model parameters through gradient-based training. The former still attends to the prefix at inference, while the latter is training-intensive and ill-suited to prefix updates. To address these issues, we propose attention-state memory, a training-free approach that externalizes the prefix into a lightweight, lookup-based memory of precomputed attention states between prefix and query tokens. On ManyICLBench with LLaMA-3.1-8B, our method improves accuracy over in-context learning at 1K-8K memory budgets while reducing attention latency by 1.36x at 8K, and surpasses full-attention RAG performance on NBA benchmark using only 20% of its memory footprint.