Masked Autoencoders for Point Cloud Self-supervised Learning
摘要
- 灵感: 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
p
p个点的点云
X
i
∈
R
p
×
3
X^{i} \in \mathbb{R}^{p \times 3}
Xi∈Rp×3,利用FPS采样出
n
n
n个点,作为点块的中心点
C
T
CT
CT。基于中心点,对于每个点块
P
P
P,KNN选择
k
k
k个最近邻点:
C
T
=
F
P
S
(
X
i
)
,
C
T
∈
R
n
×
3
P
=
K
N
N
(
X
i
,
C
T
)
,
P
∈
R
n
×
k
×
3
\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}
CT=FPS(Xi),P=KNN(Xi,CT),CT∈Rn×3P∈Rn×k×3
每个点的坐标都通过其中心点进行归一化,便于收敛。
Masking
考虑到点云块可能会重叠,分别对这些块进行mask。将mask比率设置为 m m m,masked块被表示为 P g t ∈ R m n × k × 3 P_{g t} \in \mathbb{R}^{m n \times k \times 3} Pgt∈Rmn×k×3,实验表明, m = 60 % − 80 % m=60\%-80\% m=60%−80%时,我们工作的性能更好。
Embedding
记完整的mask tokens为 T m ∈ R m n × C T_{m} \in \mathbb{R}^{m n \times C} Tm∈Rmn×C ,其中 C C C为embedding的维度。
对于unmasked point patches,利用轻量型的PointNet进行embedding。unmasked point patches P v ∈ R ( 1 − m ) n × k × 3 P_{v} \in \mathbb{R}^{(1-m) n \times k \times 3} Pv∈R(1−m)n×k×3被embedding成 visible tokens的过程可以用下式表示:
T v = Point N e t ( P v ) , T v ∈ R ( 1 − m ) n × C . T_{v}=\operatorname{Point} N e t\left(P_{v}\right), \quad T_{v} \in \mathbb{R}^{(1-m) n \times C} . Tv=PointNet(Pv),Tv∈R(1−m)n×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 T_v Tv进行编码,encoded tokens被记为 T e T_e Te。此外,每个Transformer block中都会加入positional embeddings,提供位置信息。
本文中的decoder与encoder相似,但是包含着更少的Transformer blocks,其输入包括encoded tokens
T
e
T_e
Te和masks tokens
T
m
T_m
Tm,在decoder的每个Transformer block中也加入positional embeddings,为每个tokens提供位置信息。decoder的输出为decoded mask tokens
H
m
H_m
Hm,并将
H
m
H_m
Hm输入到接下来的预测头中,encoder-decoder的结构可以表示为:
T
e
=
Encoder
(
T
v
)
,
T
e
∈
R
(
1
−
m
)
n
×
C
H
m
=
Decoder
(
concat
(
T
e
,
T
m
)
)
,
H
m
∈
R
m
n
×
C
\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}
Te=Encoder(Tv),Te∈R(1−m)n×CHm=Decoder(concat(Te,Tm)),Hm∈Rmn×C
在encoder-decoder结构中,将mask tokens放入轻量decoder中,而不是将encoder的输入放入decoder中,是有两方面的考虑的:
- 由于我们mask的比率比较大,将mask tokens转移到decoder中能够极大地减少输入tokens的数量
- 将mask tokens转移到decoder中可以过早地避免泄露位置信息,使encoder能够更好地学习潜在特征。
Prediction Head
简单地使用一层全连接作为预测头,从decoder得到输出
H
m
H_{m}
Hm,prediction head会将其投影为一个向量,这个向量的维度和点云块坐标的全部数量相同。然后再跟一个reshape操作:
P
p
r
e
=
Reshape
(
F
C
(
H
m
)
)
,
P
p
r
e
∈
R
m
n
×
k
×
3
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}
Ppre= Reshape (FC(Hm)),Ppre∈Rmn×k×3
3.3 Reconstruction Target
本文的重建目标是恢复每个被mask点云块中点的坐标。给定预测到的点云块
P
pre
P_{\text {pre }}
Ppre 和ground truth
P
g
t
P_{g t}
Pgt,通过
l
2
l_2
l2的Chamfer Distance 计算重建损失:
L
=
1
∣
P
p
r
e
∣
∑
a
∈
P
p
r
e
min
b
∈
P
g
t
∥
a
−
b
∥
2
2
+
1
∣
P
g
t
∣
∑
b
∈
P
g
t
min
a
∈
P
p
r
e
∥
a
−
b
∥
2
2
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}
L=∣Ppre∣1a∈Ppre∑b∈Pgtmin∥a−b∥22+∣Pgt∣1b∈Pgt∑a∈Ppremin∥a−b∥22
4.实验
- 在ShapeNet训练集上对模型进行预训练
- 在不同的下游任务上评估预训练模型,包括object classification, few-shot learning 和 part segmentation
- ablation study
输入点的数量 p = 1024 p=1024 p=1024,划分为 n = 64 n=64 n=64个点云块,在KNN算法中, k = 32 k=32 k=32,在autoencoder中,encoder中包含了12个Transformer block,decoder中包含了4个Transformer block,每个Transformer block有384个隐藏维度和6个头。Transformer block中的MLP比率设置为4。
4.1 Pre-training Setup
4.2 Downstream Tasks
Object Classification on Real-World Dataset
Object Classification on clean objects dataset
Few-shot Learning
Part Segmentation
4.3 Ablation Study
Masking Strategy
- block masking就是以块为单位进行mask
- random masking就是随机mask点