跳转至

ROCKET: Residual-Oriented Multi-Layer Alignment for Spatially-Aware Vision-Language-Action Models

会议: NeurIPS 2026
arXiv: 2602.17951
代码: https://github.com/CASE-Lab-UMD/ROCKET-VLA
领域: 机器人/具身智能
关键词: 具身智能, 视觉语言动作模型, 空间表征对齐, 残差流, 梯度冲突

一句话总结

ROCKET 提出了一种基于残差流的多层表征对齐框架,通过层级不变的跨层共享投影器与套娃式稀疏激活机制,消除了多层蒸馏中的梯度冲突并平衡各深度表征,仅用约 4% 的训练计算量即在 LIBERO 基准上取得 98.5% 的最先进操作成功率。

研究背景与动机

基于视觉-语言-动作(Vision-Language-Action, VLA)的模型在指令跟随与具身操作任务中取得了显著进展,然而现有主流 VLA 模型主要基于预训练的 2D 视觉-语言骨干网络构建,缺乏对物理世界稳健的三维空间几何与约束理解。当机械臂操作涉及复杂的空间几何、视角变换或细粒度相对空间关系时,传统 2D VLA 往往出现定位偏差与操作失败。为了给 VLA 注入三维空间推理能力,直接输入点云或外接深度估计器的方法通常会带来显著的感知延迟与推理计算开销。相比之下,通过表征对齐(特征蒸馏)将 VLA 的中间层视觉 token 隐式对齐到强 3D 视觉基座模型(如 VGGT)已成为一种极具前景且推理零开销的范式。

然而,现有表征对齐方法普遍采用单层特征监督策略(如 Spatial Forcing、GLaD),其操作性能高度依赖于人工选择的特定对齐层,且最优层会随任务与数据分布剧烈波动,往往需要极其低效的事后网格搜索。深层残差网络中的特征在不同深度编码着从局部细节到全局几何的互补分层空间线索,理应通过多层联合对齐实现更全面的几何知识迁移。但令人意外的是,将经典知识蒸馏直接扩展至 VLA 时——即为各个不同深度独立配置专有投影器——往往会导致操作性能崩溃,在基准测试中的成功率甚至大幅落后于单层基线。

深入分析表明,这一失败根源在于深层残差网络中各层特征在参数空间存在“锥形聚集”趋势,而独立初始化的投影器在训练过程中迅速发生参数正交化,使各层向相互冲突的方向映射,进而在浅层残差流的反向传播叠加中引发剧烈的梯度干扰与破坏性抵消。本文的切入角度是:将多层对齐重新建模为“残差流到残差流”的一致动力学映射,通过结构约束确保各层梯度协同互补。核心 idea:使用单个跨层共享投影器替代独立的层级投影器以构建统一的残差流映射消除梯度冲突,并引入套娃式(Matryoshka-style)通道稀疏激活机制自适应平衡深浅层表征对齐容量。

方法详解

整体框架

ROCKET 的核心思想是将学生端 2D VLA 的残差流与教师端 3D 视觉基座模型(如 VGGT)的残差流进行多层一致对齐。针对独立投影器引发的梯度正交化与梯度抵消问题,ROCKET 采用单一参数共享的投影器 \(p_{\Phi}\) 来统一映射不同层级的残差特征,并在投影器内部引入按层深单调递增的嵌套激活通道(套娃结构),使得浅层特征使用较少通道捕捉通用局部线索,而深层特征激活全量通道精化复杂全局几何,从而实现稳定、高效的多层联合蒸馏。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["多视角观测图像 + 任务语言指令"] --> VLA["VLA 骨干网络 (Pre-LN 残差流)"]
    In --> Teacher["冻结 3D 视觉基座 (如 VGGT 教师)"]

    VLA -->|"抽取多层视觉隐藏状态 {h_s^S}"| AlignStep["跨层共享残差流投影器<br/>统一参数映射消除梯度冲突"]
    Teacher -->|"抽取对应层几何特征 {h_tau^T}"| LossCalc["对齐损失计算"]

    AlignStep -->|"按层深动态屏蔽通道"| Matryoshka["套娃式分层稀疏激活<br/>浅层小通道 / 深层全通道"]
    Matryoshka --> LossCalc

    VLA --> ActHead["动作预测头"]
    ActHead --> ActLoss["动作模仿损失 L_action"]
    LossCalc --> AlignLoss["多层对齐正则项 L_align"]

    ActLoss & AlignLoss --> JointOpt["联合优化与反向传播<br/>前向注入3D先验 / 推理零额外开销"]

关键设计

1. 跨层共享残差流投影器:从动力学视角消除层间梯度冲突 在 Pre-LN Transformer 结构中,残差更新幅度较小,浅层表征接收到的回传梯度近似为所有后续对齐层局部蒸馏梯度的线性叠加:\(\nabla_{h_i}\mathcal{L}_{\mathrm{align}} \approx \sum_{l \geq i} \alpha_l J_{p_{\phi_l}}(h_l)^\top g_l\)。若为每个对齐层配置独立的投影器 \(p_{\phi_l}\),不同投影器的参数更新轨迹会迅速解耦并在参数空间趋近正交,导致雅可比交叉项 \(J_{p_{\phi_a}}(h_a) J_{p_{\phi_b}}(h_b)^\top\) 既不对称也非半正定,使得各层回传的梯度向量在浅层产生破坏性干涉甚至符号翻转,直接破坏下游动作学习。ROCKET 强制所有对齐层复用同一个投影器参数 \(\Phi=\{W_1, W_2\}\),从而将交叉项矩阵分解为一个公共的半正定基准算子 \(M = JJ^\top \succeq 0\) 加上受残差小增量约束的微小扰动项。在理论与实证上,这种共享映射保证了教师端正相关的误差信号能够映射为学生端同向协同的梯度更新(余弦相似度显著提升为正),从数学本质上消除了梯度内耗,实现稳定加速收敛。

2. 套娃式分层稀疏激活:平衡浅层主导与深层对齐容量 通过中心核对齐(Centered Kernel Alignment, CKA)分析可知,3D 教师模型与 2D VLA 在浅层的表征相似度极高,而在深层表征差异显著扩大。浅层的对齐任务相对简单且易于快速收敛,若使用完全相同的网络容量,多层梯度的简单累加会导致浅层监督信号主导共享投影器的参数更新,抑制深层复杂全局几何特征的学习。为此,ROCKET 借鉴 Matryoshka 表示学习的思想,将共享投影器设计为两层 MLP 结构并施加前缀掩码: $\(p_{\Phi}^{(i)}(x) = W_2 \left( g_i \odot \sigma(W_1 x) \right), \quad g_i[j] = \mathbb{I}(j \leq m_i)\)$ 其中隐藏层最大宽度为 \(m\),在第 \(i\) 个对齐层激活的有效神经元数量 \(m_i = \lceil \rho_i m \rceil\) 随层深单调线性递增(激活比例 \(\rho_i\) 从最浅层的 \(0.2\) 线性增长至最深层的 \(1.0\))。这种嵌套参数子集结构使得浅层仅占用投影器前 20% 的参数容量捕捉低阶空间对齐线索,而将更丰富的表征容量留给深层去捕捉长程交互与高阶 3D 几何,有效防止浅层对齐破坏全局几何蒸馏。

3. 无需训练的层对齐选择准则:全深度覆盖与稳健迁移 针对以往单层对齐需要耗费高昂算力遍历网格搜索最佳层的问题,ROCKET 探究了多层对齐的层选择边界。由于表征沿残差流逐渐演变,均匀分布并覆盖深度的策略能够自然捕获互补的局部与全局几何结构。ROCKET 默认采用 Even-to-Max(E2M-Last1)采样策略,在浅层至中层区域均匀选择采样点,并额外纳入最深层输出特征(例如对于 OpenVLA 与 24 层的 VGGT,分别从前半段均分选取 9 层并加上末层,组成 10 个层对)。实验表明,该策略在完全无需预先微调和打分搜索的前提下,即可取得与复杂的层级信息变换率选择准则相媲美甚至更优的控制性能,大幅降低了工程落地门槛。

损失函数 / 训练策略

ROCKET 在微调阶段固定 3D 教师基座模型(如 VGGT)的全部权重,仅优化学生端 VLA 骨干(采用 LoRA 或全量微调)及共享投影器参数。训练总目标由动作预测损失与多层对齐正则项加权组合而成: $\(\mathcal{L}_{\mathrm{ROCKET}} = \mathcal{L}_{\mathrm{action}} + \lambda \frac{1}{N} \sum_{i=1}^{N} \ell\left( p_{\Phi}^{(i)}(h_{s_i}^S),\, h_{\tau_i}^T \right)\)$ 其中 \(\mathcal{L}_{\mathrm{action}}\) 为演示轨迹上的动作模仿损失(如交叉熵或回归损失),\(\ell(a, b) = 1 - \cos(a, b)\) 为余弦距离损失,\(N=10\) 为对齐层对数量,平衡系数 \(\lambda\) 默认设为 \(0.5\)(在单一任务快速收敛场景可调为 \(0.125\))。在推理执行时,整个 3D 教师与投影器均被移除,VLA 策略仅基于自身视觉骨干直接输出动作,零推理计算开销。

实验关键数据

主实验

在代表性语言引导机器人操作基准 LIBERO 上,ROCKET 以 OpenVLA-7B 为基础骨干,与代表性 2D VLA、显式 3D 输入 VLA 以及单层表征对齐方法进行了全面评测(指标为任务成功率 SR %):

方法 类型 / 特点 Spatial Object Goal Long 平均成功率 (%)
OpenVLA 基础 2D VLA (LoRA) 84.7 88.4 79.2 53.7 76.5
PI0 基础 2D 流匹配策略 96.8 98.8 95.8 85.2 94.2
OpenVLA-OFT 连续动作回归优化 97.6 98.4 97.9 94.5 97.1
GeoVLA 显式 3D 输入增强 98.4 99.0 96.6 96.6 97.7
3D-CAVLA 显式点云 / 深度上下文 98.2 99.8 98.2 96.1 98.1
Spatial Forcing 单层对齐 (VGGT-24/24) 99.4 99.6 98.8 96.0 98.5
ROCKET (本文) 多层残差共享对齐 (Ours) 98.2 99.8 98.8 97.0 98.5

消融实验

在 LIBERO 基准上针对 ROCKET 核心机制的逐级消融对比(基线为未经对齐的微调策略):

配置 Spatial Object Goal Long 平均成功率 (%) 说明
Baseline 96.9 98.7 97.6 94.8 96.4 未注入 3D 对齐的基础微调
+ 朴素多层对齐 (独立投影器) 93.6 99.2 42.2 85.0 80.0 投影器解耦导致梯度冲突,性能暴跌 16.4%
+ 跨层共享投影器 99.0 99.8 97.0 96.8 98.2 统一映射消除梯度抵消,大幅反超基线
+ 套娃式稀疏激活 (ROCKET 全量) 98.2 99.8 98.8 97.0 98.5 平衡深浅层学习,取得最佳多任务操作表现

此外,在训练成本与收敛速度分析中: - 极佳计算效率:达到 98.5% 顶尖成功率时,先前 SOTA(如 Spatial Forcing)需要 4 个模型副本联合训练 15 万步(相对计算代价 \(24.0\times\)),而 ROCKET 仅需单副本 5 万步训练(相对代价 \(1.0\times\)),相当于仅消耗了前人约 4.1%(\(\sim 4\%\))的算力开销。 - 低资源高泛化:在仅使用 1%、5%、10% 的极小演示数据子集训练时,ROCKET 分别取得 73.6%、83.4%、84.3% 的平均成功率,展现出极强的数据利用效率;在 LIBERO-Plus 的七类扰动测试中,面对机器人外形变换与空间布局重组,ROCKET 取得了 81.7% 的最佳鲁棒成功率。

关键发现

  • 朴素独立投影器的多层蒸馏会导致严重退化:由于各层投影器参数余弦相似度趋近于 0,在 Goal 等对几何目标高度敏感的子任务上成功率从 97.6% 骤降至 42.2%,证实了多头独立映射引发的严重梯度内耗。
  • 共享投影器从根本上扭转了梯度冲突:早期与中后期的梯度余弦相似度可视化显示,共享投影器使得各层对齐梯度保持高度一致的正相关,促使不同深度的几何线索同向叠加。
  • 对齐层数存在饱和区间:对齐 8 到 16 对层表现最佳(8 层 97.4%,10 层 98.5%),当层数过度膨胀至 24 层时,由于残差微小增量假设累积误差增大,性能出现下滑(94.1%)。

亮点与洞察

  • 将表征对齐建模为残差流动力学映射:突破了传统蒸馏盲目堆叠独立投影器的惯性思维,从 Pre-LN 残差反向传播的线性叠加理论证明了独立投影器引发梯度干涉的必然性,设计极度优雅。
  • 套娃式通道渐进激活解决特征层深失衡:利用浅层易对齐、深层差异大的实际表征特性,通过简单的前缀通道掩码赋予深层更大对齐容量,用几乎零参数代价解决了特征尺度不均的优化难题。
  • 极其实用的推理零开销与算力收敛加速:不改变原有 VLA 部署结构,免除了复杂昂贵的外接传感器与深度估计器,将达到 SOTA 的训练开销压缩一个数量级,非常适合具身智能真机微调场景。

局限与展望

  • 依赖预训练 3D 视觉教师的表征质量:对齐上限受制于教师模型(如 VGGT)对特殊材质、高动态反光或复杂微小工件的 3D 几何建模精度。
  • 残差流微小扰动假设的边界:在极深层网络(如 60 层以上的大尺寸模型)中,累积残差项可能不再满足近恒等映射假设,需要进一步探索更长程的动力学校准机制。
  • 真机物理交互与触觉对齐的拓展:目前主要在仿真平台(LIBERO、RoboTwin 2.0)验证,未来可进一步拓展到多模态触觉-力觉与实机长程闭环操作的对齐中。

相关工作与启发

  • vs Spatial Forcing (SF): SF 采用单层对齐范式,必须人工或事后搜索指定将某一特定层(如第 24 层)对齐到 VGGT,无法利用全深度的层次化特征;ROCKET 实现了无痛多层对齐,在同等或更优表现下算力节省约 24 倍。
  • vs 显式 3D VLA (3D-CAVLA / GeoVLA): 后者需要在推理时引入实时深度图或点云,增加了感知传感器依赖与模型推理延迟;ROCKET 在推理阶段完全不依赖 3D 额外输入或辅助分支,保持纯 2D 原生推理效率。
  • vs 经典特征蒸馏 (ReviewKD / PatientKD): 经典蒸馏在浅层 CNN 或纯分类 NLP 模型中采用分层独立映射头;ROCKET 证明了在具身控制的大规模残差流中,独立映射头会引发严重梯度正交冲突,提出以共享映射统一残差动力学。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从残差动力学与梯度相干性视角揭示多层 VLA 蒸馏冲突机理,提出的共享套娃对齐简洁有效。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 LIBERO、LIBERO-Plus、RoboTwin 2.0,横跨 OpenVLA 与 PI0/PI0.5 架构,兼具理论推导、梯度可视化与严密消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导严谨自洽,实证分析层层递进,动机与架构设计闭环清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为解决 2D VLA 缺乏 3D 空间意识的痛点提供了高性能、极低算力开销的工业级即插即用方案。