Speculative Decoding

Duży model generuje tekst poprawnie, ale wolno. Mniejszy model może szybko zaproponować kilka kolejnych kawałków. Chcesz wykorzystać jego pracę, zachowując kontrolę dużego modelu nad ostatecznym wynikiem.

Speculative Decoding tworzy tańszą propozycję kontynuacji, a potem weryfikuje ją modelem docelowym. W klasycznym algorytmie przyjmujemy początkowy ciąg zaakceptowanych tokenów — kawałków tekstu — do pierwszego odrzucenia i wtedy wyliczamy korektę.

Duży model może ocenić znaną propozycję dla kilku pozycji równolegle. Poprawny wariant losowania używa szczególnych prawdopodobieństw akceptacji i korekty, żeby zachować docelowy rozkład. Sam test „czy brzmi sensownie” tego nie zapewnia.

Przyspieszenie zależy od kosztu propozycji, liczby akceptacji i sprzętu. Jeśli mały model często się rozmija z dużym, dodatkowa praca może zniwelować zysk. Demonstracja poniżej pokazuje zachowanie szans, a nie gwarantowany czas działania.

Źródło mechanizmu: Leviathan et al., §2, algorytm 1.

Mechanizm i szczegóły

To sposób przyspieszania Decode z zachowaniem rozkładu modelu docelowego pod warunkami algorytmu. Nie wystarcza sprawdzenie, czy propozycja „wygląda dobrze”.

Dlaczego potrzebna jest korekta

Oznaczmy docelowy rozkład przez $p$, a rozkład draft przez $q$, dla tego samego prefiksu i wspólnego zbioru tokenów. Oba uwzględniają wybrane reguły losowania, np. Temperature i Top-p Sampling. Propozycję $x\sim q$ akceptujemy z prawdopodobieństwem:

$$\min\left(1,\frac{p(x)}{q(x)}\right)$$

Wylosowana propozycja ma $q(x)>0$. Po odrzuceniu losujemy z rozkładu proporcjonalnego do $\max(0,p-q)$. Taki krok uzupełnia dokładnie brakującą masę prawdopodobieństwa. Jeśli $p=q$, nic nie jest odrzucane i korekta nie jest potrzebna. Leviathan et al., §2.3 i dowód A.1.

Własny przykład: dla tokenów A, B, C niech $p=[0{,}2;0{,}5;0{,}3]$, a $q=[0{,}6;0{,}3;0{,}1]$. Bezwarunkowa masa zaakceptowanych propozycji wynosi $[0{,}2;0{,}3;0{,}1]$. Pozostałe 0,4 trzeba rozdzielić między B i C po równo. Losowanie korekty ze zwykłego $p$ dałoby zamiast tego wynik $[0{,}28;0{,}50;0{,}22]$.

Eksperyment pokazuje dokładny bilans jednej pozycji, bez symulowania całej pętli ani czasu wykonania. Zachowany rozkład nie gwarantuje identycznego tekstu dla tego samego seeda: algorytmy zużywają losowość inaczej, a obliczenia mają ograniczoną precyzję. Omawiają to Chen et al., §4.2 i §6.1.

Przyspieszenie zależy od odsetka akceptacji, kosztu draftu i sprzętu. Zbyt drogi draft lub wiele odrzuceń może zniwelować oszczędność, nawet gdy poprawność rozkładu jest zachowana.

Wykorzystuję treści generowane przez AI jako część mojego codziennego procesu nauki.