摘要

灵感： NLP和CV中的Masked Autoencoder应用很成功
问题： 点云中存在局部信息缺失和密度不均匀的现象
方法： 使用Masked Autoencoder进行点云自监督学习
细节：
①将输入的点云划分成不规则的点云块，以较高的比率 mask 这些块
②使用基于autoencoder的标准transformer从这些未被mask的块中学习高维特征，从而重建被mask的点云块。其中该transformer具有对称设计和shifting mask tokens操作
代码： https://github.com/Pang-Yatian/Point-MAE

1.引言

NLP ——> BERT
CV ——> MAE
Point cloud ——> Point-MAE
考虑到点云的数据集相对较小，使用masked autodecoder作为自监督学习方法可以解决transformer需要训练数据量大的问题。
目前方法所存在的几个问题：
①缺少一个通用的Transformer架构
②用于mask tokens的位置embedding导致缺失局部信息
③点云具有密度不均匀的特性

Point-MAE主要包括了一个point cloud masking、embedding module和an autoencoder ，主要处理流程为：

输入点云被划分为不规则的点云块，这些点云块以较高的比例被mask
Autoencoder从unmask点云中学习高维隐式特征，进而在坐标空间中重建被mask的点云块
Autoencoder的主干网络主要是以Transformer blocks为基础，并采用对称式的encoder-decoder结构
encoder只处理unmasked点云块
轻量decoder的输入包括encoded tokens 和 mask tokens
shifting mask tokens对于轻量decoder而言很重要，既可以节省运算量，又可以避免位置信息泄露，还能够提高精确度

3. Point-MAE

3.1 Point Cloud Masking and Embedding

Point Patches Generation

通过Farthest Point Sampling(FPS) 和 K-Nearest Neighborhood (KNN) 算法将点云划分为不规则的点云块。给定一组带有 $p$ 个点的点云 $X^{i} \in \mathbb{R}^{p \times 3}$ ，利用FPS采样出 $n$ 个点，作为点块的中心点 $C T$ 。基于中心点，对于每个点块 $P$ ，KNN选择 $k$ 个最近邻点：
$\begin{aligned} C T=F P S\left(X^{i}\right), & C T \in \mathbb{R}^{n \times 3} \\ P=K N N\left(X^{i}, C T\right), & P \in \mathbb{R}^{n \times k \times 3} \end{aligned}$
每个点的坐标都通过其中心点进行归一化，便于收敛。

Masking

考虑到点云块可能会重叠，分别对这些块进行mask。将mask比率设置为 $m$ ，masked块被表示为 $P_{g t} \in \mathbb{R}^{m n \times k \times 3}$ ，实验表明， $m=60\%-80\%$ 时，我们工作的性能更好。

Embedding

记完整的mask tokens为 $T_{m} \in \mathbb{R}^{m n \times C}$ ，其中 $C$ 为embedding的维度。

对于unmasked point patches，利用轻量型的PointNet进行embedding。unmasked point patches $P_{v} \in \mathbb{R}^{(1-m) n \times k \times 3}$ 被embedding成 visible tokens的过程可以用下式表示：

$T_{v}=\operatorname{Point} N e t\left(P_{v}\right), \quad T_{v} \in \mathbb{R}^{(1-m) n \times C} .$

考虑到点云块是通过归一化坐标表示的，对于embedding token而言，提供中心点的位置信息是必要的。Position Embedding (PE)的一个简单的方法就是通过一个可学习的MLP将中心点的坐标embedding到另外一个维度上。本文在encoder和decoder都分别使用了PE。

3.2 Autoencoder’s Backbone

Autoencoder的Backbone是完全基于标准的Transformer的，并且是一个对称式的encoder-decoder设计。Autoencoder的最后一层采用了一个简单的预测头来获得重建目标。

Encoder-decoder

本文中的encoder包含着标准的Transformer blocks，仅对visible tokens $T_v$ 进行编码，encoded tokens被记为 $T_e$ 。此外，每个Transformer block中都会加入positional embeddings，提供位置信息。

本文中的decoder与encoder相似，但是包含着更少的Transformer blocks，其输入包括encoded tokens $T_e$ 和masks tokens $T_m$ ，在decoder的每个Transformer block中也加入positional embeddings，为每个tokens提供位置信息。decoder的输出为decoded mask tokens $H_m$ ，并将 $H_m$ 输入到接下来的预测头中，encoder-decoder的结构可以表示为：
$\begin{gathered} T_{e}=\operatorname{Encoder}\left(T_{v}\right), \quad T_{e} \in \mathbb{R}^{(1-m) n \times C} \\ H_{m}=\operatorname{Decoder}\left(\operatorname{concat}\left(T_{e}, T_{m}\right)\right), \quad H_{m} \in \mathbb{R}^{m n \times C} \end{gathered}$

在encoder-decoder结构中，将mask tokens放入轻量decoder中，而不是将encoder的输入放入decoder中，是有两方面的考虑的：

由于我们mask的比率比较大，将mask tokens转移到decoder中能够极大地减少输入tokens的数量
将mask tokens转移到decoder中可以过早地避免泄露位置信息，使encoder能够更好地学习潜在特征。

Prediction Head

简单地使用一层全连接作为预测头，从decoder得到输出 $H_{m}$ ，prediction head会将其投影为一个向量，这个向量的维度和点云块坐标的全部数量相同。然后再跟一个reshape操作：
$P_{p r e}=\text { Reshape }\left(F C\left(H_{m}\right)\right), \quad P_{p r e} \in \mathbb{R}^{m n \times k \times 3}$

3.3 Reconstruction Target

本文的重建目标是恢复每个被mask点云块中点的坐标。给定预测到的点云块 $P_{\text {pre }}$ 和ground truth $P_{g t}$ ，通过 $l_2$ 的Chamfer Distance 计算重建损失：
$L=\frac{1}{\left|P_{p r e}\right|} \sum_{a \in P_{p r e}} \min _{b \in P_{g t}}\|a-b\|_{2}^{2}+\frac{1}{\left|P_{g t}\right|} \sum_{b \in P_{g t}} \min _{a \in P_{p r e}}\|a-b\|_{2}^{2}$

4.实验

在ShapeNet训练集上对模型进行预训练
在不同的下游任务上评估预训练模型，包括object classification, few-shot learning 和 part segmentation
ablation study

输入点的数量 $p = 1024$ ，划分为 $n = 64$ 个点云块，在KNN算法中， $k = 32$ ，在autoencoder中，encoder中包含了12个Transformer block，decoder中包含了4个Transformer block，每个Transformer block有384个隐藏维度和6个头。Transformer block中的MLP比率设置为4。