Resumo Técnico: Compressão de Chave-Valor Consciente de Cabeças para Geração Autoregressiva Eficiente de Imagens

Declaração do Problema

Modelos de geração visual autoregressiva (AR), que representam imagens como sequências de tokens visuais, alcançaram resultados de alta resolução e fotorrealismo ao escalar o tamanho do modelo e a contagem de tokens. No entanto, essa escalada introduz uma sobrecarga substancial de memória e baixa vazão de inferência devido ao cache de Chave-Valor (KV) necessário durante a decodificação autoregressiva. Embora trabalhos recentes em processamento de linguagem natural (PLN) tenham explorado o gerenciamento de cache KV (por exemplo, seleção de tokens, evicção e quantização), aplicar diretamente essas técnicas ao domínio visual é desafiador devido à estrutura espacial inerente das imagens.

Métodos específicos para visão existentes, como o LineAR, tentam reduzir o uso de cache retendo apenas um subconjunto de tokens com base na observação de que os tokens atendem principalmente a vizinhos espacialmente adjacentes. No entanto, essas abordagens alocam um orçamento fixo e uniforme para cada cabeça de atenção. Essa estratégia ignora a heterogeneidade entre as cabeças de atenção, levando a uma alocação de memória subótima. Especificamente, algumas cabeças focam em vizinhanças locais (exigindo orçamentos pequenos), enquanto outras capturam dependências de longo alcance, como estruturas de objetos e regiões de fundo (exigindo orçamentos maiores). A alocação uniforme falha em adaptar-se a esses padrões distintos, resultando em memória desperdiçada ou qualidade de geração degradada.

Metodologia: HeadKV

Os autores propõem o HeadKV, um novo framework sem treinamento que realiza compressão de cache KV consciente de cabeças. A ideia central é que as cabeças de atenção em modelos de geração de imagens AR exibem dois padrões distintos—Local e Global—e que o comportamento de uma cabeça é consistente entre as posições dos tokens dentro da mesma camada para uma dada entrada.

O pipeline do HeadKV consiste em três etapas principais:

1. Pré-preenchimento do Cache: O modelo codifica entradas condicionais e gera uma sequência inicial de tokens visuais sem compressão. Esta etapa acumula contexto suficiente para estimar distribuições de atenção por cabeça.
2. Agrupamento de Cabeças (Classificação): Uma vez que o cache atinge um comprimento pré-definido, o framework classifica cada cabeça de atenção como "Local" ou "Global" usando um critério de acumulação de atenção baseado em limiar.
   • Para cada cabeça, as pontuações de atenção são acumuladas para trás a partir da posição do token atual até que um limiar ($\tau = 0,9$) seja atingido.
   • Se o número de tokens necessários para atingir o limiar for menor que um tamanho de janela local pré-definido, a cabeça é classificada como Local. Caso contrário, é classificada como Global.
   • Essa classificação é realizada uma única vez durante a etapa inicial de geração e reutilizada para todo o processo de decodificação, não exigindo treinamento adicional ou estatísticas em nível de conjunto de dados.
3. Compressão de Cache Consciente de Cabeças: Diferentes estratégias de compressão são aplicadas com base no tipo de cabeça:
   • Cabeças Locais: Essas cabeças atendem principalmente ao contexto próximo. O framework atribui um orçamento compacto, retendo apenas os tokens mais recentes por meio de um mecanismo de janela deslizante. Tokens históricos são omitidos completamente.
   • Cabeças Globais: Essas cabeças capturam dependências de longo alcance. Para preservar tokens distantes informativos sem incorrer em crescimento linear de memória, os autores projetam uma estratégia de Evicção de Tokens Estratificada (STE).
     • Evicção de Tokens Estratificada: Tokens históricos candidatos são particionados em dois subconjuntos disjuntos: tokens de Curto Alcance e Longo Alcance, com base em sua distância da posição atual.
     • Pontuações de importância são computadas independentemente para cada subconjunto usando características de consulta passadas.
     • Tokens com as pontuações mais baixas são evictos de cada subconjunto (por exemplo, $M/2$ de curto alcance e $M/2$ de longo alcance). Isso impede que tokens localmente dominantes ofusquem tokens distantes semanticamente relevantes, garantindo uma preservação equilibrada de detalhes locais e estrutura global.

Contribuições Principais

1. Análise da Heterogeneidade de Atenção: O artigo fornece uma análise aprofundada de modelos de geração de imagens AR, identificando dois padrões distintos de atenção (Local e Global) e revelando que os comportamentos das cabeças são invariantes entre as posições dos tokens dentro de uma camada para uma entrada fixa.
2. Framework HeadKV: Um framework de compressão novo e sem treinamento que aloca orçamentos assimétricos entre as cabeças de atenção. Atribui orçamentos compactos a cabeças locais e orçamentos maiores a cabeças globais, guiado por um mecanismo de classificação em estágio inicial.
3. Evicção de Tokens Estratificada (STE): Uma estratégia específica para cabeças globais que particiona tokens por distância para garantir a retenção de dependências de longo alcance informativas, abordando a limitação da pontuação de importância padrão que frequentemente favorece tokens locais.
4. Validação Empírica: Experimentos extensivos demonstram que o HeadKV reduz significativamente o uso de memória do cache KV e melhora a vazão de inferência em múltiplos modelos (Janus-Pro, Lumina-mGPT, LlamaGen-XL) enquanto mantém alta qualidade de imagem.

Resultados

Os autores avaliaram o HeadKV em quatro backbones de geração de imagens autoregressiva: Janus-Pro-1B, Janus-Pro-7B, Lumina-mGPT-768 e LlamaGen-XL.

• Qualidade de Imagem: O HeadKV alcança desempenho comparável à linha de base de cache completo em benchmarks como GenEval e DPG. Por exemplo, no Janus-Pro-7B com uma taxa de compressão ($\rho$) de 1/8, o HeadKV mantém uma pontuação GenEval de 0,99 (vs. 0,99 para Completo) e uma pontuação DPG de 83,78 (vs. 84,20 para Completo).
• Comparação com LineAR: Comparado ao método de alocação uniforme LineAR, o HeadKV mostra menos degradação de desempenho à medida que as taxas de compressão aumentam (por exemplo, em $\rho=1/8$). No LlamaGen-XL, o HeadKV alcança uma pontuação FID-50K melhor (2,75) em comparação com o LineAR (2,95) sob compressão extrema.
• Eficiência: O HeadKV reduz significativamente o uso de memória e melhora a vazão.
  • No LlamaGen-XL ($\rho=1/6$), o HeadKV reduz a memória em 63,60% e alcança um aceleração de 3,00× em comparação com a configuração completa, superando o LineAR (2,93×).
  • No Janus-Pro-7B ($\rho=1/6$), alcança uma aceleração de 3,68× com uma economia de memória de 52,08%.
  • No Lumina-mGPT-768, embora a latência seja ligeiramente maior que a do LineAR devido a sobrecargas de implementação no agrupamento, ainda alcança uma aceleração de 1,35× com economias significativas de memória.

Significado e Alegações

O artigo afirma que o HeadKV oferece uma solução prática para o gargalo de memória na geração de imagens autoregressiva sem exigir treinamento adicional ou supervisão externa. Ao reconhecer e explorar a heterogeneidade inerente das cabeças de atenção, o método vai além da alocação de cache "tamanho único" de trabalhos anteriores.

Os autores enfatizam que sua abordagem é generalizável entre diferentes entradas e arquiteturas de modelo porque depende de padrões de atenção intrínsecos observados durante as etapas iniciais da geração. A estratégia de Evicção de Tokens Estratificada é destacada como um componente crítico para manter a qualidade das imagens geradas sob orçamentos de memória apertados, especificamente preservando informações estruturais de longo alcance que são frequentemente perdidas em métodos de evicção padrão. O trabalho sugere que a alocação de recursos adaptativa e específica por cabeça é um paradigma mais eficiente para geração visual do que a retenção uniforme de tokens.