Reducing Peak Memory Usage for Modern Multimodal Large Language Model Pipelines

Panoramica del Problema

I Modelli Linguistici Multimodali (MLLM) moderni hanno raggiunto capacità di ragionamento visivo-testuale eccezionali scalando verso immagini ad alta risoluzione e sequenze video lunghe. Tuttavia, questa scalabilità introduce un enorme sovraccarico di memoria durante l'inferenza, in particolare nella fase di prefill (l'elaborazione iniziale del contesto completo prima della generazione).

Il collo di bottiglia principale risiede nella KV Cache (Key-Value Cache), necessaria per l'attenzione self-attention nei transformer.

• Il problema: Nelle architetture multimodali, l'input visivo (es. immagini 4K o video lunghi) genera migliaia di token visivi prima ancora che inizi la decodifica. La KV Cache cresce linearmente con il numero di token, causando picchi di memoria che spesso superano la capacità delle GPU (Out-of-Memory o OOM) durante la fase di prefill.
• Limiti delle soluzioni attuali: I metodi esistenti di compressione della KV Cache operano tipicamente dopo che tutto il contesto multimodale è stato elaborato ("process first, compress later"). Questo non risolve il picco di memoria iniziale. I metodi di "token pruning" (potatura dei token) agiscono a livello di input ma ignorano le strutture eterogenee dei layer e delle teste di attenzione, rischiando di rimuovere informazioni strutturalmente importanti.

Metodologia Proposta

Gli autori propongono un framework di compressione della KV Cache consapevole del prefill (prefill-aware) e consapevole della struttura (structure-aware), che opera in modo sequenziale sotto un budget di memoria fisso.

L'approccio si basa su tre pilastri fondamentali:

1. Elaborazione a Blocchi (Block-wise Prefill):
   Invece di costruire l'intera KV Cache e poi comprimere, il metodo divide la sequenza di input in blocchi contigui (es. 256 token). Dopo l'encoding di ogni blocco, le coppie KV vengono aggiunte alla cache e immediatamente sottoposte a potatura (eviction) per rispettare il budget di memoria $M$. Questo mantiene la dimensione della cache costante durante tutto il processo di prefill, evitando i picchi di memoria.

2. Strategie di Evizione (Compressione):
   Il framework utilizza due strategie complementari per decidere quali token mantenere o scartare:

   • Query-Aware (Consapevole della Query): Per scenari a turno singolo, si utilizzano token proxy estratti dal prompt testuale per calcolare l'attenzione incrociata con le chiavi cache. I token con punteggi di importanza bassi (meno rilevanti per la domanda) vengono rimossi. Questo preserva le regioni visivamente salienti.
   • Query-Agnostic (Indipendente dalla Query): Per scenari multi-turno o quando la query non è disponibile, si utilizza una strategia basata sulla diversità rappresentazionale (KeyDiff). Si mantiene la diversità delle chiavi preservando quelle che si discostano maggiormente dalla media (outliers), evitando confronti $O(N^2)$ e garantendo che le caratteristiche visive rare non vengano perse.

3. Allineamento Strutturale:
   I confini dei blocchi sono allineati con la struttura visiva (es. griglie spaziali per le immagini, gruppi di frame per i video). Questo garantisce che le decisioni di compressione avvengano a un livello semantico significativo, migliorando la robustezza.

Contributi Chiave

• Spostamento del paradigma: Passare da una compressione post-prefill a una compressione durante il prefill, garantendo un budget di memoria costante fin dall'inizio.
• Framework ibrido: Introduzione di strategie sia query-aware che query-agnostic per adattarsi a diversi scenari di interazione (single-turn vs multi-turn).
• Ottimizzazione strutturale: Dimostrazione che la granularità della compressione deve rispettare la tokenizzazione visiva nativa del modello (es. allineare la dimensione del blocco con la griglia di patch) per massimizzare l'efficacia.

Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli come InternVL3.5-8B/14B e Qwen2.5-VL-7B/32B su benchmark di immagini (ImageNeedleInHaystack, V*) e video (MLVU, Video-MME).

• Riduzione della Memoria: Il metodo mantiene la memoria di picco quasi costante anche all'aumentare del numero di tile di immagine o frame video, prevenendo completamente gli errori OOM che affliggono i metodi baseline (Full Cache).
• Efficienza della Compressione: È possibile raggiungere una compressione della KV Cache fino al 90% (budget di 1024 token) con una degradazione delle prestazioni minima o nulla.
  • Esempio: Su InternVL3.5-8B, con un budget di 1024, la precisione media scende solo di 0.78 punti rispetto al baseline completo.
• Trade-off Memoria-Latenza: Sebbene l'elaborazione sequenziale aumenti leggermente la latenza (Time To First Token - TTFT), il metodo rende fattibili input visivi di grandi dimensioni che altrimenti causerebbero OOM.
• Robustezza: La strategia Query-Agnostic (KeyDiff) mantiene prestazioni competitive anche senza segnali di query, rendendo il metodo adatto a contesti interattivi complessi.
• Confronto con Riduzione Input: A differenza del semplice ridurre la risoluzione dell'immagine (che degrada drasticamente le prestazioni), la compressione della KV Cache preserva le informazioni visive ad alta risoluzione controllando solo la memoria.

Significato e Implicazioni

Questo lavoro è significativo perché affronta il problema fondamentale della scalabilità degli MLLM: la gestione della memoria durante l'elaborazione di contesti lunghi e ad alta risoluzione.

• Abilitazione di scenari pratici: Permette l'inferenza di modelli multimodali su hardware con memoria limitata, rendendo possibile l'uso di immagini 4K o video lunghi senza richiedere GPU di fascia altissima.
• Nuova prospettiva: Sottolinea che l'efficienza della memoria non dipende solo dalla dimensione finale della cache, ma da come il contesto visivo viene processato durante la fase di prefill.
• Futuro: Suggerisce che l'ottimizzazione dell'inferenza può avvenire senza modificare l'addestramento del modello, offrendo una soluzione immediata per l'industria e la ricerca.

In sintesi, il paper propone una soluzione elegante ed efficace che trasforma un collo di bottiglia critico (il picco di memoria nel prefill) in un compromesso gestibile tra memoria e latenza, abilitando l'inferenza multimodale su larga scala.