Суть прорыва: обучаемое отсечение токенов
Команда из Стэнфорда и UC Berkeley представила TokenGate, инновационный модуль, интегрируемый в существующие большие языковые модели (LLM). Главная проблема длинных контекстов — квадратичная сложность механизма внимания (attention), что замедляет этап префилла (prefill). TokenGate решает это через обучаемое отсечение токенов (learned token pruning).
Модуль работает как «шлюз»: после слоя L модель оценивает важность каждого токена с помощью небольшого MLP-гейта. Оставляются только наиболее релевантные токены (Top-K), а остальные отбрасываются до того, как они попадут в слои внимания. Это позволяет обрабатывать контексты до 128k токенов без потери точности.
Техническая реализация: обход не дифференцируемости
Ключевая техническая сложность заключалась в том, что операция выбора Top-K не является дифференцируемой, что блокирует обратное распространение ошибки. TokenGate решает эту проблему, заимствуя подход из Mixture-of-Experts (MoE):
- Гейт-механизм: Скорость каждого токена становится смещением (bias) в логитах внимания для слоев выше гейта.
- Обучение: Используются обучаемые коэффициенты усиления на голову и температура для коррекции количества ключей. Обучение происходит на стандартной потере следующего токена, комбинируя полную контекстную потерю и потерю со случайным бюджетом токенов.
- Интеграция: Модуль добавляется как LoRA (r=16) к базовой модели (например, Qwen3-1.7B) и не требует переобучения модели с нуля.
Результаты бенчмарков: скорость и память
Эксперименты проводились на моделях Qwen3-1.7B и Qwen3.5-2B. TokenGate демонстрирует парето-оптимальные результаты по сравнению с 6 базовыми методами (FastKV, SlimInfer, DASH и др.).
| Метрика / Задача | Результат TokenGate | Сравнение с полным префиллом | Сравнение с лучшим бейзлайном |
|---|---|---|---|
| TTFT (время до первого токена) | 1.4x – 11.3x быстрее | До 11.3x ускорение | До 6.4x быстрее |
| Пиковая память (VRAM) | 1.5x – 3.4x меньше | До 3.42x экономия | До 2.1x экономия |
| NIAH 16k (Игла в стоге сена) | 11.3x быстрее TTFT | 3.42x меньше памяти | Лучший результат среди всех методов |
| Контекст 128k токенов | 13.0x быстрее TTFT | Экономия памяти ~3.43x | Информация недостаточна |
| 2WikiMultihopQA (95% точности) | 3.0x быстрее | 2.7x меньше памяти | Информация недостаточна |
Работа с гибридными архитектурами
TokenGate демонстрирует уникальное преимущество при работе с гибридными моделями, такими как Qwen3.5-2B, сочетающая слои Gated DeltaNet и полного внимания. В отличие от FastKV, который может отсекать только слои внимания, TokenGate успешно применяет отсечение к слоям DeltaNet. На задаче NIAH 32k это дало ускорение в 3.8x по сравнению с полным префиллом и в 2.1x быстрее, чем FastKV.
Деплой и доступность
Код и чекпоинты опубликованы на GitHub под лицензией Apache-2.0. Для работы требуется Python 3.10+ и CUDA PyTorch 2.9+. Опционально поддерживается ядро FlashAttention-2 для ускорения плотных вычислений Qwen3. Обучение модуля занимает в среднем 1.8 часа на 4 GPU A100, что сопоставимо со стандартным обучением LoRA.
Источник: MarkTechPost ↗