Attention FLOPs and Visual Token Pruning

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]