TokenGate ускоряет prefill в LLM до 11,3 раза за счёт обучаемого отсева токенов

Команда исследователей опубликовала TokenGate — небольшой подключаемый компонент, который позволяет уже обученным языковым моделям пропускать малозначимые токены на длинных входах. По заявлению авторов, время до первого токена (TTFT) сокращается в 1,4–11,3 раза, пиковая память на prefill падает в 1,5–3,4 раза, а точность на разнообразных задачах не деградирует. Против лучших существующих методов прунинга токенов выигрыш доходит до 6,4 раза по скорости prefill и до 2 раз по памяти.

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

Как обойти недифференцируемый top-K

Отсев токенов сам по себе не нов. Сложность в обучении: выбор top-K токенов — операция недифференцируемая, градиент через неё не проходит, и оценщик важности нельзя обучить сквозным способом вместе с моделью. Многие прежние подходы поэтому опираются на эвристики или отдельные этапы обучения.

TokenGate решает это механизмом token-gated attention. Он вдохновлён Mixture-of-Experts: оценки важности токенов перевзвешивают логиты внимания. Благодаря этому оценщик получает градиентный сигнал в ходе обычного обучения, и вся цепочка, включая решение о том, что оставить, обучается end-to-end.

Важная практическая деталь: это не новая архитектура с нуля. TokenGate расширяет существующие модели. Авторы показывают работу и со стандартными трансформерами (например, Qwen3), и с гибридными моделями внимания (например, Qwen3.5). Адаптация, по их описанию, лёгкая и сопоставима со стандартным parameter-efficient fine-tuning.

Что показали замеры

Цифры из статьи сведены ниже.

• TTFT: быстрее в 1,4–11,3 раза
• Пиковая память на prefill: ниже в 1,5–3,4 раза
• Точность: без деградации на протестированных задачах
• Относительно лучших прежних методов прунинга: до 6,4 раза быстрее prefill и до 2 раз меньше пиковой памяти
• Совместимость: стандартные трансформеры (Qwen3) и гибридное внимание (Qwen3.5)
• Код: открыт, репозиторий tokengate-models/TokenGate на GitHub

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

Где это пригодится и что проверить

Наибольшую выгоду метод обещает там, где prefill доминирует в задержке: кодовые агенты, которые заливают в контекст целые репозитории, RAG-системы с длинными выдержками, анализ документов и логов. Снижение пикового потребления памяти также означает, что более длинный контекст может поместиться на том же GPU, либо больше запросов обработается параллельно.

Перед внедрением стоит проверить несколько вещей:

1. Прогнать TokenGate на своих данных. «Без деградации точности» в статье относится к их набору задач, а ваши могут быть строже. 2. Отдельно протестировать сценарии, где важен точный поиск одной детали в длинном тексте: именно там агрессивный отсев токенов рискует выкинуть нужное. 3. Оценить стоимость дообучения: заявлено сопоставимое с PEFT, но стоит посчитать на вашей модели. 4. Проверить, как компонент встраивается в ваш serving-стек и поддерживается ли он вашим инференс-движком.

Ограничение очевидно: перед нами препринт, а цифры принадлежат авторам, независимой воспроизводимости пока нет. Но код открыт, а метод применим к популярным семействам Qwen, так что проверить заявления можно быстро.