Resumo Técnico: ProxyKV

Declaração do Problema

A inferência eficiente de contexto longo em Modelos de Linguagem de Grande Escala (LLMs) está atualmente limitada pelas exigências de memória do cache Chave-Valor (KV). À medida que o comprimento da sequência ($N$) aumenta, o crescimento linear do cache KV frequentemente excede a capacidade de memória da GPU, resultando em sobrecarga de E/S proibitiva e latência de inferência.

Os métodos existentes de poda de cache KV enfrentam um compromisso fundamental entre eficiência de pontuação e precisão de poda:

• Métodos Heurísticos (e.g., H2O, SnapKV): São rápidos e de baixa latência, mas dependem de janelas de observação local, frequentemente perdendo dependências semânticas globais necessárias para raciocínio complexo.
• Métodos Baseados em Reconstrução (e.g., KVZip): Alcançam alta precisão ao realizar uma passagem secundária de preenchimento prévio forçada pelo professor no modelo-alvo para reconstruir o contexto global. No entanto, isso incorre em uma sobrecarga de preenchimento prévio proibitiva, anulando os benefícios de latência da poda.

O desafio central é preencher a lacuna entre heurísticas de baixa latência e reconstrução de alta precisão sem incorrer no alto custo computacional de um preenchimento prévio secundário no modelo-alvo.

Metodologia

Os autores propõem o ProxyKV, um framework de poda por proxy cruzado entre modelos projetado para descarregar a tarefa de pontuação de importância do modelo-alvo para um "Proxy de Pequeno Modelo" leve e assíncrono.

Arquitetura do Sistema

O sistema opera por meio de um pipeline de três componentes:

1. Alvo de Grande Modelo ($M_l$): O modelo primário responsável pelo caminho crítico de preenchimento prévio e por toda a fase de decodificação.
2. Proxy de Pequeno Modelo ($M_s$): Um modelo menor intra-família (e.g., 1,5B ou 4B parâmetros) que executa assincronamente em um fluxo de execução separado (ou fluxo CUDA). Ele extrai características de atenção cruzada de cabeças ($X$) do contexto de entrada.
3. HybridAxialMapper ($\Phi_\theta$): Um módulo aprendível que mapeia as características do proxy ($X$) para pontuações de importância alinhadas ao alvo ($\hat{Y}$).

Fluxo de Trabalho:
Durante a inferência, o contexto de entrada é despachado para ambos os modelos concorrentemente. O proxy gera características de atenção, que são processadas pelo HybridAxialMapper para produzir uma máscara de poda binária ($M$) via limiarização Top-K. O modelo-alvo então realiza a decodificação usando apenas o cache KV podado. Crucialmente, o cache KV do proxy é liberado imediatamente após a fase de preenchimento prévio, tornando a sobrecarga de memória transitória.

Componentes Chave

1. Arquitetura do HybridAxialMapper

Para preencher a lacuna arquitetônica entre modelos heterogêneos (diferentes contagens de cabeças e profundidades de camadas), o mapeador emprega um pipeline de três etapas que desconecta a extração de características temporais do alinhamento cruzado de cabeças:

• Etapa 1: Extração de Características Temporais: Um Tronco Convolucional 1D projeta características brutas do proxy em um espaço latente agnóstico ao modelo, normalizando a configuração de cabeças.
• Etapa 2: Codificador de Contexto do Eixo Temporal: Um codificador Transformer processa a sequência com codificações posicionais sinusoidais para capturar o contexto global, utilizando uma abordagem de janela deslizante para sequências longas (até 170k tokens).
• Etapa 3: Atenção Cruzada no Eixo de Cabeças: Consultas-alvo aprendíveis agregam tokens de cabeças sintéticos do proxy na configuração específica de cabeças do alvo. Esta etapa permite que o mapeador escale para contagens de cabeças-alvo arbitrárias sem retreinar toda a rede.

2. Função de Perda Híbrida Multi-Granularidade

Reconhecendo que a métrica de implantação é a qualidade de retenção Top-K e não a reconstrução pontual de pontuação, os autores projetam uma função de perda composta ($L_{total}$) consistindo em cinco termos:

• $L_{bin}$ (Binário Multi-Razão): Supervisiona o limite de decisão binária através de múltiplas razões de retenção, enfatizando compressão agressiva.
• $L_{mse}$ (MSE Ponderado por Valor): Garante alinhamento de magnitude numérica para tokens salientes.
• $L_{fine}$ (Rank Intra-Top-K): Impõe ordenação relativa dentro do conjunto Top-K retido.
• $L_{global}$ (Rank Global Consciente de Massa): Otimiza a separação entre tokens Top-K e não Top-K (limiar sinal-ruído).
• $L_{cos}$ (Alinhamento de Cosseno): Alinha a orientação global do vetor de pontuação previsto com o alvo.

Este design de perda desloca o objetivo de regressão rígida para consistência de classificação relativa, o que está mais diretamente alinhado com a qualidade da poda.

Contribuições Chave

1. Framework de Proxy Assíncrono: Um framework de poda inovador que remove o gargalo de pontuação do caminho crítico do alvo ao descarregá-lo para um pequeno modelo assíncrono, alcançando até 3,21× de aceleração no preenchimento prévio em Llama-3.1-8B (dual-GPU) e sustentando acelerações até 170k tokens.
2. HybridAxialMapper e Perda Híbrida: Uma arquitetura aprendível e função de perda que preenchem com sucesso lacunas estruturais entre modelos heterogêneos (e.g., alvo 8B vs. proxy 1,5B), superando a rigidez do alinhamento cabeça-a-cabeça estático usado em trabalhos anteriores.
3. Validação Empírica: Experimentos extensivos nas famílias Llama-3.1, Qwen-2.5 e Qwen-3 (alvos de 7B–32B) demonstrando que o ProxyKV recupera ~98,7% da precisão média do oráculo KVZip enquanto reduz significativamente a latência.

Resultados

• Precisão: O ProxyKV iguala o desempenho do oráculo KVZip (a linha de base de alta precisão) nos benchmarks LongBench, SCBench e RULER. Ele recupera ~98,7% da precisão média do KVZip e rastreia dentro de 1–2 pontos na maioria das tarefas, mesmo em razões de retenção agressivas ($\rho \le 0,2$).
• Eficiência:
  • No Llama-3.1-8B, o ProxyKV entrega uma aceleração de 3,21× na latência de preenchimento prévio comparado ao KVZip (dual-GPU) e ~1,5× em uma única GPU compartilhada.
  • O próprio HybridAxialMapper contribui apenas com 4,7%–6,8% do tempo de parede de preenchimento prévio.
  • O método escala efetivamente para alvos de 32B (Qwen-3) com uma razão alvo/proxy de ~8×.
• Robustez: O framework demonstra generalização zero-shot robusta para conjuntos de dados retidos (e.g., SCBench RepoQA) e mantém desempenho em tarefas complexas de raciocínio multi-hop onde métodos heurísticos (SnapKV) falham.

Significado e Alegações

O artigo alega que o ProxyKV preenche com sucesso a "lacuna custo de pontuação–precisão" que tem limitado a inferência de LLMs de contexto longo. Ao alavancar correlações de atenção intra-família e um mapeamento cruzado axial aprendível, ele alcança a alta precisão de métodos baseados em reconstrução sem a sobrecarga proibitiva de preenchimento prévio.

Os autores destacam que, enquanto o alinhamento estático camada-a-camada (como usado em IAM ou KVZip+IAM) se torna uma restrição vinculativa à medida que os modelos-alvo escalam (e.g., para 32B), o HybridAxialMapper aprendível permanece eficaz. Isso sugere que o ProxyKV oferece um caminho escalável para inferência eficiente de contexto longo, recuperando precisão próxima à do oráculo enquanto reduz significativamente a latência e a pressão de memória durante a fase de preenchimento prévio. O método é apresentado como uma solução de tempo de implantação que requer apenas um proxy leve intra-família e um pequeno mapeador treinável, evitando a necessidade de re-treinamento completo do modelo ou modificações complexas de hardware.