Kako se implementira KV keš u inferenciji autoregresivnog transformatora?

U autoregresivnim transformatorima, tokom inferencije, važno je efikasno koristiti računanja kako se cela sekvenca ne bi ponovo računala od početka. Za to se koristi KV keš, koji čuva ključeve (Key) i vrednosti (Value) iz slojeva samostalne pažnje (self-attention).

Ideja je da smo već izračunali ključeve i vrednosti za sve prethodne tokene prilikom generisanja sledećeg tokena. Umesto da ih ponovo računamo, čuvamo ih u memoriji (keš) i prilikom računanja pažnje za novi token koristimo taj keš.

Ovo funkcioniše otprilike ovako:

• Za svaki sloj transformatora, tokom obrade tokena, ključevi i vrednosti se čuvaju u odvojenim baferima.
• Pri generisanju sledećeg tokena, računaju se samo ključevi i vrednosti za njega, a zatim se kombinuju sa već sačuvanim.
• Mehanizam pažnje koristi objedinjeni skup ključeva i vrednosti za računanje konteksta.

Ovo značajno ubrzava inferencu, smanjujući složenost iz kvadratne na linearnu u dužini sekvence.

Primer pseudokoda za jedan sloj:

# kv_cache čuva ključeve i vrednosti prethodnih tokena
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)

Tako, KV keš je mehanizam za čuvanje međuzavršnih ključeva i vrednosti za ubrzavanje sekvencijalne generacije u transformatorima.