Wie wird der KV-Cache bei der Inferenz eines autoregressiven Transformers implementiert?

In autoregressive Transformatoren ist es bei der Inferenz wichtig, Berechnungen effizient zu nutzen, um nicht die gesamte Sequenz neu berechnen zu müssen. Dafür wird der KV-Cache verwendet, der die Schlüssel (Key) und Werte (Value) aus den Attention-Schichten (Self-Attention) speichert.

Die Idee ist, dass beim Generieren des nächsten Tokens bereits die Schlüssel und Werte für alle vorherigen Tokens berechnet wurden. Anstatt sie erneut zu berechnen, werden sie im Speicher (Cache) gespeichert und beim Berechnen der Aufmerksamkeit für das neue Token verwendet.

Das funktioniert ungefähr so:

• Für jede Schicht des Transformers werden beim Verarbeiten eines Tokens die Schlüssel und Werte in separaten Puffern gespeichert.
• Beim Generieren des nächsten Tokens werden nur die Schlüssel und Werte für dieses Token berechnet und mit den bereits gespeicherten kombiniert.
• Der Attention-Mechanismus nutzt die kombinierte Menge an Schlüsseln und Werten, um den Kontext zu berechnen.

Dies beschleunigt die Inferenz erheblich, indem die rechnerische Komplexität von quadratisch in der Länge der Sequenz auf linear reduziert wird.

Beispiel für Pseudocode für eine Schicht:

# kv_cache speichert die Schlüssel und Werte der vorherigen Tokens
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)

Auf diese Weise ist der KV-Cache ein Mechanismus zum Speichern der Zwischenschlüssel und -werte, um die sequenzielle Generierung in Transformatoren zu beschleunigen.