QUOKA простыми словами: как выбрать, какую историю KV читать вниманию

Понаблюдайте за работой агента для кодинга, и вы заметите, что он редко понимает задачу по одному файлу. Он открывает несколько файлов по всему репозиторию, разбирается, как связаны части, прослеживает зависимости и только потом что-то правит. Иногда это означает отправку модели поразительно большого объёма контекста.

Отсюда вопрос. Вендоры рекламируют контекстные окна в сотни тысяч и даже миллион токенов. Что на самом деле значит обработать столько текста и откуда берётся стоимость? Одним трюком на это не ответить. Инференс с длинным контекстом затрагивает кодирование позиций, вычисление внимания, управление памятью и планирование нагрузки.

Одна из статей, которая берётся за узкий кусок этой проблемы, — QUOKA, сокращение от Query-Oriented KV Selection for Efficient LLM Prefill. Она о том, как удешевить внимание, пока модель читает длинный входной промпт. Основная идея интуитивна, но чтобы понять, почему она работает, нужно три кусочка контекста: KV-кэш, разделение на prefill и decode и chunked prefill.

KV-кэш вкратце

Внутри трансформера каждый токен порождает для внимания запрос (Q), ключ (K) и значение (V). Грубо их можно понимать так:

• Запрос спрашивает, какую информацию ищет токен.
• Ключ сообщает, какого рода информацию держит токен.
• Значение — то, что передаётся, если токен окажется релевантным.

При авторегрессионной генерации ранее вычисленные ключи и значения можно хранить в KV-кэше и переиспользовать позже, а не считать заново. В общих чертах:

Previously processed tokens
          ↓
     K,V representations
          ↓
        KV cache
          ↓
reused by later attention

Чем длиннее последовательность, тем больше хранится записей K и V. Для QUOKA это важно, а чтобы понять почему, сначала нужно разобраться с двумя фазами инференса.

Prefill: чтение вашего промпта

Когда вы отправляете модели промпт, будь то вопрос в одну строку или гора кода, входные токены нужно обработать до того, как начнётся вывод. Эта фаза называется prefill.

Input prompt
     ↓
   PREFILL
     ↓
process input tokens
     ↓
produce their K,V representations

Decode: написание ответа

Когда вход обработан, модель выдаёт ответ по одному токену за раз. Это decode.

Processed input
      ↓
generate token 1
      ↓
generate token 2
      ↓
generate token 3
      ↓
     ...

Разница важна, потому что выбрать, какие части длинной истории заслуживают внимания, довольно легко, когда новый запрос один. Совсем другая задача, когда в целом фрагменте сразу сотни или тысячи запросов. Прежде чем к этому перейти, нужно ещё одно понятие.

Chunked prefill

Обрабатывать очень длинный промпт за один проход prefill дорого. Поэтому системы инференса часто используют chunked prefill: вход делится на меньшие фрагменты, которые обрабатываются по порядку. Для промпта на 3 000 токенов:

3,000-token prompt
Chunk 0 → tokens    0–999
Chunk 1 → tokens 1000–1999
Chunk 2 → tokens 2000–2999

За каждым фрагментом после первого стоит всё больше истории:

Chunk 0 → no previous chunk
Chunk 1 → 1,000 historical KV positions
Chunk 2 → 2,000 historical KV positions
Chunk 3 → 3,000 historical KV positions
...

Деление на фрагменты делает длинный prefill терпимым, но каждый следующий фрагмент может обращаться ко всё большей куче закэшированной истории. Именно на эту кучу и нацелена QUOKA.

Наблюдение, на котором стоит метод

Представьте, что обрабатывается поздний фрагмент, а в кэше уже лежит 100 000 исторических позиций KV. Плотное внимание позволяет текущим запросам смотреть на каждую из них.

Current chunk
     ↓
K1 K2 K3 K4 K5 .............. K100000

Действительно ли текущим запросам нужны все 100 000? Скорее всего нет. Некоторые ключи будут очень важны для того, что делает этот фрагмент, а многие почти ничего не дадут. Значит, есть лазейка:

100,000 historical KV
          ↓
find the important ones
          ↓
much smaller KV subset
          ↓
attention

Вот базовая интуиция. Вместо дорогого внимания по всей истории выбрать меньший набор исторических KV, которые, скорее всего, важны для текущего фрагмента. Достаточно просто, но тут встаёт вопрос посложнее: кто решает, какие KV важны?

Один запрос может выбрать сам, а во фрагменте их тысячи

При генерации токен за токеном новая позиция запроса одна, и этот запрос по сути может сам спросить, какие исторические ключи ему релевантны.

one query
    ↓
Which historical keys
are relevant to me?

Во время chunked prefill позиций запросов сразу много:

q1
q2
q3
...
q4096

Разным запросам могут быть важны совершенно разные части истории. Можно свести их к одному усреднённому запросу, но так легко потерять необычные, но важные потребности. Допустим, большинство запросов смотрят в одну сторону:

q1  ↗
q2  ↗
q3  ↗
q4  ↗
q5  →

В среднем будут доминировать q1–q4. Если q5 — единственный запрос, которому нужен определённый ключ, усреднение похоронит это взаимодействие.

Выбираем запросы, которые не согласны с толпой

Сначала QUOKA вычисляет средний запрос текущего фрагмента, затем измеряет косинусное сходство каждого запроса с этим средним. Например:

Similarity to mean
q1 → 0.96
q2 → 0.94
q3 → 0.91
q4 → 0.35
q5 → 0.18

Вместо того чтобы оставлять запросы, ближайшие к среднему, метод отдаёт приоритет наименее похожим на него:

q1 → 0.96
q2 → 0.94
q3 → 0.91
q4 → 0.35  ← keep
q5 → 0.18  ← keep

Логика простая. Очень похожие запросы смотрят примерно в одну сторону, поэтому, оставив много таких, вы получите избыточную информацию. Запросы с низким сходством как раз говорят, что ищут что-то другое, и именно они, скорее всего, найдут ключи, которые большинство упустило бы. Так тысячи запросов превращаются в небольшой разнообразный набор:

thousands of queries
        ↓
a small set of diverse queries

Пусть эти запросы оценят историю

Теперь разнообразные запросы просматривают сохранённые ключи, и каждый спрашивает, какие части прошлого выглядят для него полезными. QUOKA снова использует косинусное сходство, на этот раз между выбранными запросами и историческими ключами. Допустим, выбраны три запроса:

qA
qB
qC

а ключей тысячи:

K1, K2, K3, ... K10000

Каждый выбранный запрос оценивает каждый ключ по косинусному сходству:

            K1      K2     K3     K4
qA          .90    .10    .20    .15
qB          .10    .85    .15    .20
qC          .20    .10    .95    .10

Это полезно, потому что разные запросы находят разные части прошлого. Здесь qA сильно интересует K1, qB — K2, а qC — K3.

Но теперь у каждого ключа несколько оценок важности, по одной на выбранный запрос, а QUOKA нужно одно число, чтобы решить, останется ли ключ.

Почему берётся максимум

Посмотрим на K1:

qA → K1 = .90
qB → K1 = .10
qC → K1 = .20

Очевидный вариант — усреднить:

(.90 + .10 + .20) / 3 = .40

Тогда K1 выглядит умеренно важным, и прячется ценная деталь: qA считает его крайне важным. Ключу не обязательно быть важным для каждого запроса, чтобы заслужить место. Он может содержать ровно то, что нужно одному запросу. Поэтому QUOKA берёт максимум по выбранным запросам:

K1 → max(.90, .10, .20) = .90
K2 → max(.10, .85, .10) = .85
K3 → max(.20, .15, .95) = .95
K4 → max(.15, .20, .10) = .20

Каждый ключ получает одну итоговую оценку, и сильные взаимодействия запрос–ключ сохраняются. Если хотя бы одному представительному запросу ключ очень нужен, остальным нельзя позволять размыть этот сигнал.

От оценок к меньшему набору KV

Получив оценки, QUOKA ранжирует ключи и оставляет только лучшие в пределах своего бюджета разреженного внимания.

Historical keys
K1 → .90  ✓
K2 → .85  ✓
K3 → .95  ✓
K4 → .20
K5 → .12
...
K10000 → .08

        ↓
      Top-K
        ↓

small set of selected keys
K1 → .90  ✓
K2 → .85  ✓
K3 → .95  ✓

Когда выбирается позиция ключа, вместе с ней берётся и соответствующее значение. Если QUOKA оставляет такие ключи:

K3, K1, K2

то оставляет и такие значения:

V3, V1, V2

В итоге во внимании участвует меньший набор пар K,V. По сути, вместо того чтобы заставлять внимание рассматривать каждую историческую позицию, QUOKA сначала строит гораздо меньший набор кандидатов, исходя из того, что действительно, похоже, нужно текущим запросам.

Нюанс: grouped-query attention

У современных моделей не обязательно столько же голов запросов, сколько голов KV, и тут появляется grouped-query attention (GQA). Допустим, у модели четыре головы запросов и две головы KV:

Q head 1 ─┐
          ├── KV head 1
Q head 2 ─┘
Q head 3 ─┐
          ├── KV head 2
Q head 4 ─┘

Вместо того чтобы у каждой головы запросов была своя голова ключей и значений, несколько голов запросов делят одну. Это касается QUOKA, потому что в модели с GQA у каждого выбранного токена несколько представлений по головам запросов. Для одного выбранного токена qA может быть так:

qA_head1
qA_head2
qA_head3
qA_head4

Первые два взаимодействуют с головой KV 1, а другие два — с головой KV 2. Оценивать каждую голову отдельно и потом усреднять выглядело бы так:

qA_head1 × (K head 1)ᵀ ─┐
                         ├── average scores
qA_head2 × (K head 1)ᵀ ─┘

qA_head3 × (K head 2)ᵀ ─┐
                         ├── average scores
qA_head4 × (K head 2)ᵀ ─┘

Это повторяет дорогие сравнения запрос–ключ с одной и той же головой KV. QUOKA пользуется простым алгебраическим фактом и сначала усредняет нормализованные представления запросов:

qA_head1 ─┐
          ├── MEAN → q̄A_1 → × (K head 1)ᵀ
qA_head2 ─┘

qA_head3 ─┐
          ├── MEAN → q̄A_2 → × (K head 2)ᵀ
qA_head4 ─┘

Это работает, потому что (Q1 K^T + Q2 K^T) / 2 равно ((Q1 + Q2) / 2) K^T. Вместо двух умножений для оценки и последующего усреднения вы сначала усредняете запросы и умножаете один раз. В статье о QUOKA этот приём называется предагрегацией (pre-aggregation).

MEAN и MAX делают разную работу

Этот MEAN легко спутать с MAX из предыдущего раздела, но они работают по разным измерениям.

MEAN действует по головам запросов, которые делят одну голову KV, и даёт более дешёвый заменитель запроса для оценки этой головы:

qA_head1 ─┐
          ├── MEAN
qA_head2 ─┘

MAX действует по разным выбранным токенам запросов и защищает ключ, если хотя бы один нетипичный запрос считает его очень релевантным:

qA → K50 = .20
qB → K50 = .91
qC → K50 = .30
                 ↓
                MAX
                 ↓
          K50 score = .91

Чтобы не путать:

MEAN across related GQA heads
          ↓
make KV selection cheaper

MAX across diverse query tokens
          ↓
preserve important query–key interactions

Важны ещё две детали. Усреднённые головы запросов служат исключительно для выбора позиций KV. Исходные головы запросов модели остаются на месте и выполняют настоящее разреженное внимание над тем, что выбрано, так что предагрегация — это сокращённый путь только для отбора.

Заметьте также, что ничего не вытесняется. Полный исторический KV-кэш остаётся, выбранные KV — лишь подмножество для текущего шага внимания, и новый фрагмент может выбрать другое подмножество по своим запросам. Экономия — в том, сколько истории участвует в вычислении внимания, а не в размере хранимого кэша.

Где искать цифры

Результаты бенчмарков приведены в трёх таблицах статьи, и именно там стоит смотреть, как метод выглядит на фоне альтернатив. Читайте её на arXiv.