大家优化视觉语言模型(Vision Language Model)推理的时候,最先看到的优化点往往是量化、KV Cache、批量推理这些“通用招数”。直到我自己动手去拆真实请求的耗时分布,才意识到一个容易被忽略的事实:图像输入会被切成一堆视觉令牌,塞进和文本一样的自注意力序列里,这部分的计算开销常常才是延迟的大头。所以今年看到“2025_NIPS_Balanced Token Pruning: Accelerating Vision Language Models Beyond Local Optimization”这个标题时,我的第一反应是:终于有工作把令牌剪枝从“局部最优”这条老路上拽出来了。这篇文章我会完整拆解它背后的技术逻辑、实验观察,以及实际部署中需要注意的细节,适合正在做多模态模型推理优化的同学,也适合对模型加速方法感兴趣的人读一读。
1. 先算一笔推理账:视觉语言模型到底慢在哪
1.1 “一视同仁”的代价
很多同学一开始会有一个直觉:视觉语言模型的文本部分是自回归生成的,要一个词一个词往外蹦,肯定慢;视觉部分不过是“看一张图”,顶多编码一次,应该快得很。我当初也是这么想的,直到把一个真实请求的耗时拆开看,才承认想反了。
问题出在“一视同仁”这四个字上。以常见的视觉编码方式为例,一张 336×336 的输入图,如果按 14×14 的补丁切分,得到的网格是 24×24,也就是 576 个视觉令牌;如果渲染到 768×768,令牌数会直接跳到 2304 附近。这些视觉令牌会和文本令牌一起进入后面的自注意力层,参与两两之间的注意力计算。
自注意力的计算量是随序列长度近似平方增长的。我们算一笔很粗略的账:单层自注意力里,对每个头都要做 Query-Key 点乘和 Attention-Value 加权,大概需要4·n²·d次浮点运算,这里n是序列长度,d是隐藏维度。假设某模型隐藏维度是 4096,输入图片带来 576 个视觉令牌,文本只有 60 个令牌,序列总长是 636。单层注意力就是4 × 636² × 4096,算出来大约 66 亿次浮点运算。如果模型有 32 层,那就奔着 210 亿次去了。再看看 2304 视觉令牌的情况,光注意力部分就会膨胀到千亿级别,视觉侧的统治力一下子就体现出来了。
更关键的是,文本令牌数量受对话长度限制,可以通过截断、摘要来控制;视觉令牌数量却和输入分辨率强绑定,分辨率稍微提一档,成本肉眼可见地上涨。
1.2 剪枝解决的究竟是哪部分开销
既然排在面前的成本大头是注意力计算,那自然想到了剪掉一部分令牌,直接从序列长度上做文章。
有人可能会问:为什么不直接降分辨率,或者用更小的视觉编码器?这样做当然简单,但它牺牲的是信息密度。图像里总有一部分区域是后面推理真正需要用到的,比如文档 OCR 里的文字区域、自动驾驶场景里的交通标志、电商图片里的商品主体,如果无差别地压缩掉,模型就失去了回答问题的依据。
令牌剪枝的思路是动态的:与其把整张图等比缩小,不如保留那些对当前任务最有信息量的区域。这张图上如果有五个物体,只保留四个,推理照样进行,但注意力计算量少了一截,同时 KV Cache 的压力也小了很多。
不过,理解了这个“为什么剪”之后,真正的问题才浮出水面:怎么决定留谁、杀谁、以及每一层是不是都应该做同样的决定?这一步远比想象中容易踩坑。
2. 局部优化的老路:看似步步最优,实际全局打转
2.1 经典局部方案的常见形态
最早一批令牌剪枝方法,大多数是走“局部贪心”的路线。它们的共同点是:在某一层,基于某个局部信号,立刻做出剪枝决定,被剪掉的令牌后续层就不再参与计算。
具体来说,常见的做法是在视觉编码器或者跨模态融合层里,计算每个令牌和全局相关令牌的注意力权重,或者统计注意力熵。权重很低、熵特别高的令牌就被认为是“没什么用”的,当场丢掉。这个过程可以每一层做一次,也可以每隔几层做一次——英文里管它叫 token dropping 或者 token merging,本质上都是局部决策。
这套做法很符合直觉:我在这一层看来不重要的东西,后面应该也不重要,早扔早省钱。大部分早期工作的实验也确实表明,在保持一定剪枝比例的前提下,模型在不少基础评测任务上不会掉太多分。所以很长一段时间里,“每层挑出 Top-K 保留”几乎成了 token pruning 的标准范式。
2.2 局部选择与全局保真度的冲突
但真要把这套局部方案推到更高剪枝率的时候,问题就来了。
麻烦在于:一个令牌在这一层不重要,不代表它在后续层里始终不重要。Transformer 每一层的注意力结果都会发生语义重组。举一个简化但很常见的例子:一个人形轮廓在浅层是由几十个边缘令牌共同组成的,单看某个边缘令牌,它和其他部位的注意力权重都不高,很容易被局部策略判定为“冗余”。但到了深层,这个轮廓需要被整体识别成“人”的语义中心,之前被剪掉的那部分边缘令牌就缺席了。模型只好把剩下的令牌强行拉出来顶替位置,结果是识别质量明显下降,或者模型开始出现幻觉成分。
这里存在一个“不可逆性”:局部方案在浅层做的决定无法在深层挽回。它等同于每一步都只盯着当前位置的小山坡往上爬,完全不管整张地形图长什么样。
在数学优化里,这类方法最典型的毛病就是“陷入局部最优解”。放在剪枝场景里,局部最优具体表现为:当前层剪枝选区让本层损失最小,但经过多层传播之后,全局信息保留度反而被破坏了。这个破坏不一定立刻出现在准确率上,可能积攒到某个下游任务上突然爆发。
2.3 只做递推的局部优化为什么难自愈
还有一个容易被忽视的细节:剪枝这个动作本身会改变后续层看到的输入分布。你在一开始剪掉的一些令牌,会导致某一层的注意力概率发生偏移;下一层看到的是已经偏移过的分布,再在此基础上做新的局部最优选择时,误差就一层层滚下去了。
这种误差累积有点像把一张 JPEG 反复压缩再解压:第一次压缩损失一点细节,视觉上还能忍;第二次、第三次叠加起来,边缘就开始糊了。局部剪枝的问题就是它允许这种“反复压缩误差”无阻碍地在层与层之间扩散。
另外,很多局部方案依赖的注意力得分,本身也不是绝对可靠的重要性指标。注意力权重高,有可能只是代表“这个令牌经常被众头关注”,并不等于“这个令牌承载了唯一不可替代的信息”。如果一张图里所有令牌都高频聚焦在背景天空上(很多视觉模型面对自然图片的时候就是这样的),局部方法就会把重要的前景细节误判成次要角色。
所以,面对“局部最优”这个坑,解决方案不应该是再设计一个更复杂的局部重要性函数,而是要把剪枝决策从“每层各自为政”改成“全局统筹分配”。
3. 平衡剪枝的全局设计:从逐层挑选到整体预算
3.1 “平衡”体现在三个维度
这篇工作标题里最值得琢磨的词,就是 Balanced,也就是“平衡”。它并不是单纯指“剪枝比例均衡”,而是指剪枝策略要在几个维度上同时做取舍:
第一是信息量与多样性的平衡。如果只按“注意力重要性”来排序,很容易出现“留下所有令牌都集中在同一个对象上”的极端情况。举例来说,给模型看一只猫站在草地上,局部重要性排序靠前的令牌几乎全是猫脸附近的区域。这时候如果不做多样性约束,草地部分就被全部剪掉,后续如果模型需要回答“猫在什么环境里”,它就失去依据了。平衡剪枝会在重要性之外,额外维持一个覆盖度约束,让保留的令牌分区分布尽可能接近原图的语义覆盖范围。
第二是层间预算的平衡。局部方法每一层都在独立剪枝,有些层可能被剪得太狠,有些层又被剪得太保守。理想的做法是设定一个全局总预算,再按深层的需求分配。比如浅层可以多留一些令牌,保证底层特征供给;深层则可以适当收紧,因为语义已经比较集中,冗余度更低。这个过程正如资源分配,把水从冗余的渠引到最缺水的田里。
第三是局部效率与全局效率的平衡。深度网络内部存在一条隐式的信息瓶颈,如果剪枝信号完全来自最后一层的最终目标,计算代价会很高,而且难以训练;如果完全来自当前层,又会陷入局部最优。“平衡剪枝”的解法是,用一个跨越多层的信号来衡量每个令牌的边际价值,既保留当前局部的实时信息,又兼顾深层语义的需要。
3.2 全局重要性信息的获取方式
实际操作里,想要拿到“全局重要性”,通常不能靠逐层前向传播里临时停下来算,它会让实现变得非常复杂。更常见的做法是引入一个轻量级的全局重要性估计器。
它的基本形态可以是:利用模型内部某个特定层(比如中层)的表示,去计算每个视觉令牌在未来层中可能发挥作用的概率。训练这个估计器时,可以用“完整模型输出”作为监督信号,比如计算各个令牌对最终答案预测的贡献度。因为是离线训练好的,所以在推理时只需要额外做一次很轻的预测,就能预先知道每个令牌对全局目标的价值,而这个价值又被同时用来做跨层预算分配。
这其实是在每个令牌上建立一个“全局效用”函数,近似替代之前那种局部贪心的排序逻辑。有了全局效用之后,只需要在这个基础上加入区块覆盖度约束,就构成了一个非常接近全局优化的求解问题。虽然严格意义上它不是完美全局最优(比如上了某些近似算法),但比起那种纯逐层贪心,它已经算翻越了局部山坡。
3.3 配合模型结构的预算分配策略
不同视觉语言模型的结构其实差别蛮大:有的模型视觉编码器和语言解码器是严格分开的,视觉令牌在跨模态注意力层融入文本;有的模型则一开始就把视觉补丁和文本令牌拼在一起,走统一的自注意力。
这也决定了剪枝的插入点应该不一样。如果是分离式结构,那么剪枝更适合放在视觉编码器末端和跨模态融合之前的区域,主要目的是削减进入后续自注意力计算的令牌总数。如果是拼接式结构,视觉令牌会和文本令牌共同参与注意力的概率分布计算,这时剪枝的粒度需要考虑得更细,比如是否允许文本令牌与视觉令牌采用不同的保留策略,以及是否在每一层重复执行剪枝。
从标题的研究倾向来看,它的思路更偏向“在整体层面确定每个令牌的保留权重”,再根据模型不同层的冗余度做动态剪裁。这样一来,剪枝决策就不再是某一层的“个人行为”,而是整个模型的一次全局资产配置。
4. 实验观察:加速收益在什么条件下才成立
4.1 衡量加速不能只看 FLOPS
另一个我特别想提醒的地方是:只盯着浮点运算次数看加速,很容易被纸面数字骗过去。
注意力层本身的计算量下降确实和序列长度直接相关,但在实际推理引擎里,决定一个请求延迟的因素还包括显存带宽、算子启动开销、并行效率、甚至输入图片的尺寸是否触发某些 kernel 的边界条件。很多剪枝方法在实验里缩短了序列长度,FLOPs 下降可观,但端到端的服务延迟却可能只下降了百分之二三十,原因就在于瓶颈不在某一块计算,而在整体管线调度。
所以当你看到一篇论文给出的加速倍数时,一定要先搞清楚它的测量口径:是纯模型的 encoder 耗时,还是包含 tokenizer、图像预处理、采样循环在内的完整端到端延迟;是并发请求下的吞吐,还是单请求的 P50 延迟。这两类数字之间的差距可能非常大。
类比的实务:修改剪枝比例时,不要只依赖实验室里的 benchmark,一定要把改动放进真实服务的链路里做一次耗时采样。我自己遇到过一次非常离谱的情况:把剪枝率从 30% 提到 50%,按理说延迟应该掉一大截,结果因为其中一个算子触发了一个更低效的 kernel 分支,反而比之前更慢了。这种问题在离线评测环境里基本发现不了。
4.2 剪枝比例与准确率的权衡曲线
关于令牌剪枝,一个经常出现在实验里的关系模式是:保留率较高时(比如保留 80% 到 90%),准确率几乎不掉,但加速效果有限;继续压到 50% 到 60% 时,会出现一段相对平缓的“软着陆”区间,速度明显上升,任务性能略有下降;一旦压到 30% 甚至以下,任务性能会突然跳水。
这个“跳水点”的位置非常关键,因为它决定了你能安全地把剪枝比例推到多高。论文里强调“超越局部优化”,本质上就是希望把软着陆区间延长,把跳水点往更高剪枝率的方向推。如果用局部贪心方法,跳水点通常来得更早,可能在保留率 50% 左右就开始崩;使用全局平衡策略后,同等保留率下掉分会更少,或者能在更低保留率下维持相近的分数。
从工程视角看,需要先在自己的业务测试集上画一条这样的权衡曲线。不同任务对剪枝的容忍度差异会很大:比如拍照识花这种“细粒度识别”任务,重要细节非常零散,剪枝的容忍度就很低;而一般的常识性问答任务,对令牌冗余的容忍度明显更高。务必用自己的数据做一张真实的“剪枝比例-业务指标”表,不要抄别人论文里的经验数值。
4.3 什么样的场景最适合这类方案
结合这个全局平衡思路的特点,我觉得真正适合它的场景至少具备两个特征:
第一,输入图像的语义有相对明显的“主次关系”或“冗余区域”。生活类图片、场景图、网络截图往往就是这样:大片背景、相同纹理区域占了很多令牌,真正派上用场的只有一小部分。这种情况下令牌剪枝的收益会非常直观。
第二,任务本身对上下文长度敏感,或者推进到高分辨率输入。我们看到不少线上模型想把输入分辨率从 336 提到 768 甚至更高,以解决小文字、小物件的识别问题,但又怕高分辨率带来巨大的延迟开销。这类场景最需要令牌剪枝:保留高分辨率的信息密度,同时把令牌总数量拉回低分辨率时代的水平。可以说,这种“高分辨率输入 + 激进剪枝”的组合,才是这个方向最理想的应用形态。
反过来,如果任务本身的图像内容非常密集,每个区域都可能影响答案(比如医疗影像、显微照片、各种布满文字的文档扫描件),剪枝的收益就很有限,强行上方案大概率会看到准确率下滑。
5. 实践建议与踩坑记录
5.1 最容易踩的坑:把剪枝当成免费午餐
我见过不少团队评估了剪枝方案之后,只取了论文里的剪枝比例,直接套到业务模型上,很快发现效果崩了。这里面的坑往往不在算法本身,而是“训练-剪枝-微调”这条链路没有对齐。
令牌剪枝不是单纯在推理时删掉几个令牌就可以的。哪怕你用的是全局重要性估计器,它对“哪些令牌重要”的判断,也是在某种训练分布下学会的。一旦业务数据分布和训练数据分布差异较大,重要性估计就会失真。这时候比较稳妥的做法是做任务内微调,让模型在有限令牌的条件下学会适应,而不仅仅是把剪枝当作后处理手段看待。用我刚才那个类比来说,训练阶段就应当预留“压缩-解压”的余量,而不是让模型在推理时被突然打一个措手不及。
更稳的做法是分步走:先在离线完整模型上做剪枝模拟,找到业务指标能接受的剪枝区间;再把这个区间内的剪枝结果和完整模型输出做对比,逐批量计算 KL 散度或任务损失差异;最后才进入微调和端到端测速。贪快直接跳到上线,大概率会花更多时间返工。
5.2 更稳妥的落地路径
结合我接触过的几个实践项目,我觉得比较可靠的落地顺序是这样的:
- 先做诊断:收集有代表性的输入样本,统计不同区域令牌的重要性分布,确认图像里是否存在大量的“低信息密度”令牌。这一步不用改任何代码,只做统计就可能已经说服大家不值得上剪枝。
- 再定指标:选择一个端到端延迟指标,明确是 P50 还是 P95,并且固定硬件环境。不要在不同的 GPU 型号之间比剪枝效果,否则变量不可控。
- 然后做小范围 A/B:在少量内部任务上跑一版剪枝,把输出和全量模型输出做逐样本对比,尤其是看一眼失败样本是不是集中在某些特定图像类别上。
- 最后再加微调:在保留的剪枝预算下,做短周期的领域微调,再复测一次业务指标。
这套流程听起来不复杂,但它能避免大多数学术方法“迁移到业务上就失灵”的尴尬。
5.3 和其他优化手段的配合
单独使用令牌剪枝的收益上限是存在的,但和其他手段搭配时,它往往能发挥出更大的价值。
最典型的搭配是量化。量化主要节省的是每个算子内部的乘加运算成本以及权重的显存占用,而令牌剪枝节省的是序列长度。两者作用维度不同,叠加起来基本算是乘法效应。实际操作里可以先剪枝、再量化,或者先量化、再剪枝,顺序不同,最终精度损失也不同,值得专门花时间对比一下。
另一个搭配是“提前退出”或“层跳过”。如果某些层对视觉令牌的计算特别冗余,那可以在保留全局预算的同时,把部分层直接跳过。这需要模型结构本身支持,但一旦支持,剪枝+跳层组合起来的端到端加速会非常可观。
KV Cache 的优化也能和剪枝协同:剪枝减少的令牌,同时也会减少 KV Cache 在长对话场景中的显存增长,对服务端并发能力的提升尤其明显。我自己在压测过程中感受最深的不是单请求变快,而是并发吞吐的提升——在很多部署形态里,吞吐能力往往比单请求延迟更值钱。
5.4 哪些时候别用
最后不吐不快:不是所有视觉语言模型都需要令牌剪枝。如果业务本身用的是低分辨率输入,图像内容极度依赖全局上下文,而且并发压力也不大,那花大量时间去调剪枝策略,性价比其实很低。还不如老老实实做量化,或者把时间花在缓存和调度优化上。
另外,如果团队内部已经有一套非常成熟且稳定的推理管线,那么引入剪枝带来的每层复杂度(尤其是全局重要性估计器的维护和校准成本)可能大于收益。有时候最简单的优化就是最好的优化,判断“不做”也是资深从业者最有价值的技能之一。
从这次“Balanced Token Pruning”的思路里,我最想带走的一点是:剪枝这类决策问题,一定要跳出“当前层最优”的思维惯性,回到全局信息流的角度来思考。局部最优点再漂亮,也架不住后续层多轮传播之后的误差放大。真正的加速方案,应当是一张从头到尾统一设计过的地图,而不是沿途路口各凭感觉的临时转向。