Inżynierowie z Google DeepMind i YouTube opublikowali pracę opisującą STATIC – framework, który przyspiesza tak zwane constrained decoding w modelach generatywnych nawet 948 razy względem dotychczasowych rozwiązań. Brzmi jak marketingowa przesada, ale liczby z wdrożenia produkcyjnego na YouTube mówią same za siebie.
Żeby zrozumieć, o co chodzi, trzeba cofnąć się o krok. Tradycyjne systemy rekomendacji szukają podobnych treści przez porównywanie wektorów w przestrzeni embeddingów – to szybkie, sprawdzone, ale mało elastyczne. Coraz więcej firm, w tym Google, przechodzi na Generative Retrieval, gdzie model językowy dosłownie „generuje” identyfikator rekomendowanego elementu, token po tokenie, jak tekst. Problem pojawia się, gdy biznes potrzebuje narzucić ograniczenia: tylko świeże filmy, tylko dostępne produkty, tylko treści zgodne z polityką platformy. Standardowe generowanie nie ma o tym pojęcia i spokojnie halucynuje identyfikatory filmów, które nie istnieją albo zostały usunięte.
Drzewo prefiksowe na TPU – coś tu nie gra
Dotychczasowe podejście: drzewo prefiksowe (trie) jako maska na każdym kroku dekodowania. Koncepcyjnie eleganckie. W praktyce – katastrofa wydajnościowa na nowoczesnych akceleratorach.
Dwa główne powody:
- Struktury oparte na wskaźnikach wymuszają losowy, nieciągły dostęp do pamięci, co zabija możliwości HBM (High-Bandwidth Memory) w TPU i GPU
- Kompilatory ML jak XLA Googla wymagają statycznych grafów obliczeń, a klasyczne trie używa dynamicznego, rekurencyjnego przechodzenia – to po prostu nie daje się skompilować bez kosztownych transferów między CPU a akceleratorem
STATIC: spłaszcz drzewo, zwektoryzuj wszystko
Rozwiązanie Google’a nosi nazwę STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding) i polega na… zrezygnowaniu z drzewa jako struktury do przechodzenia. Zamiast tego cały trie zostaje spłaszczony do macierzy rzadkiej w formacie CSR (Compressed Sparse Row), co zamienia nieregularne przejścia po drzewie w w pełni zwektoryzowane operacje macierzowe.
Architektura dekodowania działa w dwóch fazach. Dla pierwszych dwóch warstw, gdzie rozgałęzienia są najgęstsze, STATIC używa skompresowanego tensora boolowskiego z odczytem O(1). Dla głębszych warstw wchodzi kernel VNTK (Vectorized Node Transition Kernel), który wykonuje „spekulatywne cięcie” stałej liczby wpisów bez żadnych rozgałęzień warunkowych. Efekt: cały proces dekodowania pozostaje jednym, statycznym grafem obliczeń.
Piotr Wolniewicz, Redaktor Naczelny AIPORT.pl: To ciekawy przykład na to, że wąskie gardło w systemach AI często leży nie w samym modelu, ale w infrastrukturze wokół niego. Google rozwiązał tu realny problem produkcyjny i chwała im za opublikowanie pełnej pracy z wynikami wdrożenia. Z drugiej strony – warto pamiętać, że STATIC jest mocno zoptymalizowany pod ekosystem Google: TPU v6e, XLA, własna infrastruktura YouTube. Pytanie, na ile inne firmy będą mogły z tego skorzystać bez przepisywania dużej części stacku? I drugie pytanie: czy rosnąca zależność od tak niszowej, hardware-specyficznej optymalizacji nie tworzy nowej formy lock-in, tym razem nie na poziomie danych, ale architektury obliczeniowej?
Liczby, które trudno zbagatelizować
Testy przeprowadzono na TPU v6e z modelem 3-miliardowym, rozmiarem batcha 2 i beam size 70. Wyniki względem innych metod są następujące:
- STATIC: +0,033 ms na krok (0,25% całkowitego czasu inferencji)
- PPV Approximate: +1,56 ms (11,9%)
- Hash Bitmap: +12,3 ms (94%)
- CPU Trie: +31,3 ms (239%)
- PPV Exact: +34,1 ms (260%)
948-krotne przyspieszenie wobec offloadowanego na CPU trie, 1033-krotne wobec dokładnego binary search. Zużycie pamięci HBM: około 90 MB na milion ograniczeń, co dla słownika 20 milionów elementów daje górny limit ok. 1,5 GB.
YouTube już to używa. I działa.
STATIC trafił na produkcję YouTube z ograniczeniem „tylko filmy z ostatnich 7 dni”. Słownik świeżych treści: 20 milionów elementów, compliance z regułą biznesową: 100%.
Wyniki testów A/B:
- +5,1% wyświetleń świeżych filmów (do 7 dni)
- +2,9% wyświetleń filmów z ostatnich 3 dni
- +0,15% wzrost CTR
Framework rozwiązuje też tzw. problem cold-start, czyli rekomendowanie treści, których model nigdy nie widział podczas trenowania. Na benchmarkach Amazon Reviews model bez ograniczeń uzyskiwał 0,00% Recall@1 dla nowych produktów. Z STATIC-iem – wyniki istotnie niezerowe, co w kontekście systemów rekomendacyjnych to przepaść.
Pełna praca dostępna na arxiv.org, kod źródłowy na GitHubie YouTube.
