Изследователи от Stanford и UC Berkeley представят TokenGate: внимание с управление на токените ускорява префила за дълъг контекст до 11,3 пъти
Изследователи от Stanford и UC Berkeley публикуваха TokenGate — модул за предварителна обработка на дълъг контекст. Наученото от TokenGate прореждане на токени позволява на съществуващ LLM да преглежда набързо входа си. Пълна обработка получават само важните токени. Изследователският екип отчита 1,4x до 11,3x по-бързо време до първия токен (TTFT) без загуба на точност. Пиковата памет при предварителната обработка намалява 1,5x до 3,4x. Кодът и чекпойнтите са налични в GitHub под лиценза Apache-2.0.
Накратко
- MLP гейт: Малък MLP гейт и LoRA (r=16), добавени към Qwen3-1.7B или Qwen3.5-2B. Тествано с контексти до 128k токена.
- Работи на: Данните от статията са получени с 1 A100 при размер на партидата 1. Обучението е извършено с 4 A100.
- Производителност: Оптимално по Парето при 5 задачи спрямо 6 базови метода.
- Най-добър резултат: NIAH 16k с 11,3x по-бързо TTFT и 3,42x по-малко памет. Печалбата при TTFT достига 13,0x при 128k.
- Най-слаб резултат: Ускорението спада до 1,4x при някои задачи.
Какво представлява TokenGate?
TokenGate е компонент, който може да се добави директно към предварително обучени LLM и осигурява обучаемо прореждане на токени. Трансформърите прекарват всеки токен през всеки слой. Цената на attention нараства квадратично с дължината. Прореждането без обучение е евтино, но губи точност при високи нива на прореждане. Архитектурните промени често изискват предварително обучение от нулата.
Как работи наученото прореждане на токени в TokenGate?
Правият проход се разделя на 3 фази:
- Стандартен прав проход: Слоевете под слоя на гейта L обработват всички N токена по обичайния начин.
- Модул на гейта: MLP оценява всеки токен въз основа на неговия остатъчен сигнал след attention на слой L. Селекторът top-K запазва токените с най-високи оценки.
- Прав проход с прореждане: Слоевете над L обработват само тези K токена и никога не създават KV записи за премахнатите токени.
Може да се прорежда само частта от промпта, съдържаща документа. Въпросът никога не се прорежда.
Какво представлява attention с гейт на токени?
Изборът top-K не е диференцируем, което възпрепятства обучението от край до край. TokenGate заема трика с рутера от смесицата от експерти. Оценката на всеки запазен токен се превръща в отместване на логитите на attention над гейта. Това отместване осигурява път за градиентите обратно към оценяващия модул.
Обучаем коефициент за всяка глава и обучаема температура усъвършенстват този процес. Температурата коригира факта, че всеки query вижда по-малко ключове. Обучението използва стандартна загуба за следващия токен. На всяка стъпка се усредняват загуба при пълен контекст и загуба при случаен бюджет. Така един обучен гейт работи с различни бюджети.
Как се представя TokenGate при бенчмарковете?
Изследователският екип тества Qwen3-1.7B върху RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX и needle-in-a-haystack (NIAH). Базовите методи са FastKV, SlimInfer, DASH, GemFilter, FlexPrefill и MInference, върху същата LoRA основа.
- Спрямо пълната предварителна обработка: 1,4x до 11,3x по-бързо TTFT и 1,5x до 3,4x по-малко памет.
- Спрямо следващия най-добър базов метод: 1,2x до 6,4x по-бързо TTFT и до 2,1x по-малко памет.
- 2WikiMultihopQA: При 95% от точността с пълен контекст предварителната обработка е 3,0x по-бърза и използва 2,7x по-малко памет.
- По-дълги контексти: Гейт, обучен при 16k, достига 12,8x при 32k и 13,0x при 64k и 128k. Спестяването на памет остава близо до 3,43x.
Наученият селектор превъзхожда евристиката на FastKV при същия слой. Дори изборът на случаен K постига по-висок резултат върху основата на TokenGate. Изследователският екип отдава това на съвместната адаптация с основния модел.
Работи ли с модели с хибриден attention?
Да, Qwen3.5-2B комбинира слоеве Gated DeltaNet и слоеве с пълен attention. При NIAH 32k TokenGate на слой 3 достига точността на пълната предварителна обработка. Той работи 3,8x по-бързо от пълната предварителна обработка и 2,1x по-бързо от FastKV. Гейтирането на слоеве 1 или 2 достига 7,8x при точност в рамките на 6%. FastKV може да прорежда само attention слоеве.
TokenGate спрямо FastKV, SlimInfer, GemFilter и DASH
| Характеристика | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| Избор на токени | Научен MLP гейт | Attention от последните query заявки | Блоков attention, на няколко етапа | Attention на последния query на филтриращия слой | Стабилност на скритото състояние |
| Необходимо ли е обучение | Да (гейт + LoRA r=16) | Не | Не | Не | Не |
| Слой на прореждане в изпълненията на TokenGate с Qwen3-1.7B | 1 | 13 (по подразбиране) | 8, 17, 25 | 12 | 11 |
| Прорежда ли на слоеве Gated DeltaNet | Да | Не (само attention слоеве) | Не е разкрито | Не е разкрито | Не е разкрито |
| Най-добра заявена от авторите печалба при предварителната обработка (при собствена конфигурация) | 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+. Незадължително ядро за bias на ключовете във FlashAttention-2 ускорява плътния Qwen3. То струва 1,0x до 1,1x спрямо стандартното извикване. Без него инференсът преминава към FlexAttention.
Заместителят {state} обозначава документа, който може да бъде прорязан. generate приема бюджет max_kept_prefill_tokens. Чекпойнтите обхващат и 5-те задачи върху Qwen3-1.7B, както и гейт за NIAH на Qwen3.5-2B. Пример: liana313/tokengate-niah-qwen3-1.7b.
Обучението следва рецепта за TRL/PEFT. Всяка стъпка е 1,3x по-бавна от стандартното LoRA. Изпълненията са се сходими по-бързо при 4 от 5 задачи. Средното време за обучение е 1,8 спрямо 2,2 часа на A100.
Основни изводи
- TokenGate добавя научено прореждане на токени към съществуващи LLM чрез малък гейт и LoRA.
- Предварителната обработка е 1,4x до 11,3x по-бърза, а пиковата памет е 1,5x до 3,4x по-малка.
- Attention с гейт на токени прави избора top-K обучаем — по подобие на рутер в MoE.
- Той превъзхожда следващия най-добър базов метод за прореждане с до 6,4x по отношение на TTFT.
- Работи с хибридния Qwen3.5-2B, при който FastKV не може да прорежда слоевете DeltaNet.
Разгледайте статията, GitHub хранилището и страницата на статията в Hugging Face. Също така можете да ни последвате в Twitter и не забравяйте да се присъедините към нашия ML SubReddit с над 150 хил. членове и да се абонирате за нашия бюлетин. Чакайте! В Telegram ли сте? Вече можете да се присъедините към нас и в Telegram.
Имате нужда от партньор за популяризиране на вашето GitHub хранилище ИЛИ страница в Hugging Face, ИЛИ продуктова премиера, ИЛИ уебинар и т.н.? Свържете се с нас
Преведено автоматично от английски. Оригиналната статия е на връзката по-долу.