Hogyan valósul meg a KV-cache az autoregresszív transzformátor inferenciájában?

Az autoregresszív transzformátorokban az inferencia során fontos hatékonyan kihasználni a számításokat, hogy ne kelljen az egész szekvenciát újraszámolni. Ehhez használják a KV-cache-t, amely a figyelem (self-attention) rétegeiből származó kulcsokat (Key) és értékeket (Value) tárolja.

Az ötlet az, hogy a következő token generálásakor már kiszámoltuk az összes előző tokenhez tartozó kulcsokat és értékeket. Ahelyett, hogy újra kiszámolnánk őket, ezeket memóriában (cache) tároljuk, és az új token figyelem számításakor ezt a cache-t használjuk.

Ez nagyjából így működik:

• Minden transzformátor rétegben, a token feldolgozása során, a kulcsokat és értékeket külön pufferben tárolják.
• A következő token generálásakor csak az ő kulcsait és értékeit számítják ki, majd ezeket összekapcsolják a már tároltakkal.
• A figyelem mechanizmus a kombinált kulcsok és értékek halmazát használja a kontextus kiszámításához.

Ez jelentősen gyorsítja az inferenciát, mivel a számítási komplexitás a szekvencia hosszának négyzetes növekedéséről lineárisra csökken.

Egy réteg pseudokód példája:

# kv_cache tárolja az előző tokenek kulcsait és értékeit
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)

Így a KV-cache egy olyan mechanizmus, amely közbenső kulcsokat és értékeket tárol, gyorsítva a szekvenciális generálást a transzformátorokban.