ICLR 2025｜首个动态视觉-文本稀疏化框架来了，计算开销直降50%-75%
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
本文由华东师范大学和小红书联合完成，共同第一作者是华东师范大学在读硕士、小红书NLP团队实习生黄文轩和翟子杰，通讯作者是小红书NLP团队负责人曹绍升，以及华东师范大学林绍辉研究员。
多模态大模型（MLLMs）在视觉理解与推理等领域取得了显著成就。然而，随着解码（decoding）阶段不断生成新的token，推理过程的计算复杂度和GPU显存占用逐渐增加，这导致了多模态大模型推理效率的降低。现有的方法通过减少预填充（prefill）阶段的视觉token冗余来实现推理加速。遗憾的是，这种在预填充阶段实现的视觉token稀疏化所带来的加速优势，在解码阶段会逐渐减弱。当解码输出的文本token数量增多时，这些方法仍然会遇到性能瓶颈。
为了解决上述问题，团队创新性地提出了一个全新的动态视觉-文本上下文稀疏化推理加速框架——Dynamic-LLaVA。该框架针对多模态大模型在不同推理模式下（包括预填充阶段以及有无KVCache的解码阶段），设计了定制化的稀疏化推理方案，以实现多模态大模型的高效推理。实验结果表明，Dynamic-LLaVA在几乎不损失视觉理解和生成能力的前提下，能够将预填充阶段的计算开销减少约75%；在无KVCache的解码阶段，计算开销减少约50%；在有KVCache的解码阶段，GPU显存占用减少约50%。Dynamic-LLaVA为多模态大模型推理加速领域树立了新的标杆。
论文GitHub仓库：https ://github.com/Osilly/dynamic_llava
论文链接：https ://arxiv.org/abs/2412.00876
研究团队：华东师范大学、小红书NLP团队
1引言
1.1前置信息：预填充与解码
本文主要围绕以LLaVA为范式的多模态大模型展开研究。一个多模态大模型的推理过程可以分为预填充和解码两个阶段：
在预填充阶段，不同模态的特征被映射到与大语言模型（LLM）输入embedding相同的特征分布空间中。这些多模态特征与文本token会一起被大语言模型处理，以生成初始输出文本token。以图片理解场景为例，该阶段主要处理输入的图片和文本格式的问题。
在随后的解码阶段，预填充阶段生成的所有token以及后续生成的所有输出文本token，将被用于自回归生成，从而产生完整的输出。同样以图片理解场景为例，该阶段生成针对整个问题的完整回答。
1.2多模态大模型推理加速困境
图1：多模态大模型生成过程（有/无KVCache）中FLOPs（计算复杂度）和GPU显存开销的增长趋势
现有的多模态大模型大多以基于解码器架构的大语言模型（LLM）为核心，这些模型通常拥有庞大的参数规模。在生成输出文本token的过程中，模型计算负担会逐渐加重，导致对计算资源的巨大消耗。为了提升推理速度，现有模型通常会在解码过程中运用KVCache技术，通过存储并复用之前计算的KV激活值来减少重复计算。然而，如图1(B)所示，即使使用了KVCache，LLaVA在输出token不断增加时，仍会迅速面临GPU显存耗尽的问题。
与文本不同，视觉信息往往包含大量冗余。因此，许多方法尝试通过减少视觉上下文来加速多模态大模型的推理，即对预填充阶段的视觉token进行剪枝处理。但这种方法存在局限性：其主要提升了多模态大语言模型在预填充阶段的推理效率，而在解码阶段，其效率提升会逐渐减弱。
如图1(B)和(C)所示，FastV这种针对视觉token剪枝的方法，虽然相较于原始的LLaVA能够节省一定的GPU显存和计算开销（FLOPs），但当输出token数接近5K时，它仍然会遭遇计算资源瓶颈。此外，FastV和原始LLaVA的曲线斜率基本一致，这表明在长输出的解码阶段，这类方法并没有显著的推理效率优势。因此，仅通过减少预填充阶段的视觉token，在输出文本token数量远超视觉token时，难以实现整个推理效率的显著提升。
1.3迈向全阶段推理加速：Dynamic-LLaVA
针对上述问题，我们认为：为了实现真正的全阶段推理加速，不仅需要对预填充阶段的视觉token进行剪枝，还必须对解码阶段输出的文本token进行稀疏化处理，限制参与自回归运算的token数量。为此，我们提出了Dynamic-LLaVA，针对多模态大模型的视觉-语言上下文稀疏化推理加速框架。该框架能够集成到多模态大模型推理的不同阶段中，实现以下目标：
显著降低预填充阶段计算开销：通过优化视觉token的处理方式，减少不必要的计算。
提升解码阶段的推理效率：无论是否使用KVCache，都能减少计算开销，提高推理速度。
保持性能优势：在视觉理解任务上几乎不损失性能；在长文本输出场景中，生成能力也几乎不受影响。
通过这些创新，Dynamic-LLaVA为多模态大模型的高效推理提供了一种全新的解决方案。
2方法
图2：Dynamic-LLaVA整体框架
如图2所示，Dynamic-LLaVA可以集成到多模态大模型推理流程中的不同阶段。具体而言，在预填充阶段，该框架对视觉token执行精准剪枝操作，剔除冗余信息；在不使用KVCache的解码阶段，限制参与自回归运算的视觉与输出文本token数量，避免不必要的计算负担；而在使用KVCache的解码阶段，Dynamic-LLaVA则动态调控KVCache，自适应判断是否将当前输出文本token的KV激活值纳入KVCache，优化资源利用效率。为了使模型适应这种全新的稀疏化推理模式，Dynamic-LLaVA在预训练的LLaVA-1.5基础上进行了1个epoch的监督微调（SFT），确保模型能够高效地运行在稀疏化的推理路径上。
2.1预填充阶段
在预填充阶段，我们对输入的视觉token进行稀疏化操作。如图2左侧部分所示，我们引入一个可训练的轻量化的图像预测器（ImagePredictor），来判断应当丢弃哪些视觉token。该图像预测器的结构如下图：
图3：图像预测器的结构示意图
图像预测器会对每个视觉token产生“决策分数”，以决定对哪些视觉token进行保留。在端到端训练中，视觉token的剪枝通过0-1二值化的掩码操作实现（具体过程见2.4节）。在实际推理阶段中，通过保留“决策分数”前k大的视觉token（即图2左侧部分的“Yes”分支），实现视觉token数量减少，以实现推理加速。
2.2解码阶段
不使用KVCache的解码过程：
对于视觉token，采用和上一小节相同的做法，进行稀疏化处理。
对于输出的文本token，分两类进行处理：
最后一个输出的文本token（即图2中间部分的“Lastoutputtexttoken”），不进行任何处理，完整输入LLM的decoder层进行计算。这样做的目的是保证模型的输出内容是连贯的，产生新的输出文本token时，始终保证自回归运算包含上一个输出文本token。
对其他历史的输出文本token进行稀疏化操作，其形式类似于对视觉token的处理。引入一个结构如下图的输出预测器（OutputPredictor），给出每个输出文本token的“决策分数”，以决定当前产生新的输出内容时，应当包括哪些文本token进行自回归运算。图2中间部分的“Yes”分支，表明保留的输出文本token。
图4：输出预测器的结构示意图
使用KVCache的解码过程：
KVCache是节省冗余计算的一个关键推理加速技术，其思想是“用GPU显存的空间换计算时间”。显而易见的是，KVCache也并非无限大，在长输出情况下，必须丢弃一些KVCache以适应有限的GPU显存。目前在LLM领域已有大量的KVCache压缩方案，以方法为代表，这一类方法一般基于当前token和历史KVCache进行重要性分数计算，以压缩历史KVCache。
与上述方法不同的是，我们对有KVCache的解码阶段的设计，核心在于“仅判断当前新token的KV激活是否需要加入KVCache中”。如图2右侧所示，对于当前正在处理的新token（Lastoutputtexttoken），使用和上一部分结构相同的输出预测器，以决定是否加入KVCache集合中。这种“OnlineKVCache压缩”方法，判断是否保留KVCache的过程计算复杂度更低，也更加适应多模态场景。在论文附录中，我们详细讨论了我们的方法和现有的LLMKVCache压缩方法的区别。
需要特别说明的是，和不使用KVCache的解码阶段相同，无论当前处理的token是否加入KVCache，其都会输入LLMdecoder层进行计算，以保证输出的连贯性。
2.3端到端训练
图5：Dynamic-LLaVA在端到端训练过程中的示意图
Dynamic-LLaVA是一个需要训练的多模态大模型推理加速框架。我们基于LLaVA进行了一个epoch的指令微调，以实现对token动态选择的稳定性，保证最终的性能。为了保证端到端训练，在训练阶段的稀疏化操作通过0-1二值化掩码实现（在推理中的实现是直接从历史token序列中丢弃token）。如图5所示，上半部分表示训练中进行mask的过程，在得到整个token序列的重要性分数后，我们选取前k重要的token进行保留，相对应的生成掩码向量，其中0对应丢弃的冗余token（不参与注意力过程的计算），1对应保留的重要token，进一步基于掩码向量生成注意力过程的掩码矩阵。掩码矩阵用来对多头注意力机制进行掩码操作，以确保丢弃的token不参与注意力过程的计算。由于二值化操作会导致不可微问题，所以我们借助了GumbalSoftmax和梯度直通估计器（StraightThroughEstimator,STE）来保证梯度流的正确传播，以进行端到端的训练，如图5下半部分所示。
3实验
Dynamic-LLaVA基于LLaVA-1.5-7B和13B的两个版本进行了1个epoch的指令微调，训练使用的数据和LLaVA-1.5相同。
3.1视觉理解能力
我们首先评估了Dynamic-LLaVA在主要的视觉理解基准的性能，选取了目前主流的多模态大模型推理加速方法进行比较。
表1：视觉理解基准效果对比。其中，Free表示方法是否是Training-Free的。Dynamic-LLaVA的下标”I”和”I|T”分别表示仅对视觉token做稀疏化和同时对视觉和文本token都做稀疏化（该标识适用于下文所有的表格）
如表1所示，Dynamic-LLaVA在大部分视觉理解任务上取得了优越的性能。和其他对视觉内容稀疏化的方法相比，Dynamic-LLaVA在能大幅减小计算复杂度的同时，能够实现相比原始的LLaVA-1.5性能几乎不下降。此外，在SciQA、POPE、MME和MMBench上，Dynamic-LLaVA相比LLaVA-1.5甚至有一定的性能提升。例如，在SciQA任务上，Dynamic-LLaVA的7B和13B版本，相较于LLaVA-1.5实现了2.3%和0.8%的性能提升。
表2：与其他高效视觉projector的SOTA方法对比
值得一提的是，Dynamic-LLaVA并没有对LLaVA-1.5的视觉projector进行修改，就可以实现大幅降低预填充阶段计算复杂度，同时维持模型性能。在表2中，和其他针对视觉projector做高效设计（以提高推理效率）的SOTA方法进行了对比。相较于其他使用了高效的视觉projector的方法，Dynamic-LLaVA使用和LLaVA-1.5相同的MLP结构作为视觉projector，实现了更好的性能，同时也大幅降低了预填充阶段的计算复杂度。此外，Dynamic-LLaVA也可以和其他使用高效视觉projector的方法集成。例如，表2中Dynamic-LLaVA使用TokenPacker这一高效视觉projector的版本，在原始的TokenPacker方法基础上，进一步减少了视觉token。相较于其他基于TokenPacker的推理加速方法，性能损失最少。
3.2生成能力
现有的视觉理解任务中，一般只要求模型给出简短的回复，这和现实世界中多模态大模型的应用场景仍然存在不小的区别。在现实使用中，多模态大模型多数情况下会被要求生成更长、更细致的描述。为了和现实世界的场景对齐，评估Dynamic-LLaVA在更长的输出情况下的生成能力和推理效率。我们额外构建了两个评估模型生成能力的基准：
ShareGPT4V-VQA：基于ShareGPT-4V数据集，选取了caption超过300个单词的单论对话样本，平均输出token长度超过1000。
表3：生成能力基准比较。其中，解码阶段的TFLOPs和Mem.（GPU显存占用）分别在无/有KVCache的情况下测量得出。PPL越低越好，METEOR越高越好
如表3所示，相比LLaVA-1.5，只进行视觉内容稀疏化的Dynamic-LLaVA的生成流畅度（PPL）和生成质量（METEOR）几乎没有变化；同时对视觉和文本进行稀疏化的Dynamic-LLaVA，PPL仅变高了0.3，METEOR甚至略有提升，而在推理效率上，在无KVCache的解码阶段降低了～50%的TFLOPs，在有KVCache的解码阶段降低了～50%的GPU显存占用。实验结果充分表明，Dynamic-LLaVA针对视觉和文本同时进行稀疏化，几乎不影响实际生成能力，却可以实现大幅的推理效率提升。
3.3实际推理效率
表4：Dynamic-LLaVA-13B推理效率实测。其中，2K/4K表示输出的文本token数，所有结果均在一张A100(80G)上测试得出，batchsize固定为8。“”表示GPU显存耗尽
在表4中，我们测试了多模态大模型实际推理的时间和GPU显存占用。Dynamic-LLaVA实现了更快的推理速度和更低的显存占用。FastV这种对预填充阶段的视觉token进行剪枝的方法，随着输出长度的增长，推理效率也逐渐降低。而我们提出的Dynamic-LLaVA，随着输出变长，相比于FastV的推理效率优势也逐渐显现出来。
3.4实例展示
图6：Dynamic-LLaVA-13B在LVIS-VQA(single-round)上的推理结果展示。图片的白色部分表示该位置的图像块被丢弃，文字中的灰色部分表示其在稀疏化过程中被丢弃，这表示它们不参与后续的自回归解码过程，但在模型的输出中都被完整保留
图6中展示了Dynamic-LLaVA-13B在LVIS-VQA(single-round)上的推理结果，以及对视觉和文本token的稀疏化情况。可视化结果表明，视觉token部分的主要信息得以保留；文本token中，一些不影响整体语义理解的连词、介词等被丢弃。这表明Dynamic-LLaVA能够实现关键的视觉、语义信息的保留，从而保证了模型整体的性能。
4总结
针对当前多模态大模型推理效率受限的问题，团队通过分析多模态大模型推理过程中的不同阶段，针对性的设计了推理加速方案。提出了Dynamic-LLaVA——第一个同时稀疏化视觉和语言上下文的多模态大模型推理加速框架，将不同推理模式的推理效率优化集成到统一框架中。
随着多模态大模型技术的发展，尤其是其在复杂推理、长思维链领域的不断进步。我们有理由相信，Dynamic-LLaVA的应用场景正变得更加广泛，其对输出文本token进行稀疏化的模式，会在当前的更长输出、更复杂推理的场景下，体现出更明显的推理加速优势。
作者简介
黄文轩：小红书NLP团队算法实习生，现硕士就读于华东师范大学计算机科学与技术学院2023级。他在ICLR、CVPR等国际顶级会议上以第一作者身份发表了多篇学术论文，主要研究方向包括多模态大模型、大模型的高效训练与推理等。
翟子杰：小红书NLP团队算法实习生，现硕士就读于华东师范大学计算机科学与技术学院2023级。他在ICML、ICLR、EMNLP等国际顶级会议上发表过多篇学术论文，研究方向主要集中在多模态大模型、生成式搜索与推荐大模型等领域。
曹绍升：小红书NLP团队负责人，发表论文30余篇，授权专利100余项，引用近4000次，获得ICDE2023年最佳工业论文奖、CIKM2015-2020年最高引用论文、AAAI2016最具影响力论文。此外，还荣获了中国发明协会创新成果一等奖（排名1）、中国人工智能学会吴文俊科技进步二等奖（排名1），连续4年入选世界人工智能学者榜单AI-2000新星榜前100名、Elsevier中国区高被引学者，CCTV-13《新闻直播间》采访报道。
叶哲宇：硕士毕业于帝国理工学院计算机专业，小红书NLP团队算法工程师，专注于大模型算法与应用方向，开源社区DMLC成员。他在ICLR、NAACL、EMNLP等国际顶级会议上发表过多篇论文，研究领域涵盖大模型应用、多模态大模型、Agent模拟等。
林绍辉：华东师范大学计算机学院研究员，紫江青年学者，2021年扬帆计划获得者，曾获中国人工智能学会优秀博士论文提名奖、《中国科学：技术科学》最佳审稿人。在国际顶级期刊和会议发表超过50篇论文，包括TPAMI、TNNLS、TMI、CVPR、ECCV、AAAI、IJCAI等。担任CVPR2024领域主席、IJCAI2020SPC以及国际顶级期刊和会议审稿人。目前主要研究方向有计算机视觉、机器学习、图像视频理解、低层视觉等。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
