When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?

1. Declaração do Problema

Decaimento de Longo Alcance na Decodificação Especulativa:
A decodificação especulativa acelera a inferência de Modelos de Linguagem de Grande Porte (LLMs) utilizando um modelo "rascunho" (drafter) leve para propor múltiplos tokens, que são verificados por um modelo "alvo" maior. Os modelos rascunho mais avançados atuais (por exemplo, EAGLE-3, MTP) dependem da reutilização de estados ocultos, alimentando os estados ocultos internos do modelo alvo no modelo rascunho.

• O Problema: Esses sistemas sofrem de decaimento de longo alcance: a precisão do rascunho diminui progressivamente à medida que o passo especulativo $k$ aumenta.
• Explicação Existente: A comunidade atribui isso ao descompasso entre treinamento e inferência (o modelo rascunho condiciona-se aos seus próprios estados ocultos em deriva durante a inferência, ao contrário do treinamento). Embora o Treinamento no Tempo de Teste (TTT) mitigue isso, o decaimento persiste.
• A Lacuna: O artigo argumenta que o descompasso entre treinamento e inferência não é a única causa. A questão fundamental reside na preservação de informações de contexto. Os estados ocultos atuam como uma "compressão de contexto enviesada" otimizada para o próximo token imediato, descartando informações relevantes para futuros passos especulativos.

2. Hipótese Central e Análise Teórica

Os autores propõem a Hipótese de Reutilização de KV: Reutilizar o Cache KV (pares Chave-Valor) do modelo alvo, em vez de apenas os estados ocultos, pode preservar um contexto mais rico para a redação de longo alcance.

• Reutilização de Estado Oculto (Compressão Enviesada):
  • Os estados ocultos do alvo são o resultado da agregação de atenção: $h_t = \sum \alpha_i v_i$.
  • Os pesos $\alpha_i$ são otimizados para prever o token atual $x_{t+1}$.
  • Tokens fracamente relevantes para $x_{t+1}$ recebem pesos próximos de zero e são efetivamente descartados em $h_t$.
  • Consequência: Para passos especulativos posteriores ($x_{t+k}$), o modelo rascunho deve recuperar informações que foram suprimidas pelo mecanismo de atenção do alvo, um problema de "recuperação de informação" mal condicionado.
• Reutilização de KV (Contexto Explícito):
  • O cache KV retém o conjunto completo de pares $(k_i, v_i)$ antes da agregação.
  • O modelo rascunho pode realizar reatenção usando suas próprias consultas futuras estimadas contra o cache KV completo do prefixo.
  • Consequência: Nenhuma informação é descartada antes do processamento do modelo rascunho. A dificuldade desloca-se de "recuperar informações perdidas" para "estimar consultas futuras com precisão" (um problema de aproximação de função).

3. Metodologia: O Framework KVShot

Para testar essa hipótese, os autores introduzem o KVShot, um framework de diagnóstico que compara três paradigmas de reutilização sob o mesmo pipeline de Treinamento Autoregressivo no Tempo de Teste (TTT) (utilizado pelo EAGLE-3):

1. Reutilização Apenas de Oculto: Abordagem padrão do EAGLE-3/MTP (concatenando estados ocultos do alvo).
2. Reutilização Apenas de KV: O modelo rascunho recebe o cache KV do alvo via atenção cruzada, substituindo completamente os estados ocultos.
3. Reutilização Híbrida: Um mecanismo de fusão com portão onde o estado oculto fornece a âncora principal, e o cache KV fornece uma "correção delta com portão" para passos de longo alcance.

Configuração Experimental:

• Modelo Alvo: Qwen3-8B.
• Conjuntos de Dados: ShareGPT e UltraChat (escalados de 70k para 280k amostras).
• Métricas: Taxas de aceitação passo a passo ($\alpha_k$), Tokens Médios Aceitos (MAT) e proxies de aceleração de Ponta a Ponta (E2E).

4. Resultados Principais

A. Validação de Previsões (Análise Passo a Passo)

Os experimentos confirmaram três previsões teóricas:

1. Vantagem de Longo Alcance: A reutilização Apenas de KV degrada-se de forma mais graciosa do que a reutilização Apenas de Oculto à medida que $k$ aumenta. No passo $k=6$, um modelo rascunho de 4 camadas apenas de KV superou a linha de base EAGLE-3 ($\alpha_6 = 0,495$ vs. $0,469$).
2. Gargalo de Estimativa de Consulta: O desempenho Apenas de KV depende fortemente da profundidade do modelo rascunho. Um modelo rascunho de 1 camada performou mal (MAT 1,84) porque não consegue aproximar consultas complexas do alvo. Aumentar a profundidade para 2 camadas resultou em um salto massivo (MAT 2,23), mas 4 camadas ofereceram retornos decrescentes, ainda ficando atrás da linha de base EAGLE-3 de 1 camada (MAT 2,37) em eficiência geral.
3. Desvantagem de Curto Alcance: Em $k=0$ ou $1$, a reutilização de Estado Oculto supera a reutilização Apenas de KV porque os estados ocultos carregam conteúdo semântico mais rico do que os embeddings brutos usados para reatenção KV.

B. Desempenho Híbrido

O Modelo Rascunho Híbrido (combinando ambos os sinais) alcançou os melhores resultados passo a passo:

• MAT: 2,54 (vs. 2,37 para a linha de base EAGLE-3).
• Retenção: Melhorou a taxa de retenção de longo alcance de 73,5% para 77,3%.
• Mecanismo: A regra delta com portão permitiu com sucesso que o caminho KV corrigisse o caminho de estado oculto em horizontes mais longos sem perturbar a precisão de curto alcance.

C. A Verificação da Realidade de Ponta a Ponta

Apesar das melhorias passo a passo, a aceleração de Ponta a Ponta (E2E) foi marginal:

• MAT Medido no HF: Aumentou apenas de 5,01 (EAGLE-3) para 5,04 (Híbrido).
• Latência: O caminho adicional de atenção cruzada introduziu 5–10% de latência extra na redação.
• Conclusão: O pipeline atual de TTT autoregressivo não consegue explorar eficientemente o sinal KV para superar o custo computacional adicional.

5. Gargalos Estruturais Identificados

O artigo identifica três razões específicas pelas quais o pipeline atual falha em traduzir ganhos passo a passo em acelerações E2E:

1. Dificuldade de Estimativa de Consulta: Estimar consultas futuras do alvo requer replicar as transformações não lineares de múltiplas camadas do modelo alvo. Modelos rascunho rasos (1-2 camadas) carecem da capacidade de aproximar essas consultas com precisão.
2. Sinais de Gradiente Esparsos: No TTT autoregressivo, as projeções KV do lado do rascunho ($W_K, W_V$) são atualizadas apenas pelos poucos tokens gerados pelo rascunho (por exemplo, 7 tokens) por passo, enquanto o prefixo KV é copiado. Isso resulta em gradientes extremamente esparsos e desequilibrados para o caminho KV.
3. Fome de Gradiente Induzida pelo Portão: No modelo híbrido, o portão que controla a correção KV tende a colapsar para próximo de zero no início do treinamento (porque o ramo KV inicializado aleatoriamente começa como ruído). Uma vez fechado, o portão bloqueia gradientes para o ramo KV, impedindo-o de aprender a ser útil. Isso cria um "estado de baixa injeção auto-reforçador".

6. Significado e Direções Futuras

• Mudança Conceitual: O artigo reenquadra o decaimento de longo alcance não apenas como um descompasso entre treinamento e inferência, mas como um problema de preservação de informação. Ele prova que os estados ocultos são uma compressão com perdas para tarefas de longo horizonte.
• Ferramenta de Diagnóstico: O KVShot fornece um ambiente de teste fundamental para isolar efeitos de representação de artefatos de treinamento.
• Roteiro para Arquiteturas de Próxima Geração: As descobertas sugerem que o TTT Autoregressivo é insuficiente para decodificação consciente de KV. Para realizar o potencial da reutilização de KV, o campo deve avançar em direção a Paradigmas de Treinamento em Blocos (por exemplo, DFlash).
  • O treinamento em blocos gera muitos tokens de rascunho em paralelo, fornecendo sinais de gradiente densos para projeções KV.
  • Permite modelos rascunho mais profundos capazes de estimativa precisa de consultas sem penalidades de latência sequencial.

Resumo: Embora os caches KV ofereçam teoricamente um sinal superior para decodificação especulativa de longo alcance ao preservar o contexto, os pipelines de treinamento autoregressivo atuais falham em aproveitar isso devido a limites de estimativa de consulta e fome de gradiente. O caminho a seguir reside em arquiteturas de treinamento paralelo e em blocos.