Resumo Técnico: HeatKV - Compressão de Cache KV Ajustada por Cabeça para Modelagem Autoregressiva Visual

Declaração do Problema

Os modelos Autoregressivos Visuais (VAR) emergiram como uma alternativa competitiva aos modelos de difusão para geração de imagens, oferecendo velocidades de inferência rápidas e propriedades de escalabilidade favoráveis ao gerar imagens como uma sequência de mapas de tokens em múltiplas escalas. No entanto, esses modelos enfrentam restrições severas de memória devido à natureza de seu cache KV. Diferentemente dos Modelos de Linguagem de Grande Escala (LLMs), que tipicamente geram um token por etapa, os modelos VAR geram mapas inteiros de tokens em resoluções crescentes. Consequentemente, o número de tokens no cache KV cresce de forma superlinear com a escala de geração.

Os modelos VAR mais avançados, como Infinity-2B e Infinity-8B, podem exigir até 10 GB de memória de cache KV para gerar uma única imagem de 1024 × 1024. Essa alta pegada de memória limita o tamanho máximo do lote e o número de usuários simultâneos que um provedor de serviços pode suportar em aceleradores de hardware especializados. Métodos de compressão existentes para VARs, como ScaleKV e HACK, frequentemente dependem de estratégias de granularidade grosseira, como classificar cabeças ou camadas de atenção em categorias binárias (por exemplo, "orçamento alto" vs. "orçamento baixo"). Os autores argumentam que essas classificações binárias são excessivamente restritivas, pois os padrões de atenção dentro da mesma camada variam significativamente entre diferentes cabeças, levando a desperdício de memória ou poda subótima.

Metodologia: HeatKV

Os autores propõem o HeatKV (Compressão de Cache KV Ajustada por Cabeça), uma estratégia inovadora que adapta a alocação de cache na granularidade de cabeças de atenção individuais e escalas de imagem específicas. O método consiste em três componentes principais:

1. Importância em Escala de Cabeça (Calibração Offline)

Diferentemente de métodos anteriores que utilizam métricas sensíveis à variância, o HeatKV classifica as cabeças de atenção com base em sua Pontuação de Atenção Cumulativa Dependente de Escala (S-CAS).

• Mecanismo: Utilizando um pequeno conjunto de dados de calibração offline, o método agrega probabilidades de atenção de tokens em escalas futuras de volta para tokens em escalas anteriores.
• Métrica: Para uma dada camada $\ell$, cabeça $h$, e escala de origem $k$, a S-CAS mede a massa de atenção média de todas as escalas futuras ($\tau > k$) para os tokens na escala $k$.
• Resultado: Isso gera uma classificação específica de pares "cabeça-escala" para cada escala de origem, indicando o quão crítico é o cache de uma cabeça específica de uma escala anterior específica para a geração futura de tokens. Isso permite que o sistema identifique que uma cabeça específica pode precisar reter a Escala 1, mas pode descartar a Escala 2, um nível de granularidade que métodos binários não conseguem alcançar.

2. Alocação de Orçamento

O HeatKV calcula o número de cabeças ($N_k$) que devem ser podadas após gerar cada escala $k$ para garantir que o tamanho total do cache permaneça dentro de um orçamento de memória definido pelo usuário $b$.

• Restrição: O método adere ao princípio do "Streaming-LLM" de reter as primeiras $s$ escalas como um "sumidouro de atenção" (essencial para estabilidade) e nunca armazenar em cache a escala atual sendo gerada.
• Cálculo: O algoritmo determina o número mínimo de cabeças a serem podadas para satisfazer a restrição de orçamento, permitindo uma redução dinâmica do cache à medida que a geração progride.

3. Otimização da Ordem de Remoção (Poda Antecipada Gananciosa)

Um desafio crítico nos VARs é que o tamanho do cache cresce de forma superlinear dentro de uma única etapa de geração (através das camadas do transformador). Métodos padrão que podam apenas entre escalas podem causar sobre-alocações de memória transitórias que excedem os limites de hardware.

• Solução: O HeatKV introduz um algoritmo de Poda Antecipada Gananciosa (Algoritmo 1). Ele impõe a restrição de orçamento após cada camada do transformador, e não apenas entre escalas.
• Estratégia: Se o tamanho do cache exceder o orçamento após uma camada específica, o algoritmo seleciona gananciosamente pares adicionais de cabeça-escala para podar daquela camada ou de camadas anteriores. Ele prioriza a poda de cabeças com classificações S-CAS mais baixas e, em caso de empate, prioriza a poda de camadas posteriores ou escalas de origem posteriores para maximizar o alívio de memória.
• Optimalidade: Os autores provam (Proposição 1) que essa abordagem gananciosa produz um conjunto de cardinalidade mínima de cabeças podadas antecipadamente, garantindo que o orçamento nunca seja violado enquanto maximiza a utilização da memória disponível.

Contribuições Principais

1. Observação de Granularidade Fina: Os autores demonstram que os padrões de atenção nos modelos VAR variam significativamente entre diferentes escalas dentro da mesma cabeça de atenção, mas permanecem relativamente estáveis entre diferentes prompts de entrada. Isso motiva uma mudança da classificação binária de cabeças para uma alocação de orçamento contínua e dependente de escala.
2. Métrica S-CAS: A introdução da Pontuação de Atenção Cumulativa Dependente de Escala permite a construção de um cronograma de poda estático adaptado a um orçamento de memória específico, permitindo que cada cabeça descarte escalas não essenciais enquanto retém as críticas.
3. Algoritmo de Poda Otimizado: O desenvolvimento de um algoritmo de evicção ganancioso que impõe restrições de orçamento no nível da camada (em vez de apenas no nível da escala) previne transbordamentos de memória transitórios, o que é crucial para o crescimento superlinear característico dos VARs.
4. Desempenho de Estado da Arte: O método alcança um novo estado da arte para compressão de cache KV em modelos VAR, demonstrando que a alocação de granularidade fina supera significativamente as abordagens de granularidade grosseira.

Resultados Experimentais

Os autores avaliaram o HeatKV nos modelos Infinity-2B e Infinity-8B utilizando o conjunto de validação MS-COCO 2017. Eles compararam o HeatKV com linhas de base incluindo StreamingLLM, SnapKV, HACK e ScaleKV em orçamentos de memória de 4%, 10% e 20%.

• Fidelidade: O HeatKV superou consistentemente todas as linhas de base. Notavelmente, no modelo Infinity-2B, o HeatKV com um orçamento de 10% alcançou melhores pontuações de FID, LPIPS e PSNR do que o ScaleKV com um orçamento de 20%. Com um orçamento de 4%, o HeatKV manteve um FID de 2,70 (Infinity-2B), significativamente menor do que o próximo melhor método (ScaleKV com 4,92).
• Alinhamento com Prompt: Usando GenEval, o HeatKV manteve forte alinhamento com o prompt mesmo em orçamentos baixos, com pontuações comparáveis à linha de base de cache completo.
• Qualidade Perceptiva: Os resultados da Pontuação de Preferência Humana (HPSv2.1) indicaram que o HeatKV preservou alta qualidade perceptiva, superando outros métodos com orçamento de 4% no Infinity-2B.
• Eficiência: O HeatKV alcançou uma razão de compressão 2× maior na alocação de memória em comparação com métodos existentes, mantendo fidelidade de imagem similar ou superior. O método também permitiu lotes maiores, reduzindo efetivamente a latência de geração em comparação com a inferência de cache completo.

Significado e Alegações

O artigo afirma que o HeatKV representa um avanço significativo na viabilização de modelos VAR para implantação em hardware com capacidade de memória limitada, como GPUs de nível consumidor, ou para atender a um número maior de usuários simultâneos em infraestrutura existente.

Os autores enfatizam que seu método não requer alterações na arquitetura do modelo ou no processo de treinamento; ele depende de um pequeno conjunto de calibração offline (tão pequeno quanto 1–10 prompts) para gerar um cronograma de poda estático. Os resultados sugerem que a alocação de cache específica por cabeça e de granularidade fina é essencial para preservar a estrutura semântica e a qualidade de geração dos modelos VAR sob restrições agressivas de memória, superando métodos que dependem de classificações binárias grosseiras de cabeças ou camadas de atenção.

O artigo conclui observando que, embora o HeatKV otimize a memória, trabalhos futuros poderiam explorar combiná-lo com outras técnicas de otimização, como exploração de esparsidade de atenção ou quantização de modelos, para aprimorar ainda mais a eficiência.