KEEP: A KV-Cache-Centric Memory Management System for Efficient Embodied Planning

Título: KEEP: Um Sistema de Gerenciamento de Memória Centrado em KV-Cache para Planejamento Embutido Eficiente

1. Problema e Motivação

O planejamento embutido (embodied planning) envolve agentes que utilizam Modelos de Linguagem de Grande Porte (LLMs) para prever sequências de ações em cenários do mundo real. Para evitar exploração repetitiva e manter uma visão global, esses agentes dependem de memória (histórico de tarefas, estados do ambiente, status do agente).

No entanto, existem desafios críticos nas abordagens atuais:

• Latência de Pré-preenchimento (Prefill): Armazenar a memória como texto bruto gera prompts extremamente longos (dezenas de milhares de tokens), resultando em alta latência antes da geração da primeira resposta.
• Ineficiência do Cache KV: Embora seja possível armazenar e reutilizar os pares Chave-Valor (KV Cache) para evitar recálculos, a natureza dinâmica da memória em ambientes físicos (ex: um objeto muda de posição, uma tarefa é concluída) invalida frequentemente o cache.
• Limitações das Métodos Existentes:
  • Pré-fix Caching: Uma única atualização invalida todo o cache subsequente.
  • Reuso Total/Recomputação de Blocos Fixos: Blocos de tamanho fixo não se alinham com a granularidade fina e rápida das atualizações de memória, causando perda de atenção cruzada (cross-attention) entre segmentos ou recálculos excessivos.
  • Recomputação Estática: Métodos que reavaliam tokens em posições fixas falham em capturar a importância dinâmica e dependente do contexto da memória em tarefas de planejamento.

2. Metodologia: O Sistema KEEP

O KEEP propõe um sistema de gerenciamento de memória centrado no KV-Cache, composto por três inovações principais:

A. Construção de Memória Estática-Dinâmica (Static-Dynamic Memory Construction)

• Conceito: Reconhece que diferentes segmentos de memória têm frequências de atualização distintas (ex: a posição de uma mesa é estática; a posição de um copo é dinâmica).
• Mecanismo:
  • Agrupa segmentos de memória com base na frequência de mudança.
  • Grupos Estáticos: Segmentos que não mudaram por $t$ passos são agrupados e calculados com atenção cruzada completa entre eles, preservando relações contextuais ricas.
  • Grupos Dinâmicos: Segmentos atualizados recentemente são calculados individualmente para evitar que uma única mudança invalide o cache de todo o grupo.
  • Transição: Grupos mudam de estado (estático $\leftrightarrow$ dinâmico) conforme a estabilidade dos dados, recalculando o KV apenas quando necessário.

B. Recomputação de Memória Multi-hop (Multi-hop Memory Re-computation)

• Conceito: A importância de um segmento de memória depende do contexto da consulta atual e do histórico, não sendo fixa.
• Mecanismo:
  • Utiliza a propagação de importância baseada nas distribuições de atenção.
  • Passo 1: Identifica memórias diretamente importantes para a consulta.
  • Passo 2: Usa essas memórias para encontrar memórias secundárias críticas (ex: se a tarefa é "abrir a porta", a porta é importante; se a porta está trancada, a chave é importante; se a chave está na mesa, a mesa torna-se importante).
  • Recomputação Seletiva: Apenas os KV caches dos segmentos identificados como críticos através dessa propagação multi-hop são recomputados, preservando a precisão da atenção cruzada com custo computacional mínimo.
  • Granularidade: Opera em nível de segmento de memória, não de token, permitindo carregamento contíguo e eficiente, reduzindo a sobrecarga de I/O.

C. Carregamento de Memória Balanceado por Camada (Layer-balanced Memory Loading)

• Conceito: Em arquiteturas de pipeline tradicionais, o carregamento de KV (da RAM para a GPU) e o cálculo (atenção) são paralelizados. No entanto, em sistemas de recomputação, as camadas iniciais exigem menos carregamento (muita recomputação) e as camadas finais exigem muito carregamento (pouca recomputação), criando "bolhas" de ociosidade no pipeline.
• Mecanismo:
  • Explora a observação de que, se um segmento não for recomputado em uma camada inicial, ele não precisará ser recomputado em camadas subsequentes.
  • Pré-carregamento Proativo: Enquanto as camadas iniciais estão computando, o sistema de carregamento já começa a pré-carregar os KV caches das memórias que serão necessárias nas camadas futuras (garantidas como não recomputadas).
  • Isso redistribui a sobrecarga de carregamento das camadas tardias para as camadas iniciais subutilizadas, eliminando bolhas de pipeline e maximizando a utilização do hardware.

3. Resultados Experimentais

O KEEP foi avaliado nos conjuntos de dados ALFRED e WAH-NL, utilizando modelos Qwen-2.5 (14B e 32B).

• Comparação com Métodos Baseados em Texto:
  • O KEEP alcançou um aceleração de 2.68x em relação a métodos que usam memória em texto bruto, com perda de precisão negligenciável.
• Comparação com CacheBlend (Método de Recomputação KV):
  • Taxa de Sucesso (SR): Melhoria de 4.13% no dataset ALFRED.
  • Latência (TTFT - Time-to-First-Token): Redução de 1.90x no tempo para a primeira resposta.
  • No dataset WAH-NL, houve melhoria de 3.31% na taxa de sucesso de subobjetivos e redução de 1.91x no TTFT.
• Análise de Ablação:
  • A remoção da construção estática-dinâmica causou uma queda de 6.94% na taxa de sucesso e aumento de 1.54x na latência.
  • A remoção da recomputação multi-hop reduziu a precisão em 2.52%.
  • A remoção do carregamento balanceado aumentou a latência em 1.20x devido a ineficiências no pipeline.

4. Contribuições Principais

1. Análise Fundamental: Identificação das diferenças críticas entre a construção de memória para planejamento embutido (dinâmica, atualizações frequentes) e a inferência de LLMs convencional (texto estático).
2. Sistema KEEP: Proposta de um sistema integrado que combina:
   • Gerenciamento híbrido de granularidade (estático/dinâmico) para minimizar invalidação de cache.
   • Mecanismo de recomputação multi-hop para recuperar interações de memória críticas sem custo excessivo.
   • Agendamento de pipeline balanceado para otimizar o uso de hardware em cenários de recomputação assimétrica.
3. Desempenho Superior: Demonstração de que é possível alcançar alta eficiência e precisão em tarefas de planejamento de longo prazo, superando os métodos atuais de reuso de KV Cache.

5. Significado e Impacto

O trabalho KEEP é significativo porque resolve o gargalo de latência e precisão no uso de LLMs para robótica e agentes autônomos. Ao tratar a memória não como texto estático, mas como um recurso dinâmico gerenciado via KV-Cache com granularidade adaptativa, o sistema permite que agentes operem em ambientes complexos com atualizações frequentes de estado sem sacrificar a velocidade de inferência. Isso viabiliza a aplicação prática de LLMs em cenários de "longo horizonte" onde a memória precisa ser persistente, atualizável e acessível rapidamente.

O código do projeto está disponível publicamente, facilitando a reprodução e adoção pela comunidade de pesquisa em IA e robótica.