Reducing Peak Memory Usage for Modern Multimodal Large Language Model Pipelines

Probleemstelling

Multimodale grote taalmodellen (MLLM's) bereiken sterke prestaties in visueel-taalredenering door te schalen naar hoge-resolutie afbeeldingen en lange video's. Deze schaalbaarheid introduceert echter een aanzienlijke geheugentijdoverhead tijdens de inferentie, voornamelijk veroorzaakt door de groei van de Key-Value (KV) cache.

• De Prefill-fase als bottleneck: Bestaande compressiemethoden voor de KV-cache werken meestal na dat de volledige multimodale context is verwerkt (post-prefill). Dit laat echter de piekgeheugengebruik tijdens de prefill-fase (waarbij de volledige invoer wordt gecodeerd) onaangeroerd.
• Explosie van visuele tokens: Bij hoge-resolutie afbeeldingen en video's kan het aantal visuele tokens duizenden bedragen voordat de generatie begint. Dit zorgt ervoor dat de KV-cache die tijdens de prefill wordt opgebouwd, het geheugengebruik domineert en vaak leidt tot "Out-of-Memory" (OOM) fouten.
• Beperkingen van bestaande methoden:
  • KV-cache compressie: Werkt vaak te laat (na volledige invoer).
  • Token pruning: Verwijdert invoertokens op input-niveau, maar negeert de heterogene rol van tokens in verschillende lagen en attention-heads, wat het risico op het verliezen van structureel belangrijke informatie vergroot.

Methodologie

De auteurs stellen een prefill-bewust, structureel bewust KV-cache compressieframework voor. In plaats van eerst de volledige cache te bouwen en deze later te comprimeren ("first process, then compress"), past hun methode compressie incrementeel toe tijdens het encoderen van de invoer ("compress as you prefill").

1. Block-wise Prefill (Algoritme 1):
In plaats van de volledige sequentie in één keer te verwerken, wordt de invoer opgesplitst in opeenvolgende blokken (bijv. grootte 256).

• Na het coderen van elk blok worden de gegenereerde KV-paren toegevoegd aan de cache.
• Onmiddellijk wordt de cache gecontroleerd op een vast geheugenbudget ($M$).
• Als de cache het budget overschrijdt, worden overtollige tokens direct verwijderd (geëvacueerd) voordat het volgende blok wordt verwerkt.
• Dit zorgt ervoor dat het piekgeheugengebruik gedurende de hele invoerverwerking constant blijft en niet lineair groeit.

2. Evictie-strategieën:
Binnen dit raamwerk worden twee strategieën gebruikt om te bepalen welke tokens worden verwijderd:

• Query-aware Eviction (SnapKV): Voor single-turn scenario's. Hierbij worden proxy-query tokens uit de tekstuele prompt gebruikt om cross-attention te berekenen over de gecacheerde keys. Tokens met een lage belangheidsscore (minder relevant voor de vraag) worden verwijderd. Dit behoudt visueel opvallende gebieden die relevant zijn voor de taak.
• Query-agnostic Eviction (KeyDiff): Voor multi-turn scenario's waar de query nog niet beschikbaar is. Deze methode behoudt de representatieve diversiteit door keys te behouden die het sterkst afwijken van het gemiddelde (anchor vector). Dit voorkomt het verwijderen van zeldzame visuele kenmerken zonder afhankelijk te zijn van een query.

3. Structurele Alignering:
De auteurs benadrukken dat blokgrenzen moeten aligneren met de structurele organisatie van visuele data (ruimtelijke tegels voor afbeeldingen, temporale groepen voor video's). Dit verbetert de robuustheid van de compressie.

Belangrijkste Bijdragen

1. Shift in Paradigma: De eerste aanpak die KV-cache compressie toepast tijdens de prefill-fase in plaats van erna, waardoor het piekgeheugengebruik effectief wordt begrensd.
2. Vaste Geheugenbegroting: Een framework dat een strikt geheugenbudget handhaaft tijdens de invoerverwerking, waardoor OOM-fouten bij grote visuele inputs worden voorkomen.
3. Flexibiliteit: Ondersteuning voor zowel query-afhankelijke (single-turn) als query-onafhankelijke (multi-turn) scenario's.
4. Structuur-bewustzijn: Het inzien dat visuele tokens structurele regulariteiten hebben die eerder in de pijplijn kunnen worden benut dan tekstuele tokens.

Resultaten

De methode is getest op modellen zoals InternVL3.5-8B/14B en Qwen2.5-VL-7B/32B met benchmarks voor afbeeldingen (ImageNeedleInHaystack, V*) en video's (MLVU, Video-MME).

• Geheugenreductie: De methode bereikt een compressie van ongeveer 90% (bij een budget van 1024 tokens) met minimale degradatie in prestaties.
• Piekgeheugen: In tegenstelling tot de baseline (waar het geheugen lineair groeit met het aantal tegels), blijft het piekgeheugen bij de voorgestelde methode bijna constant, zelfs bij inputs van 12.800 visuele tokens. De baseline faalt hierbij vaak door OOM.
• Prestatiebehoud:
  • Bij InternVL3.5-8B met een budget van 1024 is de daling in gemiddelde nauwkeurigheid slechts 0,78% (SnapKV).
  • Bij Qwen2.5-VL-7B is de daling bij 1024 tokens ongeveer 4-11% afhankelijk van de strategie, maar blijft het model bruikbaar waar de baseline faalt.
• Trade-off: Er is een afweging tussen geheugen en latentie. De sequentiële verwerking verhoogt de "Time to First Token" (TTFT) iets, maar maakt inferentie mogelijk op hardware waar dit anders onmogelijk zou zijn.
• Block Size: De prestaties zijn het beste wanneer de blokgrootte aligneert met de native visuele tokenisatie van het model (bijv. 784 tokens voor Qwen2.5-VL-7B), wat aantoont dat structurele alignering cruciaal is.

Significantie

Dit werk is significant omdat het een praktische oplossing biedt voor het schaalprobleem van MLLM's. Het maakt inferentie mogelijk op grote visuele inputs (hoge resolutie, lange video's) op hardware met beperkt VRAM, zonder dat het model opnieuw getraind hoeft te worden.

• Het verandert de focus van "hoe comprimeren we na het laden?" naar "hoe houden we het geheugen onder controle tijdens het laden?".
• Het biedt een schaalbare route voor multimodale toepassingen in real-world systemen waar geheugen een beperkende factor is.
• Het benadrukt dat geheugenefficiëntie niet alleen afhangt van de uiteindelijke cache-grootte, maar ook van hoe de visuele context tijdens de prefill wordt verwerkt.

Kortom, de auteurs tonen aan dat door slim, structureel bewust en incrementeel comprimeren tijdens de prefill-fase, de barrière voor het verwerken van complexe multimodale data op bestaande hardware aanzienlijk kan worden verlaagd.