LinearKV: One Cached State Suffices for Position-Independent Caching in Hybrid LLMs

Resumo Técnico: LINEARKV

Definição do Problema
O atendimento (serving) de Grandes Modelos de Linguagem (LLMs) enfrenta gargalos significativos de latência e vazão devido ao custo crescente do preenchimento (prefilling) de contextos longos. O Caching Independente de Posição (PIC - Position-Independent Caching) mitiga isso armazenando e reutilizando fragmentos (chunks) de tokens preenchidos independentemente de sua posição no prompt. No entanto, os métodos de PIC existentes são projetados para arquiteturas de atenção total (full-attention), baseando-se em caches de Chave-Valor (KV) indexados por tokens que podem ser concatenados e reparados seletivamente. Os LLMs híbridos, que intercalam camadas de atenção total (FA) com camadas recorrentes (ex: Mamba-2, Gated DeltaNet), quebram essas primitivas. As camadas recorrentes comprimem o contexto em um estado de tamanho fixo em vez de manter um cache indexado por tokens. Isso cria dois desafios: (1) Não está claro como inicializar e reparar estados recorrentes a partir de fragmentos (chunks) cacheados disjuntos sem entradas indexadas por tokens para concatenar, e (2) é desconhecido se uma estratégia de PIC unificada pode se generalizar através de diferentes arquiteturas híbridas (ex: Mamba-2 de decaimento escalar vs. Gated DeltaNet de transição densa).

Metodologia: LinearKV
Os autores propõem o LinearKV, um framework livre de treinamento para caching independente de posição em LLMs híbridos. A inovação central é uma estratégia de inicialização desacoplada aplicada após o pareamento dos fragmentos cacheados:

1. Inicialização Desacoplada:

   • Camadas de Atenção Total (FA): Estas camadas procedem como no PIC padrão, concatenando entradas KV cacheadas de fragmentos correspondentes ao longo do eixo do token.
   • Camadas Recorrentes (Lineares): Em vez de tentar concatenar estados (o que é impossível, pois são resumos de tamanho fixo), o framework mapeia $K$ estados locais correspondentes $\{S^{\text{local}}_1, \dots, S^{\text{local}}_K\}$ em um único estado inicial $S_{\text{init}}$ através de uma função $f$.
   • Reuso do Seletor: Este desacoplamento permite que o LinearKV reutilize algoritmos de seleção de tokens de PIC existentes (ex: CacheBlend, EPIC, ProphetKV) sem alterações. O seletor determina um conjunto de posições de tokens $U$ para serem recomputados. O sistema então repara o cache sobrescrevendo as entradas KV de FA em $U$ e avançando o estado recorrente a partir de $S_{\text{init}}$ através dos tokens em $U$ em ordem de contexto.

2. O Inicializador de Estado Linear ($f$):
   O artigo investiga duas escolhas para $f$:

   • Composição Exata: Uma abordagem algebricamente fundamentada (adotada pelo trabalho concorrente HYPIC) que compõe todos os $K$ estados cacheados usando os operadores de transição dos fragmentos para reconstruir o estado exato do prefixo completo.
   • Inicialização de Bloco Único (Padrão do LinearKV): Uma abordagem mais simples onde $f$ seleciona um único estado cacheado (tipicamente o último fragmento correspondente, $S^{\text{local}}_K$) como o inicializador, descartando os demais.

Principais Descobertas e Contribuições
O artigo apresenta três contribuições principais, apoiadas por evidências empíricas em três modelos híbridos (Granite/Mamba-2, OLMo/GDN, Qwen/GDN) e três seletores de PIC:

1. Formalização de PIC Híbrido: Os autores formalizam o PIC híbrido como um framework desacoplado onde a única operação específica da arquitetura é a inicialização do estado linear. Isso permite que seletores de PIC de atenção total existentes sejam integrados sem modificação.
2. Fragilidade da Composição Exata: Contrariando a intuição algébrica, o artigo demonstra que a composição exata é desnecessária e fragiliza a arquitetura.
   • Nos modelos GDN (OLMo, Qwen), a composição exata e a inicialização de bloco único apresentam desempenho equivalente, ambas recuperando até 92% da qualidade total.
   • Nos modelos Mamba-2 (Granite), a composição exata colapsa. Devido à natureza de decaimento escalar do Mamba-2, os erros de fragmentos preenchidos independentemente se acumulam através da profundidade, levando a um estado inicial construído com alto erro relativo. Consequentemente, a composição exata recupera apenas ~46,6% da qualidade total sob o seletor EPIC, enquanto o inicializador de bloco único recupera 86,8%.
   • Os autores identificam que a falha decorre da agregação de $K$ operadores profundos condicionados independentemente, não da arquitetura em si, já que a abordagem de bloco único evita esse acúmulo de erro.
3. Ganhos de Eficiência: O inicializador de bloco único é computacionalmente mais barato. Ele evita o custo online de compor matrizes de transição densas exigidas pela composição exata. Isso reduz o Tempo para o Primeiro Token (TTFT) em adicionais de 5 a 17% em comparação com a composição exata, mantendo ou melhorando a precisão.

Resultados Experimentais

• Benchmarks: Avaliado em LongBench QA, RULER (8K–32K) e LongBench Summarization.
• Desempenho:
  • Granite (Mamba-2): O LinearKV (bloco único) supera significativamente a composição exata em todos os seletores. Por exemplo, sob o EPIC com orçamento de recomputação de 20%, o LinearKV atinge 86,8% da qualidade total contra 46,6% da composição exata.
  • OLMo & Qwen (GDN): Ambos os métodos apresentam desempenho similar, com o LinearKV igualando o desempenho da composição exata.
  • Robustez: A superioridade do inicializador de bloco único no Mamba-2 mantém-se em várias razões de recomputação (0,03 a 0,4) e comprimentos de contexto (8K a 32K).
• Latência: O LinearKV reduz o TTFT para 0,46×–0,62× do tempo de preenchimento total, superando a composição exata em todos os cenários medidos.

Significância
O artigo estabelece que, para LLMs híbridos, a abordagem "algebricamente fundamentada" de reconstruir o estado exato do prefixo completo a partir de fragmentos cacheados não é apenas computacionalmente cara, mas pode ser prejudicial ao desempenho do modelo em arquiteturas específicas como o Mamba-2. Ao demonstrar que um único estado cacheado é suficiente como um inicializador robusto, o LinearKV fornece uma solução simples, livre de treinamento e agnóstica à arquitetura que maximiza os benefícios do caching independente de posição, minimizando a latência e evitando o acúmulo de erro inerente à composição de múltiplos estados. Essa descoberta desafia a suposição de que uma composição de estado mais complexa é necessária para um caching eficaz em arquiteturas recorrentes.