[CVPR 2024] STTS: Unified Spatio-Temporal Token Scoring for efficient Video VLMs

TL;DR

Training and deploying Video-Language Models (VLMs) is notoriously expensive due to the massive number of tokens generated by high-frame-rate videos. STTS (Spatio-Temporal Token Scoring) solves this by introducing a lightweight, learnable module that prunes 50% of visual tokens across the entire pipeline (ViT + LLM). This results in a 62% to 2.25x speedup with negligible impact on accuracy, outperforming traditional merging or heuristic-based pruning methods.

The Bottleneck: Why "Selective Hearing" is Hard for AI

In video understanding, we face a "Double Whammy" of redundancy:

1. Spatial Redundancy: Many pixels in a frame are just background (e.g., the sky).
2. Temporal Redundancy: Consecutive frames are often nearly identical.

Previous researchers tried to solve this by pruning tokens inside the Vision Transformer (ViT) or just before the Large Language Model (LLM). However, pruning only in the LLM leaves the ViT to struggle with every single frame, while pruning only in the ViT often misses the "semantic" needs of the downstream language task.

Methodology: Unified Scoring and Efficient Packing

The core innovation of STTS lies in its Dual-Axis Scoring Mechanism. Unlike methods that require complex text-conditioning (asking the LLM "what should I look for?"), STTS learns what's important implicitly.

1. The Scorer

STTS is inserted early in the ViT (typically after layer 3). It uses a small MLP to assign a score to each patch.

• Spatial Saliency: Learned through backpropagation from the final LLM loss. The model "realizes" which tokens helped it answer the question.
• Temporal Redundancy: Enforced by an Auxiliary Loss that compares the cosine similarity of patches between adjacent frames. If a patch hasn't changed, its score is lowered.

Figure: The STTS architecture showing how scores are injected as attention bias before hard pruning.

2. The Packing Algorithm (Hardware Acceleration)

Pruning creates "ragged" tensors (each frame has a different number of tokens). Standard hardware (GPUs) hates this. STTS uses a First-Fit Descending algorithm to pack these sparse tokens into dense "bins" (Algorithm 1), ensuring that the GPU stays at 100% utilization while processing fewer total tokens.

Experiments: High Speed, Low Drag

The researchers tested STTS on 13 benchmarks using the Molmo2 backbone.

• The "Sweet Spot": At 30% pruning, the model actually improved on some tasks like NextQA, likely because pruning removed distracting "background noise."
• Massive Efficiency: As the number of input frames increases, the benefits of STTS grow quadratically. For a 256-frame input, STTS delivered a 2.25x training speedup.
• Test-Time Scaling (TTS): Because STTS makes tokens "cheaper," you can afford to feed the model more frames during inference. By doubling the frames but pruning 50%, the model outperformed the unpruned baseline by 1.1% on long-video tasks.

Figure: Throughput gains increase significantly as the pruning ratio (k) and frame count increase.

Deep Insight: STTS vs. Heuristics

Why not just use a simple rule like "prune if frames look similar"? The authors show that simple heuristics (green boxes below) often fail by pruning critical details like facial expressions or moving foreground objects just because they look "statistically" similar to previous frames. STTS (purple boxes) learns to keep the "semantic anchors" of the video.

Figure: STTS intelligently preserves the player in a game and facial expressions in a video, while heuristics prune them as "redundant."

Conclusion & Takeaways

STTS proves that we don't need complex, heavy-weight attention mechanisms to handle long videos. A simple, learnable scoring module can drastically reduce the computational footprint of VLMs.

Key Takeaways for Practitioners:

• Prune Early, Prune Often: Pruning at ViT Layer 3 provides the best balance between efficiency and feature robustness.
• Gradients are Better than Rules: Learning importance via the downstream task is vastly superior to hard-coded similarity rules.
• Scaling Potential: This method makes "Long Video" understanding viable on consumer-grade hardware by swapping spatial redundancy for temporal depth.

Limitations

While STTS is highly effective, the authors noted that pruning too early (e.g., at Layer 0) hurts performance because the model hasn't had enough time to form high-level representations of the patches. Finding the optimal "injection layer" remains a vital hyperparameter for different architectures.