ASAP: Attention-Shift-Aware Pruning for Efficient LVLM Inference

Titel: ASAP: Attention-Shift-Aware Pruning voor Efficiënte LVLM-Inferentie

Auteurs: Surendra Pathak en Bo Han (George Mason University)

1. Het Probleem

Grote Vision-Language Modellen (LVLMs), zoals LLaVA-NeXT, tonen uitzonderlijke prestaties in multimodale taken, maar kampen met een kritieke bottleneck: de kwadratische rekenkosten die gepaard gaan met het verwerken van visuele tokens.

• Token-explosie: Hoogresolutie afbeeldingen worden vaak opgesplitst in veel kleine tegels (tiles), wat leidt tot duizenden visuele tokens (bijv. ~2.304 tokens voor LLaVA-NeXT).
• Attention Shift (Aandachtverschuiving): Bestaande methoden voor token-pruning (het verwijderen van minder belangrijke tokens) vertrouwen vaak op de "attention scores" van het model. De auteurs identificeren echter een fundamenteel probleem: door de Rotary Position Embeddings (RoPE) in Transformer-architecturen krijgen visuele tokens die later in de sequentie verschijnen (vaak de onderkant van een afbeelding) onterecht hoge attention scores. Dit komt door een "distance decay" (afstandsdemping) in RoPE, niet door de semantische relevantie van de token.
• Gevolg: Naïeve pruning-methoden verwijderen belangrijke achtergrondinformatie (zoals verre objecten) en behouden irrelevante tokens, wat de modelprestaties aanzienlijk verlaagt.

2. Methodologie: ASAP

ASAP is een trainingsvrij (training-free) en KV-Cache-compatibel pruning-framework dat bestaat uit drie kerncomponenten om de attention shift te mitigeren en redundantie te elimineren.

A. Salience-Guided Bidirectional Attention (SG-BiMask)

Om het probleem van de attention shift op te lossen, verzwakt ASAP de strikte causale (unidirectionele) attention-maskering specifiek voor visuele tokens.

• In plaats van dat een token alleen naar eerdere tokens kan kijken, wordt een dynamische, bidirectionele zachte attention-mask ingevoerd.
• Dit stelt eerdere visuele tokens in staat om "te kijken" naar latere tokens, waardoor context kan worden geaggregeerd voor een holistisch begrip van het 2D-beeld.
• De maskering is salience-gestuurd: de mate waarin een toekomstige token zichtbaar is, hangt af van zijn structuur- en informatiewaarde (berekend via de pre-softmax alignment matrix). Dit voorkomt dat ruis uit achtergrondgebieden wordt doorgegeven.

B. Salience-Weighted Consolidation (Samenvoeging)

Nadat de attention shift is gecorrigeerd, blijven er vaak semantisch redundante tokens over (bijv. aangrenzende tokens van een uniforme lucht).

• ASAP introduceert een gewogen samenvoegingsmechanisme. In plaats van tokens simpelweg te verwijderen of te middelen, worden semantisch vergelijkbare tokens samengevoegd.
• De weging is gebaseerd op de attention mass (salience score) die in de vorige stap is berekend. Tokens met een hogere salience fungeren als "anker" en behouden hun kenmerken, terwijl minder belangrijke tokens daarin worden opgenomen. Dit voorkomt verdunning van belangrijke visuele features.

C. Budget Reallocation (Herverdeling)

• Tijdens het samenvoegen ontstaan er vrije plekken in het token-budget.
• ASAP vult deze gaten niet met willekeurige tokens, maar redt (salvage) de hoogst gerangschikte tokens uit de oorspronkelijk "gesnoeide" pool.
• Dit zorgt ervoor dat het eindresultaat een compacte set is die maximale semantische diversiteit en dichtheid behoudt.

3. Belangrijkste Bijdragen

1. Identificatie van Attention Shift: De auteurs attribueren de bias in attention scores systematisch aan de afstandsdempingseigenschappen van RoPE in LVLMs.
2. Structurele Mitigatie: Voor het eerst wordt de attention shift structureel opgelost door een zachte bidirectionele attention-mechanisme te gebruiken, in plaats van RoPE volledig te verwijderen (wat de ruimtelijke redenering zou beschadigen).
3. Novel Framework: ASAP combineert attention-shift-aware pruning met feature-based redundantie-eliminatie in een trainingsvrij, plug-and-play framework dat compatibel is met bestaande KV-caches en multi-turn conversaties.
4. State-of-the-Art Prestaties: Het framework behaalt bijna verliesloze compressie met aanzienlijke reducties in rekenkracht.

4. Resultaten

ASAP werd getest op modellen zoals LLaVA-1.5, LLaVA-NeXT-7B en InternVL 2.5 tegenover bestaande baselines (zoals FastV, PruMerge+, DivPrune).

• Prestatiebehoud:
  • Op LLaVA-NeXT-7B behoudt ASAP 99,02% van de originele prestaties, terwijl het het aantal FLOPs (rekenoperaties) met ~80% reduceert.
  • Op LLaVA-1.5-13B behoudt het 98,68% van de prestaties bij een FLOPs-reductie van ~66%.
• Superioriteit: ASAP presteert consequent beter dan de sterkste bestaande methoden (bijv. DivPrune), met name bij agressieve pruning (bijv. 128 tokens behouden).
• Supra-vanilla Prestaties: In sommige benchmarks (zoals MMBench en MME) overtreft ASAP zelfs het volledige, niet-gepruneerde model.
• Efficiëntie: De methode reduceert de GPU-geheugenvraag en de CUDA-latentie aanzienlijk, met slechts een minimale overhead van ~2 ms voor de dynamische selectie.

5. Betekenis en Conclusie

Dit werk toont aan dat een zorgvuldig ontworpen pruning-strategie, die rekening houdt met de specifieke architecturale beperkingen van LVLMs (zoals RoPE-bias) en de inherente redundantie van visuele data, leidt tot enorme efficiëntiewinsten zonder kwaliteitsverlies.

ASAP biedt een praktische oplossing voor de schaalbaarheid van LVLMs, waardoor het mogelijk wordt om hoogresolutie beelden te verwerken met een fractie van de huidige rekenkosten, wat essentieel is voor de adoptie van deze modellen in real-time toepassingen en op beperkte hardware.