Kā tiek īstenota KV kešatmiņa autoregresīvā transformatora inferencē?

Autoregresīvajos transformatoros, inferenču laikā, ir svarīgi efektīvi izmantot aprēķinus, lai visa secība nebūtu jāaprēķina no jauna. Tam tiek izmantots KV kešatmiņa, kas glabā pašattences (self-attention) slāņu atslēgas (Key) un vērtības (Value).

Ideja ir tāda, ka, ģenerējot nākamo tokenu, jau ir aprēķinātas visas iepriekšējo tokenu atslēgas un vērtības. Tā vietā, lai tās atkārtoti aprēķinātu, tās glabājam atmiņā (kešā) un, aprēķinot uzmanību jaunajam tokenam, izmantojam šo kešatmiņu.

Tas aptuveni darbojas šādi:

• Katram transformatora slānim, apstrādājot tokenu, atslēgas un vērtības tiek glabātas atsevišķos buferos.
• Izveidojot nākamo tokenu, tiek aprēķinātas tikai tās atslēgas un vērtības, un tās tiek apvienotas ar iepriekš saglabātajām.
• Uzmanības mehānisms izmanto apvienoto atslēgu un vērtību kopu, lai aprēķinātu kontekstu.

Tas ievērojami paātrina inferenci, samazinot aprēķinu sarežģītību no kvadrātiskās uz lineāro atkarībā no secības garuma.

Pseido-koda piemērs vienam slānim:

# kv_cache glabā iepriekšējo tokenu atslēgas un vērtības
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ādējādi, KV kešatmiņa ir mehānisms, kas glabā starpposma atslēgas un vērtības, paātrinot secīgu ģenerēšanu transformatoros.