Дослідники зі Stanford і UC Berkeley представили TokenGate: увага з керуванням на рівні токенів пришвидшує префіл довгого контексту до 11,3 раза

Дослідники зі Stanford та UC Berkeley представили TokenGate — модуль для попереднього заповнення в контексті великої довжини. TokenGate використовує навчене скорочення кількості токенів, щоб наявна LLM могла швидко переглядати вхідні дані. Повністю обробляються лише важливі токени. Дослідницька команда повідомляє про прискорення часу до першого токена (TTFT) у 1,4–11,3 раза без втрати точності. Пікове споживання пам’яті зменшується у 1,5–3,4 раза. Код і чекпоїнти доступні на GitHub за ліцензією Apache-2.0.

Коротко

MLP-гейт: Невеликий MLP-гейт і LoRA (r=16), додані до Qwen3-1.7B або Qwen3.5-2B. Тестування проводилося на контекстах обсягом до 128 тис. токенів.
Запуск: У статті наведено результати для 1 A100 із розміром батчу 1. Для навчання використовувалися 4 A100.
Продуктивність: Парето-оптимальний результат на 5 завданнях порівняно з 6 базовими методами.
Найкращий результат: NIAH 16k — у 11,3 раза швидший TTFT і у 3,42 раза менше пам’яті. Виграш за TTFT досягає 13,0 раза при 128k.
Найгірший результат: На деяких завданнях прискорення знижується до 1,4 раза.

Що таке TokenGate?

TokenGate — це компонент, який можна безпосередньо інтегрувати та який додає навчене скорочення кількості токенів до попередньо навчених LLM. Трансформери пропускають кожен токен через кожен шар. Вартість механізму уваги зростає квадратично зі збільшенням довжини. Скорочення без навчання є дешевим, але за високих коефіцієнтів скорочення призводить до втрати точності. Архітектурні зміни часто потребують попереднього навчання з нуля.

Як працює навчене скорочення кількості токенів у TokenGate?

Прямий прохід складається з 3 етапів:

Стандартний прямий прохід: шари нижче шару гейта L обробляють усі N токенів у звичайному режимі.
Модуль гейта: MLP оцінює кожен токен за його залишковим представленням після механізму уваги на шарі L. Селектор top-K залишає токени з найвищими оцінками.
Прямий прохід після скорочення: шари вище L обробляють лише ці K токенів і не створюють KV-записів для відкинутих токенів.

Скорочувати можна лише частину промпту, що містить документ. Запит ніколи не скорочується.

Що таке увага з керуванням токенами?

Вибір top-K є недиференційовним, що блокує наскрізне навчання. TokenGate запозичує трюк маршрутизатора суміші експертів. Оцінка кожного залишеного токена стає зміщенням логітів уваги вище гейта. Це зміщення забезпечує градієнтам шлях назад до оцінювача.

Це доповнюється навчуваним коефіцієнтом для кожної голови та навчуваною температурою. Температура компенсує те, що кожен запит бачить менше ключів. Для навчання використовується звичайна функція втрат наступного токена. На кожному кроці усереднюються втрати на повному контексті та втрати для випадкового бюджету. Таким чином, один навчений гейт працює з різними бюджетами.

Як TokenGate показує себе на бенчмарках?

Дослідницька команда протестувала Qwen3-1.7B на RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX і needle-in-a-haystack (NIAH). Базовими методами були FastKV, SlimInfer, DASH, GemFilter, FlexPrefill і MInference на тій самій LoRA-основі.

Порівняно з повним попереднім заповненням: TTFT у 1,4–11,3 раза швидший, а використання пам’яті у 1,5–3,4 раза менше.
Порівняно з наступним найкращим базовим методом: TTFT у 1,2–6,4 раза швидший, а використання пам’яті — до 2,1 раза менше.
2WikiMultihopQA: за точності 95% від повноконтекстної попереднє заповнення працює у 3,0 раза швидше та використовує у 2,7 раза менше пам’яті.
Довші контексти: гейт, навчений на 16k, досягає 12,8 раза на 32k і 13,0 раза на 64k та 128k. Економія пам’яті залишається близькою до 3,43 раза.

Навчений селектор перевершує евристичний метод FastKV на тому самому шарі. Навіть вибір випадкових K показує вищі результати на основі TokenGate. Дослідницька команда пояснює це коадаптацією основної моделі.

Чи працює це з моделями з гібридною увагою?

Так, Qwen3.5-2B поєднує шари Gated DeltaNet і повної уваги. На NIAH 32k TokenGate на шарі 3 досяг точності повного попереднього заповнення. Він працював у 3,8 раза швидше за повне попереднє заповнення та у 2,1 раза швидше за FastKV. Керування на шарах 1 або 2 дало прискорення у 7,8 раза за збереження точності в межах 6%. FastKV може скорочувати лише шари уваги.

TokenGate проти FastKV, SlimInfer, GemFilter і DASH

Характеристика	TokenGate	FastKV	SlimInfer	GemFilter	DASH
Вибір токенів	Навчений MLP-гейт	Увага від останніх запитів	Блокова увага, багатостадійна	Увага останнього запиту на шарі фільтра	Стабільність прихованого стану
Потрібне навчання	Так (гейт + LoRA r=16)	Ні	Ні	Ні	Ні
Шар скорочення у запусках TokenGate на Qwen3-1.7B	1	13 (за замовчуванням)	8, 17, 25	12	11
Скорочує на шарах Gated DeltaNet	Так	Ні (лише шари уваги)	Не розкрито	Не розкрито	Не розкрито
Найкраще заявлене прискорення попереднього заповнення (власна конфігурація)	11,3x TTFT, у 3,4 раза менше пам’яті	1,82x попереднього заповнення	2,53x TTFT (LLaMA3.1-8B, RTX 4090)	2,4x швидкість, на 30% менше пам’яті GPU порівняно з SOTA	Не розкрито
Код	GitHub	GitHub	GitHub	GitHub	Анонсовано

Заявлені результати отримано в різних конфігураціях, тому їх не можна безпосередньо порівнювати. Шари скорочення взято з додатка C.2 TokenGate.

Як розгорнути TokenGate?

Для репозиторію GitHub потрібні Python 3.10+ і CUDA PyTorch 2.9+. Додаткове ядро FlashAttention-2 для зміщення ключів пришвидшує роботу зі щільною моделлю Qwen3. Воно працює у 1,0–1,1 раза повільніше за звичайний виклик. Без нього під час виконання використовується FlexAttention.

Заповнювач {state} позначає документ, який можна скорочувати. Функція generate приймає бюджет max_kept_prefill_tokens. Чекпоїнти охоплюють усі 5 завдань на Qwen3-1.7B, а також гейт Qwen3.5-2B для NIAH. Приклад: liana313/tokengate-niah-qwen3-1.7b.

Навчання повторює рецепт TRL/PEFT. Кожен крок у 1,3 раза повільніший за звичайну LoRA. На 4 із 5 завдань збіжність настала швидше. Середній час навчання становив 1,8 проти 2,2 годин A100.

Основні висновки

TokenGate додає навчене скорочення кількості токенів до наявних LLM за допомогою невеликого гейта та LoRA.
Попереднє заповнення працює у 1,4–11,3 раза швидше, а пікове споживання пам’яті у 1,5–3,4 раза менше.
Увага з керуванням токенами робить вибір top-K придатним до навчання за принципом маршрутизатора MoE.
За TTFT він перевершує наступний найкращий базовий метод скорочення до 6,4 раза.
Він працює на гібридній моделі Qwen3.5-2B, де FastKV не може скорочувати шари DeltaNet.

Ознайомтеся зі статтею, репозиторієм GitHub і сторінкою статті на Hugging Face. Також підписуйтеся на нас у Twitter і не забудьте приєднатися до нашого субреддіту ML із понад 150 тисячами учасників та підписатися на нашу розсилку. Стривайте! Ви є в Telegram? Тепер ви також можете приєднатися до нас у Telegram.

Потрібно співпрацювати з нами для просування вашого репозиторію GitHub, сторінки на Hugging Face, релізу продукту, вебінару тощо? Зв’яжіться з нами

Перекладено автоматично з англійської. Оригінал статті — за посиланням нижче.