Come viene implementata la cache KV nell'inferenza di un trasformatore autoregressivo?

Nei trasformatori autoregressivi, durante l'inferenza, è importante utilizzare i calcoli in modo efficiente per non ricalcolare tutta la sequenza da zero. A tal fine, si utilizza la cache KV, che memorizza le chiavi (Key) e i valori (Value) degli strati di attenzione (self-attention).

L'idea è che, generando il token successivo, abbiamo già calcolato le chiavi e i valori per tutti i token precedenti. Invece di ricalcolarli, li memorizziamo in memoria (cache) e, nel calcolo dell'attenzione per il nuovo token, utilizziamo questa cache.

Funziona circa così:

• Per ogni strato del trasformatore, durante l'elaborazione di un token, le chiavi e i valori vengono memorizzati in buffer separati.
• Quando si genera il token successivo, vengono calcolate solo le chiavi e i valori per esso, e poi vengono combinati con quelli già memorizzati.
• Il meccanismo di attenzione utilizza l'insieme combinato di chiavi e valori per calcolare il contesto.

Questo accelera notevolmente l'inferenza, riducendo la complessità computazionale da quadratica in lunghezza di sequenza a lineare.

Esempio di pseudocodice per uno strato:

# kv_cache memorizza le chiavi e i valori dei token precedenti
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)

In questo modo, la cache KV è un meccanismo per salvare le chiavi e i valori intermedi per accelerare la generazione sequenziale nei trasformatori.