Stratified Transformer for 3D Point Cloud Segmentation论文阅读笔记

任务概述

通过注意力机制提取点云中点的特征，通过这个特征预测点云中点的类别。 

论文创新点 

1.Transformer结合U-net，融合不同尺度的点云特征。

2.提出一种分层注意力机制，较密集的近程点和较稀疏的远程点被采样在一起以形成关键字，从而实现了显著扩大的有效接受范围，而引起的额外计算可以忽略不计。

3.在Point Embedding中，传统方法是用点云中点的位置编码，这里提出了一种新的方法：Contextual Relative Position Encoding结合上下文关系的位置编码。

Stratified Transformer 分层注意力机制

Vanilla Version 初始版本

window-based self-attention

 将三维空间分成一个个不重合的window，每一个点只从同一个window下的所有点中通过注意力机制提取特征。注意力机制提取的特征过程如下面公式所示：

若每一个窗口下的点的平均数量为k，点云中点的总数量为N，则总窗口数量为N/k，计算复杂度为：

Stratified Key-sampling Strategy.

初始版本存在的问题：感受野太小，仅仅限制在一个窗口内。无法捕获远距离对象上的远程上下文依赖关系，从而导致错误预测。

如何既能捕获远距离对象上的远程上下文依赖关系，又不会增加太多计算量？

 如上图所示，将密集的点云通过farthest point sampling采样，将三维空间用较大的窗口分割，将query（绿色星星）所在的小的dense keys网格中的点和大的sparse keys网格中的点合并，作为这个query要查询的点。

小trick:如何让各个window中的点之间交流特征？作者采用shifted windows策略，如下图所示，相邻的两个Transformer bolcks中三维空间window的划分移动1/2Swin,这样可以使两次同一个window涵盖不同的点。

 Contextual Relative Position Encoding

传统的位置编码：通过MLP将位置信息映射到高维度。导致不同的key之间的位置编码都是相似的。如下图所示，仅仅使用MLP进行位置编码，因为位置相似，所以编码得到的特征相似，注意力机制得到的注意力都是相似的（如下面中间图，都是黄色的）。

 Contextual Relative Position Encoding方法：

首先将第t个窗口下的kt个点的坐标用矩阵p表示：

 定义同一窗口下点与点之间的关系矩阵r，用来表示同一窗口下第 i 个点和第 j 个点x、y、z坐标差异。

 然后将relative xyz coordinates r映射到位置编码

定义三个可学习的查询表tx,ty,tz,如下图所示

 将r的值域分成L份，如图所示：

 则该窗口中的点i和点j之间的坐标差所对应特征矩阵中的向量下标可由下面公式算得，其中m=1,2,3,对应上面tx,ty,tz三个表。

 通过下标取三个表中提取行向量，加起来，作为点i和点j之间的坐标差异编码成的特征向量用ei,j表示，如下图：

 ei,j组成的特征矩阵如下图所示，其中的一个e1,2用橙色表示出来了，便于理解。

 

 位置编码和query和key的特征向量进行点积，获得位置偏差pos_bias。

 使得注意力不仅要关注query和key的相似度，还要关注query和key的位置偏差即Contextual Relative Position Encoding。

与原始版本对比：

 更加通俗的理解：

传统的注意力机制只在输入点云Point Embedding时加入positional encoding，但是这种细粒度的位置信息可能在深层神经网络中丢失。Contextual Relative Position Encoding做的就是在每次计算注意力的时候加上一个包含queryi和keyj位置偏差的向量。

 实验结果

 论文原文

2203.14508.pdf (arxiv.org)