做过多模态模型部署的人应该都有这种体验：一个视觉语言模型跑起来，大部分时间都耗在自注意力计算上，而自注意力的复杂度是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%和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 一套简单好用的排查顺序

当剪枝后效果不对时，我一般按固定顺序排查：

1. 先关闭剪枝，确认原模型本身没问题。
2. 打开剪枝但保留率设到90%，排查决策模块是否引入了异常。
3. 逐步降低保留率，定位精度拐点出现在哪个区间。
4. 可视化每一层的token保留图，检查结构坍缩。
5. 单独调整空间平衡权重和模态平衡权重，观察对不同类型任务的分别影响。

这个顺序能帮你快速区分是评分器问题、平衡约束问题，还是单纯配置不合理。我见过很多人一上来就调保留率，结果搞了半天发现评分器根本没训练好，路径完全跑偏。

我自己的体会是，token剪枝这类工作，真正难的不是剪掉哪个token，而是让整个保留集合在几十层网络里始终保持合理的结构。Balanced Token Pruning的价值，在于把“剪掉谁”提升到“保留成什么样”的层面，让剪枝决策变得可解释、可约束、可诊断。如果你正在做VLM推理优化，建议先从可视化你自己模型里的token重要性分布开始，可能你会发现冗余比想象中大得多。最后再分享一个小技巧：部署时保留一个极低保留率的应急模式，配合自动回退逻辑，能在不影响在线体验的情况下持续收集剪枝失败样本，这些样本是后续微调评分器最值钱的素材。