★ FEATURED ARTICLE 如果你做过点云分类或者三维分割大概率体会过那种尴尬点数一多显存告急训练时间翻倍推理延迟也压不下来。Transformer 用自注意力把点云建模往前带了一大步但同时也把 N² 的计算复杂度带进了三维世界。PointMambaNeurIPS 2024正是冲着这个痛点去的——它把点云 token 序列喂给 Mamba 这种状态空间模型用线性复杂度完成全局上下文建模精度却能跟主流 Transformer 方法掰手腕。名字里的 Mamba 本来就带“轻快、灵活”的意思在点云这种动辄上万点的数据上它的表现确实配得上“凌波微步”这个比喻。这篇文章适合两类人一是想把点云骨干网络从 Transformer 换成更高效率方案的算法工程师二是想复现近期三维深度学习工作的研究生。我会从动机、网络设计、复现细节到踩坑记录都过一遍尽量把门道讲透而不是只贴摘要。1. 为什么点云处理需要 Mamba而不是更大的 Transformer1.1 点云的三个硬约束无序、稀疏、量大点云跟图像很不一样。图像有像素格子和固定的行列关系CNN 能靠卷积核在局部滑动点云只是一堆三维坐标没有顺序相邻点之间也没有固定距离。处理点云时通常要同时面对三个约束无序性、稀疏性和规模大。无序性意味着同一个物体在不同角度扫描出来的点顺序完全不一样模型不能依赖输入顺序否则换一下点的排列顺序精度就崩了。稀疏性表示一堆点里真正落在物体表面上的只有一部分传感器噪声、遮挡会让很多区域缺信息。规模大就更好理解了激光雷达的一帧可能就有几十万到上百万个点即便下采样到几千个点再做细粒度分割序列长度也远超图像 patch 数量。Transformer 在点云上的适配方案通常是先把点云划分成 patch再用自注意力建模 patch 间的关系比如 Point-BERT、Point-MAE 走的就是这个路线。自注意力确实能捕捉很远的依赖但它的计算量是所有 patch 两两之间的patch 数翻倍注意力矩阵的规模就要翻四倍。PointMamba 想解决的问题就是点云需要全局依赖但不需要为全局依赖付出二次复杂度。1.2 Transformer 的 N² 困境从显存到延迟举个例子算笔账。假设你对一个点云采样 1024 个点分成 128 个 patch每个 patch 的 token 维度是 384。Transformer 里 Q 和 K 相乘会得到一个 128×128 的注意力矩阵这一层还好。但是当点数增多、patch 数变成 512 甚至 1024 时注意力矩阵就会变成 512×512、1024×1024计算量和显存占用都在平方级往上走。更麻烦的是点云任务经常需要做逐点分割。一个室内场景的点云切成几万个点哪怕组成了几千个 tokenself-attention 依然会成为整个网络的计算瓶颈。推理延迟在实时机器人、自动驾驶场景里是没法妥协的。所以很多团队开始考虑 RNN 这类方案但传统 RNN 很难并行训练而且长距离记忆总是衰减。Mamba 的出现算是把“线性复杂度”和“可并行训练”这两个需求同时满足了。1.3 Mamba 的状态空间模型线性复杂度的秘密Mamba 本质上是一个带选择性机制的状态空间模型SSM。直觉上可以把它理解成 RNN网络维护一个固定大小的隐状态 h每读入一个新 token就更新一次 h再基于 h 输出当前结果。所有历史信息都被压缩进 h 里读下一个 token 时不需要回头去找之前所有的 token所以每一步的计算量是常数整个序列处理下来就是 O(N) 复杂度。这和 attention 那种“每一步都要和所有历史 token 做两两匹配”完全不同。但这带来一个问题固定大小的 h 怎么能装下那么多信息关键在“选择性”三个字。Mamba 让 B、C、∆ 这些参数都依赖输入也就是说模型可以根据当前输入决定哪些信息要写进状态、哪些要丢弃、步长该多大。对点云来说这意味着扫描到一个结构复杂的局部区域时模型可以更精细地记住局部细节扫描到平坦区域时则快速略过。SSM 在训练时可以通过循环扫描的算子高效实现所以它既保持了 RNN 式的高效推理又能在训练阶段把整个序列一起处理这也是 PointMamba 选择 Mamba 而不是传统 RNN 的根本原因。2. PointMamba 的核心设计拆解2.1 从原始点云到 Patch TokenFPS 和 KNN 的分组逻辑为了把点云喂给 MambaPointMamba 参考了 Point-BERT 和 Point-MAE 那套 token 化流程。先对整体点云做最远点采样FPS选出一批分布均匀的中心点然后对每个中心点用 K 近邻取邻域点组成一个 patch。为什么用 FPS 而不是随机采样随机采样容易在密集区域扎堆稀疏区域没有中心点FPS 每选一个点都会尽量远离已选点所以能保证中心点覆盖整体形状。得到 patch 之后把每个 patch 里的点送进一个小型 PointNet 网络共享权重的 MLP聚合成一个 patch token。这样点云就从 N 个原始点变成了 M 个 tokenM 通常取 128 或 256。这一步的好处是局部几何被压缩成紧凑特征后续 Mamba 不用直接面对点数爆炸也能通过局部 patch 获得类似于卷积的感受野。K 的取值一般是 16 或 32太大会让 patch 内混入不同部分的结构太小又会丢失细节。2.2 序列重排给无序点云“编个队”Mamba 模型对序列顺序非常敏感因为它是一步一步从左往右扫描的。图像上可以按光栅顺序扫描文本天然有先后顺序但点云没有。PointMamba 的核心技巧就在这里给点云的 patch token 重新排出一个合理顺序。论文里设计了一个几何感知的序列重排步骤——先选一个起始 patch然后每一次从剩下的未访问 patch 里选择与当前 patch 空间距离最近的作为下一个。你可以把它理解为旅行商问题的贪心版本最终形成一条尽量在空间上连续的扫描路径。为什么这么排有效因为 SSM 的隐状态空间是有限的扫描路径如果频繁跨越大片空间跳跃隐状态会被不相关的信息反复冲掉按空间连续顺序扫描相邻 token 在几何上本来就接近隐状态更新时只需要小幅调整模型就更容易把局部上下文串成一段连续的故事。这里需要特别说明这不是 Transformer 里那种 position embedding而是给 Mamba 这个序列模型本身提供一种合理的顺序。Position embedding 是为了弥补 attention 对位置不敏感的缺陷Mamba 天然顺序敏感所以排序告诉它的是“点云的空间连接关系”。这是 PointMamba 和其他基于 Transformer 的点云模型很本质的差别。2.3 扫描方向单向往复与双向扫读Mamba 的原始设计是单向因果扫描只看当前 token 之前的信息。图像、文本有明确的先后方向点云没有。PointMamba 的解决方案是双向扫描——把重排后的序列正着扫一遍再反着扫一遍然后把两个方向的特征拼接或求和作为整个 block 的输出。你可以想象成先把点云沿着贪心路径从左到右看一遍再从右到左看一遍这样每个 token 既能感知“前面”的上下文也能感知“后面”的上下文。这个设计和很多双向 RNN 类似但 PointMamba 不会做铺张的隐藏层拼接只是在同一组 Mamba 块里分别处理两个方向的扫描结果。想要堆叠深度时可以让每个 block 都做双向扫描也可以隔几层交换一下扫描方向。从论文里的消融实验来看双向扫描比单向扫描平均能带来更稳定的精度尤其在真实扫描数据上差距更明显。原因也很好理解真实扫描的物体有遮挡和噪声某些 patch 的信息必须通过反向扫描才能被其他 patch 正确解读。2.4 整体 Encoder 结构Mamba 模块怎么叠PointMamba 的编码器主体由多个 Mamba block 叠加而成。每个 Mamba block 的基本结构是先通过线性层把输入 token 映射到更高维度再用一个因果卷积做局部窗口混合激活函数之后送入选择性 SSM最后通过门控机制和残差连接输出。这个结构和语言模型里的 Mamba 块大体一致区别是去掉了位置编码依赖把序列交换层换成了上述点云重排层。整个网络不需要堆得很深通常 4 个 block 就能获得很强性能。原因在于 SSM 本身会在时间维度上逐步传递记忆每一个 block 内部已经做了全局信息压缩不需要像 Transformer 那样通过很多层去反复做两两交互。网络浅一点训练更快过拟合风险也小一些。分类任务接一个全局池化或分类 token再加 MLP 就能输出结果。3. 复现关键参数与训练配置3.1 数据和任务准备分类与分割基准先讲评估协议。PointMamba 在论文里主要在三个数据集上做验证ModelNet40 做物体分类ScanObjectNN 做真实场景扫描物体分类ShapeNetPart 做零件分割。ModelNet40 是 40 类 CAD 模型渲染出来的点云比较干净ScanObjectNN 是从真实场景扫描里截出来的物体包含遮挡、噪声和背景点难度明显更大ShapeNetPart 是 16 类物体的零件级分割要求每个点都被分类到一个零件标签。复现时我的建议是先跑 ModelNet40把全流程打通再上 ScanObjectNN。因为 ModelNet40 数据干净、类别均衡好收敛能快速验证模型实现有没有写错。ScanObjectNN 则能看出方法在真实数据上的鲁棒性也是和论文做对比时最关键的数字。如果你连 ModelNet40 都跑不出合理精度那问题大概率出在实现而不是数据难度上。3.2 训练超参数哪些能帮你快速收敛基于常见的实验设置PointMamba 训练时可以参考这些配置输入点数 1024patch 数量 128邻域点数 16token 维度大约 384。优化器用 AdamW初始学习率约 1e-3配合 warmup 和 cosine 衰减batch size 在常见训练卡上可以开到 64训练 300 个 epoch 左右。数据增强包括随机缩放、随机平移、随机旋转、点坐标抖动。损失函数直接走交叉熵不需要额外辅助损失。分类 head 用全局池化之后接 MLP。实际操作中有一个点对训练速度影响很大SSM 相关的算子在自定义 CUDA 实现里对 batch size 很敏感batch 太小时算子启动开销占比大训练效率上不去。建议在显存允许的情况下尽量把 batch size 顶到 64 或者 128。如果显存不够可以降 patch 数到 96精度损失很小但速度会明显改善。3.3 环境配置Mamba 库编译注意事项复现 PointMamba 不是 pip install 完就能跑的因为要用到选择性扫描的 CUDA 算子。通常需要安装配套的 mamba-ssm 和 causal-conv1d 库然后编译扩展。这一步最容易踩坑版本组合非常受限某个 PyTorch 版本对应某个 CUDA 版本gcc 版本太新也可能编译失败。我建议严格按项目提供的环境配置来操作先创建干净的虚拟环境再安装匹配的 PyTorch最后编译扩展。编译报错时优先确认三件事CUDA 版本是否和 torch 匹配、gcc 版本是否在支持范围里、有没有正确设置 CUDA_HOME 环境变量。如果一时找不到预编译好的对应包这一步可能耗掉你半天时间但值得耐心解决编译好之后后续训练基本不会再碰环境问题。3.4 与 Point-MAE 的自监督衔接PointMamba 除了做有监督训练也能接基于掩码的自监督预训练框架思路和 Point-MAE 类似把一部分 patch token 遮住用 Mamba 编码器去重建被遮住的 patch。Mamba 在这里天然适合做这个任务因为它的循环扫描在信息缺失时依然能利用上下文做重建比 Transformer 的非局部注意力更接近“按顺序推理”。预训练对 ScanObjectNN 这种真实数据特别有意义。真实扫描数据量少且噪声大先在大规模无标注合成数据上预训练再在 ScanObjectNN 上微调能明显提升精度。如果只做有监督训练也可以跳过预训练但最后的精度会差一些。我的习惯是先把有监督跑通确认模型和优化器没问题再回头加预训练流程这样定位问题更容易。4. 实验结果、效率对比与消融分析4.1 分类精度ModelNet40 和 ScanObjectNN从论文报告的实验数据来看PointMamba 在 ModelNet40 上的分类准确率落在 93% 左右的档位和 Point-MAE 等基于 Transformer 的方法基本持平在 ScanObjectNN 上PointMamba 与主流 Transformer 方法保持同一水平或略有优势。这里要强调一下PointMamba 的目标不是在精度上一骑绝尘而是用更低的复杂度达到同一档精度。方法骨干复杂度ModelNet40 AccScanObjectNN Acc显存/速度特征PointNetO(N)约 89约 68很轻但全局信息弱Point-BERTO(N²)约 93约 83Transformer显存高Point-MAEO(N²)约 94约 85Transformer MAEPointMambaO(N)约 93与 Point-MAE 同级线性复杂度显存更低上表是我基于论文数据和公开结果整理的近似参考具体数字会受训练配置影响。核心结论是PointMamba 在精度上没有明显妥协复杂度却降了一个量级。4.2 显存和吞吐量线性复杂度的实际收益这个部分最值得细看。同样输入条件下PointMamba 的显存占用随 patch 数增加基本是线性增长而 Transformer 的显存增长速度要快得多。原因是 Transformer 中间要保存 QK^T 矩阵梯度回传时也是平方级Mamba 只需要保留状态 h 和少量中间激活。吞吐量上Mamba 的推理不依赖两两注意力扫描算子非常规整所以在固定显存预算下batch size 可以开得更大线上部署时延迟也更低。如果你需要在车载、机器人端做点云实时识别这个效率优势比多一个点的精度更值钱。用大白话说同样的算力预算以前只能跑一个很小的 Transformer现在可以跑更大的点云范围或者更密的 patch 划分。4.3 消融与分析排序策略和扫描方向的影响论文对重排序策略做了消融对比随机排序、按原始坐标顺序、贪心连续排序。结果显示贪心连续排序明显优于随机排序说明“空间连续性”对 SSM 建模点云至关重要。这非常好理解隐状态容量有限如果扫描顺序乱跳模型必须频繁更新隐状态里的主题把大量容量浪费在“切换话题”上空间连续性则让信息在路径上自然流动隐状态只需要小幅更新。扫描方向也同样重要。单向扫描时靠近序列尾部的 token 能看到的信息更少双向扫描能缓解首尾不对称的问题。如果你的显存非常紧张可以只在后半段 block 做双向扫描精度下降有限速度还能再提一点。我复现时试过一次去掉双向扫描的设计ModelNet40 上掉了大概一个点ScanObjectNN 上掉得更明显所以这个设计不是可有可无的装饰。5. 常见问题与排查心得5.1 Mamba 模块编译失败怎么办编译是最常见的第一道坎。报错类型通常分三类缺少 gcc、xxd 等工具CUDA 头文件和 torch 版本不一致算子库源码里没包含当前显卡的架构。我的处理建议先升级或固定 gcc 版本到项目支持范围不要用太新的 CUDA最好卡在项目发布时对应的版本区间设置 TORCH_CUDA_ARCH_LIST 只保留你的显卡架构避免编译全部架构浪费大量时间和磁盘空间。如果实在搞不定版本匹配可以试试官方提供的预构建镜像或者发行说明里推荐的组合。另外尽量别在 Windows 裸环境下折腾这种编译很多算子没有预编译的 Windows 支持用一个 Linux 容器会省心很多。5.2 精度上不去排序策略和超参数调整复现时如果精度一直比论文低先别怀疑模型结构优先排查三件事。第一patch 划分参数对不对patch 太小局部特征不足patch 太大容易把跨区域结构混在一起。第二学习率是不是太高Mamba 类模型对学习率比较敏感太高会在后期震荡建议从一开始就用 warmup 加 cosine。第三数据增强有没有做全点云抖动和归一化特别重要。还有一个容易被忽略的细节Mamba 对输入数值尺度敏感。建议先把点坐标中心化到以原点为中心的单位球范围再做后续操作。我之前因为没做归一化直接拿原始坐标训练收敛速度明显变慢精度也掉了不少。5.3 显存和并发问题batch size 与梯度检查点PointMamba 在 token 数量非常大时显存同样可观只是比 Transformer 好很多。如果你跑分割任务patch 数上千再叠加大 batch显存也会吃紧。这时候可以考虑开梯度检查点gradient checkpointing。不过有一点要提醒Mamba 的 CUDA 算子对流过状态 h 的梯度已经比较省开检查点的收益没有 Transformer 里那么大优先调 batch size把梯度检查点当作最后手段。数据加载也是容易忽略的瓶颈。点云格式有很多种训练前要做统一预处理下采样到固定点数、坐标归一化、记录类别标签。把这些步骤写好缓存下来能省去大量训练中断的时间。否则每次启动训练都要重新读原始文件做预处理网络训练反而在等数据。结尾最后说点个人体会。PointMamba 给我最大的惊喜不是那个平均精度的提升而是它让我意识到点云建模里的“顺序问题”比想象中重要。同一个 Mamba 模型给一条更合理的扫描路径效果就能从“能用”变成“好用”。如果你手头正好在做点云识别又一直被 Transformer 的显存占用折腾得头疼把骨干网络换成 PointMamba 是性价比很高的选择——复杂度降一个量级精度不掉队。后续真想再深入我认为有两个方向值得跟进一是把重排策略从“贪心连续排序”升级成可学习的排序方式让网络自己决定扫描路径二是在更大规模的点云场景理解任务上测试它的扩展性毕竟线性和二次方的差距在数据量越大的地方体现得越明显。祝复现顺利。
阅读完成 · 觉得有帮助？
