Speculative Decoding

Speculative Decoding przyspiesza generowanie przez proponowanie kilku tokenów tańszym mechanizmem i weryfikowanie ich przez model docelowy. W klasycznym wariancie mały model draft tworzy kontynuację kolejno, a model docelowy oblicza rozkłady dla jej prefiksów równolegle. Akceptujemy początkowy ciąg propozycji do pierwszego odrzucenia, po czym losujemy korektę. Jeśli wszystkie propozycje przejdą, można dobrać dodatkowy token. Leviathan et al., §2, algorytm 1.

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.