Resumo Técnico: KV-Fold

Declaração do Problema

Transformadores pré-treinados de grande porte são limitados pelo custo quadrático de memória da atenção completa e pela janela de contexto fixa de uma única passagem direta. Embora modelos recentes suportem janelas de contexto de até centenas de milhares de tokens, cargas de trabalho do mundo real (por exemplo, código em escala de repositório, registros médicos longitudinais) frequentemente excedem esses limites. As soluções existentes enfrentam um trade-off persistente:

• Métodos de streaming (por exemplo, StreamingLLM) mantêm memória limitada retendo apenas tokens recentes e "sumidouros de atenção", mas perdem a capacidade de recuperar informações assim que elas caem fora da janela deslizante.
• Compressão de cache KV e métodos de memória recorrente frequentemente exigem modificação do modelo, ajuste fino ou introduzem aproximações com perdas que degradam a fidelidade da recuperação.
• Atenção completa é computacionalmente inviável para sequências extremamente longas em hardware padrão devido ao tamanho massivo da matriz de pontuação de atenção (por exemplo, ~1 TB para uma sequência de 128K em uma GPU de 40 GB).

O artigo questiona se um transformador pré-treinado congelado pode carregar seu estado interno através de blocos sucessivos para realizar inferência de contexto longo como uma recorrência, sem alterações arquitetônicas ou treinamento.

Metodologia: KV-Fold

Os autores propõem o KV-Fold, um protocolo de inferência sem treinamento que trata o cache Key-Value (KV) como o acumulador em um foldl (dobrar à esquerda) sobre blocos de sequência.

• Mecanismo: Uma sequência longa de comprimento $T$ é dividida em blocos de comprimento $C$. Ao processar o bloco $t$, o modelo realiza uma passagem direta onde:
  1. As consultas ($Q$) são geradas a partir dos tokens do bloco atual.
  2. As chaves ($K$) e os valores ($V$) são formados concatenando o cache KV acumulado de todos os blocos anteriores ($0$a$t-1$) com os estados KV do bloco atual.
  3. Embeddings posicionais (RoPE) são aplicados continuamente através das fronteiras dos blocos para manter o alinhamento absoluto da sequência.
  4. O cache KV ampliado é passado adiante para o bloco $t+1$.
• Estado: Os parâmetros do modelo permanecem congelados. Nenhum token de memória especial ou ajuste fino é introduzido. O cache KV cresce linearmente com o comprimento da sequência, atuando como um registro estruturado e não comprimido do cálculo passado do modelo.
• Analogia: O processo é análogo a um fold à esquerda de programação funcional, onde o cache KV serve como o estado acumulador atualizado por uma função de recorrência de passo único $F_\theta$ em cada fronteira de bloco.

Contribuições Principais

1. Protocolo KV-Fold: Uma mudança operacional simples que converte um transformador padrão em um modelo de contexto longo recorrente ao carregar o cache KV acumulado através dos blocos.
2. Estabilidade da Recorrência: Evidência empírica mostrando que a recorrência induzida não sofre de acumulação de erro ilimitado. Em vez disso, a "deriva" (diferença na Negativa Log-Verossimilhança em comparação com a atenção completa) sobe brevemente e depois satura em um platô plano que persiste através de cadeias profundas.
3. Deriva Estrutural vs. Numérica: O platô é mostrado como uma propriedade estrutural da dinâmica de atenção, e não um artefato numérico. Ele permanece estável sob um aumento de 10.000× na precisão numérica (de bf16 para fp32) e é robusto entre diferentes famílias de modelos (OLMoE, Qwen, Llama) e tamanhos de bloco.
4. Recuperação Exata de Longo Alcance: O método preserva informações factuais exatas ao longo de grandes distâncias. Em benchmarks de "agulha no palheiro", o KV-Fold alcança 100% de recuperação de correspondência exata em contextos de até 128K tokens e profundidades de cadeia de até 511, igualando o desempenho da atenção completa.
5. Trade-offs Práticos: O artigo caracteriza a posição do método no espectro memória-versus-recuperação. Embora exija crescimento linear de memória (diferente do streaming de memória limitada), ele permite recuperação em profundidades arbitrárias onde os métodos de streaming falham, tornando viável executar recuperação equivalente à atenção completa em hardware que não suporta uma única passagem direta monolítica.

Resultados Principais

• Saturação da Deriva: No Qwen2.5-7B-Instruct, a deriva satura na profundidade 7 e permanece plana até a profundidade 63, com uma mudança total de apenas -0,0003 nats entre as profundidades 15 e 60.
• Desempenho de Recuperação:
  • Agulha Única: 100% de precisão de recuperação em 152 tentativas abrangendo contextos de 16K a 128K tokens e profundidades de cadeia de até 511 no Llama-3.1-8B.
  • Múltiplas Agulhas: 99,4% de precisão de recuperação (175/176) ao recuperar múltiplos fatos distintos simultaneamente.
• Viabilidade de Hardware: Em uma GPU A100 de 40 GB, uma única passagem direta de atenção completa para uma sequência de 128K é inviável (requerendo ~1 TB para pontuações de atenção). O KV-Fold processa a mesma sequência em 171 segundos com um uso de memória pico de 35,6 GB, permanecendo dentro dos limites de hardware.
• Comparação com StreamingLLM: Em contexto de 128K, o StreamingLLM mantém memória limitada (16,6 GB) e tempo de execução mais rápido (22s), mas cai para 0% de recuperação assim que a "agulha" sai de sua janela de 1024 tokens. O KV-Fold mantém 100% de recuperação em todas as distâncias testadas, mas requer maior memória (35,6 GB) e tempo de execução mais longo.

Significado e Alegações

O artigo afirma que transformadores pré-treinados congelados já suportam uma forma estável de recorrência de cache KV. O significado principal é demonstrar que a inferência de contexto longo não exige estritamente alterações arquitetônicas, treinamento ou compressão com perdas para alcançar recuperação de alta fidelidade.

Os autores argumentam que o KV-Fold revela um "regime de atenção ligeiramente deslocado, mas estável", onde o modelo pode processar informações recursivamente sem degradar sua capacidade de acessar tokens passados específicos. Isso fornece uma rota prática para estender janelas de contexto em hardware existente, trocando o crescimento linear de memória pela capacidade de realizar recuperação exata sobre sequências que, de outra forma, seriam truncadas ou aproximadas. O artigo conclui que, embora o método não reduza os FLOPs totais ou a memória abaixo da atenção completa, ele reorganiza o cálculo em etapas tratáveis, tornando a inferência de contexto longo viável onde anteriormente era impossível. Trabalhos futuros sugeridos incluem ajuste fino modesto para reduzir a deriva residual e combinar o método com técnicas de extensão de posição para ir além dos limites de contexto nativos.