3D点云学习新架构！PointRWKV联合RWKV，刷新点云表征学习性能及FLOPs
Transformer 彻底改变了点云学习任务，但其二次复杂度阻碍了其向长序列的扩展，这给有限的计算资源带来了负担。最近出现的 RWKV 是一种新型的深度序列模型，在 NLP 任务中显示出序列建模的巨大潜力。在这项工作中提出了 PointRWKV，这是一种线性复杂度的新模型，具有 3D 点云学习任务所需的适应性。
论文链接：
https://arxiv.org/abs/2405.15214项目主页：
https://hithqd.github.io/projects/PointRWKV/背景
3D 点云分析是众多现实应用的基础，包括自动驾驶、虚拟现实和机器人技术等。与 2D 图像不同，点云的内在不规则性和稀疏性使得进行准确的点云特征学习成为一项具有挑战性的任务。并且同时平衡准确性和复杂性仍然是一个持久的问题。现有的点云特征学习方法主要是基于自注意力结构（Transformer）或者是线性时间序列结构（Mamba）的，如下图所示。
然而，基于自注意力结构的算法对扩展点标记进行全面注意力机制的部署会显著增加对计算资源的需求，这种效应直接归因于注意力计算中固有的二次复杂性，影响了计算和内存。而基于线性时间序列结构的算法尽管有效，但原始的单向建模的固有属性阻碍了它们达到卓越的性能。
本文提出了一种基于 RWKV 的算法，该算法可以在极小的线性复杂度和参数量上达到较高的效率，并且能够处理多尺度的点云输入。
方法
PointRWKV 的整体流程如上图所示，其中通过分层网络架构对点云进行编码。给定一个输入点云，首先采用多尺度掩蔽策略在不同尺度上对不同点数进行采样。
然后应用轻量级 PointNet 来嵌入点并生成 embedding 嵌入。这些点标记由块堆叠编码器（即 PRWKV 块）使用，其中每个块由两个并行分支组成，用于分层局部和全局特征聚合。每个 PRWKV 块，采用两个并行分支的处理策略来聚合局部和全局特征。上面的是综合特征调制流程，具有空间混合和通道混合，下面的是基于局部图的合并。最后，两个分支的连接用作每个块的输出。
2.1 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：
之后，分别执行线性投影和门机制。最终输出的公式如下：
2.2 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 视觉任务中的潜力。
3.1 3D点云分类
3.2 Few-shot分类
3.3 Part Segmentation
总结
在本文中，我们介绍了一种基于 RWKV 的新型点云学习架构 PointRWKV。PointRWKV 采用分层架构，通过对多尺度点云进行编码来学习生成强大的 3D 表示。为了促进局部和全局特征聚合，我们设计了并行特征合并策略。
实验结果表明，PointRWKV 在不同的点云学习数据集上表现出优于基于 transformer 和 mamba 的同类工作的性能，同时显著减少了参数和 FLOP。凭借其线性复杂性能力，我们希望 PointRWKV 将成为更多 3D 任务的高效且经济高效的基准。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧