Visual Token Pruning by Attention Mass
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:
print(prune_visual_tokens([0.5, 0.3, 0.1, 0.1], 0.75))
{'kept': [0, 1], 'num_kept': 2, 'flops_saved': 0.75}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, allattn[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. keptis returned sorted ascending;flops_savedrounded to 4 decimals.
1. Background Knowledge
Vision-Language Models (VLMs) process images by splitting them into a grid of visual tokens (e.g., 14×14 patches for a 224×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×N, leading to O(N2) 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 k2, 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
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
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.