Resumo Técnico: Transferência de Cache KV entre Modelos (Cross-Model KV Cache Transfer) em Famílias de LLMs

1. Formulação do Problema

O atendimento (serving) de Grandes Modelos de Linguagem (LLMs) em produção depende cada vez mais da orquestração multimodelo para cascateamento custo-qualidade, alternância no meio da conversa e roteamento. Esses fluxos de trabalho frequentemente alternam entre membros de tamanhos diferentes da mesma família de modelos (ex: Qwen3 14B para 32B). Embora esses modelos compartilhem escolhas arquiteturais centrais e dados de pré-treinamento, eles diferem em escala, contagem de camadas e receitas de treinamento.

O gargalo principal em tais fluxos de trabalho é o custo de prefill. Toda vez que um contexto de conversa é passado de um modelo de origem para um modelo de destino, o destino deve reprocessar todo o contexto acumulado para popular seu próprio cache de Chave-Valor (KV). Essa passagem direta (forward pass) escala com o tamanho do modelo e o comprimento do prompt, anulando os ganhos de eficiência da troca de modelos. Soluções existentes, como o cache de prefixo, funcionam apenas dentro de um único modelo, e abordagens anteriores de cross-model (ex: fusores neurais, adaptadores latentes) tipicamente exigem treinamento baseado em gradiente ou restrições arquiteturais estritas (ex: contagem de camadas idêntica).

Este artigo propõe a Transferência de Cache KV entre Modelos, um método para reutilizar o cache KV pré-preenchido do modelo de origem diretamente no modelo de destino, pulando o prefill do destino por completo. O objetivo é encontrar um mapeamento $f: C_S \to \hat{C}_T$ tal que o modelo de destino, decodificando a partir do cache mapeado $\hat{C}_T$, produza saídas equivalentes à decodificação a partir de seu próprio cache real $C_T$.

2. Metodologia

Os autores hipotetizam que as relações de KV entre modelos possuem uma estrutura linear substancial, permitindo um mapeamento de forma fechada e sem gradiente. O framework proposto opera em pares de "KV correspondentes" (matched-KV), onde o modelo de origem e o de destino compartilham o mesmo número de cabeças de KV e dimensão por cabeça, mesmo que suas contagens totais de camadas ou tamanhos de parâmetros difiram.

O mapeador é projetado com três componentes principais:

A. Regressão de Ridge por Cabeça

Em vez de uma rede neural global, o método ajusta um mapeamento linear independente para cada camada e cabeça de destino.

• Entrada: Características de KV concatenadas de um subconj conjunto selecionado de camadas de origem.
• Saída: Os vetores de Chave ($K$) e Valor ($V$) esperados do destino.
• Otimização: Os pesos são resolvidos via regressão de ridge (solução de forma fechada) em um pequeno conjunto de calibração (500 sequências de 1.024 tokens do FineWeb-Edu). A regularização de ridge ($\lambda=0.01$) é usada para lidar com a natureza quase singular da matriz de características quando as camadas de origem são correlacionadas.

B. Seleção de Origem entre Camadas (Cross-Layer)

Como os modelos de origem e destino possuem profundidades diferentes ($L_S \neq L_T$), o método deve determinar quais camadas de origem informam cada camada de destino.

• Estratégia de Seleção: Para cada camada de destino, as $k$ melhores camadas de origem são selecionadas com base em seu poder preditivo (medido por $R^2$ médio por cabeça no conjunto de calibração).
• Mecanismo: As características de KV dessas $k$ melhores camadas de origem são concatenadas para formar a entrada para a regressão de ridge. A análise empírica mostra que a informação complementar é distribuída através de múltiplas camadas de origem; uma única camada é insuficiente, mas o desempenho satura rapidamente (ex: $k=6$ captura a maior parte da variância de $k=todos$).

C. Mapeamento de Espaço de Conteúdo (Fatoração de RoPE)

LLMs padrão utilizam Embeddings de Posição Rotativos (RoPE), que aplicam rotações dependentes da posição às chaves.

• Problema: Ajustar um mapeador diretamente em chaves rotacionadas vincula os pesos à distribuição de posição específica vista durante a calibração, limitando a transferibilidade para diferentes comprimentos de contexto.
• Solução: O método remove o RoPE das chaves de origem antes do mapeamento e reaplica o Ro-PE do destino após o mapeamento.
  • $K_{stripped} = K_{source} \cdot R^{-1}_{\Theta_s}$
  • $\hat{K}_{target} = (K_{stripped} \cdot W + b) \cdot R_{\Theta_t}$
• Isso desacopla a rotação posicional do conteúdo semântico, tornando o ajuste linear independente de posição e reutilizável para diversos comprimentos de contexto (até 32k tokens).

D. Extensão Não Linear (MLP)

Para pares onde o mapeador linear de ridge falha em reter precisão suficiente, os autores propõem um substituto (drop-in replacement) usando um pequeno Perceptron de Camadas Múltiplas (MLP). Este MLP redistribui o erro residual para longe de subespaços sensíveis à atenção, recuperando o desempenho em pares difíceis.

3. Principais Contribuições

1. Framework de Mapeamento de Forma Fechada: Um framework de regressão de ridge por cabeça, sem gradiente, que permite a transferência de cache KV entre modelos sem necessidade de treinamento. Baseia-se na seleção de $k$ camadas de origem e no mapeamento de conteúdo com remoção de RoPE.
2. Validação Multi-Família: Validação em três famílias de modelos (Qwen3, Llama 3.1, Ministral 3) e seis pares de KV correspondentes.
3. Recuperação Não Linear: Demonstração de que um MLP leve pode recuperar o desempenho em pares "difíceis" onde o mapeamento linear falha, especificamente ao redistribuir o erro para longe de dimensões críticas de atenção.
4. Métrica Diagnóstica: Identificação de que a similaridade de cosseno da saída de atenção é um preditor melhor da retenção de transferência downstream do que as métricas de reconstrução padrão ($R^2$). Os autores mostram uma correlação de Pearson de $r=+0,57$ entre o cosseno da saída de atenção e a retenção do HellaSwag, comparado a $r=-0,20$ para $R^2$.

4. Resultados Experimentais

O framework foi avaliado em seis pares de KV correspondentes em três famílias usando cinco benchmarks de acurácia (ARC-C, HellaSwag, WinoGrande, MMLU, GSM8K) e perplexidade do WikiText-2.

• Retenção de Acurácia:
  • Nível 1 (Sucesso): Quatro pares (Qwen3 14B→32B, 8B→32B; Llama 3.1 8B→70B; Ministral 3B→8B) retiveram 73–98% da acurácia autônoma do modelo de destino.
  • Nível 2 (Falha): Dois pares (Ministral 3B→14B, 8B→14B) degradaram significativamente (retenção de 42–44%). No entanto, substituir o mapeador linear por um MLP recuperou a retenção do HellaSwag em +24,3 a +36,8 pontos percentuais, elevando esses pares acima de 90%.
• Latência:
  • O mapeador executa de 2,7 a 25× mais rápido do que o re-prefill do modelo de destino.
  • Os ganhos de velocidade são mais pronunciados em transferências de pequeno-para-grande e comprimentos de contexto maiores (ex: 25× de aceleração em 32k tokens para Qwen3 14B→32B).
• Estabilidade Multi-turno:
  • Em cenários de entrega (handoff) de múltiplos turnos (CoQA), o desvio permanece pequeno. O desvio de pequeno-para-grande aumenta apenas 1,7 pp em 10 turnos; o desvio de grande-para-pequeno cresce linearmente a 0,33 pp/turno, sugerindo estabilidade para sessões típicas.
• Eficiência de Calibração:
  • O mapeador requer apenas 500 sequências para calibração.
  • O ajuste leva cerca de 47–87 minutos por par em um único nó 8×H100, sem treinamento baseado em gradiente.

5. Significância e Alegações

O artigo afirma que a transferência de cache KV entre modelos dentro de uma família de LLMs admite um ajuste de forma fechada e sem treinamento. Os autores argumentam que a relação de KV entre modelos é amplamente linear, desafiando a suposição de que adaptadores neurais complexos são necessários para tais transferências.

A significância deste trabalho reside em:

1. Praticidade: Permite o cascateamento custo-qualidade e a alternância no meio da conversa sem o custo proibitivo de re-prefill, tornando a orquestração dinâmica de modelos viável em produção.
2. Mudança na Avaliação: As descobertas sugerem que os critérios de avaliação para mapeadores cross-model devem mudar de métricas de reconstrução bruta ($R^2$) para diagnósticos conscientes de subespaço (como o cosseno da saída de atenção), pois estes predizem melhor o desempenho de tarefas downstream.
3. Redistribuição de Erro: O trabalho destaca que a falha de transferência ocorre frequentemente devido a onde o erro cai (em subespaços sensíveis à atenção), e não pela magnitude do erro, motivando futuros objetivos alinhados à atenção.

Os autores mantêm a modéstia quanto às limitações, observando que a abordagem é atualmente restrita a modelos de atenção total densa dentro da mesma família com configurações de KV correspondentes, e que a calibração é realizada em um único domínio (FineWeb-Edu). Eles deixam a transferência entre diferentes famílias e condições de KV não correspondentes como questões abertas para trabalhos futuros.