面向视觉Token剪枝的AI4AI自动算法设计框架 An AI4AI Framework for Visual Token Pruning
用LLM在残差式DSL空间中自动搜索视觉token剪枝策略,94.4%剪枝下保留99%以上性能
前置知识
视觉Token剪枝(Visual Token Pruning)
多模态大模型把一张图像编码成数百到数千个视觉token(如LLaVA-1.5-7B每图576个、LLaVA-NeXT-7B每图2880个),这些token在语言模型的注意力计算和KV缓存中占大头。剪枝就是在送入语言模型前,按某种打分标准删除或合并冗余视觉token,只保留最信息量最大的子集,属于training-free的推理加速手段。
本文的搜索对象就是剪枝策略本身,理解预算、打分、筛选、重组这套流程是读懂TPDSL四个组件的必要背景。
MLLM推理流水线与Prefill/KV Cache
MLLM推理分两阶段:prefill阶段一次性处理全部输入token(视觉+文本)并生成KV缓存,计算量与token数的平方相关;decode阶段逐token生成,延迟受KV缓存大小影响。视觉token越多,FLOPs、prefill延迟和显存占用都越高。
论文用FLOPs降低9.9倍、prefill加速6.4倍、KV缓存从1440MB降到160MB来衡量剪枝收益,需要理解这些指标才能看懂效率实验。
注意力重要性剪枝(如FastV)
FastV等方法利用语言模型对视觉token的注意力权重作为重要性分数,删除低注意力token。这类方法假设注意力高的token对回答更重要,是现有剪枝方法最常见的信号之一,但注意力并不总是可靠的代理指标。
论文的打分组件S把attention-proxy saliency作为五种信号之一,且实验显示单一注意力信号不如多信号融合,理解它有助于把握消融实验的动机。
评估器在环的LLM算法搜索(FunSearch范式)
让LLM反复提出候选程序/算法,用自动评估器打分,把分数和失败原因反馈给LLM进入下一轮提议,形成『提议→验证→评估→反馈』的迭代搜索闭环。FunSearch、AlphaEvolve、EvoPrompting都属于这一范式。
AutoPrune继承了这个范式,但把搜索对象从完整程序换成结构化残差状态,理解原范式才能体会本文在表示层面的改进。
残差式搜索状态表示
不从头生成完整方案,而是把候选表示为『强基线方案+一个有界修正量』。类似深度学习中的残差连接:$S_\alpha = (\pi_P, \alpha)$,基座策略$\pi_P$提供可靠初值,残差$\alpha$只修改少量不确定的决策。这缩小了搜索空间并保底性能。
这是全文最核心的思想,TPDSL的每个设计(参考锚定、交换配额、最低保留数)都是围绕它展开的。
DPP式多样性选择
行列式点过程(DPP)是一类在子集选择中同时考虑质量和多样性的概率模型:选出的子集内部相似度越低、行列式越大。视觉token剪枝用它避免选出一堆视觉上重复的token,CDPruner和DivPrune都基于类似思想。
TPDSL的约束组件C包含diversity-aware selection原子,消融显示去掉它MME下降12.66分,理解多样性约束才能明白候选池构建的意图。
研究动机
多模态大模型(MLLM)把图像编码成大量视觉token:LLaVA-1.5-7B每图576个,LLaVA-NeXT-7B每图2880个,导致推理延迟、显存和部署成本大幅上升。视觉token剪枝能显著缓解这个问题,但现有方法几乎都依赖人工设计的固定启发式:FastV和PyramidDrop用注意力重要性,PruMerge和VisionZip做token合并,DART做冗余估计,DivPrune做多样性最大化,CDPruner做指令条件相关性选择。设计这些策略需要大量领域专业知识和昂贵的试错。更麻烦的是,随着剪枝目标(精度优先还是延迟优先)、token预算(16到640不等)和模型架构(LLaVA、Qwen2.5-VL等视觉token组织方式完全不同)日益多样化,人工逐个遍历这个组合空间变得越来越不可行。
本文的目标是本文想回答一个自然的问题:能否让大语言模型自动设计有效的视觉token剪枝算法,即构建一个AI4AI(AI for AI)框架。具体目标是提出一个training-free的框架AutoPrune:不对目标MLLM做任何微调,由一个预训练LLM在结构化空间中搜索剪枝策略,要求搜索到的策略在激进攻剪枝(如删94.4%视觉token)下保留99%以上全token性能,同时大幅降低FLOPs和prefill延迟,并且能跨token预算、跨MLLM骨干迁移而无需重新搜索。
与已有工作不同的是,已有LLM驱动算法设计的工作(FunSearch、AlphaEvolve、EvoPrompting、LLaMEA等)都把候选表示成完整的独立Python程序,这直接用在剪枝任务上有三个问题:设计空间过大导致LLM要重新发明已知机制;剪枝必须同时满足预算、张量形状、索引、数值稳定性等严格约束,自由生成代码极易违反;每个候选都是从零写的,无法复用已有强基线的可靠结构。本文的独特切入点是:核心不在LLM多强,而在搜索状态的表示方式——把每个候选表示为对强基座策略的结构化残差修改,用领域专用语言TPDSL(131个原子)编码,从而把LLM的通用算法知识和剪枝任务的结构性约束连接起来。
核心方法
直觉上,AutoPrune不让LLM从零写剪枝代码,而是给它一个乐高积木库和一个可靠的默认方案,让它只做『微调式』的创新。形式化地,给定图像token序列$T_v$和指令$T_q$,剪枝策略$\pi$在预算$K \ll N$下选出保留子集$T_v^* = \pi(T_v, X_q; K), |T_v^*| = K$,随后$Y_q = f_\phi([T_q; T_v^*])$送入语言模型。框架分三部分:一是TPDSL搜索状态表示,把每个候选表示为对基座策略$\pi_P$的残差修改$S_\alpha = (\pi_P, \alpha)$,其中$\alpha = (B, S, C, R)$分别对应预算控制、token打分、选择约束和token重组;二是LLM驱动的策略设计,预训练LLM在评估器反馈下迭代提议、验证、评估这些残差状态;三是先验引导的残差精化,运行时保留基座策略的绝大多数选择,只替换有界数量的不确定token。
核心创新是搜索状态表示本身:TPDSL不把候选表示为完整策略,而是表示为基座策略$\pi_P$之上的残差修改$\alpha = (B, S, C, R)$。TPDSL库含131个实例化原子,组织成9个功能组:$\mathcal{L} = \bigcup_{g=1}^{9} \mathcal{L}^{(g)}, |\mathcal{L}| = 131$,覆盖token来源与基座策略、预算控制、打分、相似度核、分数融合、候选池构建、选择约束、重组与token处理、安全检查。每个原子有预定义语义和标准化类型化运行时接口,迫使LLM只能组合与预算、形状、索引要求兼容的操作。这与自由代码生成的本质区别在于三点:残差公式保留了可靠的先验结构(基座策略的绝大多数选择不被触碰),有效缩小搜索空间并把LLM注意力引向对性能最关键的部分;组合式空间天然可执行、可解释;策略结构与运行时依赖字段解耦,使搜索结果无需重新搜索即可跨预算跨骨干迁移。
方法步骤详情
搜索阶段:以CDPruner为基座$\pi_P$、MME为评估器$\mathcal{E}$、源预算$K_0=32$,在第$t$轮LLM收到TPDSL库$\mathcal{L}$、基座策略和累积搜索历史$H_{t-1}$,提议$A_t = \mathrm{LLMPropose}(\mathcal{L}, \pi_P, H_{t-1}; K_0)$,共10轮、每轮5个候选。每个候选先经$\pi_\alpha^{(K)} = \mathrm{Materialize}(\alpha, \pi_P, K)$绑定运行时实现并解析预算相关字段,再经$d_\alpha^{(K)} = \mathrm{SafetyCheck}(\pi_\alpha^{(K)}, K)$验证预算一致性、索引有效性、数值稳定性、输出形状兼容性和确定性执行,无效候选直接淘汰。有效候选得分$e_{t,i} = \mathcal{E}(\pi_\alpha^{(K_0)})$汇入全局池$\mathcal{V}$,历史$H_t = \mathrm{Summarize}(H_{t-1}, O_t)$压缩反馈高分组合与常见失败。$T$轮后选$\alpha^* = \arg\max e_{t,i}$。运行时精化:基座策略先选出$T_P$,打分组件$S^*$对每个视觉token计算五种信号$s_i = (s_i^{rel}, s_i^{attn}, s_i^{spa}, s_i^{red}, s_i^{con})$(指令相关性、注意力代理显著性、空间中心性、冗余密度、局部对比度),加权积融合$g_i = \prod_m w_m \bar{s}_i^{(m)}$;约束组件构建候选池$T_S$并执行有界残差交换$(T_D, T_A) = \mathrm{RE}(T_P, T_S, g; q_e, r_{min})$,精化集$\tilde{T}_v = (T_P \setminus T_D) \cup T_A$,约束$|T_D| = |T_A| \le q_e$、$|T_P \setminus T_D| \ge r_{min}$(默认$q_e=2, r_{min}=30$);最后重组组件恢复token顺序并可融入被删token的残差信息。迁移时仅重实例化预算和模型相关字段:$\pi_\alpha^*(K_n) = \mathrm{Materialize}(\alpha^*, \pi_{P,n}, K_n)$。
技术新颖性
技术新颖性体现在四个层面。第一,问题形式化新颖:把『LLM能否设计剪枝算法』落成AI4AI问题,并论证表示设计是关键瓶颈,这在视觉token剪枝领域是首次系统尝试。第二,TPDSL作为约束化组合搜索空间:131个原子带类型化接口和四大安全检查(预算一致、索引有效、数值稳定、确定性),把非法解在评估前就过滤掉,实验显示4个随机种子下50个候选全部有效,验证了该空间能稳定产出合法策略。第三,残差交换机制:搜索到的状态每图最多替换2个基座token($q_e=2$)、至少保留30个($r_{min}=30$),这种『极少改动』的设计被消融证明是收益来源——允许全部替换反而比基座低155.12分。第四,结构-运行时解耦:选出的状态$\alpha^*$只是结构描述,预算和骨干相关字段在部署时重新绑定,因此从LLaVA-1.5-7B搜出的策略可直接用于Qwen2.5-VL-7B(视觉token组织完全不同)且无需任何LLM调用。与OpenEvolve和EvoPrompting在相同搜索预算下的对比(1413.46 vs 1407.23/1409.21)也直接支撑了表示设计的价值。
实验结果
主实验(Table 1,94.4%剪枝):LLaVA-1.5-7B从576token压到32,AutoPrune聚合分63.2、保留全token性能99.7%,超CDPruner 3.2分(60.0/94.3%)、超PruMerge+ 8.9分;LLaVA-NeXT-7B从2880压到160,聚合分65.2、保留99.9%,超CDPruner 2.4分、超PruMerge+ 7.9分。最引人注目的是MMBench的异常大幅提升:32token下MMBench-EN从全模型的64.70升到70.55、MMBench-CN从58.10升到65.90,双双超过全token模型,且该优势在各预算和Qwen2.5-VL上持续存在,作者归因于剪枝去除了视觉干扰。Figure 3显示在77.8%/88.9%/94.4%三档削减率下AutoPrune均最优且预算越紧优势越大。效率(Table 2,LLaVA-NeXT-7B保留320token,RTX 3090):FLOPs从41.7T降到4.2T(9.9倍),prefill延迟从815降到128 ms/token(6.4倍),KV缓存从1440MB降到160MB,显存从17.2GB降到13.8GB,同时MME 1457.9反超同预算CDPruner的1453.0。跨骨干迁移(Table 3,Qwen2.5-VL-7B):128token下相对性能从CDPruner的74.9%升到81.0%(+6.1个百分点),TextVQA从23.24升到41.42,四个预算(256/128/64/32)全部领先。消融(Table 4):去掉参考锚定MME从1413.46暴跌到1223.29(-190.17),去掉多信号融合降到1391.70,去掉多样性选择降到1400.80,完整TPDSL超基座40.46分。交换配额(Table 6)呈倒U形:$q_e=2$最优,全替换($q_e=32$)只剩1217.88。基座泛化(Table 7):对五种策略提升22.46到145.40分,PruMerge+获益最大。搜索策略对比(Table 13):超OpenEvolve 6.23分、超EvoPrompting 4.25分。稳定性(Table 14):4个种子全部收敛到相同的最佳1413.46。跨预算迁移(Table 12):16token仍保留96.6%,64token为100.9%,128token达102.3%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 多模态综合能力(LLaVA-1.5-7B,576→32 token,削减94.4%) | 14基准聚合分 Acc. / 相对全模型 Rel. | AutoPrune 63.2(Rel. 99.7%) | CDPruner 60.0(94.3%);PruMerge+ 54.3(86.1%) | 超CDPruner +3.2分,超PruMerge+ +8.9分,性能保留率提升5.4个百分点 |
| 多模态综合能力(LLaVA-NeXT-7B,2880→160 token,削减94.4%) | 聚合分 Acc. / Rel. | AutoPrune 65.2(Rel. 99.9%) | CDPruner 62.8(96.0%);PruMerge+ 57.3(87.7%) | 超CDPruner +2.4分,超PruMerge+ +7.9分 |
| 推理效率(LLaVA-NeXT-7B保留320 token,单张RTX 3090) | FLOPs / prefill延迟 / KV缓存 | 4.2T(×9.9)/ 128 ms/token(×6.4)/ 160.0 MB,MME 1457.9 | 全模型:41.7T / 815 ms/token / 1440.0 MB;CDPruner MME 1453.0 | FLOPs降9.9倍、prefill加速6.4倍、KV缓存降9倍,MME还反超CDPruner 4.9分 |
| 跨骨干迁移(Qwen2.5-VL-7B,保留128 token,削减90.1%) | Rel. 相对全token性能 | AutoPrune 81.0%(TextVQA 41.42) | CDPruner复现版 74.9%(TextVQA 23.24) | +6.1个百分点,TextVQA提升18.18分 |
| 基座策略泛化(LLaVA-1.5-7B,MME,32 token) | 相对各策略MME增益 | 五种策略全部提升,增益22.46~145.40 | Strategy only(统一接口下的各策略本身) | PruMerge+ +145.40、VisionZip +127.80、DivPrune +34.66、TRIM +54.60、DART +22.46 |
| 与其它LLM搜索策略对比(相同10×5搜索预算与评估器) | MME(相对CDPruner基座) | AutoPrune 1413.46(+40.46) | OpenEvolve 1407.23(+34.23);EvoPrompting 1409.21(+36.21) | 分别领先6.23分和4.25分 |
局限与改进
作者明确承认两点局限:AutoPrune依赖TPDSL搜索空间的表达能力,以及任务评估器的可靠性——搜索空间里没有的机制永远搜不出来,评估器有偏则搜索目标有偏。我自己的观察还有几点:第一,搜索只用MME一个粗粒度基准做评估器,虽然迁移结果显示没有严重过拟合,但MMBench-CN上异常大的提升(32token下65.90远超全模型58.10)缺乏机理解释,不排除某些基准的评分方式对特定token子集敏感;第二,131个原子仍全部由人工设计,所谓AI4AI目前是『AI在人类划定的空间里做AI设计』,真正的自动化还有距离;第三,对基座策略CDPruner的依赖较强,参考锚定消融(去掉后-190.17分)说明残差框架的性能下限几乎完全由基座决定,若基座本身有系统性偏差,残差修正幅度(每图2个token)可能不足以纠正;第四,实验仅覆盖图像基准和7B模型,没有视频MLLM、多图高分辨率场景,也没有批处理/连续批处理等真实服务栈下的端到端吞吐数据;第五,POPE等个别基准上AutoPrune(87.67)实际略低于CDPruner(87.90),说明残差修正并非全面占优。
独立分析的弱点
第一,评估器单薄:整条搜索流水线的优化目标只有MME(一个二元问答式基准),对OCR密集、细粒度空间推理等能力没有显式监督,改进方向是多任务复合评估器或按目标基准族加权的评估器,代价是搜索成本上升,可用早停和代理基准缓解。第二,残差交换量是全局固定超参($q_e=2$),对简单图像可能一个token都不用换、对困难图像可能2个不够,改进方向是按图像难度或基座token置信度自适应确定$q_e$,例如置信度门控的动态配额。第三,搜索空间封闭:LLM不能定义新原子,遇到TPDSL未覆盖的机制(如跨层token交互、动态分辨率选择)无能为力,改进方向是分层DSL——允许LLM提议新原子但必须通过安全验证和最小性能门槛才能入库。第四,性能收益高度集中于多语言MMBench而其余基准提升幅度小(多数1分以内),说明残差策略学到的可能是较窄的信号融合方式,改进方向是在更多样化评估器下搜索并分析学到的原子组合在不同能力维度上的贡献。第五,缺乏与轻量学习式剪枝(如训练小型打分网络)的对比,无法判断LLM搜索相对直接梯度优化的性价比优势。每个弱点都可以通过把TPDSL原子库当作可版本化、可扩展的社区资产来长期改进。
未来方向
作者提出三个方向:更丰富的搜索状态表示(超越当前四元组残差,如引入跨层、时序结构)、更自适应的评估策略(多任务、多粒度评估器或在环主动学习)、在更多MLLM架构和多模态任务上扩展验证(视频、音频、具身智能)。基于本文成果还可以延伸:一是把残差搜索范式推广到MLLM效率的其他环节,如KV cache逐层驱逐策略、动态分辨率选择、视觉投影器结构设计,TPDSL思路是通用的;二是基座策略与残差的协同搜索——当前CDPruner是固定的,可以让搜索同时微调基座超参,或迭代交替『搜基座、搜残差』;三是多目标搜索,用Pareto前沿代替单点最优,同时优化精度、延迟、显存;四是在线自适应,利用搜索历史作为先验,在推理时对分布外的图像类型做轻量在线残差调整;五是把搜索稳定性现象(4个种子收敛到同一最优)作为理论切入点,研究约束化组合空间中LLM搜索的收敛性质;六是反向输出——AutoPrune搜出的多信号加权积打分公式本身可反哺为可解释的人工剪枝新基线。
复现评估
复现条件相当友好。代码已开源(github.com/AIM-ResearchLab/AutoPrune),TPDSL的131个原子提供机器可读定义、参数范围和运行时接口。框架training-free、MLLM全程冻结:搜索阶段仅需LLM驱动的10轮×5候选共50次提议(用API模型Qwen-Plus即可,成本很低),评估用MME单基准;运行阶段无任何LLM调用。所有14个基准均为公开数据集,遵循官方划分、协议和指标;三个骨干(LLaVA-1.5-7B、LLaVA-NeXT-7B、Qwen2.5-VL-7B)都是公开权重,效率实验单张RTX 3090即可完成。需要注意的坑:Qwen2.5-VL上CDPruner无官方实现,论文是自己按原文协议复现的,复现对比数字时可能有小幅出入;搜索历史摘要和多重子稳定性依赖的随机种子细节在补充材料有4种子报告(全部50个候选有效、最佳一致),复现时建议同样做多种子验证。总体难度中等:有MLLM推理工程经验的小团队在一到两周内应能复现主表结果,搜索部分因依赖外部LLM API的版本变化(Qwen-Plus会迭代)可能存在轻微不可复现性,但论文的种子稳定性实验表明结果对这类扰动不敏感。
论文图表