Attention FLOPs and Visual Token Pruning

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

Problem Statement

Count the FLOPs of one attention layer over n visual tokens, then recount after pruning to a fraction of them, and report the reduction.

Background

Attention has two cost terms with different growth rates. Use this standard count (one multiply-accumulate = 2 FLOPs):

Projections - Q, K, V and the output projection, each an n x d by d x d matmul: 4 * 2 * n * d * d = 8 n d^2
Attention itself - Q K^T at 2 n^2 d and A V at another 2 n^2 d: 4 n^2 d

$\text{FLOPs}(n) = 8 n d^2 + 4 n^2 d$

The projection term is linear in n; the attention term is quadratic. For a language prompt of 30 tokens the linear term dominates completely, but a single AnyRes image is 2000+ visual tokens, and there the quadratic term takes over. That crossover is why visual-token pruning exists at all: dropping tokens buys a superlinear saving, whereas making the model narrower only buys a linear one.

Pruning keeps floor(n * keep_ratio) tokens, and the reduction is:

$\text{reduction \%} = 100 \left(1 - \frac{\text{FLOPs}(n_{\text{kept}})}{\text{FLOPs}(n)}\right)$

Halving the tokens therefore saves strictly more than 50% - and the more quadratic the regime, the closer to 75%.

Your Task

Implement:

def attention_flops(num_tokens, dim, keep_ratio=1.0):

Return [flops_full, tokens_kept, flops_pruned, reduction_percent] - three ints and a float rounded to 4 decimals.

Input Format

num_tokens - positive integer n
dim - positive integer model width d
keep_ratio - float in (0, 1]

Output Format

[int, int, int, float]

Sample

print(attention_flops(576, 1024, 0.5))

Output:

[6190792704, 288, 2755657728, 55.4878]

Example

Input

print(attention_flops(576, 1024, 0.5))

Output

[6190792704, 288, 2755657728, 55.4878]

Reasoning

Full cost is 85761024^2 = 4831838208 projections plus 4576^21024 = 1358954496 attention, totalling 6190792704. At 288 tokens both terms shrink - the first by half, the second by four - giving 2755657728, a 55.4878% reduction from halving the tokens.

Constraints

Use exactly FLOPs(n) = 8 * n * d * d + 4 * n * n * d
tokens_kept = floor(num_tokens * keep_ratio) and is at least 1 for the given inputs
Compute the counts with integer arithmetic so no precision is lost at billions of FLOPs
reduction_percent is relative to the full count, rounded to 4 decimals
A keep_ratio of 1.0 must give a reduction of exactly 0.0

Learn this in

Benchmarks, Deployment & the FrontierVLM Study Plan Vision EncodersVLM Study Plan

More from VLM 1: Vision-Language Models

Visual Token Pruning by Attention Masshard CHAIR-i and CHAIR-s Hallucination Ratesmedium Merge a LoRA Update into the Base Weightmedium QLoRA Memory Footprint Estimatemedium POPE Precision, Recall, and F1medium

Concept primerPREVIEW

1. Background Knowledge

In Vision-Language Models (VLMs), attention mechanisms are the primary computational bottleneck, especially when processing high-resolution images. An image is typically split into patches, each represented as a visual token. If an image yields $n$ tokens and the model dimension is $d$ , the standard multi-head attention layer involves two distinct types of operations: projections and attention scoring.

The projection operations involve computing Query ( $Q$ ), Key ( $K$ ), Value ( $V$ ) matrices and the final output projection. These are matrix multiplications of shape $(n \times d)$ by $(d \times d)$ . Since there are four such projections, and each multiply-accumulate operation counts as 2 FLOPs, the total cost for projections is linear with respect to the number of tokens: $8 n d^2$ .

The attention scoring operations involve computing the similarity matrix $Q K^T$ (shape $n \times n$ ) and then multiplying by $V$ ( $A V$ ). These operations depend on the pairwise interactions between all tokens, making them quadratic in $n$ . The cost is $4 n^2 d$ . For large $n$ (e.g., 2000+ tokens from high-res images), the quadratic term dominates, making the model computationally expensive. Visual token pruning reduces $n$ to a smaller subset, leveraging the quadratic nature of attention to achieve superlinear reductions in FLOPs.

2. Algorithm Approach

The problem requires implementing a direct mathematical evaluation based on the provided FLOPs formula. The approach is analytical rather than iterative or algorithmic in the traditional sense.

Define the FLOPs Function: Implement the formula $\text{FLOPs}(n) = 8 n d^2 + 4 n^2 d$ .
Calculate Pruned Tokens: Determine the number of tokens kept after pruning using floor(n * keep_ratio).
Compute Pruned FLOPs: Apply the FLOPs formula to the new token count.
Calculate Reduction: Compute the percentage reduction using the formula provided.

This is a straightforward application of arithmetic operations. The key is to ensure integer arithmetic is used for FLOPs counts (as they are discrete operations) and floating-point precision is handled correctly for the reduction percentage.

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