🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

Este artigo introduz o Looped Latent Attention (LLA), um codec de pós-treinamento que explora a estrutura de baixo posto dos caches KV indexados por loop em Transformers com pesos compartilhados para alcançar compressão extrema (até 32x) e aumentar significativamente a capacidade de lote, mantendo um desempenho quase sem perdas através de inicialização por SVD e destilação.

Autores originais: James O' Neill, Fergal Reid

Publicado 2026-07-20

📖 6 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: James O' Neill, Fergal Reid

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

O Problema da Memória nos Cérebros de IA

Imagine que você está tentando ensinar um robô a escrever uma história. O robô tem um cérebro feito de milhões de pequenos interruptores (parâmetros) que sabem como as palavras se encaixam. Mas aqui está o detalhe: para escrever uma história longa, o robô não pode apenas olhar para o seu cérebro; ele precisa de um bloco de notas. Cada vez que ele escreve uma nova palavra, ele tem que se lembrar de tudo o que acabou de escrever para não se repetir ou perder o fio da meada. No mundo da Inteligência Artificial, esse bloco de notas é chamado de KV cache (cache de Chave-Valor). É um caderno privado para cada conversa individual, e ele cresce a cada palavra que o robô fala.

Agora, imagine um tipo especial de robô projetado para ser super eficiente. Em vez de ter um cérebro enorme com muitas salas diferentes (camadas) para diferentes tarefas, este robô tem uma sala pequena e inteligente que ele visita repetidamente. Ele reutiliza o mesmo conjunto de instruções, fazendo loops através delas para pensar mais profundamente. Isso é chamado de Looped Transformer (Transformer em Loop). Isso economiza espaço no próprio cérebro, o que é ótimo. Mas há um problema sorrateiro: embora o robô reutilize a mesma sala, ele ainda escreve uma nova página em seu bloco de notas toda vez que faz o loop. Se o robô fizer quatro loops para resolver um problema de matemática, ele acaba com quatro páginas separadas de notas para o mesmo momento da história. Isso significa que o robô "em loop" ainda precisa de uma quantidade massiva de memória para funcionar, derrotando o propósito de ser eficiente. A grande questão para os cientistas é: podemos encolher esse bloco de notas sem fazer o robô esquecer como pensar?

A Descoberta do Artigo: Um Atalho Secreto nas Notas

Este artigo apresenta um truque inteligente chamado Looped Latent Attention (LLA) para resolver esse problema de memória. Os pesquisadores descobriram que as notas do robô não são tão bagunçadas quanto parecem. Quando o robô percorre seu processo de pensamento em loop, as notas que ele escreve para a mesma palavra não mudam drasticamente; elas seguem um caminho suave e previsível, como uma bola de gude rolando por uma colina suave. Em vez de salvar cada página individual do bloco de notas do robô, o novo método salva apenas um pequeno "resumo" desse caminho e um conjunto de instruções sobre como reconstruir as páginas específicas apenas quando o robô realmente precisar lê-las.

Pense assim: Imagine que você está assistindo a um filme onde o personagem principal caminha por um corredor. Em um computador normal, você salvaria uma foto completa, de alta definição, do personagem em cada passo que ele dá. Isso ocupa um enorme espaço de armazenamento. Mas com o LLA, você percebe que o personagem está apenas caminhando em linha reta. Então, em vez de salvar milhares de fotos, você salva uma foto do personagem e uma nota simples dizendo: "Mova 1 polegada para frente no passo 2, 2 polegadas no passo 3". Quando o computador precisa ver o personagem no passo 3, ele desenha rapidamente aquela imagem usando a nota. É exatamente isso que o artigo faz pela IA: ele comprime a parte do "loop" da memória, transformando uma pilha de quatro páginas em um único resumo minúsculo que pode ser expandido de volta para as páginas completas instantaneamente.

Os pesquisadores testaram isso em vários modelos de IA, incluindo um chamado Ouro-1.4B. Eles descobriram que este método funciona incrivelmente bem. Na verdade, em um chip de computador poderoso chamado H200, eles conseguiram encaixar 768 conversas diferentes na memória de uma só vez, enquanto antes conseguiam encaixar apenas 32. Esse é um salto massivo na capacidade de quantas pessoas podem usar a IA ao mesmo tempo. Eles também mostraram que esse truque funciona mesmo quando a IA está resolvendo problemas matemáticos muito longos e difíceis, embora precise de um pouco de treinamento extra para garantir que não se confunda ao escrever sua própria história em vez de apenas copiar um professor.

O Que Este Artigo Diz "Não"

É importante saber o que este método não é. Os pesquisadores testaram uma ideia muito simples primeiro: "E se jogarmos fora as três primeiras páginas de notas e mantivermos apenas a última?". Eles pensaram que talvez o robô se estabilize em uma resposta final após alguns loops, então as notas anteriores não importariam. Eles tentaram isso, e foi um desastre. A IA esqueceu completamente como fazer matemática, pontuando zero nos testes. Isso prova que o "caminho" que as notas percorrem é importante; você não pode simplesmente pegar o ponto final. Você precisa capturar a jornada.

Eles também compararam seu método com outras formas de reduzir a memória, como esmagar as notas juntas (quantização) ou compartilhar notas entre diferentes partes do cérebro (compressão de eixo de cabeça). Embora esses métodos ajudem um pouco, o artigo mostra que comprimir o eixo do "loop" é o ingrediente secreto. Se você tentar comprimir as outras partes em vez dos loops, a IA perde sua capacidade de raciocínio. O artigo é muito claro: o loop é a parte mais redundante da memória, e é o único lugar onde você deve apertar.

O Quão Certo Eles Estão?

Os autores estão bastante confiantes nesses resultados porque os mediram diretamente. Eles não apenas adivinharam; eles rodaram a IA em tarefas reais, como resolver problemas de matemática (GSM8K) e escrever código (HumanEval). Eles mostraram que, com uma compressão de 4x (tornando a memória 4 vezes menor), a IA ainda performa quase tão bem quanto a versão original, não comprimida. Quando levaram ao extremo da compressão (21.3x), a IA ainda conseguia lidar com muitas tarefas, embora começasse a ter dificuldades com os passos de raciocínio mais longos e complexos.

Eles também testaram isso em um tipo diferente de modelo de IA chamado Huginn-3.5B, que possui uma estrutura de cérebro diferente e faz loops de até 32 vezes. Mesmo lá, o método funcionou, provando que isso não é apenas uma coincidência para um robô específico. O artigo sugere que essa "trajetória de baixo posto" (o caminho suave das notas) é uma propriedade fundamental de como esses robôs em loop pensam. Embora tenham provado que funciona para economizar memória e aumentar a capacidade, eles observam que, para o melhor desempenho em tarefas muito longas, a IA precisa de uma segunda rodada de treinamento, onde ela pratica com suas próprias notas comprimidas, em vez de apenas copiar um professor. Este treinamento "on-policy" corrige os pequenos erros que acontecem quando o robô começa a escrever sua própria história.

Em resumo, este artigo mostra que podemos tornar os modelos de IA em loop muito mais eficientes ao perceber que sua memória é repetitiva de uma forma específica. Ao armazenar um pequeno resumo dessa repetição em vez das notas completas, podemos encaixar centenas de vezes mais conversas em um único chip de computador, tornando a IA poderosa disponível para mais pessoas sem a necessidade de hardware super caro.

Resumo Técnico: Looped Latent Attention (LLA)

Declaração do Problema

Em Transformers com pesos compartilhados e em loop (Looped, weight-tied Transformers), um único bloco é aplicado repetidamente (recorrência) para reduzir a contagem de parâmetros. Embora isso comprima efetivamente os pesos do modelo, o gargalo de memória na inferência permanece o cache Key/Value (KV). Na inferência em loop padrão, o modelo recomputa a atenção em cada etapa de recorrência, armazenando um cache KV separado para cada etapa, camada e cabeça. Consequentemente, um modelo com $T$ loops sobre $D$ camadas carrega um cache KV equivalente a um decodificador de $T \times D$ camadas, anulando os benefícios de memória da partilha de pesos.

O artigo postula que este cache indexado pelo loop é altamente estruturado. Para um token, camada e cabeça fixos, os vetores K/V através das etapas de recorrência traçam uma trajetória curta e de baixo posto (low-rank), em vez de explorar um espaço de alta dimensão. A redundância é específica ao eixo do loop; os espectros ao longo dos eixos de cabeça e camada permanecem significativamente mais planos (posto mais alto). O desafio é comprimir este eixo de recorrência sem colapsar a trajetória em um único ponto fixo, o que degradaria o desempenho.

Metodologia: Looped Latent Attention (LLA)

Os autores introduzem a Looped Latent Attention (LLA), um codec de cache pós-treinamento que substitui os $T$ vetores por loop K/V por uma representação latente compacta e mapas de reconstrução específicos para cada loop.

Mecanismo Central

Armazenamento Latente: Em vez de armazenar $T$ $T$ vetores ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ ) para um token, camada e cabeça específicos, a LLA armazena um único vetor latente centralizado de baixo posto $c \in \mathbb{R}^r$ $c \in R^{r}$ .
- $c = (x - \mu)W_{\text{down}}$
- Aqui, $\mu$ é a média de um conjunto de calibração, e $W_{\text{down}}$ é uma matriz de projeção descendente (down-projection).
Reconstrução Específica do Loop: O vetor K/V específico para a etapa de recorrência $t$ $t$ é reconstruído usando uma projeção ascendente (up-projection) específica do loop:
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
Ranks Assimétricos: O método utiliza ranks separados para Keys ( $r_k$ ) e Values ( $r_v$ ). Uma análise empírica mostra que os vetores de Value estabilizam mais lentamente entre os loops do que os vetores de Key, motivando $r_v > r_k$ .
Variante LLA-2D: Para compressão extrema, a LLA-2D dobra o eixo da cabeça (head axis) no latente, armazenando um único latente por token e camada para reconstruir todas as $H$ cabeças. Isso reduz o tamanho do cache por um fator de $H$ , mas acopla as cabeças.

Treinamento e Inicialização

Inicialização via SVD: O codec é inicializado via Decomposição em Valores Singulares (SVD) das ativações do professor em um conjunto de calibração. A projeção descendente é definida pelos $r$ vetores singulares à direita de maior valor. Esta inicialização é crítica; a inicialização aleatória resulta em uma divergência KL significativamente maior.
Destilação: O codec é refinado congelando o modelo professor e minimizando a divergência KL frontal entre a distribuição do professor e a distribuição do modelo com o codec aplicado, aumentada por um termo de correspondência de saída de atenção.
Refinamento On-Policy: Para abordar o viés de exposição (exposure bias) em rollouts autorregressivos longos (onde o modelo condiciona em seus próprios prefixos gerados em vez dos do professor), uma segunda etapa de destilação on-policy é empregada. O codec é atualizado usando prefixos amostrados do próprio modelo estudante, estabilizando o desempenho de geração longa.

Principais Contribuições

Definição de LLA: Um codec de K/V latente trans-loop para Transformers com pesos compartilhados que alcança uma razão de compressão $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ . Ele armazena $(r_k + r_v)H$ escalares por token por camada em vez de $2Td_{\text{head}}$ .
Compressão Específica de Eixo: O artigo demonstra que o eixo de recorrência (loop) é a dimensão mais compressível em modelos em loop (Ouro-1.4B, Ouro-2.6B-Thinking e Huginn-3.5B). Comprimir o eixo do loop supera a compressão do eixo da cabeça (MLA), do eixo da camada (compartilhamento entre camadas) ou a aplicação de quantização com orçamentos de cache equivalentes.
Rejeição do Colapso de Ponto Fixo: Os autores mostram que a trajetória do loop não pode ser substituída por um único estado (o último loop). Um controle de zero parâmetros que reutiliza apenas o K/V do último loop colapsa o desempenho de geração (a precisão no GSM8K cai para 0.000), provando que a "trajetória de baixo posto curta" é necessária.
Impacto no Deployment: A redução exata de escalares traduz-se diretamente no aumento da capacidade de serviço. Em uma única GPU H200, a LLA aumenta o batch size máximo para o Ouro-1.4B com contexto de 4k de 32 para 768 sequências (com 21.3× compressão).

Resultados Experimentais

Compressão vs. Precisão

Desempenho no GSM8K: No Ouro-1.4B, a compressão leve (1.33× a 2×) iguala a precisão do professor (aprox. 0.79–0.80). Com 4× compressão, a LLA mantém 0.59 de precisão, superando significativamente a compressão do eixo da cabeça (0.52) e o reuso do último loop (0.00).
Generalização de Escala: No Ouro-2.6B-Thinking, a LLA permanece quase sem perdas (0.81–0.83) entre 4× e 10.67× de compressão, enquanto as bases de eixo da cabeça perdem precisão substancial.
Transferência para Huginn: O método transfere-se para o Huginn-3.5B (32 etapas de recorrência). Um codec SVD livre de treinamento alcança fidelidade quase sem perdas até 32× em métricas independentes do decodificador.

Estabilidade de Geração Longa

MATH-500: Rollouts matemáticos longos são sensíveis ao desvio de prefixo (prefix drift). O treinamento off-policy (usando prefixos do professor) leva à degradação da precisão e altas taxas de "sem resposta" conforme o comprimento da geração aumenta.
Ganho On-Policy: Aplicar o refinamento on-policy eleva a precisão do MATH-500 de 0.43 para 0.66 com 4× compressão e reduz significativamente as gerações sem resposta.

Capacidade de Serviço

Redução de Memória: A redução do cache é exata. Com 21.3× de compressão, o cache KV para o Ouro-1.4B com contexto de 1M cai de 768 GiB para 36 GiB.
Throughput: O "caminho de latente de armazenamento" (armazenando apenas latentes) é limitado por computação (compute-bound), mas permite batches e contextos anteriormente impossíveis devido a OOM (Out of Memory). O "caminho rápido" (reconstruindo e armazenando K/V completo) preserva o throughput, mas não realiza a economia de memória.

Significância e Alegações

O artigo afirma que a recorrência é o eixo de cache mais compressível em Transformers com pesos compartilhados. A significância reside em:

Insight Mecanístico: Identifica que o eixo do loop é espectralmente distinto (baixo posto) dos eixos de cabeça e camada, permitindo uma compressão eficiente sem colapsar o comportamento dinâmico do modelo em um ponto fixo estático.
Deployment Prático: Fornece uma solução pós-treinamento que aumenta diretamente a capacidade de batch size e comprimento de contexto em hardware existente sem retreinar a arquitetura base do modelo.
Direção Futura: Os resultados sugerem que arquiteturas recorrentes nativas poderiam se beneficiar de caches latentes desacoplados, onde a recorrência computa conteúdo compartilhado mais deltas baratos por loop, em vez de rematerializar estados K/V completos.

Os autores enfatizam que o método não depende de implementações de atenção absorvida para suas alegações de precisão; o resultado central é a validade da suposição de trajetória de baixo posto e a eficácia do codec em preservar o comportamento do professor sob restrições rigorosas de memória.