Como é que a cache KV é implementada na inferência de um transformador autoregressivo?

Nos transformadores autorregressivos, durante a inferência, é importante usar os cálculos de forma eficiente para não recalcular toda a sequência do zero. Para isso, é utilizado o cache KV, que armazena as chaves (Key) e os valores (Value) das camadas de atenção (autoatenção).

A ideia é que, ao gerar o próximo token, já calculámos as chaves e os valores para todos os tokens anteriores. Em vez de recalculá-los, guardamo-los na memória (cache) e, ao calcular a atenção para o novo token, usamos esse cache.

Funciona aproximadamente assim:

• Para cada camada do transformador, ao processar um token, as chaves e os valores são armazenados em buffers separados.
• Ao gerar o próximo token, apenas a chave e o valor para ele são calculados, e depois combinados com os já armazenados.
• O mecanismo de atenção usa o conjunto combinado de chaves e valores para calcular o contexto.

Isto acelera significativamente a inferência, reduzindo a complexidade computacional de quadrática em comprimento de sequência para linear.

Exemplo de pseudocódigo para uma camada:

# kv_cache armazena as chaves e valores dos tokens anteriores
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)

Assim, a cache KV é um mecanismo para guardar as chaves e valores intermediários para acelerar a geração sequencial em transformadores.