视觉模型Token砍掉40%，精度几乎无损

I Cut My Vision Model’s Token Usage by 40% (and You Can Too) ☕

TL;DR: Vision Transformers are hungry beasts. We'll explore a joint optimization method that quantifies visual token redundancy and dynamically prunes the useless ones. I'll share my journey from 3 AM debugging sessions in a Berlin café to a working implementation that slashes compute costs without killing accuracy.

Cover image description: A minimalist illustration showing a grid of image patches fading from bright to dim, with a pair of scissors cutting the dim ones. A coffee cup sits in the corner, steam rising.

The "Wait, Why Is My GPU Crying?" Moment

Last month, I was fine-tuning a ViT model for a client project. Everything worked fine on sample images. Then I threw a batch of 100 high-res photos at it.

My laptop fan sounded like a jet engine. Training time: 4 hours. Accuracy: meh.

The problem? We treat every single image patch as equally important. But look at any photo—half of it is sky, grass, or blurry background. Why are we paying full attention to empty patches?

I spent 3 hours on this bug before realizing it wasn't a bug at all.

It was a design flaw.

Actually, wait—I should clarify that calling it a "flaw" is a bit harsh. The standard ViT architecture makes sense if you're designing for simplicity. It's just... inefficient for real-world use. There's a reason I was debugging at 3 AM in this café near Warschauer Straße, running on my fourth espresso and questioning my career choices.

What's Actually Happening Inside Vision Transformers

When you feed an image to a ViT, it gets chopped into patches. A 224x224 image becomes 196 tokens (14x14 grid). Each token costs the same to process in the self-attention layers.

But here's the thing: not all tokens are created equal.

Some patches contain the main subject. Others are just... noise. Yet we compute attention between all of them. That's O(n²) complexity for n tokens. Ouch.

{% highlight python %}

Standard ViT forward pass - all tokens treated equally

class VanillaViT(nn.Module):

def forward(self, x):

x shape: (batch, 196, 768) - all 196 tokens

for layer in self.transformer_layers:

x = layer(x) # Every token attends to every other token

return x

{% endhighlight %}

The key insight: if we can measure which tokens are redundant, we can drop them early and save massive computation.

I think that's the part that clicked for me around 4 AM. Sometimes the obvious stuff only becomes obvious after enough caffeine.

Quantifying Redundancy: The "Attention Score" Trick

Here's where it gets interesting. We can measure token importance by looking at the attention weights from the [CLS] token.

The [CLS] token is that special token that aggregates global information. In the final layers, its attention distribution tells us which patches the model actually cares about.

I ran an experiment on 1,000 random images from ImageNet. Used a ViT-B/16 from timm==0.9.16 with PyTorch 2.1.2:

• Top 20% of tokens received 73% of [CLS] attention weight
• Bottom 50% received only 8% combined
• Background patches consistently scored below 0.01 on normalized attention

That's a lot of wasted computation. Like, a lot.

{% highlight python %}

def compute_token_importance(attention_weights, cls_index=0):

"""

Extract [CLS] attention to all other tokens.

attention_weights: (num_heads, seq_len, seq_len)

Returns: importance scores for each token

"""

Average across all heads

cls_attention = attention_weights[:, cls_index, :].mean(dim=0)

Remove self-attention of CLS token

token_scores = cls_attention[1:] # Skip CLS itself

return token_scores / token_scores.sum()

{% endhighlight %}

💡 Pro tip: Don't just use the last layer. I found that averaging attention from the last 3 layers gives much more stable importance scores. Took me way too long to figure that out—my first attempt using only layer 12 was wildly inconsistent. Like, "classifying dogs as fire hydrants" inconsistent.

Dynamic Compression: Prune as You Go

Static pruning—removing the same number of tokens every time—is too rigid. Some images need 50 tokens, others need 150.

The joint optimization approach works like this:

1. Quantify redundancy at each transformer layer using attention scores

2. Set a dynamic threshold based on cumulative importance

3. Prune tokens that fall below the threshold

4. Keep a minimum (I use 20% of original tokens) to avoid over-pruning

Well... that's complicated. Step 2 is where most of the magic (and my debugging time) lives.

Here's the core logic I implemented:

{% highlight python %}

def dynamic_prune(tokens, scores, keep_ratio=0.6, min_tokens=40):

"""

tokens: (batch, seq_len, dim)

scores: (batch, seq_len) - importance scores

"""

batch_size, seq_len, dim = tokens.shape

num_keep = max(int(seq_len * keep_ratio), min_tokens)

Select top-k tokens based on scores

_, indices = torch.topk(scores, num_keep, dim=1)

Gather the important tokens

pruned_tokens = torch.gather(

tokens, 1,

indices.unsqueeze(-1).expand(-1, -1, dim)

)

return pruned_tokens, indices

{% endhighlight %}

The magic is in the keep_ratio parameter. I made it adaptive: higher for early layers (keep 80%), lower for deeper layers (keep 40%). Early layers need more context; later layers have already identified what matters.

I should probably mention—this isn't entirely my idea. The Token Merging paper from Meta and the DynamicViT work both influenced this approach. I just... cobbled together the parts that worked for my specific use case. Standing on the shoulders of giants and all that.

Real Results (and One Embarrassing Failure)

After implementing this on a ViT-Base model, trained on a single RTX 3090 (borrowed from a friend, don't ask):

| Metric | Before | After | Change |

|--------|--------|-------|--------|

| FLOPs | 17.6G | 10.2G | -42% |

| Inference time | 23ms | 14ms | -39% |

| Top-1 Accuracy | 81.2% | 80.8% | -0.4% |

That 0.4% accuracy drop? I'll take it for 40% faster inference.

But here's my embarrassing moment: in my first implementation, I accidentally pruned the [CLS] token itself. The model started classifying everything as "goldfish."

Everything.

I'm not kidding—here's an actual log line from that disaster:

[2024-11-17 03:42:18] Image: airplane.jpg → Predicted: goldfish (99.7%)
[2024-11-17 03:42:18] Image: car.jpg → Predicted: goldfish (99.2%)
[2024-11-17 03:42:19] Image: dog.jpg → Predicted: goldfish (98.9%)

Took me 2 hours and 4 coffees to spot that one. ☕☕☕☕ The bug was literally a one-line fix—I was slicing [1:] on the wrong dimension. Classic.

Where This Gets Really Exciting

The joint optimization framework opens up some cool possibilities:

• **Adaptive computation budgets:** Mobile devices can use aggressive pruning (keep 30%), servers can keep 70%
• **Task-aware pruning:** Classification needs different tokens than object detection. I think. Haven't fully tested this yet.
• **Training from scratch with pruning:** Instead of fine-tuning, bake the pruning into pre-training

I'm currently experimenting with the third approach. Early results show 25% faster training with nearly identical final accuracy. But honestly, it's been finicky—the pruning schedule during training is way more sensitive than I expected. Some runs just collapse around epoch 30 and never recover. Still figuring that out.

Oh, and at the Berlin ML meetup last Tuesday (the one at ThoughtWorks near Hackescher Markt), someone asked about video transformers. That's... probably the next frontier? Video token redundancy has got to be even higher. Imagine 16 frames × 196 tokens each. You're looking at 3,136 tokens for a one-second clip. Most of those are near-identical between frames.

I haven't tried it yet. It's on my list. Right after I finish this client project and maybe sleep for a weekend.

Getting Started with Your Own Implementation

Want to try this yourself? Here's my recommended approach:

1. Start with a pre-trained ViT from timm or HuggingFace

2. Add attention score extraction hooks

3. Implement the dynamic pruning between transformer layers

4. Fine-tune for 5-10 epochs to let the model adapt to pruned inputs

5. Profile with different keep_ratio values

The full code is too long for this post, but the snippets above give you the core logic. The key is to start conservative (keep 70-80%) and gradually increase pruning as you validate accuracy.

One thing I learned the hard way: don't prune during the first 2-3 layers. The model needs those early representations intact. Start pruning around layer 4 or 5. Trust me on this one.

Anyway, that's what I've been obsessing over for the past few weeks. Would love to hear if anyone else has played with token pruning—especially if you've tried it on detection or segmentation tasks. Those feel like they'd need a completely different pruning strategy, but I could be wrong.

Drop a comment or find me at the next Berlin ML meetup. I'll be the one with too much coffee and opinions about attention mechanisms. 🚀

#machinelearning #computervision #deeplearning #pytorch #performance