Kronecker Product Attention comprime o KV-cache em até 10 vezes e alcança o MLA em qualidade

A memória do KV-cache há muito tempo é o principal limitador da geração autorregressiva: ela cresce linearmente com o comprimento do contexto, o tamanho do batch e a profundidade do modelo. Um novo preprint no arXiv propõe o mecanismo Kronecker Product Attention (KPA), que, segundo os autores, comprime esse cache em até 10 vezes sem perder em qualidade para MHA, MQA, GQA, MLA e TPA. A comparação foi feita em quatro escalas de modelo e nove benchmarks downstream, em condições idênticas de treinamento.

Como funciona o KPA

A ideia central: chaves e valores não são armazenados por inteiro, mas decompostos em uma soma de produtos de Kronecker de fatores matriciais compactos. Os próprios fatores dependem do token, ou seja, são calculados a partir da entrada e não fixados nos pesos. Apenas esses fatores permanecem no cache, e os tensores completos são reconstruídos no momento do cálculo da atenção.

As abordagens existentes têm compromissos diferentes. MQA e GQA economizam memória ao agrupar cabeças e perdem expressividade. O MLA comprime o cache por meio de uma representação latente, mas exige um ramo «desacoplado» separado para a codificação posicional RoPE. O TPA usa fatoração por produto externo de vetores, ou seja, uma estrutura rígida de posto um. O KPA, segundo os autores, generaliza estritamente o TPA: o produto externo vetorial surge como caso degenerado da decomposição de Kronecker, e fatores matriciais não degenerados conseguem descrever padrões de bloco mais ricos.

Um resultado técnico à parte diz respeito às posições. Os autores introduzem um operador rotacional estruturado, compatível com a representação de Kronecker. Ele preserva a propriedade de posição relativa do RoPE, mas rotaciona apenas os fatores compactos. Por isso, não é necessário um ramo adicional para a codificação posicional, como no MLA.

O que o artigo afirma

• Compressão do KV-cache: até 10x.
• Comparação com MHA, MQA, GQA, MLA e TPA: resultado igual ou superior.
• Escalas: quatro tamanhos de modelo, nove benchmarks downstream.
• Condições: treinamento idêntico para todas as arquiteturas.
• Codificação posicional: operador compatível com RoPE sobre fatores compactos, sem ramo separado.
• Ajuste: a «divisão de cabeças de Kronecker» (Kronecker head split) é definida como hiperparâmetro.

O último item talvez seja o mais prático. A divisão define o compromisso entre grau de compressão e precisão, e pode ser alterada na fase de projeto sem trocar a arquitetura. A equipe não precisa reescrever o modelo para passar de uma compressão mais suave para uma mais agressiva. Basta treinar de novo com outro valor.

O que fazer como engenheiro e onde ter cautela

Por enquanto, é um preprint v1 sem verificação independente. O resumo não informa os tamanhos exatos dos modelos, nem os comprimentos de contexto, nem a velocidade de decodificação em hardware real. O ganho de memória nem sempre equivale a ganho de latência: reconstruir chaves e valores a partir dos fatores adiciona computação, e o custo dela depende dos kernels e da GPU. O «até 10x» declarado refere-se ao melhor cenário. Em qual configuração ele é alcançado e a que custo em qualidade só ficará claro com o texto completo e o código, caso os autores o publiquem.

Para equipes que montam inferência de contexto longo ou atendem batches grandes, vale manter o KPA na lista de candidatos ao lado do MLA. Há três pontos a verificar: quanta memória é liberada nos seus comprimentos de contexto, como a qualidade se comporta com diferentes divisões e qual throughput se obtém com os seus kernels de atenção. Para quem já treinou um modelo com GQA ou MLA, a migração exige novo treinamento. Por isso, faz mais sentido incorporar o KPA ao próximo ciclo de pre-training do que tentar trocar o mecanismo em pesos já prontos.