Unified Spatio-Temporal Token Scoring for Efficient Video VLMs

Probleemstelling

Video-Visuele Taalmodellen (Video VLMs) staan voor een aanzienlijke rekenkundige uitdaging. Het verwerken van video vereist het coderen van een groot aantal frames, waarbij elk frame door een Vision Transformer (ViT) wordt opgesplitst in honderden patch-tokens. Omdat de aandachtsmechanisme (attention) kwadratisch schaalt met de lengte van de tokenreeks, leidt een toename in het aantal frames tot een explosie in geheugengebruik, een verlaagde trainingsdoorvoer en een hogere inferentielatentie.

Bestaande oplossingen voor token-pruning (het verwijderen van niet-informatieve tokens) zijn vaak onvolledig:

1. Pre-/In-ViT-methoden: Deze werken alleen binnen de ViT en zijn vaak ontworpen voor unimodale taken (zoals objectdetectie), zonder rekening te houden met temporele redundantie in video's of de behoeften van downstream VLM-taken.
2. Post-ViT-methoden: Deze verwijderen tokens pas na de ViT, voordat ze de Large Language Model (LLM) bereiken. Hierdoor blijft de ViT echter alle frames verwerken, wat de grootste rekenkundige bottleneck blijft, vooral bij lange video's. Veel van deze methoden vereisen bovendien complexe, tekst-geconditioneerde selectiemechanismen of token-merging.

Er is dus behoefte aan een holistische oplossing die de rekenlast verlaagt in zowel de ViT als de LLM, zonder de redeneercapaciteiten van het model te schaden.

Methodologie: Spatio-Temporal Token Scoring (STTS)

De auteurs introduceren STTS, een lichtgewicht, end-to-end trainable module die visuele tokens selectief verwijdert over de hele architectuur (zowel ViT als LLM) zonder tekst-geconditioneerde selectie of complexe merging-algoritmen.

De werking van STTS omvat drie gecoördineerde stappen:

1. Dual-Axis Scoring Mechanisme:

   • Ruimtelijke Saliëntie: Een scorer (bestaande uit een Token Pooler en een 3-laags MLP) leert welke tokens binnen een frame belangrijk zijn. Dit wordt geleerd via gradiënten van de downstream VLM-taak (zonder expliciete tekst-input).
   • Temporele Redundantie: De scorer neemt context van het vorige frame mee. Een auxiliary loss (hulpverlies) wordt gebruikt om de scorer te trainen om hoge overeenkomst (cosine similarity) tussen aangrenzende frames te koppelen aan lage belangrijkheidsscores. Dit dwingt het model om statische, redundante achtergrondinformatie te verwijderen.

2. Bias Injectie en Hard Pruning:

   • De scores worden omgezet in een bias die wordt geïnjecteerd in de attention-matrix van de volgende ViT-laag. Dit maakt de module trainbaar via backpropagation.
   • Na deze laag worden de tokens met de laagste scores (onder een bepaalde drempel $k\%$) permanent verwijderd ("hard pruning").

3. Efficiënte Packing Algorithm:

   • Omdat het aantal verwijderde tokens per frame varieert (bijv. minder beweging = meer pruning), ontstaan er "ragged" (onregelmatige) tensors die inefficiënt zijn voor GPU-berekeningen.
   • STTS gebruikt een packing-algoritme (first-fit descending) om de overgebleven tokens van meerdere frames te consolideren in compacte, dichte tensors. Dit maximaliseert de rekenkundige doorvoer en zorgt voor echte versnelling, in plaats van alleen masking.

Belangrijkste Bijdragen

1. Unificatie van Token Pruning: STTS is de eerste methode die naadloos tokens pruneert in zowel de ViT als de LLM, waardoor de totale rekenlast van de pipeline wordt verlaagd.
2. Dual-Axis Scoring: Het combineert ruimtelijke relevantie (geleerd via LLM-gradiënten) en temporele redundantie (geleerd via een cosine-similarity loss) in één eenvoudige module.
3. Significante Efficiëntiewinst: De methode kan veilig 50% van de visuele tokens verwijderen, wat leidt tot een tot 62% hogere efficiëntie tijdens training en inferentie.
4. Schalbaarheid: De efficiëntiewinst neemt toe bij het verwerken van meer frames per video. Bovendien kan STTS worden gecombineerd met Test-Time Scaling (TTS) om bij lange video's meer frames te verwerken binnen hetzelfde rekenbudget, wat de prestaties verder verbetert.

Resultaten

De auteurs hebben STTS geëvalueerd op 13 korte en lange video-QA-benchmarks (zoals VideoMME, MLVU, LongVideo).

• Prestatiebehoud: Bij een pruning-ratio van 50% daalt de gemiddelde prestatie over alle taken slechts met 0,7%. Op sommige benchmarks (zoals NextQA) zelfs een lichte verbetering ten opzichte van de baseline.
• Efficiëntie:
  • Bij 128 frames: 1,62x snellere training en 1,61x snellere inferentie.
  • Bij 256 frames (zwaardere last): De versnelling loopt op tot 2,25x voor training en 2,22x voor inferentie, dankzij de kwadratische complexiteit van de attention-mechanismen.
• Vergelijking met Baselines: STTS presteert aanzienlijk beter dan willekeurige pruning (Random) en heuristische methoden (gebaseerd op cosine similarity zonder leren). Vooral bij lange video's is het verschil groot, omdat STTS semantisch belangrijke voorgrondobjecten behoudt terwijl heuristische methoden deze soms per ongeluk verwijderen.
• Test-Time Scaling (TTS): Door bij inferentie meer frames te sample (bijv. van 64 naar 128) terwijl er 50% wordt gepruned, behaalt het model een prestatieverbetering van 0,5% tot 1,1% ten opzichte van de niet-gepruned baseline.

Betekenis en Conclusie

STTS biedt een eenvoudige, interpreteerbare en krachtige oplossing voor de schaalbaarheid van Video VLMs. Door de rekenlast in de ViT te verminderen, wordt de bottleneck voor de LLM ook opgelost. Het bewijst dat men de token-lengte drastisch kan verkorten zonder de redeneercapaciteit van het model te verliezen, zolang de pruning gebaseerd is op een leerbaar mechanisme dat zowel ruimtelijke als temporele context begrijpt.

De methode maakt het mogelijk om Video VLMs in te zetten in omgevingen met beperkt geheugen of lage latentie-eisen, en opent de deur voor het verwerken van veel langere video-sequenties dan tot nu toe haalbaar was.