QKV Projections Require a Fraction of Their Memory

O artigo propõe a técnica PAMM, que comprime as ativações das projeções Q, K e V em camadas de atenção, reduzindo seu consumo de memória em até 512 vezes sem prejudicar a perplexidade final, tornando-se um método complementar e eficaz para o treinamento eficiente de LLMs.

O essencial

Imagine que você está organizando uma festa gigante para treinar um "cérebro" digital (uma Inteligência Artificial chamada LLM). Para que esse cérebro aprenda, ele precisa ler milhões de frases e lembrar de tudo o que leu.

O problema é que, para processar essas frases, o computador precisa guardar uma quantidade absurda de "bilhetes de memória" (chamados de ativations) na sua memória RAM. É como se, para cada frase que o cérebro lê, ele tivesse que escrever cada palavra em um post-it e colar na parede. Com frases longas e muitas pessoas (frases) na festa, a parede enche, a memória acaba e o computador trava.

A maioria das pesquisas tenta encurtar as frases ou apagar post-its que parecem inúteis. Mas os autores deste paper olharam para o processo e perceberam algo curioso: muitos desses post-its são quase idênticos.

A Grande Descoberta: A Festa Repetitiva

Imagine que na sua festa, 500 pessoas estão falando sobre "café". Elas usam palavras diferentes, mas a ideia é a mesma. Em vez de escrever 500 post-its diferentes, você poderia escrever apenas um post-it com a ideia principal de "café" e, para as outras 499 pessoas, apenas um pequeno bilhete dizendo: "Você é igual ao post-it de café, só que um pouquinho mais forte" ou "Você é igual, só que mais fraco".

Isso é o que o PAMM (Multiplicação de Matriz Aproximada por Pontos) faz.

Como o PAMM Funciona (A Analogia do "Chefe e os Estagiários")

O método funciona em duas etapas simples:

1. Escolha dos "Chefes" (Geradores): O computador olha para todas as frases (os dados) e escolhe aleatoriamente um pequeno grupo de "representantes" ou "chefes". Digamos que ele escolha apenas 10 pessoas de um grupo de 5.000.
2. Os "Estagiários" (Aproximação): Para as outras 4.990 pessoas, o sistema não guarda o post-it completo. Ele guarda apenas:
   • Quem é o "chefe" que essa pessoa se parece mais (ex: "Pessoa 42 se parece com o Chefe 3").
   • Um multiplicador (ex: "Pessoa 42 é 1,5 vezes mais intensa que o Chefe 3").

O Resultado? Em vez de guardar 5.000 post-its gigantes, você guarda 10 post-its completos e 5.000 pequenos bilhetes de anotação.

Por que isso é incrível?

• Economia Extrema: O paper mostra que você pode reduzir o uso de memória em 512 vezes. É como transformar uma biblioteca inteira de livros em uma única caixa de sapatos cheia de resumos.
• Sem Perda de Qualidade: O mais surpreendente é que, mesmo com essa compressão absurda, o "cérebro" da IA aprende exatamente da mesma forma (ou até melhor!). Isso acontece porque, na verdade, a maioria das frases que a IA lê é repetitiva e redundante. O sistema estava gastando energia guardando coisas que já sabia.
• Funciona com Tudo: O PAMM é como um "plug-and-play". Você pode usá-lo junto com outras tecnologias modernas de IA (como FlashAttention) sem quebrar nada.

A Metáfora Final: O Mapa da Cidade

Pense no treinamento da IA como um turista tentando memorizar o mapa de uma cidade enorme.

• O método antigo: O turista tenta desenhar cada rua, cada poste de luz e cada árvore em um mapa gigante. O mapa fica tão pesado que ele não consegue carregá-lo.
• O método PAMM: O turista percebe que a cidade tem bairros inteiros que são idênticos (todos têm uma padaria, uma praça e um posto de gasolina). Então, ele desenha apenas um bairro (o "Chefe") e faz uma lista: "O Bairro 2 é igual ao Bairro 1, só que virado para a esquerda. O Bairro 3 é igual ao Bairro 1, mas com casas azuis."

Ele consegue carregar o mapa no bolso, economiza 99% do espaço e, quando precisa navegar, consegue reconstruir a cidade inteira na mente perfeitamente.

Conclusão

Os autores criaram uma técnica que "limpa a mesa" da memória do computador durante o treinamento de IAs. Eles provaram que não precisamos guardar tudo o que a IA vê; basta guardar os "exemplos principais" e dizer aos outros dados como se relacionar com eles. Isso permite treinar IAs maiores, mais rápido e com computadores mais baratos, sem perder inteligência.

Resumo técnico

Título: QKV Projections Require a Fraction of Their Memory

Autores: Malik Khalaf, Yara Shamshoum, Nitzan Hodos, Yuval Sieradzki, Assaf Schuster (Technion, Israel).

1. O Problema

O mecanismo de Multi-Head Attention é central para o funcionamento dos Grandes Modelos de Linguagem (LLMs). Embora a pesquisa recente tenha focado intensamente na eficiência computacional e de memória da operação de scaled dot-product attention (como FlashAttention), a memória consumida pelas projeções lineares que geram os tensores $Q$ (Query), $K$ (Key) e $V$ (Value) a partir da entrada $x$ foi frequentemente negligenciada.

• O Gargalo: Durante o treinamento, as ativações de entrada ($x$) para as camadas de projeção Q, K e V devem ser salvas no forward pass para serem usadas no backward pass (cálculo dos gradientes).
• Impacto: Essas ativações acumulam-se rapidamente, consumindo até 20% da memória de pico da GPU necessária para os blocos de atenção.
• Limitação das Soluções Atuais: Técnicas recentes de compressão focam na dimensão oculta (rank baixo) ou no estado do otimizador. No entanto, os tensores de ativação em modelos Transformer exibem redundância significativa na dimensão da sequência (tokens repetidos, padding, similaridade contextual), que não é explorada por métodos existentes.

2. Metodologia: PAMM (Point-Approximate Matrix Multiplication)

Os autores propõem o PAMM, uma técnica de compressão de tensores que reduz drasticamente a memória das ativações Q, K e V sem degradar o desempenho do modelo.

Conceito Central

O PAMM trata as linhas do tensor de entrada $X$ (onde cada linha é um token) como pontos em um espaço de alta dimensão. Em vez de armazenar todas as $b$ linhas (onde $b$ é o tamanho do lote total), o método:

1. Seleciona um pequeno subconjunto de pontos geradores ($C \in \mathbb{R}^{k \times n}$), onde $k \ll b$.
2. Aproxima as linhas omitidas como projeções escalares desses geradores.

Algoritmo em Duas Etapas

1. Compressão (Forward Pass):

   • Amostra-se aleatoriamente $k$ linhas de $X$ para formar a matriz de geradores $C$.
   • Para cada linha $A_i$ de $X$, encontra-se o gerador $C_j$ que maximiza a similaridade cosseno (o ponto mais próximo na linha spanned pelo gerador).
   • Armazena-se apenas:
     • A matriz $C$ (os geradores).
     • Um vetor de índices $f$ (qual gerador representa qual linha).
     • Um vetor de escalares $\alpha$ (o fator de escala para a projeção).
   • Otimização: É introduzido um parâmetro de tolerância $\epsilon$. Se a melhor aproximação não estiver dentro de uma vizinhança $\epsilon$, a linha é descartada (representada pelo vetor zero), o que é compensado por um fator de correção $\beta$ no cálculo do gradiente.

2. Multiplicação Aproximada (Backward Pass):

   • Para calcular o gradiente $\nabla W = X^\top \nabla Z$, o PAMM não reconstrói $X$.
   • Em vez disso, ele contrai o gradiente $\nabla Z$ usando os índices e escalares para criar uma matriz intermediária $\tilde{B} \in \mathbb{R}^{k \times m}$.
   • O produto final é calculado como $\tilde{O} = C^\top \tilde{B}$, que é muito mais barato computacionalmente do que a multiplicação original $X^\top \nabla Z$.

Garantias Teóricas

• O artigo prova que, sob amostragem uniforme, o número de geradores $k$ necessário para cobrir a distribuição de dados cresce apenas logaritmicamente com o tamanho do lote ($b$), devido à alta densidade de redundância nos dados de linguagem.
• Isso permite compressões extremas (ex: $k = b/512$) mantendo a qualidade.

3. Contribuições Chave

1. Identificação de Redundância: Demonstra que a redundância na dimensão da sequência (tokens) é uma fonte de compressão muito mais significativa do que a redundância na dimensão oculta para as projeções de atenção.
2. Técnica PAMM: Apresenta um método simples, eficaz e composável que reduz a memória das ativações Q, K, V em até 512x.
3. Compatibilidade: O PAMM é totalmente compatível com técnicas de eficiência existentes, como FlashAttention, Gradient Checkpointing e LoRA (Low-Rank Adaptation), sem modificar os pesos do modelo ou o comportamento de inferência.
4. Simplicidade: Diferente de clustering complexo (que é $O(n^2)$), o PAMM usa amostragem aleatória simples e projeções lineares, tornando-o computacionalmente viável.

4. Resultados Experimentais

Os autores avaliaram o PAMM em pré-treinamento e ajuste fino (fine-tuning) em várias arquiteturas (LLaMA, RoBERTa, Pixtral).

• Redução de Memória:
  • Redução de >97% na memória de pico das ativações Q, K, V.
  • Em modelos LLaMA-1B e 7B, a memória das ativações caiu de GBs para MBs (ex: de 3GB para 24MB no LLaMA-1B com $r=1/512$).
• Desempenho (Perplexidade/Acurácia):
  • Pré-treinamento: Em modelos LLaMA (60M a 7B), o PAMM manteve a perplexidade igual ou até melhorou ligeiramente em comparação ao baseline (sem compressão), sugerindo que a remoção de linhas redundantes pode até ajudar a estabilidade do treinamento.
  • Ajuste Fino (GLUE): No RoBERTa-base, o PAMM manteve a performance competitiva em todas as tarefas do benchmark GLUE, mesmo com compressão de 1/256.
  • Multimodal: Testado no Pixtral-12B (VLM) com LoRA, mantendo a acurácia F1 enquanto economizava >97% da memória de ativação.
• Throughput (Velocidade):
  • A sobrecarga computacional é mínima. Para modelos maiores (1B+), a degradação no throughput é inferior a 2.7%.
  • O custo adicional é negligenciável no contexto do treinamento completo do modelo.
• Comparação com Outros Métodos:
  • O PAMM superou significativamente métodos como CompAct (projeção aleatória Gaussiana) e Uniform-CRS (amostragem simples), que sofreram degradação severa de perplexidade em taxas de compressão altas.
  • A configuração $\epsilon = \infty$ (sem restrição de vizinhança, permitindo que todos os pontos sejam representados) mostrou-se a melhor opção, indicando que a estrutura de clusters nos dados de atenção é robusta.

5. Significado e Conclusão

O trabalho demonstra que a memória consumida pelas projeções lineares de atenção, muitas vezes ignorada, é um alvo viável para compressão extrema. O PAMM oferece uma solução prática que:

• Permite treinar modelos maiores ou com batches maiores em hardware limitado.
• É "plug-and-play", podendo ser integrado a pipelines de treinamento modernos sem reescrever a lógica de atenção.
• Sugere que a redundância na dimensão da sequência é uma propriedade fundamental dos LLMs que pode ser explorada para eficiência, desafiando a noção de que todas as ativações de tokens são igualmente informativas durante o treinamento.

Em suma, o PAMM "apaga" a pegada de memória das projeções Q, K, V, permitindo um treinamento de LLMs mais eficiente sem sacrificar a qualidade do modelo final.