Hoe wordt de KV-cache geïmplementeerd bij de inferentie van een autoregressieve transformer?

In autoregressieve transformers is het tijdens inferentie belangrijk om berekeningen efficiënt te gebruiken, zodat de hele sequentie niet opnieuw berekend hoeft te worden. Hiervoor wordt de KV-cache gebruikt, die de sleutels (Key) en waarden (Value) uit de aandachtlagen (self-attention) opslaat.

Het idee is dat bij het genereren van de volgende token, de sleutels en waarden voor alle voorgaande tokens al berekend zijn. In plaats van ze opnieuw te berekenen, worden ze in het geheugen (cache) opgeslagen en bij het berekenen van de aandacht voor de nieuwe token wordt deze cache gebruikt.

Dit werkt ongeveer als volgt:

• Voor elke laag van de transformer worden bij het verwerken van een token de sleutels en waarden in aparte buffers opgeslagen.
• Bij het genereren van de volgende token worden alleen de sleutels en waarden voor dat token berekend en gecombineerd met de eerder opgeslagen.
• Het aandachtmechanisme gebruikt de gecombineerde set van sleutels en waarden om de context te berekenen.

Dit versnelt de inferentie aanzienlijk, doordat de computationele complexiteit van kwadratisch in de lengte van de sequentie naar lineair wordt verminderd.

Voorbeeld van pseudocode voor één laag:

# kv_cache slaat de sleutels en waarden van eerdere tokens op
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)

Op deze manier is de KV-cache een mechanisme om tussenliggende sleutels en waarden op te slaan, waardoor sequentiële generatie in transformers versneld wordt.