🤖 machine learning

Regularize or Localize: When Training-Time KV-Cache Geometry Pays Under Quantization

Este artigo demonstra que, embora a aplicação do objetivo anti-colapso LeJEPA apenas aos estados ocultos falhe em melhorar a quantização do KV-cache, regularizar diretamente a geometria do KV-cache durante o treinamento melhora significamente o desempenho sob esquemas de quantização grosseiros por canal, embora essa vantagem diminua quando configurações de quantização mais sofisticadas como KIVI são empregadas.

Autores originais: Libo Sun, Po-Wei Harn, Zewei Zhang, Peixiong He, Xiao Qin

Publicado 2026-07-21

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Libo Sun, Po-Wei Harn, Zewei Zhang, Peixiong He, Xiao Qin

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você esteja tentando ensinar um robô a escrever histórias. Para fazer isso, você lhe dá uma biblioteca massiva de livros e o deixa aprender prevendo a próxima palavra em uma frase, repetidamente. É assim que os modelos de linguagem de IA modernos funcionam. Mas há um detalhe: conforme o robô lê, ele constrói uma "memória" da história até o momento, chamada KV-cache (cache de Chave-Valor). Pense neste cache como uma pilha de post-its que o robô mantém em sua mesa. Cada vez que ele escreve uma nova palavra, ele adiciona um bilhete. Se o robô tiver que lembrar de uma história longa, essa pilha fica enorme, e armazená-la consome muita memória do computador.

Para economizar espaço, os engenheiros tentam encolher esses post-its usando um processo chamado quantização. É como pegar uma foto de alta definição e comprimi-la em uma miniatura minúscula. Geralmente, isso funciona bem, mas às vezes os "bilhetes" na memória do robô são escritos de uma forma estranha e desequilibrada. Eles se agrupam todos em uma direção, como uma pilha de lápis todos inclinados para o mesmo lado. Isso os torna difíceis de comprimir sem perder detalhes importantes. A grande questão para os cientistas da computação é: Podemos ensinar o robô a escrever seus bilhetes de uma forma mais equilibrada e "redonda" desde o início, para que sejam mais fáceis de encolher depois? Este artigo investiga exatamente isso, usando um truque de treinamento especial para ver se ele altera a forma da memória do robô e ajuda a encaixá-la em uma caixa menor.

O Experimento: Moldando a Memória do Robô

Os pesquisadores deste artigo decidiram testar um truque de treinamento específico chamado SIGReg. Imagine que o cérebro do robô é uma fábrica gigante de múltiplas camadas. Normalmente, a fábrica apenas tenta obter a resposta correta (prever a próxima palavra). O SIGReg adiciona uma segunda regra: "Ei, certifique-se de que as formas dos seus pensamentos não estejam todas inclinadas na mesma direção". Ele empurra as representações internas do robô para serem mais como uma nuvem de dados perfeita e redonda do que como uma panqueca plana e desequilibrada.

A equipe treinou um robô de tamanho médio (110 milhões de parâmetros) em um conjunto de dados massivo de 10 bilhões de palavras. Eles testaram três ideias principais, agindo como detetives tentando resolver um mistério sobre de onde vem a forma da memória.

1. O truque funciona nos "pensamentos" do robô?
Sim, funciona. Quando aplicaram o SIGReg aos pensamentos ocultos do robô (as etapas de processamento interno), o "desequilíbrio" desses pensamentos caiu 38%. O robô não ficou muito pior ao escrever histórias; seu índice de confusão (perplexidade) aumentou menos de 0,35%, o que é quase imperceptível. Assim, o truque conseguiu remodelar a geometria interna do robô sem quebrar sua capacidade de falar.

2. O truque conserta os "post-its" (o KV-cache)?
É aqui que a trama muda. Os pesquisadores descobriram que apenas consertar os pensamentos do robô não consertou automaticamente os post-its (o cache). O cache permaneceu desequilibrado, exatamente como antes.

O Conserto Falho: Eles tentaram "adaptar" o robô congelando seu cére de forma a apenas adicionar novas cabeças para ler os bilhetes. Isso não funcionou.
O Conserto Real: Para realmente mudar a forma dos post-its, eles tiveram que manter o céreão do robô destravado e aplicar o truque SIGReg diretamente aos próprios bilhetes enquanto o robô continuava aprendendo. Quando fizeram isso, o desequilíbrio do cache caiu drasticamente 94%.
A Lição: Você não pode apenas consertar o cérebro e esperar que a memória o siga. Você tem que treinar a memória diretamente.

3. Um cache mais redondo gera miniaturas melhores (quantização)?
Esta é a parte mais emocionante. A equipe tentou encolher os post-its em miniaturas de 3 bits (um tamanho muito pequeno) para ver o quanto a qualidade da história sofreria.

A Boa Notícia: Para o robô treinado com o truque direto de cache, o "dano" causado pelo encolhimento dos bilhetes foi de 4,3 a 7,9 vezes menor do que para o robô normal. Na verdade, o robô treinado foi o único que preferiu uma forma específica de medir os bilhetes (escalonamento por canal) que funcionou muito melhor para a compressão.
A Má Notícia (O Porém): Essa enorme vantagem só apareceu quando usaram uma forma de encolhimento simples e "grossa" (coarse). Quando usaram uma configuração mais avançada e realista (chamada de configuração "estilo KIVI", que mistura métodos de escalonamento e adiciona pontos de zero), a vantagem desapareceu. Nesse cenário complexo, o robô treinado e o robô normal tiveram um desempenho quase idêntico.

O Veredito

Então, o que eles aprenderam? O artigo sugere que você pode treinar um modelo de linguagem para ter uma memória (KV-cache) mais "redonda" e equilibrada, aplicando um truque de regularização específico diretamente a essa memória durante o treinamento. Isso torna a memória muito mais fácil de comprimir se você estiver usando métodos de compressão simples e grosseiros.

No entanto, o artigo é muito cuidadoso ao dizer que isso não é uma solução mágica para toda compressão. Se você usar as técnicas de compressão sofisticadas e modernas que os sistemas do mundo real utilizam (com escalonamento misto e pontos de zero), o benefício deste truque de treinamento desaparece. A "redondeza" da memória não ajudou nesses cenários complexos.

Em resumo: os pesquisadores encontraram uma maneira de remodelar a memória do robô, e isso funciona muito bem para tarefas de encolhimento simples. Mas para as tarefas de encolhimento de alta tecnologia usadas na vida real, o treinamento extra não trouxe nenhum ganho extra de velocidade ou de espaço. A vantagem é real, mas tem um limite muito específico: ela só ajuda quando as escalas de compressão são "grossas" e não sobrevive à complexidade da quantização de configuração mista moderna.

Resumo Técnico: Regularizar ou Localizar: Quando a Geometria do Cache K/V em Tempo de Treinamento Paga sob Quantização

Declaração do Problema

Modelos de linguagem Transformer treinados com entropia cruzada (CE) tipicamente desenvolvem uma geometria de estado oculto anisotrópica, onde as representações se concentram ao longo de alguns poucos direcionais compartilhados. Embora as Arquiteturas Preditivas de Embedding Conjunto (JEPA) tenham introduzido objetivos como o SIGReg (um teste de normalidade Epps–Pulley esquematizado) para empurrar os embeddings em direção à isotropia ( $N(0, I)$ ), tentativas anteriores de aplicar o SIGReg a Grandes Modelos de Linguagem (LLMs) produziram resultados nulos. Essas falhas ocorreram quando o SIGReg foi aplicado como um auxiliar a fluxos não padronizados ou apenas aos estados da camada final durante o ajuste fino (fine-tuning).

Uma lacuna crítica permanece: o SIGReg remodela a geometria do treinamento autorregressivo (AR) padrão e, mais importante, essa remodelação beneficia o cache Key/Value (K/V) durante a quantização pós-hoc? O trabalho existente foca em compressão pós-hoc ou mudanças arquiteturais, deixando a questão de se a regularização distributiva em tempo de treinamento da geometria padrão do cache K/V pode melhorar a robustez da quantização amplamente inexplorada.

Metodologia

Os autores empregam um design experimental de três estágios usando modelos do tipo Llama com 110M de parâmetros pré-treinados em 10B de tokens FineWeb.

1. Fase 1: Dose-Resposta e Caracterização do Estado Oculto

Objetivo: Determinar se a adição de SIGReg à perda CE durante o pré-treinamento reduz a anisotropia do estado oculto.
Configuração: Os autores realizam um varredura (sweep) do peso de regularização $\lambda$ ($0, 0.001, 0.01, 0.1, 1.0$).
Métrica: Anisotropia de cosseno pareado e perplexidade.
Descoberta: $\lambda=0.01$ é identificado como o ponto de operação ideal, reduzindo significativamente a anisotropia com custo mínimo de perplexidade.

2. Fase 2: Sondas de Ponte (Adaptação vs. Retrofit)

Para determinar se a mudança de geometria se propaga para o cache K/V e como alcançá-la, os autores testam três sondas específicas partindo de um checkpoint $\lambda=0.01$ :

A1 (Adaptação): Descongelar o tronco e continuar o treinamento com SIGReg aplicado diretamente tanto aos estados ocultos ( $h$ ) quanto às projeções K/V ( $K_\ell, V_\ell$ ) em todas as camadas.
A2 (Predição Congelada): Treinar preditores por camada em um tronco congelado para ver se a geometria K/V é recuperável apenas a partir dos estados ocultos.
A3 (Decodificação Congelada): Treinar decodificadores para reconstruir o K/V a partir do estado oculto final do tronco congelado.
Controle: Uma "continuação correspondente" onde o tronco é continuada pelo mesmo orçamento de tokens, mas sem o termo de regularização K/V ( $\lambda_{KV}=0$ ).

3. Fase 3: Avaliação de Quantização

Os autores avaliam os modelos resultantes sob quantização fake pós-hoc (quantizar $\to$ desquantizar) do cache K/V.

Esquemas Testados: Quantização simétrica livre de grupo (por token vs. por canal), grades assimétricas (com pontos de zero) e arranjos mistos (estilo KIVI: por canal K, por token V).
Métrica: $\Delta$ NLL (Negative Log-Likelihood) em relação à própria referência de precisão total do modelo.

Principais Resultados

1. Geometria do Estado Oculto

Em $\lambda=0.01$ , o SIGReg reduz a anisotropia de cosseno pareado do estado oculto em 38% através de três sementes pareadas. A perplexidade aumenta em menos de 0.35%, e as capacidades de zero-shot (LAMBADA, HellaSwag, etc.) permanecem limpas, sem sinais de degradação consistente.

2. Geometria do Cache K/V e a Falha do "Retrofit"

Não Propagação Automática: A redução da anisotropia nos estados ocultos não se propaga automaticamente para o cache K/V. As sondas de tronco congelado (A2 e A3) falham em recuperar a geometria, indicando que a anisotropia K/V não é meramente uma função do estado oculto final.
Regularização Direta é Necessária: Apenas A1 (treinamento contínuo com regularização direta de K/V) reduziu com sucesso a anisotropia do cache, alcançando uma redução de 94% na anisotropia média do cache.
Isolamento de Mecanismo: A continuação correspondente sem o termo K/V ( $\lambda_{KV}=0$ ) falhou em reproduzir esse efeito, isolando o termo K/V como a causa necessária.

3. Ganho de Quantização

Os benefícios da geometria regularizada dependem fortemente do esquema de quantização:

Escalas Grossas (Livre de Grupo): Sob quantização simétrica livre de grupo não transformada, o modelo adaptado A1 é o único que prefere escalonamento por canal em vez de escalonamento por token. Sob quantização simétrica de 3 bits por canal, o baseline sofre de 4.4–7.9× mais dano de $\Delta$ NLL do que o modelo A1.
Escalas Finas (Estilo KIVI): Sob uma configuração completa simulada do estilo KIVI (arranjo misto, pontos de zero e escalas agrupadas com $g=32$ ), a vantagem desaparece. Todos os modelos, incluindo o baseline, atingem paridade de desempenho.
Sobrecarga de Armazenamento: Mesmo ao considerar a sobrecarga de metadados dos pontos de zero agrupados (tornando o armazenamento agrupado de 3 bits comparável ao de 4 bits livre de grupo), a vantagem do tempo de treinamento desaparece sob a configuração de escalonamento misto e agrupamento local de tokens.

Significância e Alegações

O artigo afirma ser o primeiro a avaliar a regularização distributiva de tempo de treinamento da geometria do cache K/V padrão contra a quantização de cache pós-hoc.

Controle Geométrico: Os autores demonstram que a geometria do cache K/V pode ser explicitamente controlada durante o treinamento via regularização direta, mas esse efeito não surge apenas da regularização do estado oculto e não pode ser "reajustado" (retrofit) via sondas de tronco congelado.
Utilidade Condicional: A intervenção de treinamento fornece um ganho tangível apenas quando os escalonadores de quantização são grosseiros (abrangendo extensões longas de tokens). A vantagem é específica da interação entre a geometria regularizada e o esquema de quantização; ela não se generaliza para configurações modernas e finas (como o estilo KIVI com pontos de zero e agrupamento local).
Interpretação de Métrica: O artigo esclarece que a métrica de anisotropia de cosseno pareado mede o alinhamento de direção média (dominado por médias compartilhadas) em vez de branqueamento de covariância. Os resultados mostram uma dissociação: enquanto o alinhamento da média cai, as propriedades espectrais (como o share top-1) comportam-se de forma diferente para as camadas K e V, sugerindo que o regularizador atua primariamente na média/direções erráticas.

Limitações

Os autores delimitam explicitamente suas alegações:

Regime: Os resultados são limitados a 110M de parâmetros e ao corpus FineWeb; modelos maiores com diferentes layouts de cache (GQA/MQA) não foram testados.
Simulação: A quantização é simulada (fake quantization) sem medições de tempo de execução ou memória.
Escopo da Sonda: Os resultados negativos para sondas congeladas baseiam-se em capacidades e orçamentos específicos; um "meio termo" (ex: adaptação em escala LoRA) permanece não testado.
Não Alegação de "Primeira Isotropia": Os autores não afirmam ser os primeiros a regularizar para isotropia (citando CosReg, SimCTG, etc.), mas sim os primeiros a visar especificamente a geometria do cache K/V para robustez de quantização.