Jak je implementována cache KV při inferenci autoregresivního transformátoru?

V autoregresivních transformátorech je při inferenci důležité efektivně využívat výpočty, aby se znovu nepočítala celá sekvence od začátku. K tomu slouží KV cache, která ukládá klíče (Key) a hodnoty (Value) z vrstev self-attention.

Myšlenka je, že při generování dalšího tokenu již máme vypočtené klíče a hodnoty pro všechny předchozí tokeny. Místo jejich opětovného výpočtu je ukládáme do paměti (cache) a při výpočtu pozornosti pro nový token používáme tuto cache.

Toto funguje přibližně takto:

• Pro každý vrstvu transformátoru, při zpracování tokenu, jsou klíče a hodnoty uloženy do samostatných bufferů.
• Při generování dalšího tokenu jsou vypočteny pouze klíče a hodnoty pro něj, a poté jsou sloučeny s již uloženými.
• Mechanismus pozornosti používá sloučenou sadu klíčů a hodnot k výpočtu kontextu.

Toto výrazně urychluje inferenci, snižuje výpočetní složitost z kvadratické na lineární vzhledem k délce sekvence.

Příklad pseudokódu pro jednu vrstvu:

# kv_cache ukládá klíče a hodnoty z předchozích tokenů
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)

Tímto způsobem je KV cache mechanismus pro ukládání mezivýsledných klíčů a hodnot, který urychluje sekvenční generaci v transformátorech.