TokenGate: Stanford i Berkeley uczą model, które tokeny pominąć przy czytaniu długiego kontekstu

Prefill, czyli etap, w którym model wczytuje cały prompt przed wygenerowaniem pierwszego słowa, przy długich kontekstach staje się wąskim gardłem. Koszt uwagi rośnie kwadratowo z długością wejścia. Zespół ze Stanfordu i UC Berkeley proponuje na to TokenGate: niewielki, uczony moduł, który dokleja się do gotowego modelu i decyduje, które tokeny dokumentu w ogóle warto przepuszczać przez wyższe warstwy. O pracy pisze też MarkTechPost.

Jak to działa

Przebieg dzieli się na trzy fazy. Warstwy poniżej wybranej warstwy bramki (gate) przetwarzają wszystkie N tokenów jak zwykle. Potem mały MLP ocenia każdy token na podstawie jego reprezentacji po warstwie uwagi, a selektor top-K zostawia tylko najwyżej oceniane. Warstwy powyżej bramki widzą już wyłącznie K tokenów i nie tworzą wpisów KV-cache dla pozostałych. Przycinana jest tylko część promptu będąca dokumentem; pytanie zostaje nietknięte.

Główny problem: wybór top-K nie jest różniczkowalny, więc nie da się nauczyć bramki end-to-end. Autorzy zapożyczyli sztuczkę z routerów w architekturach Mixture-of-Experts. Wynik oceny zachowanego tokenu trafia jako bias do logitów uwagi w wyższych warstwach, dzięki czemu gradient wraca do scorera. Uzupełniają to uczony współczynnik na głowicę i uczona temperatura, która koryguje fakt, że każde zapytanie widzi mniej kluczy. Według autorów usunięcie któregokolwiek z tych elementów obniża dokładność na BioDEX.

Trening używa zwykłej straty next-token, uśredniając stratę przy pełnym kontekście i przy losowym budżecie K. Dzięki temu jedna bramka działa dla różnych budżetów. Do bazowego modelu dochodzi bramka MLP i adapter LoRA (r=16).

Wyniki

Testy objęły Qwen3-1.7B na pięciu zadaniach: RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX i needle-in-a-haystack (NIAH). Porównano je z sześcioma bazowymi metodami: FastKV, SlimInfer, DASH, GemFilter, FlexPrefill i MInference, na tym samym backbone z LoRA. Czasy w pracy mierzono na jednej karcie A100 przy batch size 1; trening wymagał czterech A100.

• W porównaniu z pełnym prefillem: TTFT szybszy 1,4–11,3 raza, szczytowa pamięć mniejsza 1,5–3,4 raza – przy braku utraty dokładności, jak twierdzą autorzy.
• W porównaniu z najlepszą z metod bazowych: TTFT szybszy 1,2–6,4 raza, pamięć do 2,1 raza mniejsza.
• 2WikiMultihopQA: przy 95% dokładności pełnego kontekstu prefill jest 3,0 raza szybszy i zużywa 2,7 raza mniej pamięci.
• NIAH przy 16k tokenów: najlepszy wynik, czyli 11,3 raza szybciej i 3,42 raza mniej pamięci.

Bramka wytrenowana na kontekście 16k przenosi się na dłuższe wejścia: przyspieszenie rośnie do 12,8 raza przy 32k i 13,0 raza przy 64k oraz 128k, a oszczędność pamięci utrzymuje się w okolicach 3,43 raza. Najsłabszy wynik to 1,4 raza na części zadań. Autorzy zauważają też, że uczony selektor wygrywa z heurystyką FastKV na tej samej warstwie, a nawet losowy wybór K tokenów działa lepiej na backbone’ie TokenGate. Przypisują to współadaptacji backbone’u do przycinania.

Metoda działa również na modelu hybrydowym Qwen3.5-2B, który miesza warstwy Gated DeltaNet z pełną uwagą. Na NIAH 32k bramka na warstwie 3 dorównała dokładnością pełnemu prefillowi, była 3,8 raza szybsza od niego i 2,1 raza szybsza od FastKV. Bramka na warstwach 1–2 dała 7,8 raza przyspieszenia kosztem około 6% dokładności. FastKV potrafi przycinać wyłącznie warstwy uwagi, TokenGate – także DeltaNet.

Wdrożenie i zastrzeżenia

Kod jest na GitHubie na licencji Apache-2.0. Wymaga Pythona 3.10+ i PyTorcha 2.9+ z CUDA. Opcjonalny kernel FlashAttention-2 z key-bias przyspiesza gęste Qwen3 (koszt 1,0–1,1 raza zwykłego wywołania), a bez niego inferencja przechodzi na FlexAttention. Dokument do przycięcia oznacza się w prompcie placeholderem {state}, a w generate podaje się budżet max_kept_prefill_tokens. Checkpointy obejmują wszystkie pięć zadań na Qwen3-1.7B oraz bramkę NIAH dla Qwen3.5-2B; przykładowy to tokengate-niah-qwen3-1.7b na Hugging Face (bramka na warstwie 1 z adapterem LoRA).

Metoda nie jest darmowa: trzeba trenować bramkę i LoRA, a pojedynczy krok jest około 1,3 raza wolniejszy niż w zwykłym LoRA. Według autorów trening zbiegał jednak szybciej na czterech z pięciu zadań, ze średnim czasem 1,8 zamiast 2,2 godziny A100.

Warto pamiętać o ograniczeniach tych danych. Wszystkie liczby pochodzą od samych autorów, a testy dotyczą modeli o rozmiarze 1,7–2 mld parametrów. Przyspieszenia konkurencyjnych metod, które podają ich twórcy, wynikają z innych konfiguracji sprzętowych i modeli, więc nie da się ich bezpośrednio porównać z TokenGate. Rozrzut od 1,4 do 11,3 raza pokazuje też, że zysk silnie zależy od zadania: najlepiej wypada NIAH, gdzie większość kontekstu rzeczywiście jest zbędna.