MemDecay: Region-Aware KV Cache Eviction for Efficient LLM Agent Inference

Resumo Técnico: MemDecay

Definição do Problema

Modelos de Linguagem de Grande Escala (LLMs) estão sendo cada vez mais implantados como agentes que executam fluxos de trabalho de longo horizonte envolvendo interações de múltiplos turnos, uso de ferramentas e planejamento iterativo. Diferente das tarefas estáticas de contexto longo, os rastros desses agentes acumulam componentes heterogêneos: instruções do sistema, restrições de segurança, planos em evolução, turnos do usuário, documentos recuperados, saídas de ferramentas e estados de raciocínio intermediários.

O cache de Chave-Valor (KV) necessário para a decodificação autorregressiva cresce linearmente com esse rastro, criando uma sobrecarga significativa de memória, demanda de largura de banda e latência. As políticas de evasão de cache KV existentes geralmente dependem de sinais centrados no modelo (ex: atenção acumulada, janelas de recência ou sumidouros de atenção) e aplicam regras de retenção uniformes para toda a sequência. Essa abordagem falha em aproveitar a estrutura semântica inerente dos prompts orquestrados de agentes. Em cargas de trabalho de agentes, diferentes componentes possuem papéis funcionais e tempos de vida esperados vastamente distintos; por exemplo, as instruções do sistema devem persistir durante todo o episódio, enquanto o scratchpad de raciocínio ou saídas obsoletas de ferramentas podem ser relevantes apenas por curtos períodos. Tratar todos os tokens com o mesmo tempo de vida esperado ignora essa estrutura, levando a um gerenciamento de memória subótimo.

Metodologia: MemDecay

Os autores propõem o MemDecay, um framework de evasão de cache KV consciente de regiões e livre de treinamento, projetado para agentes de LLM de longo horizonte. O método opera sob um orçamento fixo de cache e integra priors estruturais da camada de orquestração com a importância derivada da atenção.

Componentes Principais

1. Rótulos de Região: A camada de orquestração fornece rótulos de região para extensões contíguas de prompt (ex: system, plan, user, tool_in, tool_out, retrieval, scratchpad).
2. Pontuação de Retenção: Cada token $i$ recebe uma pontuação de retenção $s_i(t)$ no tempo $t$, calculada como:
   $$s_i(t) = b_{r(i)} \exp(-\lambda_{r(i)} \Delta t_i(t)) + \alpha a_i(t)$$
   • Prior Estrutural: $b_{r(i)}$ é uma prioridade base específica para a região do token.
   • Decaimento Temporal: $\lambda_{r(i)}$ é uma taxa de decaimento específica da região. O termo $\Delta t_i(t)$ representa o tempo desde que o token foi "atualizado" pela última vez.
   • Importância da Atenção: $a_i(t)$ é uma média móvel exponencialmente ponderada (EWMA) da massa de atenção recebida pelo token.
   • Mecanismo de Atualização (Refresh): Se a atenção a um token exceder um limiar $\tau$, seu relógio de decaimento é resetado ($\Delta t_i$ torna-se 0), estendendo efetivamente sua vida com base no uso contínuo.
3. Fixação (Pinning): Regiões do sistema (e opcionalmente segmentos de plano) podem ser "fixadas", o que significa que nunca são evacuadas, independentemente da pontuação.
4. Política de Evasão: O cache é gerenciado na granularidade de página (ex: 16 tokens por página). Quando o orçamento é excedido, as páginas são pontuadas pelo valor médio da pontuação de retenção de seus tokens constituintes. As páginas com as menores pontuações são evacuadas primeiro.

Calibração

O MemDecay é livre de treinamento em relação aos pesos do modelo, mas emprega uma calibração estatística leve para ajustar as taxas de decaimento ($\lambda_r$) para cada região. Este processo estima os tempos de vida da atenção condicionada à região a partir de rastros de agentes, permitindo que o sistema se adapte a características específicas de carga de trabalho sem o retreinamento do LLM.

Contribuições Principais

1. Formulação: O artigo formula o esquecimento seletivo consciente de região como um problema de gerenciamento de memória para agentes de LLM, fundamentando a solução na estrutura semântica de prompts orquestrados.
2. Algoritmo: Introduz o MemDecay, uma política que combina prioridades base específicas de região, decaimento temporal e atualização disparada pela atenção em uma única pontuação de retenção interpretável. Ele opera em layouts de granularidade de página compatíveis com stacks de serviço existentes (ex: PagedAttention).
3. Protocolo de Medição: Os autores especificam um protocolo para estimar os tempos de vida da atenção condicionada à região e calibrar as taxas de decaimento. Eles executam isso através de dois comprimentos de contexto (~450 e ~1.700 tokens) e dois tamanhos de modelo (Qwen2.5-1.5B e 3B), totalizando 4.320 sondagens de evacuação pontuadas.

Resultados Experimentais

O estudo avalia o MemDecay contra baselines incluindo evacuação aleatória, retenção baseada em recência (streaming) e retenção de atenção acumulada (estilo H2O).

Principais Descobertas

• Hierarquia de Tempo de Vida de Região: Os tempos de vida da atenção medidos diferem em uma ordem de magnitude entre as regiões. Em todas as configurações, os tokens de system têm meias-vidas de 148–189 passos de decodificação, enquanto os tokens de scratchpad decaem rapidamente (14–16 passos). Os tokens de Retrieval foram identificados como a região móvel de vida mais longa, contradizendo suposições padrão de que decairiam rapidamente.
• Eficácia da Fixação (Pinning): A fixação de regiões do sistema garante a sobrevivência da instrução. O MemDecay preservou 21–24 de 24 sondagens de sistema em todas as configurações no limite máximo de cache, enquanto os baselines baseados em recência preservaram 0–2.
• Recência vs. Estrutura: A retenção baseada em recência (streaming) colapsa conforme o comprimento do contexto aumenta, performando próximo de zero em níveis longos. Em contraste, o prior estrutural no MemDecay permanece eficaz.
• Atenção Acumulada: Em conteúdo não fixado, a atenção acumulada (estilo H2O) fortalece-se com a escala do modelo, superando o MemDecay em recall não fixado na configuração do modelo 3B.
• Limitação do MemDecay: O MemDecay mostrou uma perda específica de recall não fixado para fatos antigos (ex: fatos da região do usuário no turno inicial) em comparação com baselines baseados em atenção. Um estudo de ablação revelou que o termo de atenção ($\alpha a_i$) era pequeno demais em magnitude para neutralizar o decaimento estrutural para esses tokens. Os autores identificam que o termo de atenção requer normalização de magnitude (ex: relativa a níveis de atenção uniforme) para ser eficaz, em vez de apenas ajuste de peso.

Significância e Alegações

O artigo posiciona o MemDecay como uma solução prática e livre de treinamento que alinha a retenção de cache com a estrutura semântica de fluxos de trabalho de agentes. Sua significância reside em:

• Interpretabilidade: A política é impulsionada por rótulos de região explícitos e taxas de decaimento, tornando as decisões de retenção transparentes.
• Compatibilidade: Opera em caches de granularidade de página sem exigir retreinamento de modelo ou módulos auxiliares, integrando-se a stacks de serviço existentes.
• Estabilidade Garantida: Ao fixar regiões críticas, garante que as instruções do sistema e as restrições de segurança nunca sejam perdidas devido à pressão de memória, uma garantia que métodos baseados em atenção ou recência não fornecem.
• Adaptabilidade: O protocolo de calibração demonstra que os tempos de vida das regiões dependem da carga de trabalho e podem ser medidos para otimizar as taxas de decaimento, desafiando o uso de políticas de evacuação estáticas e universais.

Os autores concluem que, embora o MemDecay consiga preservar os priors estruturais e as instruções fixadas, trabalhos futuros devem abordar a normalização de magnitude do termo de atenção para melhorar o recall não fixado e explorar extensões em camadas (ex: descarregamento para a memória host) para lidar com a evacuação irreversível.