The Better You Learn, The Smarter You Prune: Towards Efficient Vision-language-action Models via Differentiable Token Pruning
作者: Titong Jiang, Xuefeng Jiang, Yuan Ma, Xin Wen, Bailin Li, Kun Zhan, Peng Jia, Yahui Liu, Sheng Sun, Xianpeng Lang
分类: cs.RO, cs.CL, cs.CV
发布日期: 2025-09-16 (更新: 2025-09-21)
备注: Under review. Project site: https://liauto-research.github.io/LightVLA
💡 一句话要点
LightVLA:通过可微Token剪枝提升视觉-语言-动作模型的效率与性能
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 视觉-语言-动作模型 Token剪枝 可微学习 机器人 效率优化 Gumbel Softmax
📋 核心要点
- VLA模型计算量大,难以在资源受限平台部署,主要瓶颈在于视觉Token上的注意力计算。
- LightVLA通过动态查询评估Token重要性,使用Gumbel softmax实现可微剪枝,保留关键Token。
- 实验表明,LightVLA在LIBERO上优于其他VLA模型和剪枝方法,显著降低计算开销并提升成功率。
📝 摘要(中文)
本文提出了一种简单而有效的可微Token剪枝框架LightVLA,用于提升视觉-语言-动作(VLA)模型的效率。VLA模型在执行现实世界机器人任务中表现出令人印象深刻的能力,但其在资源受限平台上的部署常常受到基于大量视觉Token的繁重注意力计算的限制。LightVLA通过自适应的、性能驱动的视觉Token剪枝来解决这一挑战:它生成动态查询来评估视觉Token的重要性,并采用Gumbel softmax来实现可微的Token选择。通过微调,LightVLA学会保留信息量最大的视觉Token,同时剪枝对任务执行没有贡献的Token,从而同时提高效率和性能。值得注意的是,LightVLA不需要启发式魔法数字,也不引入额外的可训练参数,使其与现代推理框架兼容。实验结果表明,LightVLA在LIBERO基准测试的各种任务中优于不同的VLA模型和现有的Token剪枝方法,以显著降低的计算开销实现了更高的成功率。具体而言,LightVLA将FLOPs和延迟分别降低了59.1%和38.2%,同时任务成功率提高了2.6%。此外,我们还研究了基于可学习查询的Token剪枝方法LightVLA*,它也取得了令人满意的性能。我们的工作表明,随着VLA追求最佳性能,LightVLA自发地学习从性能驱动的角度剪枝Token。据我们所知,LightVLA是第一个将自适应视觉Token剪枝应用于VLA任务,同时兼顾效率和性能的工作,标志着朝着更高效、更强大和更实用的实时机器人系统迈出了重要一步。
🔬 方法详解
问题定义:VLA模型在机器人任务中表现出色,但其计算复杂度高,尤其是在处理大量视觉Token时,导致难以在资源受限的平台上部署。现有方法要么依赖启发式规则,要么引入额外的参数,限制了其通用性和效率。
核心思路:LightVLA的核心思路是自适应地剪枝对任务执行贡献小的视觉Token,从而降低计算量,同时保留关键信息以维持甚至提升性能。通过学习Token的重要性,并使用可微的方式进行选择,模型能够自动优化Token选择策略。
技术框架:LightVLA框架主要包含以下几个阶段:1) Token嵌入:将输入的视觉信息转换为Token嵌入表示。2) 动态查询生成:生成动态查询,用于评估每个视觉Token的重要性。3) Token重要性评估:使用查询与Token嵌入进行交互,计算每个Token的重要性得分。4) 可微Token选择:使用Gumbel softmax函数,根据重要性得分对Token进行概率选择,实现可微的剪枝过程。5) 任务执行:使用选择后的Token嵌入执行VLA任务。
关键创新:LightVLA的关键创新在于其自适应和可微的Token剪枝方法。与传统的启发式剪枝方法不同,LightVLA通过学习的方式确定Token的重要性,并使用Gumbel softmax实现可微的选择过程,使得模型能够端到端地优化Token选择策略。此外,LightVLA无需额外的可训练参数,保持了模型的简洁性。
关键设计:LightVLA的关键设计包括:1) 动态查询:使用与输入相关的动态查询,而非静态查询,能够更准确地评估Token的重要性。2) Gumbel Softmax:使用Gumbel softmax函数实现可微的Token选择,使得模型能够通过梯度下降优化Token选择策略。3) 无额外参数:LightVLA不引入额外的可训练参数,避免了增加模型复杂度和训练难度。损失函数通常是任务相关的损失函数,例如,机器人任务的成功率或动作预测的准确率。
🖼️ 关键图片
📊 实验亮点
LightVLA在LIBERO基准测试中表现出色,相较于现有VLA模型和Token剪枝方法,在降低计算开销的同时提升了任务成功率。具体而言,LightVLA将FLOPs降低了59.1%,延迟降低了38.2%,同时任务成功率提高了2.6%。LightVLA*虽然引入了额外的可训练参数,但也取得了令人满意的性能,验证了可学习查询的有效性。
🎯 应用场景
LightVLA在机器人领域具有广泛的应用前景,尤其是在资源受限的移动机器人、无人机等平台上。通过降低VLA模型的计算复杂度,LightVLA使得这些平台能够实时执行复杂的视觉-语言-动作任务,例如目标导航、物体操作等。此外,该技术还可以应用于其他需要处理大量视觉信息的任务,例如视频理解、图像搜索等。
📄 摘要(原文)
We present LightVLA, a simple yet effective differentiable token pruning framework for vision-language-action (VLA) models. While VLA models have shown impressive capability in executing real-world robotic tasks, their deployment on resource-constrained platforms is often bottlenecked by the heavy attention-based computation over large sets of visual tokens. LightVLA addresses this challenge through adaptive, performance-driven pruning of visual tokens: It generates dynamic queries to evaluate visual token importance, and adopts Gumbel softmax to enable differentiable token selection. Through fine-tuning, LightVLA learns to preserve the most informative visual tokens while pruning tokens which do not contribute to task execution, thereby improving efficiency and performance simultaneously. Notably, LightVLA requires no heuristic magic numbers and introduces no additional trainable parameters, making it compatible with modern inference frameworks. Experimental results demonstrate that LightVLA outperforms different VLA models and existing token pruning methods across diverse tasks on the LIBERO benchmark, achieving higher success rates with substantially reduced computational overhead. Specifically, LightVLA reduces FLOPs and latency by 59.1% and 38.2% respectively, with a 2.6% improvement in task success rate. Meanwhile, we also investigate the learnable query-based token pruning method LightVLA* with additional trainable parameters, which also achieves satisfactory performance. Our work reveals that as VLA pursues optimal performance, LightVLA spontaneously learns to prune tokens from a performance-driven perspective. To the best of our knowledge, LightVLA is the first work to apply adaptive visual token pruning to VLA tasks with the collateral goals of efficiency and performance, marking a significant step toward more efficient, powerful and practical real-time robotic systems.