Kuidas rakendatakse KV-vahemälu autoregressiivse transformeri inference'is?

Autoregressiivsetes transformaatoreid kasutades on inferentsi ajal oluline tõhusalt arvutada, et kogu järjestust uuesti arvutada ei peaks. Selleks kasutatakse KV-vahemälu, mis salvestab self-attention kihtide võtmed (Key) ja väärtused (Value).

Idee on, et järgmise tokeni genereerimisel on juba arvutatud kõikide eelnevate tokenite võtmed ja väärtused. Selle asemel, et neid uuesti arvutada, salvestame need mällu (vahemällu) ja kasutame seda vahemälu uue tokeni tähelepanu arvutamisel.

See töötab umbes nii:

• Iga transformatsiooni kihi puhul salvestatakse tokeni töötlemise ajal võtmed ja väärtused eraldi pufferitesse.
• Järgmisel tokenil arvutatakse ainult selle võtmed ja väärtused ning need ühendatakse juba salvestatutega.
• Tähelepanu mehhanism kasutab ühendatud võtmete ja väärtuste komplekti konteksti arvutamiseks.

See kiirendab märkimisväärselt inferentsi, vähendades arvutuskompleksust kvadratsilisest lineaarseks seeria pikkuse järgi.

Näide pseudokoodist ühe kihi jaoks:

# kv_cache salvestab eelnevate tokenite võtmed ja väärtused
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)

Seega on KV-vahemälu mehhanism, mis salvestab vahepealseid võtmeid ja väärtusi, kiirendades järjestuslikku genereerimist transformaatoreites.