Visual Token Pruning by Attention Mass

hardPROEfficient Adaptation & Hallucination MetricsVision-Language ModelsVLM 1: Vision-Language Models

Problem Statement

Inference-time VLM speedups (FastV, and friends) drop the visual tokens the language model barely attends to. Given the LM's attention over visual tokens, keep the smallest set of tokens covering a target fraction of the attention mass and report the FLOPs saved.

Background

Let attn be the total attention weight each visual token receives (summed over the query tokens that look at it). Normalize it to a distribution, sort tokens by mass descending (ties broken by smaller original index), and keep tokens one by one until their cumulative mass first reaches keep_fraction. The rest are pruned.

Because self-attention cost scales with the number of tokens squared, dropping from N to k visual tokens cuts the visual-attention FLOPs by **1 - (k/N)2.

Your Task

Implement:

def prune_visual_tokens(attn, keep_fraction):

Return a dict:

"kept": sorted list of original indices of the kept tokens.
"num_kept": how many were kept.
"flops_saved": **1 - (num_kept/N)2, rounded to 4 decimals.

Input Format

attn: list of N non-negative attention weights.
keep_fraction (float) in (0, 1].

Output Format

A dict with the three keys above.

Sample

print(prune_visual_tokens([0.5, 0.3, 0.1, 0.1], 0.75))

Output:

{'kept': [0, 1], 'num_kept': 2, 'flops_saved': 0.75}

Example

Input

print(prune_visual_tokens([0.5, 0.3, 0.1, 0.1], 0.75))

Output

{'kept': [0, 1], 'num_kept': 2, 'flops_saved': 0.75}

Reasoning

Normalized mass is [0.5,0.3,0.1,0.1]. Take token 0 (0.5), then token 1 (cum 0.8 >= 0.75) and stop. Kept {0,1}; 2 of 4 tokens; FLOPs saved 1-(2/4)^2 = 0.75.

Constraints

1 <= N <= 100000, all attn[i] >= 0, not all zero.
Normalize to a distribution; sort by mass descending, ties to smaller index.
Keep until cumulative mass first reaches keep_fraction.
kept is returned sorted ascending; flops_saved rounded to 4 decimals.

Learn this in

Benchmarks, Deployment & the FrontierVLM Study Plan

More from VLM 1: Vision-Language Models

CHAIR-i and CHAIR-s Hallucination Ratesmedium Merge a LoRA Update into the Base Weightmedium QLoRA Memory Footprint Estimatemedium POPE Precision, Recall, and F1medium POPE Object Hallucination Accuracyeasy

Concept primerPREVIEW

1. Background Knowledge

Vision-Language Models (VLMs) process images by splitting them into a grid of visual tokens (e.g., $14 \times 14$ patches for a $224 \times 224$ image). These tokens are concatenated with text tokens and passed through a Transformer architecture. The computational bottleneck in Transformers is self-attention, which computes pairwise interactions between all tokens. If there are $N$ tokens, the attention matrix is $N \times N$ , leading to $O(N^2)$ time and space complexity. This quadratic scaling makes long sequences (many visual tokens) expensive, motivating token pruning strategies.

Token Pruning (e.g., FastV, LLaVA-PruMerge) exploits the observation that not all visual tokens are equally important. The attention mass received by a visual token—summed over all query positions—indicates its relevance. Tokens with low attention mass contribute little to the output and can be safely dropped. By keeping only the top- $k$ tokens by attention mass, we reduce the sequence length from $N$ to $k$ , cutting the visual-attention FLOPs by $1 - (k/N)^2$ . This is a greedy approximation: we keep tokens in descending order of importance until a target fraction of total attention mass is covered.

The FLOPs saved metric reflects the reduction in quadratic attention cost. If we keep $k$ out of $N$ tokens, the new attention cost is proportional to $k^2$ , so the fraction saved is $1 - (k/N)^2$ . This assumes the pruned tokens are removed before the attention computation, which is the standard inference-time pruning setup.

2. Algorithm Approach

This is a greedy selection problem with a cumulative sum stopping condition. The core pattern is:

Normalize the attention weights to form a probability distribution (sum to 1).
Sort tokens by normalized mass in descending order, breaking ties by ascending original index.
Iterate through the sorted list, accumulating mass until the cumulative sum first reaches or exceeds keep_fraction.
Collect the original indices of the kept tokens, sort them, and compute the FLOPs saved.

The key insight is that sorting by mass descending ensures we always pick the most important token first, which is optimal for minimizing the number of tokens needed to cover a target fraction of mass.

3. Step-by-Step Strategy

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results