视觉表征新范式 | ToBo用Bottleneck token压缩场景,分割指标提升12.3%

大模型机器学习数据库

点击下方名片,关注「集智书童」公众号


picture.image

精简阅读版本

本文主要解决了什么问题

动态场景中视觉表征学习的局限性 :传统自监督学习方法(如MAE)主要针对静态图像进行设计,缺乏对时间演化信息的有效建模,难以应用于需要时序理解的任务(如机器人操作、视频标签传播等)。

序列场景理解任务中对关键信息的保守编码需求 :许多任务(如机器人策略学习)要求模型能够从当前场景中提取关键信息,并保留时间线索以预测未来状态。

现有方法在真实世界部署中的泛化性与鲁棒性不足 :如SiamMAE等方法虽然尝试引入时序信息,但在真实机器人任务中表现欠佳,缺乏在物理环境中的验证。

本文的核心创新是什么

提出Token Bottleneck(ToBo)机制 :通过将参考场景压缩为一个紧凑的Bottleneck token,并使用少量目标Patch作为Prompt,引导模型预测后续场景,从而隐式建模时间依赖性。

强调对场景的保守总结与时间线索嵌入 :鼓励模型在压缩阶段保留关键视觉信息,并在扩展阶段利用这些信息进行时间推理,提升序列任务理解能力。

设计基于自注意力的解码器结构 :与以往依赖交叉注意力机制的方法不同,ToBo采用自注意力结构,确保解码过程仅关注输入信息,提升模型对Bottleneck token的依赖性。

结果相较于以前的方法有哪些提升

在模拟环境中显著优于Baseline方法

  • • 在Franka Kitchen、RLBench、CortexBench等多个模拟机器人任务中,ToBo的成功率显著高于MAE、SiamMAE、RSP等方法,例如在Franka Kitchen任务中提升超过20%。

在真实机器人操作任务中表现优异

  • • 在柜子开启、抽屉关闭、杯子堆叠等任务中,ToBo相比SiamMAE、RSP等方法分别提升了40%、10%和25%,展示了其在现实环境中的鲁棒性和可迁移性。

在视频标签传播任务中表现领先

  • • 在DAVIS、VIP、JHMDB等数据集上,ToBo在目标分割、部件分割和姿态跟踪任务中均优于现有方法,体现了其捕捉时间演化信息的能力。

具备良好的可扩展性

  • • 在不同模型规模(ViT-S/16、ViT-B/16、ViT-L/16)下,ToBo均保持性能领先,说明其适用于不同层级的视觉主干网络。

局限性总结

依赖大规模预训练数据 :尽管作者在Kinetics-400上进行了预训练,但ToBo的性能是否能在更小规模或更少数据下保持仍有待验证。

对Mask比例的敏感性 :极高Mask率虽然增强了对Bottleneck token的依赖,但可能对某些任务或数据集造成信息丢失,影响泛化能力。

缺乏对多模态任务的验证 :目前的实验主要集中在视觉任务和机器人操作,未涉及语言-视觉或多模态任务,限制了其在更广泛任务中的评估。

未与其他高级RRL方法全面比较 :尽管在部分任务上超越了如Theia、MVP、Voltron等模型,但在语言引导任务等场景下仍缺乏直接对比。

深入阅读版本

导读

从动态场景中提取紧凑且具有时间感知能力的视觉表征对于成功执行视觉追踪和机器人操作等序列场景理解任务至关重要。在本文中,作者引入了Token Bottleneck(ToBo),这是一种简单直观的自监督学习流程,它将场景压缩成一个 Bottleneck token,并使用少量 Patch 作为线索来预测后续场景。ToBo流程通过在挤压步骤中保守地将参考场景编码成一个紧凑的 Bottleneck token来促进序列场景表征的学习。在扩展步骤中,作者引导模型通过使用 Bottleneck token和少量目标 Patch 作为线索来预测目标场景,从而捕捉时间动态。这种设计鼓励视觉主干网络嵌入时间依赖性,从而实现跨场景动态转换的理解。在包括视频标签传播和模拟环境中机器人操作在内的多种序列任务中的大量实验表明,ToBo优于 Baseline 模型。此外,将作者预训练的模型部署到物理机器人上证实了其在真实环境中的鲁棒性和有效性。作者进一步验证了ToBo在不同模型尺度上的可扩展性。

1 引言

随着机器在现实世界环境中部署的兴趣日益增加,确保机器与其周围环境实现无缝感知和交互已成为一项关键挑战。这些操作本质上具有序列性,需要具备根据当前和最近的过去观测来追踪物体(例如视觉追踪)以及预测未来行为(例如操作)的能力。对周围环境的这种理解主要依赖于视觉 Backbone 网络。因此,一个强大且鲁棒的 Backbone 网络,能够跨不同任务和环境进行泛化,对于有效的序列场景理解至关重要。

自监督学习(SSL)在视觉领域被视为关键研究,预训练模型被广泛用于高效 Backbone 网络的部署。一系列研究提出了无需标签数据学习图像[17, 6, 9, 3, 4, 7, 21]和视频[46, 38]表征的有效方法。然而,这些研究主要关注对整个场景或视频的理解,这限制了序列场景理解,因为序列场景理解需要捕捉连续场景间的时序变化,并保守地编码观察到的场景视觉状态。

为应对这些挑战,一系列研究[19, 13, 25]尝试将对应学习融入MAE[21]框架,旨在保留其强大的定位能力,同时使模型能够匹配连续场景中的对应区域。然而,作者观察到这些额外的考虑对场景表征质量的影响有限,并可能导致在序列场景理解任务(如机器人操作§3.2)中的性能欠佳。这一局限性源于仅识别时序变化是不够的;

这些任务要求能够从每个场景中总结出关键信息而不丢失,同时在总结的表示中保留时间线索。

本文介绍了一种名为Token Bottleneck(ToBo)的简单而有效的SSL方法,该方法直观地促进了对观察场景的保守性总结,同时能够在总结的表示中有效识别时间演变。如图1a所示。ToBo将参考场景压缩成一个 Bottleneck token,然后仅使用一组最小的 Patch 作为 Prompt 来预测后续的目标场景。这种设计强化了对 Bottleneck token的依赖,鼓励视觉主干网络捕获场景的关键信息。此外,从 Bottleneck token预测目标场景隐式地嵌入了时间依赖性,指导视觉主干网络生成能够捕捉连续场景之间动态转换的表示。

picture.image

作者进行了全面的实验,以评估作者的预训练流程与现有自监督学习方法的有效性。作者在多种序列理解任务上评估了Token-Bottleneck,包括模拟环境中的操作任务和视频标签传播任务,显著超越了 Baseline [6, 9, 4, 21, 51, 19, 25, 13](见图1b)。此外,作者将预训练模型部署在真实机器人上,展示了在未见过物理环境中的强大泛化性能。最后,作者通过观察在不同模型规模下的一致性能提升,验证了Token-Bottleneck的可扩展性。

2 相关工作

静态场景上的自监督学习。自监督学习(SSL)方法在图像领域得到了广泛探索。对比学习方法[6, 20, 8, 9, 3]旨在通过强数据增强从静态场景中推导出的正样本对之间最大化相似性来学习有用表示。尽管这些方法在促进图像的连贯理解方面表现出色,但它们在定位能力方面存在局限性[30],这对于机器人动作预测至关重要。另一方面, Mask 图像建模(MIM)[2, 21, 52, 1, 30]最近因其通过预测学习学习视觉表示的潜力而受到关注。受Transformer中 Mask 语言建模(MLM)[11]的启发,BEiT[2]将MLM扩展到视觉领域,采用外部离线分词器。MAE[21]和SimMIM[52]通过直接重建 Mask 输入像素而不使用任何分词器展示了高效的MIM。然而,这些方法在预训练过程中没有包含捕捉时间进展的机制。

动态场景的自监督学习。近期研究集中于提升动态转换的识别能力。SiamMAE [19] 提出了利用动态场景的视觉表征学习方法。CropMAE [13] 引入一种简单的增强策略,即使从单一静态图像也能生成动态场景。另一方面,RSP 采用随机帧预测任务与 Mask 自编码相结合。一些研究还探索将这些技术应用于具身Agent和机器人操作。例如,VC-1 [35]、MVP [42] 和 Dasari 等人 [10] 采用 MAE 目标进行视觉预训练,而 STP [53] 在 SiamMAE 的基础上引入参考 Mask 策略。另一方面,一些早期工作研究了带标注监督的表征学习。Theia [44] 将从大规模预训练教师网络(其中部分网络使用标注监督训练)提炼的表征迁移到学生模型。MPI [27]、Voltron [28] 和

[36] 探索语言驱动的表征学习,利用手动标注数据提供的辅助文本指导。相比之下,作者专注于直接从原始动态场景进行自监督学习,不依赖任何标注指导。

3 方法

3.1 初步

Mask 自编码。给定一张场景图像,作者将图像分割成

个不重叠的

大小的块

,其中

。作者随机选择一个 Mask 块集

,其比例为

,且

。剩余的块

输入编码器

,生成空间表示

,其中

为编码器维度。请注意,一个可学习的 CLS Token

也作为编码过程的一部分与空间表示一同编码。编码后的 Token 通过将 Mask 位置替换为 Mask Token

扩展为

个 Token ,即对于

。解码器

接收

作为输入,并使用编码后的 Token 预测 Mask 图像块

3.2 动机

在本节中,作者从序列场景理解的角度讨论了先前自监督学习(SSL)方法的优缺点,这些讨论促使作者提出了当前方法。

MAE的时间演化感知能力有限。MAE [21]因其强大的定位能力而受到认可,促使多项研究 [35, 10, 42] 采用该方法。这源于其设计要求自动编码器基于可用的先验信息(即可见块)预测缺失信息。该流程隐式地鼓励编码器促进剩余Sparse Token 之间的交互,从而增强其定位能力。然而,由于MAE在单个静态场景上进行预测学习,编码器并未得到明确训练以处理时间上的动态变化,导致其在序列场景理解任务中的表现有限(图2)。此外,近期研究表明,MAE在更广泛上下文学习方面存在不足 [30],导致其生成的表征对观察场景的理解缺乏整体性。这些局限性进一步限制了其有效理解序列场景的潜力。

picture.image

SiamMAE在序列场景理解中的次优影响。为缓解基于静态场景的SSL方法的固有局限性,SiamMAE [19]通过从序列数据中随机采样两个动态场景构建一个非平凡的对应匹配问题。其核心原理涉及将参考场景中的块传播到目标场景中的对应位置。通过基于交叉注意力层的解码器应用这一指导,鼓励目标块与参考块之间实现细粒度的块级相似性。这个过程最终迫使编码器为对应块生成相似的表示。然而,尽管SiamMAE能够捕捉连续场景之间的对应关系,尽管其基于MAE框架,但在某些基于序列场景的任务中(图2),其对MAE的影响微不足道甚至为负。在这些任务中,由于动作是通过基于对观察到的和立即过去的场景估计的视觉状态的政策网络进行预测的,这表明仅考虑时间演化识别不足以实现序列场景理解,对观察到的场景进行保守的总结至关重要。

3.3 所Token-Bottleneck - Token Bottleneck (ToBo)

作者的主张。作者的目标是实现针对解决序列场景任务进行优化的表示。根据

中的讨论,作者将重点从单纯识别时间演变扩展;作者考虑以有效将时间动态嵌入总结表示中的方式对观察到的场景进行保守总结。

为此,作者提出了Token Bottleneck(ToBo),一种通过token Bottleneck 机制实现自监督视觉表征学习流程。ToBo包含两个关键步骤:将场景压缩成一个token,作者称之为 Bottleneck token,以及从该token扩展信息。假设给定一个参考场景和一个目标场景。在压缩步骤中,参考场景的视觉信息被紧凑地编码到 Bottleneck token中。随后,在扩展步骤中,作者引导模型使用 Bottleneck token预测目标场景,仅提供目标场景的一小部分块作为线索。在这种情况下,模型仅凭有限的线索无法精确重建目标场景,从而增强了扩展步骤对 Bottleneck token的依赖性。这种设计带来了两个优势:(1) Bottleneck token应保留参考场景的关键信息,(2) 这种信息应以某种方式编码,以便在与其他目标场景线索交织时能够识别时间动态性。最终,通过优化Token Bottleneck流程的目标函数,作者可以实现作者的目标。作者的流程整体描述如图3所示。

picture.image

整体流程公式化。假设作者采样一个参考场景

和一个目标场景

,两者之间存在时间间隔

。作者将

切分为

个不重叠的块

。参考场景块

被输入编码器

,生成空间表示

。作者使用该编码过程输出的 CLS Token 作为 Bottleneck Token

,该 Token 将指导对参考场景进行紧凑总结。目标场景

被以极高比例

进行 Mask 处理,其中

。未被 Mask 的目标块

由相同的编码器

处理,生成目标场景的

。随后,作者将 Bottleneck Token

与目标表示

以及用于缺失区域的填充 Mask Token

进行拼接,并将这些输入解码器

。解码器

利用

预测被 Mask 的图像块

。由于目标场景被应用了极高的 Mask 比例,解码器

主动依赖

,这使得编码器

能够以一种便于时间推理的方式对参考场景进行保守总结,与目标 Prompt 相比,这种方式更有利于时间推理。

作者通过以下方式在整个训练过程中最小化重建损失:

其中

是一个距离函数;作者使用余弦距离进行预训练。

解码器结构。动态SSL[19, 25, 13]中的先前方法将交叉注意力层作为学习时间演变 Aware 的核心组件,并将其置于解码器中以指导编码器学习能够有效捕捉对应关系的表示。这些方法利用交叉注意力层、自注意力层和多层感知器(MLP)层的混合结构。相比之下,ToBo采用自注意力层以确保解码器在扩展步骤中仅关注给定信息,并使用MLP层实现从表示嵌入空间到像素空间的渐进式转换。

4 实验

在本节中,作者着重通过公平的比较来展示作者预训练流程的有效性。为此,作者在序列任务上评估Token-Bottleneck,包括视频标签传播任务[40, 57, 26]以及跨多种模拟环境的基于视觉的机器人操作和运动策略学习[18, 24, 35]。作者通过将预训练模型部署到物理机器人上,扩展到真实世界环境,展示了其可迁移性。作者进一步研究了Token-Bottleneck的可扩展性。在附录中,作者验证了关于极高 Mask 率对目标场景重要性的主张,呈现了与 Baseline 方法的操作过程定性比较,并展示了真实世界操作任务的演示。

4.1 实验设置

实现细节。对于模拟环境中的视频标签传播和基于视觉的策略学习,作者遵循Jang等人[25]的评价协议。为确保与 Baseline 的公平比较,作者还使用Kinetics-400对Token-Bottleneck进行预训练400个epoch。预训练和评价的详细说明在附录中提供。

Baseline 方法。作者将Token-Bottleneck性能与用于视觉表征的传统自监督学习(SSL)方法进行比较,包括SimCLR [6]、MoCo v3 [9]、DINO [4]和MAE [21]。作者还考虑了先前的动态场景SSL方法,即SiamMAE [19]、RSP [25]和CropMAE [13]。作者验证了在这些方法中显式学习状态表征的影响。

4.2 基于视觉的机器人策略学习在模拟环境中

作者通过模仿学习在机器人操作和运动任务中,在各种模拟环境中评估Token-Bottleneck。具体而言,作者从Franka Kitchen和RLBench基准中评估了五个任务。此外,作者从CortexBench基准的Adroit [43]、MetaWorld [55]、DeepMind Control Suite (DMC) [45]和TriFinger [48]中分别考虑了两个、五个、五个和两个任务。

Franka厨房。作者在表1中展示了Token-Bottleneck与 Baseline 方法在Franka厨房环境下的基于视觉的机器人策略学习方面的比较。结果表明,Token-Bottleneck在所有任务中都显著优于所有 Baseline 方法。值得注意的是,除了Light on任务外,Token-Bottleneck在所有任务的成功率上均实现了超过20%的改进。这突显了显式编码视觉状态表示对于基于视觉的机器人策略学习的有效性。

picture.image

CortexBench。作者在Adroit、MetaWorld、DeepMind Control(DMC)和Trifinger环境中,针对基于视觉的机器人操作和运动任务,将Token-Bottleneck与 Baseline 方法进行了比较,如表2所示。结果表明,Token-Bottleneck在所有任务中都取得了优于 Baseline 方法的性能。特别是,Token-Bottleneck在DMC上以11.9%的成功率提升超越了第二优性能,在Adroit上以10.4%的成功率提升超越了第二优性能。

picture.image

RLBench 表3展示了在RLBench环境中五个演示任务上的机器人操作性能。值得注意的是,Token-Bottleneck在五个任务中始终超越所有 Baseline 。

picture.image

此外,MAE和SiamMAE性能的下降进一步凸显了状态表示学习对机器人 Backbone 网络的重要性。

4.3 基于视觉的机器人策略学习在现实环境中

定量比较。为验证Token-Bottleneck在真实环境中的鲁棒性,作者进一步研究了SSL方法在真实机器人操作任务上的表现。具体而言,作者设计了三个演示任务:柜子开启、抽屉关闭和杯子堆叠。对于每个任务,作者收集50个演示片段用于训练和10个演示片段用于评估模仿学习。遵循模拟环境中的训练协议,作者使用标准的行为克隆损失训练策略网络。每个单独任务的实验结果报告在表4中。作者首先观察到,Token-Bottleneck在三个任务上均优于SiamMAE [19]、RSP [25]和CropMAE [13]。具体而言,Token-Bottleneck在柜子开启、抽屉关闭和杯子堆叠任务上分别比 Baseline 方法提高了40%、10%和25%。虽然之前的动态场景SSL方法在需要相对较高精度的任务(如柜子开启任务)上存在困难,但Token-Bottleneck甚至以相当的成功率成功执行了该任务。这表明通过Token-Bottleneck预训练的模型能够稳健地迁移到真实环境中。

picture.image

定性比较。为了展示实际的操作过程,作者在图4中呈现了三个真实世界操作任务的成功演示中的机器人轨迹。具体来说,左侧场景展示了物理机器人的初始状态,而右侧场景展示了演示的最终状态。中间场景则说明了演示的中间状态。Token-Bottleneck在所有任务中都取得了成功。作者还在附录中与 Baseline 进行了轨迹比较。

picture.image

4.4 视频标签传播

作者对视频标签传播任务进行了比较分析。作者考虑了来自DAVIS [40]、VIP [57]和JHMDB [26]的视频目标分割、视频部件分割和姿态跟踪任务。作者遵循Jang等人 [25]中的评估协议。定量评估结果如表5所示。Token-Bottleneck在所有视频标签传播任务中均优于所有 Baseline 方法。作者还提供了定性结果,如图5所示,其中Token-Bottleneck有效地追踪了各种视频标签传播任务中的视觉外观。这些可视化结果突显了Token-Bottleneck在保持鲁棒的目标身份、部件一致性和姿态连续性方面的能力。定量和定性评估中的优异性能进一步证明了Token-Bottleneck在捕捉连续场景中视觉外观时间演化的有效性。

picture.image

picture.image

5 讨论

可扩展性。作者通过在Kinetics-400 [29]上预训练ViTB/16和ViT-L/16各100个epoch,来研究作者的ToBo的可扩展性,超越了ViT-S/16。作者使用三个不同的种子,在Franka Kitchen [18]上的基于视觉的机器人策略学习任务中评估预训练模型。作者将Token-Bottleneck与MAE、SiamMAE和RSP进行比较。表6展示了所有种子的平均值和标准差。作者观察到,使用ToBo预训练的模型在所有五个任务中始终获得最佳性能,相较于第二好的结果有显著提升。这证明了Token-Bottleneck的可扩展性。

picture.image

与机器人表示学习模型的比较

作者进一步将Token-Bottleneck与近期机器人表示学习(RRL)模型进行比较,这些模型根据其监督类型分为三类:自监督学习[10, 35]、基于基础模型输出的监督[44]以及辅助语言标注的监督[36, 42, 28, 27]。表7展示了RRL模型在多个模拟机器人操作基准测试[43, 55, 18]中的性能表现。其中,Token-Bottleneck基于在Kinetics-400上训练400个周期的ViT-Small架构。值得注意的是,尽管参数数量最少、训练数据量第二少且未使用基于标注的监督,Token-Bottleneck在MetaWorld上取得了最高分数。具体而言,Theia通过从五个大规模基础模型(CLIP large[41]、Depth Anything large[54]、DINOv2 large[37]、Segment Anything huge[31]和ViT huge[47],这些模型共同在14.3亿个标注样本上进行训练)中提炼知识进行训练,并在评估时采用基于卷积的压缩层。在这样不平衡的训练和评估设置下超越Theia值得注意。此外,带有和不带有语言指导的R3M之间的性能差距突显了辅助语言监督的显著优势。尽管训练设置存在这种不公平性,Token-Bottleneck在MetaWorld上仍优于R3M、MVP、Voltron和MPI。在Franka Kitchen上,尽管训练数据和模型规模存在显著差异,Token-Bottleneck也超越了R3M。与自监督RRL模型相比,Token-Bottleneck超越了所有模型。尽管训练数据量显著更少,Token-Bottleneck仍超越了VC-1和数据4robotics等更大规模的模型。鉴于参数数量和训练规模的极小,这些结果证明了Token-Bottleneck在机器人操作任务中的有效性和效率。

picture.image

参考

[1]. Token Bottleneck: One Token to Remember Dynamics

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
火山引擎大规模机器学习平台架构设计与应用实践
围绕数据加速、模型分布式训练框架建设、大规模异构集群调度、模型开发过程标准化等AI工程化实践,全面分享如何以开发者的极致体验为核心,进行机器学习平台的设计与实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论