Ako sa implementuje KV cache pri inferencii autoregresívneho transformátora?

V autoregresívnych transformátoroch je počas inferencie dôležité efektívne využívať výpočty, aby sa celá sekvencia nemusela prepočítavať od začiatku. Na tento účel sa používa KV cache, ktorá uchováva kľúče (Key) a hodnoty (Value) zo vrstiev self-attention.

Myšlienka je, že pri generovaní ďalšieho tokenu sme už vypočítali kľúče a hodnoty pre všetky predchádzajúce tokeny. Namiesto ich opätovného výpočtu ich ukladáme do pamäte (cache) a pri výpočte pozornosti pre nový token používame tento cache.

Toto funguje približne takto:

• Pre každý vrstvu transformátora, pri spracovaní tokenu, sa kľúče a hodnoty ukladajú do samostatných bufferov.
• Pri generovaní ďalšieho tokenu sa vypočítajú len kľúče a hodnoty pre neho, a potom sa spojí s tými už uloženými.
• Mechanizmus pozornosti používa zjednotený súbor kľúčov a hodnôt na výpočet kontextu.

Toto výrazne zrýchľuje inferenciu, znižujúc zložitosť z kvadratickej na lineárnu vzhľadom na dĺžku sekvencie.

Príklad pseudokódu pre jeden vrstvu:

# kv_cache ukladá kľúče a hodnoty z predchádzajúcich tokenov
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 spôsobom je KV cache mechanizmus na ukladanie medzičlánkových kľúčov a hodnôt, ktorý zrýchľuje sekvenčnú generáciu v transformátoroch.