Jak jest implementowana pamięć KV podczas inferencji autoregresywnego transformera?

W autoregresyjnych transformatorach podczas inferencji ważne jest efektywne wykorzystywanie obliczeń, aby nie przeliczać całej sekwencji od nowa. W tym celu stosuje się pamięć podręczną KV, która przechowuje klucze (Key) i wartości (Value) z warstw uwagi (self-attention).

Pomysł polega na tym, że podczas generowania kolejnego tokena już obliczyliśmy klucze i wartości dla wszystkich poprzednich tokenów. Zamiast je ponownie obliczać, zapisujemy je w pamięci (pamięć podręczna) i podczas obliczania uwagi dla nowego tokena korzystamy z tej pamięci.

Działa to mniej więcej tak:

• Dla każdej warstwy transformatora, podczas przetwarzania tokena, klucze i wartości są zapisywane w oddzielnych buforach.
• Przy generowaniu kolejnego tokena, obliczane są tylko klucze i wartości dla niego, a następnie łączone z wcześniej zapisanymi.
• Mechanizm uwagi korzysta z połączonego zestawu kluczy i wartości do obliczenia kontekstu.

To znacznie przyspiesza inferencję, zmniejszając złożoność obliczeniową z kwadratowej do liniowej w długości sekwencji.

Przykład pseudokodu dla jednej warstwy:

# kv_cache przechowuje klucze i wartości dla poprzednich tokenów
new_key, new_value = compute_kv(new_token)
kv_cache.keys = concatenate(kv_cache.keys, new_key)
kv_cache.values = concatenate(kv_cache.values, new_value)
output = attention(query=new_key, keys=kv_cache.keys, values=kv_cache.values)

W ten sposób, KV-cache to mechanizm przechowywania pośrednich kluczy i wartości, przyspieszający sekwencyjną generację w transformatorach.