Cum se implementează cache-ul KV în inferența unui transformator autoregresiv?

În transformatoarele autoregresive, în timpul inferenței, este important să se utilizeze calculele în mod eficient pentru a nu recalcula întreaga secvență de fiecare dată. Pentru aceasta, se folosește cache-ul KV, care stochează cheile (Key) și valorile (Value) din straturile de atenție (auto-atenție).

Ideea este că, generând următorul token, deja am calculat cheile și valorile pentru toți tokenii anteriori. În loc să le recalculăm, le stocăm în memorie (cache) și, la calculul atenției pentru noul token, folosim acest cache.

Aceasta funcționează aproximativ astfel:

• Pentru fiecare strat al transformatorului, în timpul procesării unui token, cheile și valorile sunt stocate în buffers separate.
• La generarea următorului token, se calculează doar cheile și valorile pentru acesta, apoi se combină cu cele deja stocate.
• Mecanismul de atenție utilizează setul combinat de chei și valori pentru a calcula contextul.

Acest lucru accelerează semnificativ inferența, reducând complexitatea computațională de la cuadratică în lungimea secvenței la liniară.

Exemplu de pseudocod pentru un strat:

# kv_cache stochează cheile și valorile pentru tokenii anteriori
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)

Astfel, cache-ul KV este un mecanism de stocare a cheilor și valorilor intermediare pentru accelerarea generării secvențiale în transformatoare.