AI Research 技术调研

SARA-RT: Scaling up Robotics Transformers with Self-Adaptive Robust Attention

17分钟阅读

一句话定位

SARA-RT（Self-Adaptive Robust Attention for Robotics Transformers）提出一种叫 up-training 的微调方法：把已经预训练/微调好的 Robotics Transformer 策略（包括十亿参数级的 RT-2 VLA）中二次复杂度的常规 softmax 注意力，替换成可学习的线性注意力（SARA 映射），再继续训练至质量恢复；在 5B 参数的 RT-2 上做到 TPU 前向延迟从 53.2ms 降到 45.7ms（14% 提速）且精度不降，叠加历史帧与新动作表征后精度反而提升（65.8%→76.4% 均值，约 10.6 个百分点），在 Point Cloud Transformer 抓取策略上抓取奖励从 0.64 提升到 0.75（+11pp）且点云编码器提速 2 倍以上；该工作获得 ICRA 2024 Best Robotic Manipulation Paper Award。

背景与定位

论文 2023-12-04 提交 arXiv（2312.01990），作者 Isabel Leal、Krzysztof Choromanski、Deepali Jain、Avinava Dubey、Jake Varley（并列一作）等 14 人，全部来自 Google（DeepMind 机器人团队）。2024-01-04 与 AutoRT、rt-trajectory 一起作为 DeepMind”下一代机器人三块基石”官方博客发布。

它要解决的问题是 Robotics Transformer 家族——从 35M 参数的 RT-1（open-x-embodiment 数据训练）到十亿参数级的 rt-2（PaLI-X VLM backbone）——共享的一个实际部署瓶颈：Transformer 注意力模块的二次时空复杂度，意味着传感器分辨率翻倍（或历史帧数增加）会让计算量翻两番，而 RT-1 即便体量不大，实测控制频率也不过 3Hz。作者没有走”设计全新高效架构”或”从零用线性注意力预训练”的路线，而是提出一种微调层面的转换手术：对已经训练好的 Transformer，把注意力换成线性变体后继续做少量步数的微调（作者称之为 up-training），观察到模型质量能快速自适应恢复,由此把”是否需要昂贵的线性注意力预训练”这个前提假设去掉。

方法论上直接建立在 FAVOR+ 的核化注意力（kernelized attention）理论之上——SARA 可以理解为”让 Performer 式随机特征映射中的随机投影矩阵变成可学习参数”这一思路在机器人 Transformer 上的具体落地和实证。论文同时给出配套的数学分析（无偏估计定理、误差界定理），解释为何一个只有对数级可训练参数量的线性化模块就能几乎无损地逼近原始注意力矩阵。

模型架构

核心机制——SARA 映射（up-training 的可学习线性化）：常规线性注意力用随机特征 $\phi_f(z)=(f(z_1),…,f(z_{d}))^\top$（$f$=ReLU 或 exp）来把核函数 $K(x,y)=\mathbb{E}[\phi(x)^\top\phi(y)]$ 双线性化，从而把注意力从 $O(MN)$ 降到 $O(M+N)$；但纯随机版本要 $m$（特征维度）达到 4K+ 才在实践中有优势，且相较 softmax 有精度损失。SARA 把预处理矩阵从随机高斯矩阵 $G$ 换成可学习参数：

$$\phi^{SARA}{f,1}(z)=v\odot f(G_Q z),\quad \phi^{SARA}{f,2}(z)=v\odot f(G_K z)$$

其中 $v\in\mathbb{R}^m$、$G_Q,G_K\in\mathbb{R}^{m\times d}$ 可训练，作用在 Transformer 注意力模块的原始 embedding（而非已经过 $W_Q,W_K$ 投影的 query/key）上。关键性质：可以取 $m=d_{QK}$（不需要 4K+），因而对 RT-2 这类 $M,N<200$ 的短序列也实用（论文点云实验里 $M,N>1K$，但 RT-2 提速部分序列长度只有 $L=196$）。训练时把 SARA 预处理等价看作是在线性注意力语境下微调 Transformer 原有的 $W_Q,W_K$ 投影。

数学结果（Sec. III）：(1) Lemma 3.1——若输入定长（如 L2-归一化的 CLIP embedding），随机高斯投影 + exp 核构成的线性特征映射是 softmax 核的无偏估计（差一个常数 $m\cdot\exp(r^2)$）；(2) Lemma 3.2 给出该估计量的集中不等式；(3) Theorem 3.3（主定理）——对给定误差容限 $\delta$，存在 $v,G_1,G_2,f$（构造性证明,取 $f=\exp$）使近似注意力矩阵 $\hat A$ 满足 $|A-\hat A|_\infty\le\delta$，且所需特征维度 $m$ 只需 $O(\log(MN))$ 量级——即可训练参数量相对注意力矩阵大小是对数级的。

两个落地架构：

  1. RT-2 / PaLI-X ViT 编码器：backbone 沿用 rt-2 的 PaLI-X VLM，本文选用其中更适合端上部署的 5B 变体（而非 55B）。SARA 只替换 ViT 编码器塔中的自注意力（$f=$ReLU，$v$ 取全 1 向量），不改 fuser 中的跨帧/跨模态自注意力（作者明确留作未来工作）。动作表示两种：(a) 传统方式——6-DoF 位姿+旋转位移、夹爪伸展量、终止指令，每个连续维度量化进 256 个 bin 再 tokenize（action tokens）；(b) 本文新提出的向量表示——每个连续维度四舍五入到小数点后 4 位、当作字符串走普通文本 tokenizer,发现能带来更高质量的模型。
  2. Point Cloud Transformer (PCT)：编码器为 2 层 Transformer、16 维隐层 embedding，输入是 Realsense 相机采集并经聚类分割出的单物体点云（$N\times 3$），SARA 用 ReLU 核、$v$ 全 1、$m=d$（点云 embedding 维度本身)。

辅助案例研究（建立直觉，非核心 VLA 结果）：论文先用一个零样本 VR 导航实验（CLIP ViT-B, $d_{QK}=512$，Matterport/HM3D 环境）展示 SARA 思路——常规 ReLU/exp 随机特征线性注意力导致导航失败或分心，Gaussian 预处理($m=2048$)可用但注意力分布偏平,只有 SARA（$m=d=512$，$G$ 训练成模仿 softmax）才能复现与常规 softmax 完全一致（视觉上不可区分）的尖峰注意力模式与轨迹。

数据

RT-2 侧：预训练的 PaLI-X VLM 在 rt-2 论文所述的机器人数据（源自 RT-1 数据集，[20]）上微调,该数据集由移动操作机器人上专家遥操作采集,每条演示映射到以下 7 类自然语言指令之一：pick object、knock object over、open/close drawer、place object into receptacle、place object upright、move object near object、pick object from receptacle and place on counter（具体规模数字见 RT-1/RT-2 原文，本文未重复给出）。up-training（即引入 SARA 后的继续微调）本身用的数据规模、步数论文未披露。

PCT 侧：仿真训练,agent 训练时只见过 k=5 个物体（可乐罐、水瓶、黑板擦、香蕉、章鱼玩具）；测试集包含若干训练时未见过的形状。真实机器人 AB 测试用 N=200 个随机物体摆放配置（不同物体数量与形状组合）。未观察到明显的 sim-to-real gap,策略未做额外的机器人上微调即可部署。

训练方法

Up-training 流程（本文核心贡献）：取一个已完成常规预训练/微调的 Transformer 策略,把其中的 softmax 注意力模块替换为 SARA 线性注意力,继续用下游机器人任务数据微调（而非重新预训练）。论文在 PCT 实验（Fig. 4）中展示,ReLU/exp/sqrt 三种 $f$ 变体的 SARA 都能在替换后几乎立即（几步内）适应,精度快速逼近甚至超过原始 PCT。基于此,RT-2 与 PCT 的主实验都统一选用最简单的 ReLU（论文称其为 exp 变体的”驯化版”）。

PCT 策略训练目标：不是标准监督学习,而是黑盒优化（blackbox optimization）强化学习,采用 BGS（Blackbox Gradient Sampling）变体,超参数：扰动方向数 $l=50$,高斯平滑参数 $\sigma=0.02$,步长 $\eta=0.02$,取 top $\tau=30%$ 方向;常规 PCT 先在仿真中训练至收敛（$s=1000$ 迭代）,再做 SARA up-training。真实抓取用 Kuka IIWA 机械臂 + Weiss 夹爪,开环执行。

RT-2 训练：预训练 VLM 在机器人数据上做标准的行为克隆式微调（沿用 rt-2 的训练框架）,up-training 阶段的具体 batch size、学习率、优化器、步数论文未披露。

动作 tokenization：见”模型架构”节的两种表示（256-bin 离散 token vs. 4 位小数向量文本表示）;本文发现向量表示效果更好,但两者都可与 SARA 组合。

超参数披露程度：BGS 优化超参数（$l,\sigma,\eta,\tau$）与 PCT 训练迭代数 $s=1000$ 已披露;RT-2/PaLI 微调与 up-training 的 batch size、学习率、优化器、训练步数均未披露。

Infra（训练 / 推理工程）

训练：论文全文未披露 up-training 所用 GPU/TPU 型号、数量、总训练时长或分布式并行策略（仅点云侧的 BGS 强化学习超参数被披露,见上）。

推理：

  • RT-2 / SARA-RT-2（PaLI-5B ViT 编码器,$L=196$ vision tokens,单帧无历史）：TPU 前向延迟,常规注意力 53.2ms → SARA 线性注意力 45.7ms（提速 14%）。
  • 叠加 3 帧历史 + 新动作 tokenizer 的第三配置计算量更大,论文明确说明”更难直接上机部署”,未给出该配置的具体延迟数字。
  • Point Cloud 编码器（2 层、16 维,PCT vs SARA-PCT）速度测试（Fig. 6,10 个随机种子均值）：SARA-PCT 提供近似恒定的 ~100ms 推理延迟（与点云大小基本无关）,常规 PCT 的注意力模块随点云增大成为明显瓶颈,在点云规模 >2K 时 SARA 侧提速 2 倍以上（该测试的具体硬件型号论文未注明）。
  • PaLI-ViT 编码器在不同图像分辨率（200–800,默认 16×16 patch）下的 CPU 速度测试（Fig. 8）：常规注意力延迟随分辨率快速上升、在高分辨率下不可行,SARA 保持可行,但论文未给出具体延迟数值表。
  • 边缘硬件（机器人端上实际部署用的计算平台型号）未披露。

评测 benchmark

RT-2 / SARA-RT-2 真机评测（Table I,7 类任务,固定摄像头输入,单帧或 3 帧历史）：

任务RT-2（无历史+动作token）SARA-RT-2（无历史+动作token）SARA-RT-2（3帧历史+向量表示）试验次数
pick81%83%100%36
knock86%91%91%35
open/close drawer67%78%89%18
drawer place39%31%56%36
upright57%46%51%35
move98%79%81%48
diverse pick（泛化测试）33%48%67%21
均值65.8%65.1%76.4%

前两列（同为单帧、同为动作 token,仅注意力模块不同）均值几乎相同（65.8% vs 65.1%）,验证”SARA 线性注意力不掉点”；SARA 一侧在 diverse pick 泛化任务上明显更强（48% vs 33%）。第三列叠加历史帧+新动作表示后均值达到 76.4%（较第一列 +10.6pp,论文原文称”12%+”）,但作者明确说明该配置计算量更大、更难直接部署,因此单独作为”精度上限”参考,不与 14% 提速的结论混为一谈。

Point Cloud Transformer 真机 AB 测试：N=200 个随机物体配置,每次随机分配给常规 PCT 或 SARA-PCT 策略执行抓取,二元奖励（成功/失败）。常规 PCT 平均奖励 $r_{ave}=0.64$,SARA-PCT 平均奖励 $r_{ave}=0.75$（+11 个百分点）。

推理速度（见上”Infra”节数字,不重复列出）。

论文未包含与其他线性注意力基线（如原始随机特征 Performer、cosFormer 等）在真机任务上的直接精度对比,只在零样本 VR 导航案例研究中做了定性对比。

创新点与影响

创新点：(1) 提出 up-training——把”用线性注意力替换 Transformer 注意力”从”需要昂贵重新预训练”的假设中解放出来,证明只需在已训练好的模型上做少量步数的继续微调即可恢复质量;(2) SARA 映射把随机特征线性注意力中的随机投影矩阵变为可学习参数,配合理论分析（Theorem 3.3）证明所需可训练参数量相对原始注意力矩阵是对数级,使线性注意力对 RT-2 这类中短序列（$L\approx196$）也变得实用（而非只在 $\ge$4K 长序列上有意义）;(3) 首次在十亿参数级 VLA（RT-2/PaLI-5B）与点云 Transformer 两类不同模态的机器人策略上,同时验证了”提速且不掉点甚至涨点”的线性注意力转换,并在真实机器人上（而非仅离线指标）完成 A/B 评测;(4) 方法论上足够通用,作者强调不局限于机器人领域,可推广到任意预训练 Transformer 的线性化改造,且不需要额外代码（可直接调用已开源的线性注意力实现）。该工作获 ICRA 2024 Best Robotic Manipulation Paper Award。

论文自述的局限性：(1) RT-2 实验中,SARA 只替换了 PaLI-X 的 ViT 编码器自注意力,fuser 中融合跨帧/文本 token 的自注意力模块未做线性化改造,作者明确留作未来工作;(2) 叠加历史帧与新动作表示带来的精度提升（76.4%）伴随更高计算量,“对直接机上部署更具挑战性”,论文未解决这一权衡,只呈现为”精度上限”参考;(3) 论文提出计划在未来工作中利用 SARA 对高分辨率图像输入保持可行的特性,探索比当前 RT-2 更高分辨率的输入（暗示当前工作尚未在真机上验证高分辨率场景）。

原始链接

一手源存档（sources/）

关系图谱

反向链接