VLM推理加速:平衡Token剪枝如何超越局部优化 做过多模态模型部署的人应该都有这种体验一个视觉语言模型跑起来大部分时间都耗在自注意力计算上而自注意力的复杂度是token数量的平方级。图像侧一次性塞进来成百上千个patch token文本侧每生成一个token就会新增一个越到后面序列越长整条推理链路被token数量死死卡住。我最近在给内部多模态助手做推理加速试过量化、换小模型、做蒸馏最后发现最直接的杠杆其实在token数量上如果把视觉侧的token砍掉一半整条链路的速度提升立竿见影。token剪枝不是新话题CV里很早就有基于注意力权重的token合并与丢弃方法。但在VLM场景下直接套老办法问题很多。最典型的现象是某个图像区域明明包含关键信息但因为局部显著性估计有偏差那一整片token被当成“不重要”剪掉了或者文本侧的指令token被过度保留视觉token被削得太狠模型直接变成“瞎猜”。这些问题的根源其实是剪枝决策只做了局部贪心优化没有考虑全局的平衡性。这篇文章想聊的是我在复现2025 NeurIPS上Balanced Token Pruning相关工作时的一些理解与经验。它的核心主张是把token剪枝从“逐个打分、按分砍掉”的局部优化推进到“考虑保留集合整体质量”的全局优化也就是标题里点明的Beyond Local Optimization。适合正在做VLM推理加速落地、或者对token效率这个方向感兴趣的人你可以把它当成一份踩过坑之后的复现笔记来看。1. 为什么VLM推理被token数量卡住1.1 先算一笔账要理解token剪枝为什么重要先看VLM推理到底是怎么花钱的。以典型的主流开源VLM架构为例输入图像先被切成固定大小的patch经过视觉编码器变成一组patch token。假设一个常见输入分辨率对应396个patch token再加上表示图像全局信息的class token视觉侧大概就有400个token指令文本一般也有几十个token进入LLM主干的第一层时序列长度已经超过450。真正的问题出在生成阶段。LLM是自回归的每生成一个新token完整的序列都要重新过一遍所有层。生成20个token时序列长度约470生成100个token时已经到550。自注意力计算量是序列长度的平方所以仅仅是从450个token涨到550个token单层注意力计算量就提高了接近50%。如果batch size是16或者32这个开销还要乘上batch维。再加上KV cache要存储每个token的键值中间结果序列越长显存占用也越高。换句话说在VLM里token数量同时卡着两件事计算延迟和显存占用。这与纯文本LLM不一样——文本侧token是逐步增加的模型只能被动承受而视觉侧token在开始时就是全部存在的冗余比例通常很高是可以主动优化的。相邻patch之间大量信息重叠很多情况下砍掉一半信息损失相当有限。1.2 为什么剪token是目前性价比最高的加速手段对比其他推理加速方法token剪枝有它独特的优势。量化和蒸馏都需要重新校准或重新训练模型对精度的影响是全局性的。而token剪枝是数据依赖的动态行为简单样本可以多剪复杂样本少剪甚至不剪这种动态性是量化做不到的。剪枝的另一层价值在于它和现有手段正交。你可以先量化再剪枝两个收益直接相乘。我做过一组粗糙对比仅量化能加速约20%~40%仅剪枝按50%保留率能加速约40%两者结合总加速比接近翻倍。具体数值要看硬件和实现但正交性这个判断我比较确信。从工程角度看token剪枝的接入成本也低。只要在Transformer层前加一个可选的mask模块不需要改动注意力实现本身比做算子融合、写专用kernel要友好得多。这也是这个方向能快速从论文走向实际部署的根本原因。2. 局部贪心剪枝为什么不够2.1 重要性分数是“近视眼”现有的token剪枝大多走同一个套路用一个打分函数计算每个token的重要性然后按分数从低到高剪。常用打分方式包括注意力权重和、与特殊token的相关性、以及模型自己学出来的显著性。在单层、单任务场景里这种方法有效因为关键信息往往集中在少数token上。但放到VLM里情况就变了。输入是双路的图像patch之间有强烈的空间相关性文本token之间有语法和语义依赖跨模态的注意力交互还会让token重要性在不同层之间剧烈变化。一个token在第三层看起来毫无价值到了第五层可能因为某个文本token建立联系变成关键路径的一部分。局部打分是静态的它看不到这种动态依赖所以经常“剪错了”。我实际遇到过的一个案例是对一个包含远处小目标物体的图像做视觉问答普通贪心剪枝在浅层把远目标对应token全剪了结果模型回答“图里没有该物体”而原始模型能正确指出来。如果打分器能看到后续层中该区域token的聚合效用这种错误完全可以避免。2.2 独立假设在注意力世界里不成立贪心剪枝还隐含了一个假设每个token的效用是独立可加的剪掉A不会影响B的效用。在视觉特征空间里这个假设基本不成立。相邻patch token的信息高度冗余但注意力建模又可能让一个token成为多个其他token的“聚合中心”。如果只看个体分数把聚合中心剪了后面所有依赖它的token都会跟着受损这种损失是级联扩散的。我用一个生活化类比整理书架时如果只挑“最旧的书”扔不考虑某本书可能是一套丛书的索引局部看它确实独立放着全局看没了它整套书的检索都乱了。在Transformer里一个attention head可能专门负责聚合某个区域的空间信息这个区域的token被剪掉后该head在深层几乎失效但这种效应在剪枝当时是看不到的。这也是为什么“重要性分数加阈值剪枝”在浅层网络里比在深层网络里更可靠深层特征高度抽象化单个token的独立性更弱局部决策更容易踩空。这个观察对我后续设计平衡模块很有帮助。2.3 “剪得快”不等于“整体最快”还有一个容易被忽视的点局部贪心剪枝虽然让每一步的决策变快但可能导致剪枝后的token集合在后续层里依然低效。因为模型要对一个结构不均衡的集合继续做注意力计算大量计算被浪费在残缺信息的互相补充上极端情况下整体效率不升反降。举个具体例子。某个批次的图像中文本token因为分数高被“公平地”保留而视觉token因为总体分数偏低几乎被剪到最低限度。后续层为了回答任务需要反复从残余visual token里猜测图像内容注意力分散、无谓计算多生成质量也差。这种“集合级”的浪费单看每一步贪心是看不到的只能从最终端到端的时延和精度来评估。这也是Balanced Token Pruning强调Beyond Local Optimization的直接原因。3. Balanced Token Pruning的思路拆解3.1 从“剪掉谁”变成“保留成什么样”这个方向最让我觉得有意思的地方是把优化目标从“每个token值不值得留”换成了“剪完之后这个集合是否平衡”。平衡在这里有多层含义空间上保留的visual token不能全挤在图像的一个角落模态上visual token和text token的保留比例要匹配当前任务的需求层级上不同Transformer层的剪枝力度应当协同而不是每层各自为政。这个视角转换的直接好处是决策稳健性大幅提升。即使某个token的显著性估计有误差只要保留集合在整体结构上仍然覆盖了关键信息后面的层依然能完成信息提取。这就像一个团队里某个成员状态不佳只要团队整体的角色搭配没有崩项目依然能推进。具体实现上平衡剪枝通常包含两个组件一个不错的token显著性估计器和一个平衡性约束函数。估计器决定“哪个token值得留”约束函数决定“保留集合在各种意义下是否合理”。两者互相制衡缺一不可。3.2 平衡性度量空间、模态、层级三个维度我把平衡性度量拆成三个可计算的部分。空间平衡性把图像token按patch网格组织用划分窗口统计每个区域内保留token的数量和显著性总和。剪枝后如果某个区域的保留密度过低就触发惩罚。这里推荐用软约束而不是硬顶硬性固定每个区域的token数量会让自定义形状的token失去灵活性对不规则的输入分辨率也不友好。模态平衡性VLM里文本token短但语义密度高视觉token多但冗余大。一个简单有效的做法是为每种模态设置保留率的上下界把剪枝变成带约束的选择问题。实测下来对指令跟随类任务文本token保留率一旦低于某个水平任务完成度会断崖式下跌就像给一个聪明人戴上耳塞指令都听不清了。层级协同性不同层的token重要性变化幅度不同浅层可以多剪深层要谨慎。我习惯的配置是让保留率随层深按曲线递减浅层设一个比较大的初始保留率到中层收窄到深层保持一个最低阈值避免深层信息被过度破坏。3.3 全局决策怎么把约束变成可学习的有了平衡性度量之后剩下问题是怎么做全局决策。纯贪心加约束纠正也可以但效果一般。更合理的做法是把选择过程参数化用一个可学习的评分网络估计token显著性同时把平衡性约束作为正则项加入训练目标端到端更新评分网络。这样评分器在训练时就见过“什么是失衡的剪枝方案”它的输出自然会往全局有利的方向靠。技术实现上这个思想和经典选择模型的思路相通把离散的token选择松弛成连续的软mask用Gumbel-Softmax或者直通估计器让梯度穿过选择过程。被放弃token的信息不是彻底丢弃而是通过聚合操作并入相邻token或者由注意力机制自然遗忘。这个松弛很重要否则0/1离散选择没有梯度评分网络根本训不动。一个可选的替代方案是参考最优传输的框架把保留集合的选择建模成带约束的匹配问题每次迭代寻找全局最优分配。这种方案的好处是约束可以直接满足坏处是计算开销大在线推理时不一定扛得住。我当时对比下来软mask加正则项的路线在速度和效果之间更均衡。4. 实操接入与关键配置4.1 接入位置与整体流程我在实验中采用的接入方式是在Transformer层的attention之前设置剪枝决策点。给定该层输入的token序列visual加text先算显著性分数和平衡性分数再产出保留或丢弃的mask。保留token进入attention计算丢弃token从序列中移除。为了保持模块边界清晰我只在部分层启用剪枝浅层和深层各留一些“无剪枝安全层”防止信息过早丢失或深层过度损失。整体流程分三步走第一步冻结原VLM主干训练显著性评分网络目标很简单让评分器拟合原始模型注意力中的重要性分布。第二步加入平衡性正则联合微调评分网络和少量附加层参数平衡正则的权重从0开始逐步增大同时盯住下游任务效果。第三步推理期只做前向不更新参数也可以在这个阶段用一个蒸馏出来的小模型替换评分器进一步降低剪枝决策本身的开销。需要提醒的是评分网络本身也有计算成本。如果设计得太大剪枝省下的时间又被评分吃回去了。我建议评分器控制在两个线性层以内输入取token embedding在若干层上的均值池化输出就是每个token的显著性分数。这个开销相对序列长度是线性的可接受。4.2 关键超参怎么设我实测下来有几个超参最影响最终效果整理成一张表方便对照超参数建议范围影响表现整体保留率70%~30%太高加速不明显太低精度崩平衡正则权重0.01~0.1太大导致高显著token被误剪启用剪枝层位置中浅层为主深层过度剪枝风险高评分器规模不超过两层线性层太大吃掉剪枝收益保留率是整个方案的核心旋钮。我建议先在70%和50%两个档位各跑一轮确定评分器和平衡模块本身没有引入明显问题再往30%压。平衡剪枝相对普通贪心的优势恰恰是在低保留率区域才显现出来我用同一套主干对比平衡剪枝大约能多压5到10个百分点而精度损失持平。低于30%之后无论什么方法精度下降都会加速不太建议生产环境贸然使用。平衡正则权重是个需要小心的值。权重太大模型为了凑“平衡”会把高显著性token也剪掉等于为了形式牺牲内容权重太小约束形同虚设。我的经验是从0.01起步观察验证集损失再逐步上调一个比较合理的参考区间就是0.01到0.1。具体数值最好按任务单独调不要跨任务直接套。训练数据的选择也会明显影响结果。评分网络最好用混合了指令跟随、视觉问答、图文理解任务的数据做训练。单任务数据训练出来的评分器换到其他任务上会偏科比如只在图文描述数据上训练放到OCR类场景里显著性估计会明显失真。4.3 实测收益与对比在我复现的基准上保留率50%时推理时延大约降低40%精度下降约1个百分点保留率30%时时延降低约55%精度下降接近3个百分点。相比普通贪心剪枝平衡剪枝在中低保留率下平均高1.5个百分点左右尤其在意视觉推理这类对空间覆盖要求高的任务上优势更明显。这些数字仅供参考。性能受任务、模型、硬件影响很大换架构就要重新测。我这边是在某主流开源VLM架构上加的如果你用的是更大的模型或者多分辨率输入收益会有差异。建议按自己的业务场景做一组从90%到30%保留率的全量扫描画出时延-精度曲线再决定最终配置而不是直接照搬论文里的某个数字。5. 踩坑指南与排查手册5.1 精度塌方先查结构坍缩遇到过最典型的坑是评分器loss很低但下游任务精度崩了。排查之后发现剪枝后的visual token在空间上集中到了图像中心区域边缘信息全丢也就是结构坍缩。普通视觉任务可能看不出来但凡是需要“读完整张图”的任务比如OCR、图表理解、空间关系判断全部翻车。这就是平衡约束存在的意义它把“集合结构坍缩”这个隐形问题排掉了。我现在训练完一定会做一件事把每一层的保留token分布可视化看是不是出现了结构性偏斜确认没问题再考虑上线。这个可视化检查比任何损失曲线都直观建议你也养成这个习惯。5.2 与KV cache和批量推理的交互动态token剪枝会改变序列长度导致KV cache的形状不确定。批量推理时不同样本保留的token数不同如果不做padding对齐加速卡利用率会很难看。我的折中方案是保留率按batch统一设定保证同一batch内的token数一致或者用长度分桶把不同长度样本分到不同batch。attention mask也要跟着改。因为剪枝后token序列的顺序可能重排直接复用原来的causal mask会出问题需要重新生成。这个坑新手特别容易踩跑起来不报错但结果会悄悄变差而且很难定位。建议在实现里对mask生成和token索引映射做单测覆盖。5.3 评分器过拟合与应急回退评分网络在训练分布上表现好换到另一个数据分布比如从自然图像换成文档扫描图、医疗影像显著性估计会明显失效。我的建议是两块训练数据尽量覆盖目标业务场景部署侧加一个检测机制衡量评分器输出的置信度分布发现异常就回退到不剪枝模式。这个回退开关成本很低但能避免大量线上事故。一个容易被忽略的细节是回退不该只在评分器置信度低时触发还应该在剪枝后生成的回答与不剪枝时出现显著分歧时触发。这个可以通过对比生成序列的困惑度来实现实现成本不高但能拦住相当比例的风险case。5.4 一套简单好用的排查顺序当剪枝后效果不对时我一般按固定顺序排查先关闭剪枝确认原模型本身没问题。打开剪枝但保留率设到90%排查决策模块是否引入了异常。逐步降低保留率定位精度拐点出现在哪个区间。可视化每一层的token保留图检查结构坍缩。单独调整空间平衡权重和模态平衡权重观察对不同类型任务的分别影响。这个顺序能帮你快速区分是评分器问题、平衡约束问题还是单纯配置不合理。我见过很多人一上来就调保留率结果搞了半天发现评分器根本没训练好路径完全跑偏。我自己的体会是token剪枝这类工作真正难的不是剪掉哪个token而是让整个保留集合在几十层网络里始终保持合理的结构。Balanced Token Pruning的价值在于把“剪掉谁”提升到“保留成什么样”的层面让剪枝决策变得可解释、可约束、可诊断。如果你正在做VLM推理优化建议先从可视化你自己模型里的token重要性分布开始可能你会发现冗余比想象中大得多。最后再分享一个小技巧部署时保留一个极低保留率的应急模式配合自动回退逻辑能在不影响在线体验的情况下持续收集剪枝失败样本这些样本是后续微调评分器最值钱的素材。