【论文笔记】LION: Linear Group RNN for 3D Object Detection in Point Clouds

原文链接：https://arxiv.org/abs/2407.18232

简介：Transformer在3D点云感知任务中有二次复杂度，难以进行长距离关系建模。线性RNN则计算复杂度较低，适合进行长距离关系建模。本文提出基于窗口的网络线性组RNN（即对分组的特征使用线性RNN；简称LION），用于3D目标检测，可实现在比Transformer网络更大的分组下进行充分的特征交互。为了解决线性组RNN处理空间建模的缺陷，本文引入3D空间特征描述器，整合到线性组RNN中，增强空间特征，而非增加体素序列的排序数。为进一步处理稀疏体素，本文提出3D体素生成策略，将线性组RNN用于自回归，以密集化前景特征。实验表明，LION在不同的线性组RNN（如Mamba、RWKV、RetNet）下均有效，且LION-Mamba能在多个数据集上达到SotA性能。

与组内体素数量有限的DSVT不同，LION可在数千个体素特征构成的组内进行特征交互，以建立长距离关系。

DSVT

LION

1. 概述

LION的流程如图所示。LION包含3D主干、BEV主干和检测头，其中3D主干基于线性组RNN，包含$N$个用于长距离特征交互的LION块，$N$个增强稀疏点云特征表达的体素生成操作和$N$个用于高度下采样的体素融合操作。

3D稀疏窗口分割：将输入点云体素化得到$L$个非空体素后，按照X或Y坐标排序体素，并将排序的体素分割为大小为$K$的组。由于线性组RNN的线性复杂度，组的大小$K$可以比基于Transformer的方法大很多。

2. LION块

LION块包括用于长距离特征交互的LION层、用于捕捉3D空间信息的3D空间特征描述器、用于下采样的体素融合操作和用于上采样的体素扩张操作。此外，LION有提取多尺度特征的层次结构。

LION层：LION层包括两个线性组RNN，前者基于X轴窗口分割，后者基于Y轴窗口分割。两种窗口分割可使特征进行更充分的交互。

3D空间特征描述器：由于体素排序过程会丢失空间信息（如图所示，实际很近的体素在排序后变得很远），一些方法增加排序方式，但其属于手工设计的方法，且复杂度会随着排序方式的增多而增加。本文使用3D空间特征描述器，包括3D子流形卷积、层归一化和GELU激活函数，为LION层提供丰富的3D局部位置感知的信息。此外，将3D空间特征描述器放置于体素融合前，以减小体素融合过程中的空间信息损失。

体素融合和体素扩张：为获取多尺度特征，LION使用层次特征提取结构。本文使用体素融合和体素扩张进行下采样和上采样。体素融合过程中会计算索引映射，而体素扩张会根据逆索引映射上采样体素。

3. 体素生成

为处理高度稀疏点云特征提取的挑战和体素融合中的信息损失，本文利用RNN的自回归能力，提出体素生成策略。

无监督的前景体素区分：首先需要确定体素生成的区域。有工作使用监督方法获取BEV下的前景区域，但本文注意到，主干中特征在通道维度上的响应较高的位置通常为前景。因此本文计算第$i$个LION块的特征响应$F_i^{\ast }$：
$$F_i^{\ast }=\frac{1}{C}\sum _{j=0}^C{F}_i^j$$

随后，选择前$m=r\ast L$个体素作为前景（$r$为前景比例）：
$$F_m={\text{Top}}_m(F_i^{\ast })$$

其中，$F_m$为选择的前景特征，会用于后续的体素生成。

自回归属性的体素生成：传统方法使用KNN的方式，基于KNN特征生成体素特征，但这一方法是次优的，因其存在冗余特征和有限的感受野。本文利用线性RNN的自回归能力，通过与组内其它体素特征交互，生成新的体素特征。具体来说，本文定义前景特征$F_m$相应的位置为$P_m$。首先以4个不同的偏移量（$[-1,-1,0],[1,1,0],[-1,1,0],[1,-1,0]$）扩散体素，然后初始化扩散的体素为0。将第$i$个LION块的输出特征$F_i$与上述初始化特征一起输入到第$i+1$个LION块，从而生成扩散的体素。

上述过程可表示如下：
$$\begin{gathered} F_p=F_i\oplus F_{[-1,-1,0]}\oplus F_{[1,1,0]}\oplus F_{[1,-1,0]}\oplus F_{[-1,1,0]} \\ {F}_p^{\prime }=\text{Block}(F_p) \end{gathered}$$