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

Исследователи из Стэнфорда и Калифорнийского университета в Беркли представили TokenGate — модуль для префилла в контекстах большой длины. Обучаемое TokenGate прореживание токенов позволяет существующей LLM бегло просматривать входные данные. Полную обработку получают только значимые токены. Исследовательская команда сообщает об ускорении time-to-first-token (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 тыс. токенов.
Запуск: В статье приведены замеры на одной A100 при размере батча 1. Для обучения использовались 4 A100.
Производительность: Парето-оптимальный результат по 5 задачам в сравнении с 6 базовыми методами.
Лучший результат: NIAH 16k — ускорение TTFT в 11,3 раза и снижение потребления памяти в 3,42 раза. При 128k выигрыш по TTFT достигает 13,0 раза.
Худший результат: На некоторых задачах ускорение снижается до 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,4x памяти	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 и не забудьте присоединиться к нашему сабреддиту о машинном обучении с более чем 150 тыс. участников и подписаться на нашу рассылку. Подождите! Вы есть в Telegram? теперь вы также можете присоединиться к нам в Telegram.

Хотите сотрудничать с нами для продвижения вашего репозитория на GitHub, страницы на Hugging Face, выпуска продукта, вебинара и т. д.? Свяжитесь с нами

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