Compressive Transformer
Jak działa
Compressive Transformer rozszerza pętlę pamięci Transformer-XL o trzeci poziom hierarchii: (1) Bieżący segment (T tokenów) — pełne attention, queries+keys+values aktywne. (2) Krótkoterminowa pamięć M (M tokenów) — hidden states ostatnich N segmentów, jak w XL, działa jako keys/values dla bieżącego segmentu. (3) Długoterminowa pamięć skompresowana M_c (M_c tokenów) — gdy najstarszy segment ma zostać wypchnięty z bufora M, zamiast go wyrzucić, kompresuje się go funkcją f_compress: R^{c×d} → R^{1×d}, gdzie c to compression rate (typowo 3 lub 4). Skompresowane tokeny lądują w buforze M_c, który również działa jako keys/values dla bieżącego attention. Funkcje f_compress testowane w pracy: (a) 1D mean pooling — średnia c kolejnych wektorów, (b) 1D max pooling — pointwise max, (c) 1D conv (kernel size c, stride c) — uczone, najlepsze empiryczne wyniki, (d) dilated conv — szersze receptive field, (e) most-used — zachowanie c tokenów z najwyższym kumulatywnym attention z poprzednich queries. Trening kompresorów: oprócz standardowej cross-entropy, autorzy wprowadzają attention-reconstruction loss — pierwotne hidden states i ich skompresowane warianty powinny dawać podobne wzorce attention dla zachowanych queries. To dodatkowo motywuje kompresor do zachowania informacji ważnej dla attention. Pozycja: relatywne PE jak w Transformer-XL, ale skompresowane tokeny dostają specjalny offset pozycji proporcjonalny do c.
Rozwiązany problem
Transformer-XL trzyma hidden states ostatnich N segmentów w buforze FIFO i wyrzuca starsze. To proste, ale rozrzutne — odrzucone informacje są tracone bezpowrotnie, mimo że dla zadań takich jak modelowanie książek (PG-19) odległe odniesienia są kluczowe. Naiwne zwiększenie M (memory length) skaluje VRAM liniowo z M, co jest niepraktyczne dla bardzo długich sekwencji. Compressive Transformer rozwiązuje to: zamiast wyrzucać, KOMPRESUJE — c tokenów zamienia w 1 token w buforze długoterminowym. Daje to logarytmiczną hierarchię pamięci (świeże → krótkoterminowe → skompresowane) przy stałym sumarycznym koszcie pamięci.
Komponenty
Bufor FIFO trzymający hidden states ostatnich N segmentów. Identyczna semantyka jak w Transformer-XL — różnica polega na tym, że ZAMIAST być wyrzucanym, najstarszy segment trafia do kompresora.
Funkcja zamieniająca c kolejnych hidden states w jeden. Wywoływana, gdy najstarszy segment z M ma zostać wypchnięty. Może być nauczona (1D conv) lub deterministyczna (pooling, most-used).
Oficjalna
Drugi bufor FIFO trzymający M_c skompresowanych tokenów. Każdy reprezentuje c oryginalnych tokenów. Razem z M tworzy hierarchię pamięci: świeże → krótkoterminowe → skompresowane.
Implementacja
Uczona kompresja (1D conv) bez auxiliary loss'u attention-reconstruction degeneruje do identity — model się uczy, że łatwiej ignorować skompresowane tokeny niż je wykorzystywać.
Pełen backprop przez wszystkie kroki kompresji (np. 100 segmentów wstecz) jest nierealny pamięciowo i niestabilny.
Skompresowany token reprezentuje c oryginalnych — jeśli używamy relatywnego PE jak dla zwykłych tokenów, model myśli że one są bezpośrednio za sobą, co konfunduje attention.
Ewolucja
Dai et al. wprowadzają cache hidden states i relatywne PE. Bezpośrednia podstawa Compressive Transformer — jednowartstwowa pamięć FIFO.
Niezależna ścieżka long-context: deterministyczne sparse wzorce. Compressive Transformer i Sparse Transformer powstały w tym samym roku, jako równorzędne odpowiedzi na ten sam problem.
Rae, Potapenko, Jayakumar, Hillier, Lillicrap publikują Compressive Transformer (arXiv:1911.05507). Dwupoziomowa pamięć: M (FIFO) + M_c (skompresowana). Wprowadzają również datasety PG-19 — pierwszy systematyczny benchmark long-range LM na książkach.
Longformer i BigBird (oba 2020) oferują prostszy long-context bez sekwencyjności i kompresorów. Compressive Transformer zostaje pracą ważną teoretycznie, ale rzadziej wdrażaną w produkcji.
Rozszerzenie idei skompresowanej pamięci o pamięć NIESKOŃCZONĄ — kNN lookup w gigantycznej zewnętrznej bazie hidden states. Bezpośredni dziedzic Compressive Transformer.
Google publikuje Infini-attention — pamięć skompresowana wbudowana bezpośrednio w warstwę attention bez osobnego bufora. Mamba i RWKV z kolei realizują kompresję poprzez stan ukryty SSM. Wszystkie te podejścia są ideowo bliskie Compressive Transformer.
Hiperparametry (konfigurowalne osie)
Liczba oryginalnych tokenów reprezentowanych przez 1 skompresowany token. Praca testuje c=2, 3, 4. Większe c = dłuższy efektywny kontekst, ale mniejsza precyzja zachowanych informacji.
Liczba skompresowanych tokenów w pamięci długoterminowej. Efektywny zasięg wstecz: M + c·M_c. Praca używa M=512–1024, M_c=512–1024, c=3 — co daje efektywny kontekst rzędu 4000–5000 tokenów.
Wybór mechanizmu kompresującego c tokenów do 1. Praca testuje 5 wariantów: mean pool, max pool, 1D conv, dilated conv, most-used. Najlepsze empirycznie: 1D conv z attention-reconstruction loss.
Auxiliary loss wymuszający, by attention computed na skompresowanych keys/values był podobny do attention na oryginałach. Kluczowy dla jakości kompresji opartej o uczone funkcje (1D conv, dilated conv).
Złożoność obliczeniowa
Złożoność czasowa: O(T · (T + M + M_c) · d) per segment + O(c · M · d²) per kompresję (rzadko). Złożoność przestrzenna: O((M + M_c) · d · L) cache + O(T · d · L) bieżąca aktywacja.
Wąskie gardło obliczeniowe
Wąskim gardłem treningu jest dodatkowy forward kompresora i obliczenie auxiliary loss'u. Inferencja jest tańsza — kompresja wykonywana tylko raz na N segmentów, a ostatecznie M_c jest gotowe i tylko serwowane jako keys/values.
Paradygmat wykonania
Wariant „most-used" wprowadza miękką selektywność — tokeny z najwyższym attention z poprzednich queries są zachowywane preferencyjnie. To prosta forma uczenia, jakie informacje warto zachować.
Brak uczonego routingu między poziomami pamięci. Wszystkie keys/values (bieżące + M + M_c) wchodzą do tego samego attention.
Równoległość
Trening kompresora wymaga jednego kroku forward na M tokenów co k-ty segment + dodatkowego attention-reconstruction loss — narzut rzędu 5–15% względem Transformer-XL.
Wymagania sprzętowe
Oryginalne eksperymenty DeepMind były trenowane na TPU. 1D conv kompresor mapuje się świetnie na TPU systolic array.
Standardowe attention + 1D conv kompresor działają dobrze na GPU. Dwupoziomowa pamięć wymaga dodatkowego VRAM (M·d + M_c·d) — niewielki narzut na nowoczesnych GPU.
Algorytm jest standardowym attention + dodatkowymi kerneli kompresji — przenośny, ale realna efektywność wymaga GPU/TPU.