Transformer彻底改变了点云学习任务，但其二次复杂度阻碍了其向长序列的扩展，这给有限的计算资源带来了负担。最近出现的 RWKV 是一种新型的深度序列模型，在 NLP 任务中显示出序列建模的巨大潜力。在这项工作中提出了PointRWKV，这是一种线性复杂度的新模型，具有 3D 点云学习任务所需的适应性。通过对不同点云学习任务的大量实验表明，所提出的 PointRWKV 优于基于 transformer 和 mamba 的同类网络，同时显著节省了约 42% 的 FLOPs，展示了构建基础3D点云表征学习模型的优越性。
论文：
https://arxiv.org/abs/2405.15214
主页：
hithqd.github.io/projects/PointRWKV/
代码：
https://github.com/hithqd/PointRWKV
背景
3D 点云分析是众多现实应用的基础，包括自动驾驶、虚拟现实和机器人技术等。与 2D 图像不同，点云的内在不规则性和稀疏性使得进行准确的点云特征学习成为一项具有挑战性的任务。并且同时平衡准确性和复杂性仍然是一个持久的问题。现有的点云特征学习方法主要是基于自注意力结构（Transformer）或者是线性时间序列结构（Mamba）的，如下图所示。然而，基于自注意力结构的算法对扩展点标记进行全面注意力机制的部署会显著增加对计算资源的需求，这种效应直接归因于注意力计算中固有的二次复杂性，影响了计算和内存。而基于线性时间序列结构的算法尽管有效，但原始的单向建模的固有属性阻碍了它们达到卓越的性能。本文提出了一种基于RWKV的算法，该算法可以在极小的线性复杂度和参数量上达到较高的效率，并且能够处理多尺度的点云输入。
方法
PointRWKV 的整体流程如上图所示，其中通过分层网络架构对点云进行编码。给定一个输入点云，首先采用多尺度掩蔽策略在不同尺度上对不同点数进行采样。然后应用轻量级 PointNet来嵌入点并生成embedding嵌入。这些点标记由块堆叠编码器（即 PRWKV 块）使用，其中每个块由两个并行分支组成，用于分层局部和全局特征聚合。每个PRWKV块，采用两个并行分支的处理策略来聚合局部和全局特征。上面的是综合特征调制流程，具有空间混合和通道混合，下面的是基于局部图的合并。最后，两个分支的连接用作每个块的输出。
Integrative Feature Modulation (IFM)
综合特征调制分支由空间混合模块和通道混合模块组成。空间混合模块作为一种注意力机制，执行线性复杂度的全局注意力计算，而通道混合模块则作为前馈网络（FFN）运行，促进沿通道维度的特征融合。
空间混合模块：经过一个前置的LayerNorm 之后，输入特征的token 首先通过双向二次展开 (BQE) 函数进行移位，然后输入到四个并行的线性层中，以获得多头向量：
其中，BQE的计算为：
BQE函数使注意力机制能够在不同通道上自然地关注相邻的token，而无需显著增加FLOPs。这一过程还扩展了每个token的感受野，从而显著提升了标记在后续层中的覆盖范围。此外，通过以下公式计算出一个新的时变衰减w：
然后，将K_S和V_S传递以使用新的衰减参数w计算全局注意力结果wkv。在这里，我们引入了具有线性复杂度的双向注意力机制，并进行了两项修改：（1）衰减参数独立变化，以动态方式依赖于数据，（2）在求和公式中，将原始RWKV注意力的上限从当前标记t扩展到最后一个标记T-1，以确保在每个结果的计算中所有标记都是相互可见的。对于第t个标记，注意力结果通过以下公式计算：
最终的概率输出为：
通道混合模块：来自空间混合模块的token进一步传递到通道混合模块。同样地，使用前置的LayerNorm，并在BQE操作后获得R_C和K_C：
之后，分别执行线性投影和门机制。最终输出的公式如下：
Local Graph-based Merging (LGM)
局部几何特征已被证明对点云特征学习至关重要，但RWKV结构的全局感受野无法全面捕捉局部点几何，限制了其学习细粒度特征的能力。因此我们将点云直接编码为图，使用点作为图的顶点。图的边连接在设定半径内的相邻点，允许这些点之间传递特征信息。这种图表示可以适应点云的结构，而无需对其进行规则化。此外，为了最小化局部图中的平移方差，引入了图稳定器机制。该机制允许点根据其独特特征对齐其坐标，从而提高网络的整体有效性。
通常，我们可以通过在图神经网络中沿着边聚合特征来优化顶点特征。在点云的场景中，我们旨在包含顶点所属对象的局部信息。因此，在第 (t + 1) 次迭代中，我们使用邻居的相对坐标进行边特征提取，这可以表示为：
为了减少这种平移方差，本文进一步提出基于结构特征对邻近坐标进行对齐，而不是依赖中心顶点的坐标。由于中心顶点已经包含了上一迭代中的一些结构特征，它可以用来估计对齐偏移，这促使本文设计了一个图稳定器机制。上述公式可以重写为：
实验结果
如上图所示，在 ShapeNet上进行自监督预训练后，PointRWKV 在 ScanObjectNN上实现了 93.63% (+4.66%) 的整体准确率，在 ModelNet40上实现了 96.89% (+1.79%) 的分类准确率，在 ShapeNetPart上实现了 90.26% (+3.16%) 的实例 mIoU，在预训练模型中创下了新的最先进 (SoTA)。同时，与基于 transformer 和 mamba 的同类工作相比，PointRWKV 的参数减少了 13%，FLOP 减少了 42%，展示了 RWKV 在 3D 视觉任务中的潜力。
3D点云分类
Few-shot分类
Part Segmentation
总结
在本文中，我们介绍了一种基于 RWKV 的新型点云学习架构 PointRWKV。PointRWKV 采用分层架构，通过对多尺度点云进行编码来学习生成强大的 3D 表示。为了促进局部和全局特征聚合，我们设计了并行特征合并策略。实验结果表明，PointRWKV 在不同的点云学习数据集上表现出优于基于 transformer 和 mamba 的同类工作的性能，同时显著减少了参数和 FLOP。凭借其线性复杂性能力，我们希望 PointRWKV 将成为更多 3D 任务的高效且经济高效的基准。
