SARA-RT：利用自适应鲁棒注意机制扩展机器人 Transformer

23年112月来自谷歌的论文“SARA-RT: Scaling up Robotics Transformers with Self-Adaptive Robust Attention”。

自适应鲁棒注意机制的机器人 Transformer (SARA-RT)，是一种扩展机器人 Transformer (RT) 在机器人上部署挑战的新范式。SARA-RT 依赖于一种称为向上训练（up-training）的新微调方法。它将预训练或已经微调基于 Transformer 的二次时间复杂度机器人策略（包括海量十亿参数的视觉-语言-动作或 VLA模型）转换为保持高质量的高效线性注意对应法。加速以下模型证明 SARA-RT 的有效性：(a) RT-2 模型，这是在互联网规模数据上预训练的 VLA 机器人策略，以及 (b) 在大型点云上运行的点云 Transformer (PCT) 机器人策略。

Transformers（[3]、[4]、[5]、[6]、[7]、[8]、[9]、[10]）为机器人技术提供了前所未有的语义推理能力，催生了机器学习的新领域，即探索基于 Transformer 的模型，这些模型通常在互联网规模的数据上进行预训练，用于机器人控制器。这些控制器能够从多模态输入中进行抽象推理，并已在机器人技术领域取得多项突破，包括：使用大语言模型 (LLM) 进行高级规划（[11]、[12]、[13]、[14]）、将自然语言命令转换为机器人可执行代码的策略（[15]、[12]）、多模态传感器融合 [16]，最后是第一个由大规模视觉语言模型驱动的视觉-语言-动作机器人操纵 [1]。有趣的是，即使没有经过微调，在海量网络语料库上训练的 Transformer 模型，似乎也能学习结构强化学习先验，这些先验可以通过上下文学习 ([17]、[18]) 和精心设计的提示，进行某些控制任务的轨迹优化，正如最近在 [19] 中展示的那样。

Transformer 驱动的机器人革命，是通往 AGI 的直接途径吗？剩下就是训练更大、更深的模型，并在更大规模的数据集上进行预训练吗？ 上述模型的空间和时间复杂度极其昂贵，这是尚未解决但在机器人技术中具有至关重要的实际意义的挑战之一。例如，35M 大小的参数 RT-1 模型 [20] 与非 Transformer 方法相比，它提供了显着的泛化改进，标志着机器人 Transformer  (RT) 架构类的诞生，尽管规模适中，但已经以最高 3Hz 的频率运行。对于更大的 1B+ 参数模型（例如 RT-2 [1]），问题只会变得更加严重。

如图所示：通过自适应鲁棒注意 (SARA) 获得的机器人Transformer策略，适用于三种不同的模态：视觉、语言和点云以及不同的序列长度（从 ∼ 200 到几千）。左图：来自 RT-2 类 [1] 的 5B 视觉-语言-动作模型（序列长度：L = 196）。操作策略以文本指令为条件。右图：用于抓取的点云Transformer (PCT) ([2])（点云大小 L ∈ [800, 4000]）和示例点云输入。这两种策略都使用通过 SARA 训练的线性注意，并在保持高质量的同时提供计算加速。

考虑一种纯粹的零样本基于注意的控制机制，其中智体对应于特定目标 ti (i = 1, ..., M ) 的动作 ai 取决于潜嵌入（K）、目标图像（或者对应的文本命令）的嵌入（Q）、注意得分、核函数（相似度）和基动作（左、右、前、后）。

如图所示：通过 Matterport 环境中的视觉语言注意模型进行 VR 导航 ([21])。场景的自上而下的视图位于左下角。智体的初始视图与索引图像块位于正上方。对于每个视图，突出显示三个注意得分最高的patches，并展示注意得分在所有patches上的分布。在任何给定的时间点，智体都会以与分数成比例的概率选择其中一个。执行轨迹的初始部分是可视化的。轨迹也包含在自上而下的视图中。智体收到一个文本命令：“走向桌子”（或相应的图像）。ReLU 变型（粉色块）直接导致与墙壁的碰撞。对于 exp 变型，智体可以导航，但会暂时被初始视图右下角的无关物体分散注意。两种采用高斯矩阵预处理的变型（蓝色和棕色块）都应用 m = 2048（带有 ViT-B CLIP patch/文本的嵌入维数为 d = 512），并可实现高效导航。不过，ReLU-Gaussian 的注意分数分布仍然平坦。m = d 的 SARA 变型，其中矩阵 G 经过训练模仿每个新视图的常规 softmax 注意，产生难以区分的（尖峰）注意分数并可实现高效导航（黑色块）。在涉及 Transformers 的实际应用中，SARA 训练在所谓的向上训练（up-training）过程中只进行一次。

现在考虑目标集很大甚至与 N 相当的假设情况。在这种情况下，二次空间和时间复杂度 O(M N ) 可能会变得非常昂贵。如果核函数 K 允许双线性化，则可以通过算法解决这个问题，也就是说，K 可以在新输入空间中重写为线性（点积）核，如 K(x, y) = E[ φ^T(x)φ(y)]。

线性注意机制是一个活跃的研究领域（[22]，[23]，[24]，[25]，[26]，[27]，[28]，[29]，[30]），但大多数理论重点是开发低方差变型，无偏地近似原始softmax核注意（由于与脉冲注意模式的组合性质，其具有特别强的表现力）。它们通过随机高斯投影应用随机φ。这种计算开销使它们仅在M，N足够大（通常为4K +）时才具有实际吸引力。此外，与暴力 softmax 相比，线性注意通常会产生一些性能差距。

然而，有一个非常简单的技巧可以改进 φf (其中 f 选择为 ReLU 或者 exp)，使智体成为更高效的零样本导航器。只需用随机矩阵预处理 φf 的输入即可。更具体地说，将 φf 的随机版本定义为：φrand(z) = f(Gz)，其中 G 是一个高斯矩阵，其entries 从 N (0, 1) 中独立同分布（采样一次并用于所有输入 z），并且 f 逐元素应用。正如在上图（蓝色和棕色边框）中所示，这种修改使 ReLU 和 exp 变型都能毫无干扰地达到目标，而且已经导致 exp 变型的尖峰注意模式。

上图中的随机映射 φf/rand 适用于 m = 2048，因此，它可能与 M、N < 1K 的任务无关。自适应鲁棒注意背后的核心思想是矩阵 G 不是高斯矩阵，而是可学习的。定义 SARA 的映射 φSARA，作用于原始 d-维嵌入：目标/tokens的(xi )，(yj ) （而不是查询/键 (qi)/(kj)），表现为和可学习v、f(Gqz) 和 f(Gkz) 的Hadamard积。

VL 导航是一种轻松的“宏观”案例研究，但主要目标是机器人 Transformer 架构，其中对整个图像或文本指令的查询被 Transformer 编码策略生成的图像patch、文本token甚至点云 (PC) 单点的对应项所取代。

因此，提出其注意模块的自适应过程，称之为向上训练（up-training），其实现方式是将常规的 softmax 注意替换为其 φf 编码变型，并根据来自下游机器人任务的数据对其进行微调。可学习的预处理在这里对应于来自 Transformers 注意模块的微调矩阵 WQ 和 WK ([3])，但是在线性注意环境中。

实验中抓取策略从 Realsense 摄像头接收点云（PC）。一个直通滤波器会移除桌面物体以外的所有点。这些点按层次聚类到物体中。然后将单个物体点云传递到 PCT 策略中，产生偏离物体点云中心的抓取姿势。观察空间有 3 个组成部分：(1) 云：N × 3 点云，工作空间原点位于物体云的平均处；(2) 中心：(x,y,z) 工作空间参考系中的云中心；(3) 长轴：(x, y, z) 表示物体框架中点的长轴。抓取姿势动作由相对于物体中心的指尖位置、工作空间框架中的接近方向矢量、和机械臂腕的滚动角度表示。然后使用 Kuka IIWA 机械臂和 Weiss 夹持器以开环方式执行抓取。