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.