PaperlayerView exact revision on arXiv

Source-linked AI summary

Unified Spatio-Temporal Token Scoring for Efficient Video VLMs

Jianrui Zhang, Yue Yang, Rohun Tripathi, Winson Han, Ranjay Krishna, Christopher Clark, Yong Jae Lee, Sangho Lee

arXiv:2603.18004v1cs.CVcs.AIcs.LG

TL;DR

Video VLMs incur high computation from long visual token sequences, while prior pruning approaches separately target the ViT or LLM and do not provide a unified solution. STTS introduces end-to-end spatio-temporal scoring across both components, reducing visual tokens by 50% and improving efficiency while largely preserving performance. Test-time scaling further improves long-video QA results.

  • Problem

    Long video token sequences increase attention, memory, training, and inference costs, while existing approaches do not unify pruning across the ViT and LLM.

  • Method

    STTS uses downstream gradients for spatial scoring, an auxiliary loss for temporal scoring, and token packing to prune vision tokens across the ViT and LLM.

  • Results

    50% token pruning improves training and inference efficiency by up to 62% with negligible performance loss, while test-time scaling improves long-video QA by 0.5–1%.

  • Takeaways & Limitations

    STTS provides unified architecture-wide video token pruning that reduces computation while preserving video reasoning performance across diverse QA benchmarks.

  • Takeaways & Limitations

    Post-ViT pruning methods leave the ViT processing every frame, whereas STTS’s pruning-layer choice is constrained by the need to preserve robust patch representations before discarding information.

Abstract

from arXiv · show

Token pruning is essential for enhancing the computational efficiency of vision-language models (VLMs), particularly for video-based tasks where temporal redundancy is prevalent. Prior approaches typically prune tokens either (1) within the vision transformer (ViT) exclusively for unimodal perception tasks such as action recognition and object segmentation, without adapting to downstream vision-language tasks; or (2) only within the LLM while leaving the ViT output intact, often requiring complex text-conditioned token selection mechanisms. In this paper, we introduce Spatio-Temporal Token Scoring (STTS), a simple and lightweight module that prunes vision tokens across both the ViT and the LLM without text conditioning or token merging, and is fully compatible with end-to-end training. By learning how to score temporally via an auxiliary loss and spatially via LLM downstream gradients, aided by our efficient packing algorithm, STTS prunes 50% of vision tokens throughout the entire architecture, resulting in a 62% improvement in efficiency during both training and inference with only a 0.7% drop in average performance across 13 short and long video QA tasks. Efficiency gains increase with more sampled frames per video. Applying test-time scaling for long-video QA further yields performance gains of 0.5-1% compared to the baseline. Overall, STTS represents a novel, simple yet effective technique for unified, architecture-wide vision token pruning.

1 Introduction

Video VLMs face rapidly increasing computation from long visual token sequences, while prior pruning methods address only the ViT or LLM and miss a unified multimodal solution. STTS prunes tokens across both components using spatial and temporal scoring, achieving substantial efficiency gains with limited performance loss.

  • Motivation: Video token sequences become quadratically expensive under attention as frame counts increase, raising memory use and reducing training and inference efficiency.The burden affects both the ViT encoder and the LLM consuming its output.
  • Limitations of Prior Work: Prior methods prune within the ViT for unimodal perception or after the ViT for LLM input, leaving no holistic solution for scalable video VLMs.Post-ViT methods leave the ViT processing every frame, while pre- or in-ViT methods do not explicitly address multimodal objectives and temporal redundancy.
  • STTS: STTS is a lightweight, end-to-end trainable module that prunes visual tokens across both the ViT and LLM without text-conditioned selection or complex merging.It is designed to reduce the visual token burden throughout the VLM pipeline with minimal architectural modification.
  • STTS: STTS scores intra-frame spatial saliency through downstream multimodal objectives and inter-frame temporal redundancy through an auxiliary loss.The dual-axis scoring mechanism jointly targets spatially unimportant and temporally redundant tokens.
  • Results: 50% of visual tokens can be dropped, improving training throughput and inference efficiency by up to 62% with negligible performance loss.The reported efficiency gain applies to both training and inference.
  • Results: Sampling more video frames further increases efficiency gains, while test-time scaling produces consistent 0.5–1% improvements on long-video QA benchmarks.The test-time scaling result is reported relative to the baseline.

2 Related Works

Related pruning methods primarily target spatial redundancy in images, vision-only tasks, or tokens after the ViT. STTS instead applies simple, merge-free spatial and temporal pruning within the ViT and evaluates it on downstream video-LLM tasks.

  • Within-ViT Pruning: Image-ViT methods such as SPViT, FastViT, and ToMe reduce spatial redundancy through token packaging, mixing, or matching but do not address temporal redundancy in video.These approaches primarily target static-image processing.
  • Within-ViT Pruning: Other methods use early exiting, pruning decoders, or pre-ViT temporal redundancy detection, but are typically demonstrated on vision-only tasks rather than downstream VLMs.The cited examples include DToP, VLTP, and Run-Length Tokenization.
  • STTS: STTS is designed as a simple, merge-free module that performs spatial and temporal pruning within the ViT and is evaluated on downstream video-LLM tasks.This design distinguishes it from methods focused only on static images, vision-only tasks, or post-ViT processing.
  • Post-ViT Pruning: Post-ViT approaches aggregate or merge vision tokens between the encoder and LLM using spatial, temporal, or text-conditioned mechanisms.Examples include FreeVA, PruneVid, STTM, HoliTom, FastVid, and LLaVA-PruMerge.
  • Post-ViT Pruning: Post-ViT pruning leaves the ViT processing every input frame, preserving a major bottleneck for long videos and often requiring complex merging or text-conditioned modules.STTS addresses these limitations by pruning from within the ViT using a simple, merge-free scoring mechanism.

3 Spatio-Temporal Token Scoring (STTS)

STTS is a lightweight module that scores and prunes vision tokens across the ViT and LLM while preserving end-to-end training. It combines spatial and temporal scoring with packing to achieve computational savings from non-uniform video-token pruning.

  • Architecture: STTS prunes vision tokens across both the ViT and LLM without architecture-specific constraints, text-conditioned selection, or complex merging algorithms.The module is inserted into the ViT, so reduced token counts propagate through the remainder of the VLM.
  • Scoring: STTS scores tokens along spatial saliency and inter-frame temporal redundancy using a scorer, downstream attention gradients, and an auxiliary temporal loss.The scorer uses pooled features from neighboring frames, while score bias injection allows task-loss gradients to train spatial saliency.
  • Scorer Architecture: The scorer pools ViT features by width w, concatenates each frame with its predecessor, and keeps the first video frame intact because it lacks temporal context.With w = 3, the spatial dimension changes from N to N/w^2 before scoring.
  • Token Pruning and Packing: After layer l + 1, STTS removes the bottom-k% tokens, creating non-uniform frame sequences that require packing for dense batched computation.The packing algorithm sorts frames by valid-token count, consolidates sequences into compact tensor entries, and uses an attention mask to preserve source-frame attention boundaries.
  • Temporal Scoring: The temporal auxiliary loss aligns predicted scores with one minus adjacent-frame patch cosine similarity, making higher redundancy correspond to lower importance.The final objective combines the primary task loss with the averaged MSE auxiliary loss, while the first frame receives zero auxiliary loss because it is never pruned.

4 Experiments

STTS is evaluated across short- and long-video QA tasks, with pruning preserving performance while substantially improving training and inference efficiency. Efficiency gains increase with pruning strength and sampled frame count.

  • Evaluation Setup: STTS is evaluated on standard short- and long-video QA benchmarks using the Molmo2 training recipe and architecture.The setup uses a SigLIP 2 ViT connected to a Qwen3-4B LLM and trains on the video QA subset for 6,250 steps.
  • Video QA Results: 30% pruning maintains or exceeds baseline performance on tasks including NextQA and VideoMME.The scorer preserves task-essential tokens while cosine similarity targets redundant background information.
  • Video QA Results: 50% pruning causes only a 0.7% average performance decline across diverse video tasks.On VideoMME, the decrease is 0.4 points, indicating robustness when half of the visual context is discarded.
  • Efficiency Gains: 50% pruning achieves 1.62x training and 1.61x inference speedups in the 128-frame setting.This setting reduces the token load by approximately 33%.
  • Efficiency Gains: 50% pruning achieves 2.25x training and 2.22x inference speedups in the 256-frame setting.The larger gains reflect the increasing computational savings from token reduction as sequence length grows under quadratic attention complexity.
  • Efficiency Gains: Inference speedups are marginally lower than training speedups because inference handles dynamic sequence lengths during prefill.Training benefits from torch.compile and uniformly padded static sequences.

5 Ablation Studies

Ablations show that learned scoring outperforms simpler pruning baselines, later ViT injection layers are preferable, and test-time scaling improves long-video QA without retraining. Visualizations illustrate how STTS preserves salient content while discarding redundant tokens.

  • Scorer Versus Heuristic: Random pruning underperforms heuristic and scorer-based methods by approximately 1%, while STTS leads the heuristic by 0.5% on long videos.The scorer is especially useful when long videos contain weak temporal redundancy after uniform frame sampling.
  • Scorer Versus Heuristic: The learned scorer uses downstream spatial signals to compensate for weak temporal cues in long-video inputs.The no-auxiliary-loss variant performs worse than Random, supporting the role of temporal regularization.
  • Selecting Pruning Layer Depth (l): Pruning at ViT layers l = 0 or l = 1 hurts performance, while l = 2 is marginally weaker than l = 3.The comparison indicates that premature pruning can discard information before robust patch representations form.
  • Test-Time Scaling: 30% + TTS reaches a Long QA average of 60.1, outperforming the baseline by 1.1%.TTS increases the number of sampled frames during inference without retraining.
  • Test-Time Scaling: 50% + TTS reaches a Long QA average of 59.4, surpassing the baseline by roughly 0.5%.Across configurations, TTS improves pruned models by roughly 1% relative to their pre-TTS counterparts.
  • Pruning Visualizations: STTS preserves dynamic foreground content while heuristic pruning can discard semantically important faces or moving objects.The visualizations contrast STTS’s learned semantic filtering with heuristic selection based on shallow feature similarity.
  • Aggressive Pruning: At 80% pruning, STTS achieves a QA Average of 59.8% versus Random’s 57.5%, a 2.3% absolute improvement.The performance gap widens as the retained visual-token budget becomes increasingly constrained.

6 Conclusion

STTS unifies end-to-end token pruning across the vision encoder and LLM, using downstream task gradients and an auxiliary temporal loss. It reduces visual tokens while preserving foreground semantics, accelerates video VLMs, and supports longer temporal contexts with test-time scaling.

  • STTS unifies token pruning across both the vision encoder and the LLM in an end-to-end trainable framework.
  • 50% visual-token reduction accelerates both training and inference by over 60% across 13 diverse video QA benchmarks with negligible performance degradation.
  • Downstream task gradients and an auxiliary temporal loss filter redundant background noise while preserving critical semantic foregrounds.
  • STTS pairs with test-time scaling to process substantially longer temporal contexts under strict computational constraints.

Appendix

The appendix reports image-benchmark results and video-QA comparisons at a 50% pruning rate. STTS preserves image-only accuracy and outperforms the listed inference-only and fully trained ToMe baselines on video QA.

  • STTS pruning on videos does not degrade overall image performance when evaluated on image benchmarks.
  • At a 50% pruning rate, STTS significantly outperforms Spatial Heuristic, ToMe, and fully trained ToMe on video QA tasks.

A Image Results

STTS maintains image-QA accuracy despite video-token pruning and shows a one-point improvement on multi-image QA. The authors attribute this improvement to a possible transfer effect from video-trained temporal reasoning.

  • STTS preserves image-QA accuracy compared with the differently trained Molmo2 version used as the comparison model.
  • 1-point improvement on multi-image QA is hypothesized to result from transfer of temporal reasoning skills learned from video data.

B Detailed Throughput Tables

The detailed throughput analysis separates training and inference results: Table 6 reports training, while Table 7 reports inference on MLVU.

  • Table 6 reports the training throughput analysis.
  • Table 7 reports inference throughput on MLVU.

C Comparison with ViT-Only Pruning Baselines

STTS is evaluated against ViT-centric and inference-only pruning baselines, with results favoring training on pruned sequences and joint cross-modal pruning. Its packing algorithm compresses sparse tokens into denser tensors while preserving attention structure.

  • Efficiency: Training and inference speed comparisons report fewer tokens per instance and increased throughput and speedup as pruning setup k increases.Efficiency gains grow with the maximum number of sampled frames during training, with identical trends observed during inference.
  • Baseline comparison: STTS outperforms both inference-only baselines and fully trained Token Merging, showing that pruning requires training on pruned input sequences.The comparison uses heuristic STTS, Token Merging, and a fully trained Token Merging baseline.
  • Baseline comparison: 62.3 QA Average is achieved by STTS at 50% pruning, mitigating the roughly 1-point drop observed for trained Token Merging across Short QA and Long QA.The passage attributes this margin to avoiding limitations of image-based pruning through cross-modal pruning.
  • Token packing: The packing algorithm sorts frames by valid-token count and assigns each frame to the first new frame with sufficient remaining capacity.This procedure uses an O(T^2) search to find the best fitting packed frame.
  • Token packing: The algorithm scatters tokens into a denser tensor and builds a block-diagonal attention mask before returning the packed representation.The outputs are Xpacked and Maskpacked.

E Detailed Performance Degradation Tables

Table 8 complements Figure 1 by comparing Random and STTS across k = 50 to k = 90. STTS consistently outperforms Random, while the text-only baseline at k = 100 provides a lower bound.

  • Table 8 complements the right subfigure of Figure 1.
  • STTS consistently outperforms Random from k = 50 to k = 90.
  • The text-only baseline at k = 100 is provided as a lower bound.
Loading 2603.18004v1…