MoLe-VLA:通过混合层实现的动态跳层视觉-语言-动作模型实现高效机器人操作

发布于:2025-04-03 ⋅ 阅读:(15) ⋅ 点赞:(0)

25年3月来自南京大学、香港理工、北大和香港科技大学的论文“MoLe-VLA: Dynamic Layer-skipping Vision Language Action Model via Mixture-of-Layers for Efficient Robot Manipulation”。

多模态大语言模型 (MLLM) 在理解复杂语言和视觉数据方面表现出色,使通用机器人系统能够解释指令并执行具体任务。尽管如此,它们在现实世界中的部署仍受到大量计算和存储需求的阻碍。最近对 LLM 层中同质模式的洞察启发了稀疏化技术来解决这些挑战,例如提前退出(early exit)和 token 修剪。然而,这些方法往往忽略了最终层的关键作用,这些层对与下游机器人任务最相关的语义信息进行编码。根据神经科学领域浅脑假说 (SBH,Shallow Brain Hypothesis) 的最新突破以及模型稀疏化中混合专家,将每个 LLM 层概念化为一个专家,并提出一种混合层视觉-语言-动作模型 (MoLe-VLA 或简称 MoLe) 架构,用于动态 LLM 层激活。为 MoLe 引入时-空-觉察路由器 (STAR),可根据机器人的当前状态选择性地激活部分层,模拟大脑专门用于认知和因果推理的不同信号通路。此外,为了弥补 MoLe 中 LLM 认知能力的缺失,设计一种认知自-我知识蒸馏 (CogKD),以增强对任务需求的理解,并利用认知特征生成与任务相关的动作序列。在 RLBench 模拟和真实环境中进行的大量实验,证明MoLe-VLA 在效率和性能方面的优越性。

多模态大语言模型 (MLLM) [2、3、17、27、32] 的快速发展,已证明其能够整合复杂的语言和视觉表征,从而激发能够理解视觉语言、与人交互并在操作任务中灵活解决问题的通用机器人和具身智体的发展。初步的视觉-语言-动作 (VLA) 模型 [16、19、21、25],例如 RT-2 [6] 和 OpenVLA [16],已证明使用 MLLM 进行端到端机器人控制的可行性,从而实现了稳健的策略和新兴能力,包括泛化到未见过的物体和理解新命令。然而,在现实世界的机器人系统中部署 MLLM 面临着巨大的挑战,因为它们对计算的需求很高,包括大量的内存使用、功耗和时间延迟,这与机器人平台资源受限和实时要求相冲突。例如,在 RTX 4090 等商用级 GPU 上运行的 7B VLA 模型通常可实现约 5-12 Hz 的推理频率,这远远低于 Franka 机械臂所需的 50-1000 Hz 控制频率。

最近的研究 [34, 43] 发现 LLM 层存在大量冗余,特别是在机器人任务中,跨层的同质模式会导致高计算成本而性能提升有限。例如,DeeR [43] 证明,与使用六层相比,使用 Flamingo [21] 模型的所有 24 层只能将任务成功率提高 3.2%,而 Calvin LH-MTLC [28] 的计算成本增加 4 倍。类似地,在如图 (a) 中对 OpenVLA [16] 和 RLBench [13] 的分析表明,连续层输出之间的余弦相似度超过 90%,而第一层和最后一层的特征差异很大。这表明跳过相邻层可以减少计算量,但也凸显了提前退出策略 [10, 43] 的局限性,如图 (b) 所示,丢弃较深的层可能会丢失关键的语义信息。受浅脑假说 (SBH) [37] 的启发,该假说认为,大脑通过浅层、并行的皮层-皮层下(cortico-subcortical)环路,平衡深层层次结构,以实现认知和因果推理,为此本文提出一种 VLA 模型中的选择性层激活策略。如图 © 所示,该方法反映大脑的动态深度并行平衡,仅激活与任务相关的层以提高效率和适应性,体现VLA 模型设计中 SBH 的原则。本文引入一种混合层视觉-语言-动作模型(MoLe-VLA),该模型在 LLM 的输入阶段加入了一种层-选择路由器,以实现其稀疏性。

请添加图片描述

混合专家

MoE 范式通过条件计算增强模型容量,同时保持计算效率。对于输入 x,标准 MoE 层定义为:

请添加图片描述

其中门控权重的计算是:

请添加图片描述

为了提高效率,通常采用具有 top-k 选择的稀疏门控。为了解决负载不平衡问题(太多输入被路由到少数专家),引入负载平衡损失 L_lb:

请添加图片描述

这种损失鼓励平衡专家利用并提高计算效率。

混合层:MoLe-VLA

视觉-语言-动作模型。机器人负责执行长度为 L 的语言指令 l,在时间步长 t 接收来自传感器的观察值 o_t(例如,来自相机的 RGB 图像),以预测具有 7 个自由度 (DoF) 的夹持器动作空间以执行任务。

基本 VLA 模型主要由视觉编码器 E、MLLM π 和动作模块 A 组成。视觉编码器 E 包括 DINO-v2 [30] 和 Siglip [44],它将输入图像 o_t 编码为一系列信息 token v_t。对于多模态融合,在视觉编码器 E 生成的视觉表征之上建立一个 MLLM,它充当有效的多模态特征提取器 π。其输出 f_t 表示来自时间步 t 的 MLLM 最后一层的隐状态序列,对应于认知 token。这为后续动作模块解释和得出所需动作提供条件。按照 CogAct [19],动作模块 A 将从输出特征 f_t 中提取的认知特征 e^c_t 作为输入,并预测最终动作 a∗_t 。

该视觉、语言和动作模块通过最小化动作模块预测噪声与真值噪声之间的均方误差进行端到端训练。

通过 MoLe 路由器实现跳层机制。本文提出 MoLe-VLA 来提高 LLM 在机器人任务中的效率,由于机器人任务的推理需求更简单,许多 Transformer 层未得到充分利用。MoLe 采用轻量级路由器在推理过程中自适应地跳过非必要的 Transformer 层,在保持性能的同时降低计算成本。

如图所示,对于具有 K 层的给定 MLLM π,MoLe 路由器处理输入嵌入 x_k 并生成二元门控向量 G_mol(x) = {G_k},其中 G_k ∈ [0, 1]。为了确保效率,只有 G_mol(x) 中的前 k 个值设置为 1,确定哪些层 π_k 与隐藏特征 h_k 一起执行,而其余层则被跳过。

请添加图片描述

与将 token 分配给专家的传统 MoE 路由器不同,MoLe 路由器跳过整个层,避免冗余计算。这提高了推理效率和响应能力,使 MoLe 特别适合需要轻量级和自适应处理的实时机器人任务,如操纵和导航。MoLe 的完整伪代码在算法 1 中提供。

请添加图片描述

时空-觉察路由器

本文提出一种路由机制,该机制协同利用视觉输入的空间结构和语言输入中的时间依赖性,动态地为 VLA 任务选择合适的 LLM 层。给定视觉特征 v_t 和文本特征 l,使用可学习矩阵 W_p 将两种模态投影到共享潜空间中。

从 h_img 计算空间路由权重 S 以捕获空间特征,同时使用 Transformer 模块从 h_text 得出时间路由权重 T,然后进行平均池化。一个根据 h_text 的 [CLS] token 计算的动态温度因子 α ∈ [0, 1] 可调节路由锐度(sharpness)。

最终的专家门控权重 G 结合 S 和 T,按 α 尺度化,并通过 Gumbel-Softmax 计算以进行可微分选择。通过整合空间和时间信息,本文方法使路由器能够选择 LLM 层,自适应地优化 VLA 任务的性能。该方法非常高效,每个文本样本仅需要 O(N_e (d_2 + N_text^2 )) FLOP,而标准 MoE 框架则需要 O(N_e * d),其中 d ≫ N_text, d_2。这种设计确保高适应性和计算效率。

认知自-我知识蒸馏

在实现高效跳层机制的同时,本文还设计一种自我蒸馏策略来补偿稀疏 LLM 中的认知损失,如图所示。这里,将原始模型作为老师,将 MoLe 模型作为学生。

请添加图片描述

为了蒸馏 token,一种常见的方法是逐个 token 地模仿张量 [7, 41, 49]。正式地,利用教师和学生网络的 token f^(t) 和 f^(s),可以通过 token 重构来实现模仿如下:

请添加图片描述

但是,上述公式对每个 token 都进行平等的处理和蒸馏,这是不恰当的。例如,与文本描述相关的视觉 token 应该得到更多的关注 [42, 48]。

因此,引入一个可学习的嵌入 e_tc,称为认知 token,以进行自适应蒸馏。具体来说,它被插入到底层,有效地整合视觉 token 和语言指令,以更好地理解任务要求并生成与任务相关的动作序列。教师和学生模型分别有自己的 e_tc,(t) 和 e_t^c,(s)。在蒸馏过程中,通过计算认知 token 和学生 token 之间的相似度来获得兴趣 token (ToIs) M。

接下来,利用教师和学生认知 token 生成的 ToI 的交集来决定每个 token 的提炼程度,其中 M = M(t) ⊙ M(s),因为蒸馏 token 应该包含对教师和学生都很重要的 token。故修正的蒸馏公式如下:

请添加图片描述

此外,引入 Reverse-KL [12] 与认知 token 配对,作为之前获得 L_cog-reversekl 的方​​式,以增强分布约束:

请添加图片描述

最终 CogKD 的损失是上述二个损失的结合。

模拟和实际部署。为了评估方法并展示其泛化能力,在 CoppeliaSim 模拟器和实际环境中对 RLBench [13] 进行了实验:

  1. RLBench 包括使用 Franka Panda 机器人和前视摄像头执行的 10 个不同的桌面任务。这些任务范围从目标操作到环境交互,例如:关闭盒子、关闭笔记本电脑盖、放下马桶座圈、将垃圾放入垃圾桶、扫入簸箕、关闭冰箱、将手机放在底座上、从支架上取下雨伞、从衣架上取下框架和更换时钟。任务数据是使用预定义路径点和开放运动规划库 [36] 生成的。根据先前的研究 [14],每个任务包括使用基于帧方法采样的 100 条训练轨迹,并在训练工作区内对每个任务进行 25 次试验评估。

  2. 在配备 3D 打印 UMI 夹持器 [8] 的 Franka Research 3 (FR3) 机器人上针对三项任务评估真实世界部署。安装在手腕上的 GoPro 9 摄像头,可捕捉真实世界的视觉观察。在定义的工作空间范围内使用手持式 UMI 夹持器为每个任务收集 50 个演示,包括拆卸充电器、拉抽屉和倒水。单个智体在所有任务上进行训练,并在训练工作区内每个任务进行 10 次试验评估。成功率通过人工评估确定,并作为评估指标。

所有基线都使用相同的任务配置进行训练,以便进行公平比较。每个方法的预训练参数都按照各自的训练设置进行加载。对于 MoLe-VLA,单视图 RGB 输入的大小调整为 224 × 224,机器人状态与预测动作(7-DOF 末端执行器姿势)对齐。该模型的训练批次大小为 64,每个样本有 8 个扩散步,使用预先训练的视觉和语言模块权重。视觉模块结合 DINO-v2 和 SigLIP,而语言模块 LLAMA-2 和动作模块 DiT-Base 则以 2 × 10−5 的恒定学习率进行端到端训练,进行 1k 次迭代。使用 PyTorch 的全分片数据并行 (FSDP) 框架,在 8 个 NVIDIA A800 GPU 上进行训练,耗时约 1.5 小时。