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

1. Il Problema

Il compito di pianificazione incarnata (embodied planning) richiede agli agenti di prevedere sequenze di azioni per raggiungere obiettivi a lungo termine in scenari reali. I Modelli Linguistici di Grande Dimensione (LLM) sono fondamentali per questo compito, utilizzando una memoria strutturata (stato dell'agente, stato degli oggetti, storico delle azioni) per mantenere una visione globale ed evitare esplorazioni ripetute.

Tuttavia, esistono due sfide principali nell'implementazione efficiente di questi sistemi:

• Latenza di Prefill eccessiva: La memoria è spesso memorizzata come testo grezzo. Man mano che le esperienze si accumulano, i prompt diventano estremamente lunghi (decine di migliaia di token), causando una latenza di prefilling (TTFT - Time To First Token) che diventa il collo di bottiglia principale.
• Invalidazione della KV Cache: Sebbene sia possibile riutilizzare le chiavi e i valori (KV cache) calcolati in precedenza per evitare ridondanza computazionale, gli scenari di pianificazione incarnata sono dinamici. Ogni azione dell'agente (es. "prendere il latte") modifica lo stato della memoria. Questo invalida la KV cache per tutti i token successivi alla modifica, costringendo a ricalcoli frequenti che annullano i benefici della caching.

I metodi esistenti falliscono perché:

1. Usano blocchi di dimensione fissa che non si adattano alla granularità fine degli aggiornamenti della memoria.
2. Ignorano le interazioni di cross-attention tra blocchi di memoria diversi quando riutilizzano la cache.
3. Utilizzano strategie di ricalcolo statiche che non tengono conto dell'importanza contestuale della memoria rispetto alla query corrente.

2. Metodologia: KEEP

Il sistema KEEP propone una gestione della memoria incentrata sulla KV cache, basata su tre innovazioni chiave per bilanciare efficienza e accuratezza:

A. Costruzione della Memoria Statico-Dinamica (Static-Dynamic Memory Construction)

Invece di trattare tutta la memoria allo stesso modo, KEEP classifica i segmenti di memoria in base alla loro frequenza di aggiornamento:

• Gruppi Statici: Segmenti che non cambiano frequentemente (es. storico dei compiti completati). Per questi, la KV cache viene calcolata con cross-attention completa all'interno del gruppo per preservare le relazioni contestuali.
• Gruppi Dinamici: Segmenti che cambiano spesso (es. stato degli oggetti ambientali). Per questi, la KV cache viene calcolata a livello di singolo segmento.
• Vantaggio: Questo approccio riduce drasticamente l'invalidazione della cache. Un aggiornamento in un segmento dinamico non invalida l'intera cache del gruppo, a differenza dei metodi a blocchi fissi.

B. Ricalcolo della Memoria Multi-hop (Multi-hop Memory Re-computation)

Per recuperare le interazioni critiche tra i diversi gruppi di memoria (che altrimenti verrebbero perse), KEEP introduce un meccanismo dinamico di selezione:

• Propagazione dell'Importanza: Invece di ricalcolare token in posizioni fisse, il sistema valuta l'importanza dei segmenti di memoria basandosi sulla loro relazione con la query corrente e il contesto. Utilizza le distribuzioni di attenzione per identificare i segmenti chiave e poi "propaga" questa importanza ai segmenti correlati (es. se "chiave" è importante per "serratura", e "serratura" è su "tavolo", allora "tavolo" diventa rilevante).
• Ricalcolo Selettivo: Solo i KV cache dei segmenti identificati come critici attraverso questa propagazione multi-hop vengono ricalcolati.
• Granularità: Il ricalcolo avviene a livello di segmento di memoria (non di singolo token), permettendo un caricamento I/O più efficiente e continuo.

C. Caricamento della Memoria Bilanciato per Livello (Layer-balanced Memory Loading)

Poiché la KV cache viene archiviata su memoria lenta (RAM CPU) e caricata su GPU solo quando necessario, si creano "bolle" di inattività nella pipeline di calcolo.

• Osservazione: Se un segmento di memoria non deve essere ricalcolato in uno strato (layer) iniziale, non lo sarà mai negli strati successivi (poiché lo stato nascosto corrispondente viene scartato).
• Strategia: KEEP orchestra il caricamento in modo proattivo. Negli strati iniziali, se il caricamento per lo strato successivo è completato, il motore di caricamento inizia immediatamente a precaricare i KV cache dei segmenti che non verranno mai ricalcolati negli strati futuri.
• Risultato: Questo ridistribuisce il carico di I/O dagli strati tardivi (dove il carico è alto) agli strati iniziali (dove c'è tempo libero), eliminando le bolle di pipeline e massimizzando l'utilizzo dell'hardware.

3. Risultati Sperimentali

Il sistema è stato valutato sui benchmark ALFRED e WAH-NL utilizzando modelli Qwen-2.5 (14B e 32B).

• Prestazioni rispetto ai metodi basati su testo: KEEP ottiene un speedup di 2.68x rispetto ai metodi che memorizzano la memoria come testo grezzo, con una perdita di accuratezza trascurabile.
• Prestazioni rispetto a CacheBlend (stato dell'arte per KV recomputation):
  • ALFRED: Miglioramento del 4.13% nel Success Rate (SR) e una riduzione del 1.90x nel TTFT (Time To First Token).
  • WAH-NL: Miglioramento del 3.31-3.37% nel Sub-SR e riduzione del 1.51-1.91x nel TTFT.
• Scalabilità: Con l'aumento del numero di segmenti di memoria recuperati (fino a 40), KEEP mantiene una crescita della latenza molto più lenta rispetto ai metodi di prefilling completo o caching dei prefissi, grazie alla gestione intelligente dell'invalidazione.

4. Contributi Chiave

1. Analisi delle differenze fondamentali: Identificazione delle divergenze critiche tra l'inferenza LLM tradizionale (testo statico) e la pianificazione incarnata (memoria dinamica e aggiornata).
2. Nuovo paradigma di gestione: Introduzione di una strategia ibrida statico-dinamica per minimizzare l'invalidazione della cache senza sacrificare le relazioni contestuali.
3. Meccanismo di ricalcolo adattivo: Sviluppo di un algoritmo di ricalcolo multi-hop basato sulla propagazione dell'importanza, che supera i limiti delle strategie statiche o basate su posizioni fisse.
4. Ottimizzazione dell'hardware: Progettazione di uno scheduler di caricamento bilanciato per livello che risolve i colli di bottiglia di I/O nelle pipeline di inferenza con ricalcolo KV.

5. Significato

KEEP rappresenta un passo avanti significativo per l'efficienza degli agenti LLM in ambienti reali. Dimostra che è possibile gestire grandi quantità di memoria contestuale dinamica senza incorrere nei costi computazionali proibitivi del ricalcolo completo o nelle perdite di accuratezza del semplice riutilizzo della cache. La capacità di mantenere un'alta accuratezza (Success Rate) riducendo drasticamente la latenza (TTFT) rende KEEP una soluzione pratica per l'implementazione di agenti robotici autonomi in scenari complessi e a lungo termine, dove la velocità di decisione è cruciale.