Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling
Point-BERT:基于掩码建模的点云自注意力模型预训练
摘要
我们提出了Point-BERT,一个学习自注意力的新范式,将BERT[8]的概念推广到三维点云。受BERT的启发,我们设计了一个掩蔽点建模(MPM)任务来预训练点云自注意力。具体来说,我们首先将点云划分为几个局部的点补丁,并设计了一个带有离散变异自动编码器(dVAE)的点云标记器,以生成包含有意义的局部信息的离散点标记。然后,我们随机屏蔽掉一些输入点云的斑块,并将其送入骨干变形器。预训练的目的是在标记器获得的点标记的监督下,恢复被遮蔽位置的原始点标记。大量的实验表明,所提出的BERT式预训练策略大大改善了标准点云自注意力的性能。装备了我们的预训练策略,我们表明一个纯粹的自注意力架构在ModelNet40上达到了93.8%的准确率,在ScanObjectNN的最难设置上达到了83.1%的准确率,超过了精心设计的点云模型,而手工设计的数量要少得多。我们还证明,Point-BERT学到的表征可以很好地转移到新的任务和领域,我们的模型在很大程度上推动了少样本点云分类任务的最先进水平。
1. 引言
与传统的手工设计的特征提取方法相比,卷积神经网络(CNN)[21]对先验知识的依赖要少得多。Transformer[57]进一步推动了这一趋势,因为它是朝着没有归纳偏见和最小的人为假设迈出的一步,例如CNN中的翻译等价性或定位性。最近,标准Transformer的结构优势和多功能性在语言[3, 8, 19, 29, 42]和图像任务[2, 6, 9, 53, 63, 76]中都得到了证明,通过启用更多的参数、更多的数据[9]和更长的训练计划,减少归纳偏差的能力也得到了证明。尽管Transformer在自然语言处理(NLP)和图像处理中产生了惊人的结果,但它在三维社区中的研究还不充分。现有的基于Transformer的点云模型[11,72]从局部特征聚合[72]和邻居嵌入[11]中带来了某些不可避免的归纳偏见,使它们偏离了标准Transformer的主流。为此,我们的目标是在点云上直接应用标准Transformer,并将其归纳偏差降到最低,作为建立整齐统一的三维表征学习模型的踏脚石。
显然,直接采用Transformer并不能在点云任务上取得令人满意的表现(见图5)。这个令人沮丧的结果部分归因于有限的注释的三维数据,因为没有归纳偏见的纯Transformer需要大量的训练数据。例如,ViT[9]使用ImageNet[21](14M图像)和JFT[50](303M图像)来训练视觉变形器。相比之下,准确的注释点云则相对不足。尽管随着现代扫描设备的普及,三维数据的获取变得越来越容易,但标注点云仍然很耗时,容易出错,甚至在一些极端的真实世界场景中是不可行的。这种困难促使人们开始研究从无标签的三维数据中学习。因此,自我监督的预训练成为一种可行的技术,以释放Transformer在三维点云表示学习中的可扩展性和通用性。
在所有基于Transformer的预训练模型中,BERT[8]在其发布时取得了最先进的性能,在NLP领域树立了一个里程碑。受BERT[8]的启发,我们试图探索BERT式的预训练,用于三维点云的理解。然而,由于缺乏预先存在的词汇,在点云上直接采用BERT是具有挑战性的。相比之下,语言词汇已经有了很好的定义(例如,[8]中的WordPiece),并且是现成的,可用于模型预训练。就点云Transformer而言,对点云没有预先定义的词汇表。一个天真的想法是把每个点当作一个 "词",模仿BERT[8]来预测被掩盖的点的坐标。这样的逐点回归任务随着标记数量的增加而使计算成本呈四倍增长。此外,一句话中的一个词包含基本的上下文语义信息,而点云中的一个点几乎不包含语义。
然而,从整体点云中分割出来的局部斑块包含丰富的几何信息,可以被当作一个组成单元。如果我们建立一个词汇表,其中不同的标记代表了输入单元的不同几何模式,那会怎么样呢?在这一点上,我们可以将点云表示为这种标记的序列。现在,我们可以有利地采用BERT和它的高效实现,几乎是开箱即用。我们假设,弥合这一差距是将成功的Transformer和BERT扩展到三维领域的关键。
在上述分析的推动下,我们提出了Point-BERT,一种用于学习点云Transformer的新方案。我们设想了两个基本组成部分。
- 1) 点标记化。通过基于dVAE[46]的点云重建,设计了一个点云标记器,点云可以根据学习到的词汇转换成离散的点标记;我们期望点标记应该意味着局部的几何模式,而学习到的词汇应该涵盖不同的几何模式,这样的标记序列可以代表任何点云(甚至以前从未见过)。
- 2)遮蔽点建模。一个 "遮蔽点建模"(MPM)任务被用来预训练Transformers,它遮蔽一部分输入点云,并学习重建被遮蔽区域的缺失点标记;我们希望我们的模型能够推理出点云不同斑块之间的几何关系,为点云的理解捕捉有意义的几何特征。
这两种设计在我们的实验中都得到了实现和论证。我们在图2中可视化了合成(ShapeNet[5])和真实世界(ScanObjectNN[55])数据集的重建结果。我们观察到,我们的Point-BERT模型正确地预测了被掩盖的点标记,并通过我们的dVAE解码器推导出多样化、整体性的重建。结果表明,所提出的模型已经学会了三维点云的固有和通用知识,即几何模式或语义。更重要的是,我们的模型是在ShapeNet上训练的,ScanObjectNN上的遮蔽点预测结果反映了它在具有未见物体和领域空白的挑战性场景中的卓越表现。
我们的Point-BERT采用纯转化器架构和BERT式的预训练技术,在ModelNet40上达到了93.8%的准确率,在ScanObjectNN的复杂设置上达到了83.1%的准确率,超过了精心设计的点云模型,而人类的先验因素要少得多。我们还表明,通过Point-BERT学到的表征可以很好地转移到新的任务和领域,我们的模型在很大程度上推进了少样本点云分类任务的最先进水平。我们希望一个跨越图像和点云的整齐统一的Transformer架构可以促进这两个领域的发展,因为它可以对二维和三维视觉信号进行联合建模。
2. 相关工作
我们的点云预训练模型建立在自监督学习、变形器和BERT风格预训练的基础上,以3D点云为应用。在本节中,我们试图对上述领域的最新进展做一个概述。
自我监督学习(SSL)。自监督学习是一种无监督学习,监督信号可以从数据本身产生[15]。自监督学习的核心思想是去掉一个前提任务,如图像中的拼图[35]、着色[22]和光流[33]。最近,一些研究建议使用SSL技术来理解点云[10,14,23,36,45,47,48,51,58,62,65] 。三维借口任务的例子包括方向估计[38]、变形重建[1]、几何结构线索[52]和空间线索[34, 49]。受图像中的拼图的启发[35],[48]提出从随机重新排列的部分重建点云。DepthContrast[71]提出了一个对比性学习框架,从深度扫描中学习表征。最近,OcCo[58]描述了一个编码器-解码器机制来重建被遮挡的点云。与这些研究不同,我们试图在成功的Transformers[57]之后探索一个点云SSL模型。
Transformers。Transformers[57]最近已经成为NLP[3, 8, 19, 29, 42]中的主导框架,因为其突出的优点,包括大规模并行计算、长距离特性和最小的感应偏差。它吸引了各种视觉任务[12, 20],如物体分类[6, 9]、检测[4, 76]、分割[59, 73]、姿势估计[17, 27]、图像生成[37]、视频处理[69, 75]和多模式学习[41]。然而,它在三维点云的应用仍然有限。一些初步的探索已经实现[11,67,72]。例如,[72]应用矢量化的自我关注机制来构建一个用于三维点云学习的点变换器层。作为与[72]平行的工作,[11]使用一个更典型的Transformer架构来学习点云,包括一个位置编码、一个偏移注意模块和一个邻居嵌入。尽管如此,之前基于Transformer的点云模型的努力或多或少都涉及到一些归纳性的偏见,使得它们与标准的Transformer不相符合。在这项工作中,我们试图延续标准变形器的成功,并将其扩展到三维点云学习中,将归纳偏差降到最低。
BERT式的预训练。BERT[8]的主要架构建立在一个多层变形金刚编码器之上,它首先被设计为在一个自我监督的方案中从未标记的文本中预训练双向表示。帮助BERT脱颖而出并取得令人印象深刻的性能的主要成分是掩码语言建模(MLM)的借口,它首先随机掩码,然后恢复输入标记的序列。MLM策略也激发了很多预训练任务[2, 7, 19, 29, 54]。以BEiT[2]为例,它首先将输入图像标记为离散的视觉标记。之后,它随机掩盖一些图像斑块,并将损坏的图像送入转化器主干。该模型被训练来恢复被遮蔽的斑块的视觉标记。最近,MAE[13]提出了一种用于图像表示学习的遮蔽自动编码器策略。它首先掩盖了输入图像的随机斑块,然后鼓励模型重建那些缺失的像素。我们的工作受到BEiT[2]的极大启发,它将图像编码为离散的视觉标记,这样一个Transformer骨干就可以直接应用于这些视觉标记。然而,由于点云的非结构化性质,为点云获取标记更具挑战性,这随后阻碍了BERT在点云上的直接使用。
3. Point-BERT
这项工作的总体目标是将BERT风格的预训练策略扩展到点云Transformer。为了实现这一目标,我们首先学习一个标记器,为每个输入点云获得离散的点标记。模仿BERT[8]中的 "MLM "策略,我们设计了一个 "遮蔽点建模"(MPM)任务,在这些离散点标记的帮助下对Transformer进行预训练。图3说明了我们方法的总体思路。
3.1. 点标记化
在本节中,我们将介绍如何将点云标记化为一组离散的标记。
点嵌入。一个天真的方法将每个点视为一个标记。然而,由于Transformer中自我注意的二次复杂性,这样的逐点重建任务往往会导致难以承受的计算成本。受Vision Transformers[9]中补丁嵌入策略的启发,我们提出了一个简单而高效的实现方法,将每个点云分组为几个局部补丁(子云)。具体来说,给定一个输入点云p∈R N×3,我们首先通过最远点采样(FPS)从整体点云p中采样g个中心点。然后用k-近邻(kNN)算法为每个中心点选择最近的邻居点,将g个局部斑块(子云){pi} g i=1分组。然后,我们通过减去其中心坐标使这些局部斑块成为无偏的,解除局部斑块的结构模式和空间坐标。这些无偏的子云可以被当作NLP中的词或视觉领域中的图像斑块。我们进一步采用mini-PointNet[39],将这些子云投射到点嵌入中。按照NLP和二维视觉任务中Transformer的做法,我们将点云表示为一连串的点嵌入{fi} g i=1,它们可以作为标准Transformer的输入被接收。
点标记器。点标记器将点嵌入{fi} g i=1作为输入并将其转换为离散的点标记。具体来说,标记器Qφ(z|f)将点嵌入{fi} g i=1映射为离散的点标记z = [z1, z2, ...., zg] ∈ V 1,其中V是总长度为N的学习词汇。在这一步,子云{pi} g i=1可以被标记为点标记{zi} g i=1,与有效的局部几何模式有关。在我们的实验中,DGCNN[60]被用作我们的标记化网络。
点云重构。dVAE的解码器Pj(p|z)接收点标记{zi} g i=1作为输入,学习重建相应的子云{pi} g i=1。由于局部几何结构过于复杂,无法用有限的N种情况来表示。我们采用DGCNN[60]来建立与相邻点标记的关系,这可以增强离散点标记对不同局部结构的表示能力。之后,使用FoldingNet[65]来重构子云。总体重建目标可以写成Ez∼Qφ(z|p)[ log Pj(p|z)],重建过程可以看作是最大化对数可能性Pθ(p|p˜)的证据下限(ELB)[43]:
其中p表示原始点云,p˜表示重建后的点云。由于潜伏点标记是离散的,我们不能应用重新参数化梯度来训练dVAE。继[43]之后,我们在dVAE训练中使用Gumbelsoftmax松弛[18]和统一先验。关于dVAE的结构及其在点云重建上的实现,我们参考了补充资料。
3.2. 变形器主干
我们在实验中采用了标准的变形器[57],由多头的自注意层和FFN块组成。对于每个输入的点云,我们首先将其分为g个局部斑块,以
为中心点。然后,这些局部斑块通过mini-PointNet[39]被投射到点嵌入
,它只包括MLP层和全局maxpool操作。我们进一步通过对每个补丁的中心点应用MLP来获得其位置嵌入。形式上,我们将输入嵌入定义为
,它是点嵌入和位置嵌入的组合。然后,我们将输入嵌入送入骨干变换器。按照[8],我们在输入序列上附加一个类标记。因此,Transformer的输入序列可以表示为
。在我们的模型中,Transformer块有L层,最后一层的输出代表全局特征,同时还有输入子云的编码表示。
3.3. 掩码点建模
在BERT[8]和BEiT[2]的启发下,我们将屏蔽建模策略扩展到点云学习,并设计了一个屏蔽点建模(MPM)任务来预训练我们的Point-BERT。
屏蔽序列生成。与BERT[8]和MAE[13]中使用的随机遮蔽不同,我们采用了像[2]一样的块状遮蔽策略。具体来说,我们首先选择一个中心点ci和它的子云pi,然后找到它的m个邻居子云,形成一个连续的局部区域。我们把这个区域中的所有局部斑块都屏蔽掉,生成屏蔽的点云。在实践中,我们直接将这种块状屏蔽策略(如[2])应用于变换器的输入。形式上,我们将被遮蔽的位置标记为M∈{1, - - , g} b rgc,其中r是遮蔽率。接下来,我们用相同的可学习的预定义掩码嵌入物E[M]替换所有被掩码的点嵌入物,同时保持其位置嵌入物不变。最后,被破坏的输入嵌入XM = {xi : i /∈M}g i=1 ∪ {E[M] + posi : i∈M}g i=1被送入Transformer编码器。
前言任务定义。我们的MPM任务的目标是使模型能够根据剩余的部分推断出缺失部分的几何结构。预训练的dVAE(见第3.1节)将每个局部补丁编码为离散的点标记,代表几何模式。因此,我们可以直接应用这些信息标记作为我们的代理监督信号来预训练转化器。
点贴片混合。受CutMix[68,70]技术的启发,我们另外设计了一个整齐的混合标记预测任务,作为一个辅助的前置任务,以增加我们的Point-BERT预训练的难度,称为 "点补丁混合"。由于每个子云的绝对位置信息已经被归一化排除,我们可以通过简单地混合两组子云来创建新的虚拟样本,而不需要在不同的补丁之间采用任何繁琐的对齐技术,如最优传输[70]。在预训练期间,我们还强制虚拟样本预测原始子云产生的相应标记,以执行MPM任务。在我们的实现中,我们生成了与真实样本相同数量的虚拟样本,以使预训练任务更具挑战性,这有助于在有限的数据下提高变形金刚的训练效果,正如在[53]中观察到的那样。
优化目标。MPM任务的目标是恢复与被屏蔽位置相对应的点标记。预训练的目标可以被形式化为:给定被掩盖的输入嵌入XM,使正确的点标记zi的对数可能性最大化。
MPM任务鼓励模型预测点云的屏蔽几何结构。只用MPM任务训练Transformer会导致对点云高层次语义的理解不尽人意,这也是最近在二维领域的工作[74]所指出的。因此,我们采用广泛使用的对比学习方法MoCo[14]作为工具,帮助变形器更好地学习高层语义。通过我们的点片混合技术,对比性损失的优化鼓励模型关注点云的高层语义,使虚拟样本的特征尽可能地接近原始样本的相应特征。设q为混合样本的特征,该样本来自另外两个样本,其特征为k1+和k2+({ki}由动量特征编码器[14]提取)。假设混合比例为r,对比损失可以写成:
其中τ是温度超参数,K是记忆库的大小。将MPM目标和对比性损失结合起来,使我们的Point-BERT能够同时捕获局部几何结构和高级语义模式,这在点云表征学习中至关重要。
小样本学习。我们遵循以前的工作[49],在小样本学习设置下评估我们的模型。一个典型的设置是 "K-way N-shot",首先随机选择K个类,然后为每个类取样(N+20)个对象[49]。该模型在K×N样本(支持集)上进行训练,并在剩余的20K样本(查询集)上进行评估。我们将Point-BERT与OcCo[58]进行比较,后者在这项任务上取得了最先进的性能。在我们的实验中,我们测试了在 "5路10个镜头"、"5路20个镜头"、"10路10个镜头 "和 "10路20个镜头 "下的性能。我们在每种设置下进行了10次独立的实验,并报告了10次运行中的平均性能和标准偏差。我们还在同样的条件下重现了DGCNN-rand和DGCNN-OcCo,以进行公平的比较。如表2所示,Point-BERT在小样本的学习任务中取得了最佳性能。它比基线获得了6.8%、3.0%、6.4%、3.3%的绝对改善,比基于OcCo的方法在四个设置上获得了0.6%、0.4%、1.6%、0.3%的改善。强有力的结果表明,由Point-BERT预训练的模型学习了更多的通用知识,可以在有限的数据下快速转移到新的任务中。