Kaip įgyvendinama KV talpykla autoregresyvaus transformatoriaus inferencijoje?

Autoregresiniai transformatoriai, inferencijos metu, svarbu efektyviai naudoti skaičiavimus, kad nereikėtų iš naujo skaičiuoti visos sekos nuo pradžių. Tam naudojamas KV cache, kuris saugo savarankiško dėmesio (self-attention) sluoksnių raktus (Key) ir reikšmes (Value).

Idėja ta, kad generuojant kitą žymeklį, jau yra apskaičiuoti visi ankstesnių žymeklių raktai ir reikšmės. Vietoj jų pakartotinio skaičiavimo, juos saugome atmintyje (cache) ir skaičiuojant dėmesį naujam žymekliui, naudojame šį cache.

Tai maždaug veikia taip:

• Kiekviename transformatoriaus sluoksnyje, apdorojant žymeklį, raktai ir reikšmės saugomi atskiruose buferiuose.
• Generuojant naują žymeklį, apskaičiuojami tik jo raktai ir reikšmės, ir jie sujungiami su anksčiau išsaugotais.
• Dėmesio mechanizmas naudoja sujungtą raktų ir reikšmių rinkinį konteksto skaičiavimui.

Tai žymiai pagreitina inferenciją, sumažindama skaičiavimų sudėtingumą nuo kvadratinio iki linijinio sekos ilgiui.

Pavyzdys pseudokodo vienam sluoksniui:

# kv_cache saugo ankstesnių žymeklių raktus ir reikšmes
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)

Taip, KV cache yra mechanizmas, leidžiantis saugoti tarpinio žymeklių raktus ir reikšmes, pagreitinant sekos generavimą transformatoriuose.