An AI4AI Framework for Visual Token Pruning
Zhen Liu, Wenli Huang, Wei Song, Yuhan Liu, Zhiqin Yang, Jingwen Fu
cs.LG, cs.CV
2026-08-07
AutoPrune uses an LLM to design pruning policies as small residual fixes on a strong base policy; dropping 94% of visual tokens keeps 99%+ accuracy and 9.9x fewer FLOPs.
A big chunk of MLLM inference cost sits in visual tokens. A vision encoder turns one image into hundreds or thousands of tokens that flow straight into the LLM: 576 for LLaVA-1.5, 2880 for LLaVA-NeXT, and they dominate prefill compute and the KV cache. Pruning redundant visual tokens is the standard fix, with a shelf of methods: FastV, VisionZip, DivPrune, CDPruner.
But all of these are handcrafted heuristics. Which tokens to keep, how to score them, how many to retain, every choice rests on expert intuition and trial and error. The design space is exploding as pruning objectives, token budgets, and model architectures all diversify, each combination demanding fresh human tuning. The paper asks the obvious question: LLMs already know a lot of algorithms and reason well, so can one design a visual-token pruning method on its own?
Using an LLM as an algorithm designer is not new (FunSearch and that line of work did it). The hard part is connecting broad algorithmic knowledge to one specific, tightly structured task.
AutoPrune is training-free: the MLLM stays frozen throughout, used only for policy search and inference. The loop is an LLM proposing candidate policies, scoring them on a task evaluator, and keeping the best, iterated over rounds.
The crux is not the loop but how to represent a policy. The answer is TPDSL (Token Pruning Domain-Specific Language), a DSL built for visual-token pruning with 131 reusable atoms across nine groups:
A candidate policy picks a few atoms and parameters; a materializer compiles it into a deterministic, executable pruning program.
The real trick is the last step: the residual formulation. AutoPrune starts from a strong existing base policy (CDPruner, NeurIPS25, in the experiments) and represents each candidate as a small residual fix on top of it. Most token decisions are left to the base through reference anchoring; the LLM-searched part only overturns the base's uncertain calls, via a bounded residual exchange that swaps at most qe low-confidence base tokens for higher-scoring candidates while forcing at least rmin base tokens to stay.
It is LoRA for token selection: do not rewrite the network, learn a small patch. The search space collapses, and the LLM's attention lands on the decisions that actually move performance. An ablation shows this is the linchpin: drop reference anchoring and MME falls from 1373 to 1223 (down 150); widen the exchange quota from the optimal qe=2 to qe=32 (allow full replacement) and the score collapses to 1217. The residual has to stay small.
Search runs once: on LLaVA-1.5-7B, with MME as evaluator and K=32, for 10 rounds of 5 candidates each, Qwen-Plus as proposer. The winning TPDSL state is then re-instantiated across token budgets, tasks, and even other backbones without re-searching.
14 multimodal benchmarks, 3 backbones (LLaVA-1.5-7B, LLaVA-NeXT-7B, Qwen2.5-VL-7B).
At the most aggressive setting, dropping 94.4% of visual tokens, AutoPrune barely loses anything:
| Model | Tokens | Accuracy | Retained | vs CDPruner |
| LLaVA-1.5-7B | 576 to 32 | 63.2 (full 63.4) | 99.7% | +3.2 |
| LLaVA-NeXT-7B | 2880 to 160 | 65.2 (full 65.2) | 99.9% | +2.4 |
On LLaVA-NeXT, pruned to 160 tokens, the aggregate accuracy ties the full-token model exactly. The biggest gains land on the Chinese multimodal benchmark MMBench-CN (65.9 vs CDPruner's 49.6 on LLaVA-1.5, up 16.3) and MME (1413 vs 1373).
Efficiency (LLaVA-NeXT-7B, 320 tokens, single RTX 3090): FLOPs down 9.9x (41.7T to 4.2T), prefill latency down 6.4x (815 to 128 ms/token), KV cache from 1440MB to 160MB. Against CDPruner at the same budget, FLOPs, memory, and decoding latency are identical, prefill is marginally faster, and MME is slightly higher (1457.9 vs 1453.0). The searched policy is free performance with no added inference cost.
Cross-backbone transfer is the other selling point. The policy searched on LLaVA-1.5 is moved directly to Qwen2.5-VL-7B, whose visual-token architecture is structurally different, with no re-search (CDPruner here is reproduced by the authors). At 128 tokens, relative performance rises from CDPruner's 74.9% to 81.0% (up 6.1 points), with the largest jumps on text-rich tasks (TextVQA, ChartQA). The searched state is not overfit to the MME evaluator; it captures reusable token-selection principles.
At moderate pruning ratios AutoPrune sometimes beats the full-token model, because removing redundant tokens reduces visual interference. Swapping the proposer (Qwen-Max, Qwen-Plus, DeepSeek-V4-Flash) moves results by under 3.3 MME points, so it is insensitive to which LLM does the proposing.
MLLM inference is expensive because of visual tokens, and pruning is the mandatory path to cutting cost. This paper pays off on two levels.
The first is immediate: a pruning policy that transfers across models and beats handcrafted SOTA, with no added inference cost, usable out of the box with concrete cost numbers for anyone deploying MLLMs.
The second is deeper. What the paper really validates is that when an LLM designs an algorithm, the make-or-break choice is how the search space is bounded. Expressed as a constrained residual on a strong base, the LLM reliably produces good policies; as free-form code, it fails. That conclusion probably reaches beyond token pruning, anywhere there is a decent existing baseline you want to auto-improve.
The authors flag two: dependence on TPDSL's expressiveness (it cannot search for operations absent from the atom library) and reliance on the task evaluator (the search signal is the single MME metric).
A few more reservations stand out. AutoPrune's policy is, structurally, CDPruner plus a learned residual, so the baseline it beats is its own starting point, and "outperforming CDPruner" is partly mechanical; the fairer claim is that the residual adds something beyond a strong fixed base. The entire search runs once on one backbone and one evaluator, and every cross-architecture transfer claim rests on that single run. The CDPruner numbers on Qwen2.5-VL are reproduced by the authors (the official implementation is unavailable), so that baseline is not the original authors' result. The optimal exchange quota qe=2 is tuned, and how robust it is across budgets and backbones is not explored.