TokenGate: как заставить LLM читать только важные токены — ускорение до 11,3x без потери точности

Исследователи из Стэнфорда и Калифорнийского университета в Беркли представили TokenGate — модуль для предварительной обработки длинных контекстов в LLM. Он добавляет обучаемое прореживание токенов: модель «просматривает» вход, а полную обработку получают только значимые токены. На тестах с контекстом до 128k токенов время до первого токена (TTFT) ускорилось от 1,4x до 11,3x, а пиковое потребление памяти снизилось в 1,5–3,4 раза — без потери точности.

Как это работает? Прямой проход разделён на три фазы: сначала все N токенов проходят через слои ниже «шлюзового» слоя L. Затем небольшой MLP-классификатор оценивает каждый токен по его residual-представлению после внимания, и top-K селектор оставляет только самые важные. Слои выше L обрабатывают только эти K токенов, не создавая KV-записи для отброшенных. При этом отбрасывать можно только часть промпта с документом — вопрос никогда не трогается.

Проблема недифференцируемости top-K решена трюком из Mixture-of-Experts: оценка каждого оставленного токена становится смещением (bias) в логитах внимания выше шлюза, что даёт градиентам путь обратно к классификатору. Дополнительно используются обучаемые per-head gain и температура, корректирующая разное число ключей на запрос. Обучение идёт на обычном loss следующего токена, усредняя full-context и random-budget loss — один обученный шлюз работает при любом бюджете.

Код и чекпоинты TokenGate опубликованы на GitHub под лицензией Apache-2.0. Модуль протестирован на Qwen3-1.7B и Qwen3.5-2B с контекстом до 128k токенов. Наилучший результат — задача NIAH с 16k контекстом: ускорение TTFT в 11,3x и снижение памяти в 3,42x. При 128k токенах ускорение достигает 13,0x. Худший случай — 1,4x на некоторых задачах. В целом TokenGate показывает Парето-оптимальные результаты на 5 задачах против 6 бейзлайнов.