Mask-aware inference with State-Space Models

本文提出了 Partial Vision Mamba (PVM)，这是一种将部分卷积的掩码感知重归一化机制引入 Mamba 骨干网络的新型架构组件，旨在解决状态空间模型在处理任意形状缺失数据时的推理难题，并在深度补全、图像修复及含无效数据分类等任务中展现了优异性能。

要点

这篇论文介绍了一种名为 PVM (Partial Vision Mamba) 的新技术，它能让一种非常先进的 AI 模型（叫 Mamba）学会“忽略”坏数据，只关注好数据。

为了让你更容易理解，我们可以把 AI 模型想象成一个正在拼图的超级工匠，而这张论文要解决的核心问题就是：当拼图缺了一块，或者混进了一些错误的碎片时，工匠该怎么办？

1. 背景：当世界不完美时

在现实生活中，AI 看到的图像往往是不完美的：

• 深度图（Depth）：就像激光雷达扫描，可能因为灰尘或距离太远，有些区域是黑的（缺失数据）。
• 图片修复（Inpainting）：照片里有人脸被路人挡住了，或者被水印遮住了。
• 分类任务：一张猫的照片，被一只大黑手遮住了一半。

以前的 AI 模型（比如 CNN 或早期的 Mamba）就像是一个死板的工匠。如果你给它一张缺角的拼图，它会把缺失的部分（通常用黑色或零填充）也当成真实的拼图块来处理。这就像工匠试图把一块黑色的“空气”强行拼进画里，结果导致整个画面扭曲、出错。

2. 旧方案：Partial Convolutions (部分卷积)

在传统的 AI 模型（CNN）中，人们发明了一种叫“部分卷积”的方法。

• 比喻：这就像给工匠戴上了一副智能眼镜。当他看到缺失的碎片时，眼镜会告诉他：“嘿，这块是坏的，别管它，只计算旁边好的碎片。”
• 结果：工匠能拼出不错的图，但这种方法只适用于旧式的 CNN 模型。

3. 新挑战：Mamba 模型的困境

最近，一种叫 Mamba (Vision Mamba) 的新模型火了。它处理图像的速度极快，而且能理解全局（就像工匠能一眼看穿整幅画的逻辑，而不仅仅是局部）。

• 问题：但是，Mamba 这种新模型没有那副“智能眼镜”。它处理数据是一连串地“读”过去的。如果它读到一块坏数据（比如被遮挡的像素），它会把这种“坏”的信息传递下去，污染后面所有的数据。就像工匠读到了错误的指令，后面拼的所有东西都错了。

4. 解决方案：PVM (Partial Vision Mamba)

这篇论文的作者给 Mamba 模型装上了这副“智能眼镜”，并创造了一个新组件叫 PVM。

核心创意比喻：

想象 Mamba 模型是一个在传送带上工作的流水线。

• 以前的情况：传送带上混进了几个“空盒子”（坏数据/缺失数据）。流水线机械地把空盒子也传下去，导致后面的机器误以为空盒子是零件，最后组装出一堆废品。
• PVM 的做法：
  1. 智能分拣（Partial Patch Embedding）：在数据进入流水线之前，PVM 会先检查。如果一块区域是坏的，它不会直接扔掉，而是给它贴上一个特殊的**“隐形标签”**（Learned Masked Token）。
  2. 特殊指令：这个标签告诉流水线：“这个盒子是空的，跳过它，不要把它算进计算里，但也不要让它卡住传送带。”
  3. 动态更新：随着数据在模型里流动，PVM 会不断更新“哪里是好的，哪里是坏的”这张地图。如果两个好数据拼在一起，结果就是好的；如果混进了坏数据，结果就标记为坏。

5. 他们做了什么实验？

作者把这套“智能眼镜”装到了三个不同的任务中，效果惊人：

1. 补全深度图 (Depth Completion)：

   • 场景：给一张只有零星几个点的激光雷达图，让它补全成完整的 3D 地形。
   • 结果：用了 PVM 的模型，比没用的模型准确率高了 23%。就像工匠能完美猜出被遮挡的山峰形状，而旧模型会把山填成平地。

2. 图片修复 (Image Inpainting)：

   • 场景：把照片里被涂抹掉的人脸补回来。
   • 结果：PVM 补出来的人脸更自然、细节更清晰（比如鼻子、头发），而旧模型补出来的脸要么模糊，要么有奇怪的线条。

3. 带遮挡的分类 (Classification)：

   • 场景：给一张被遮住一半的猫的照片，让 AI 猜这是什么。
   • 结果：旧模型看到一半猫一半黑块，就懵了（准确率很低）。用了 PVM 的模型，因为它知道“忽略黑块，只看猫的部分”，准确率提升了 36%。

6. 总结

这篇论文的核心贡献在于：
它证明了Mamba 这种强大的新模型，只要加上“忽略坏数据”的机制（PVM），就能在处理残缺、遮挡、不完整的现实世界数据时，表现得比传统模型更聪明、更精准。

一句话概括：
以前的 AI 看到残缺的图片会“瞎猜”导致出错；现在的 PVM 让 AI 学会了“视而不见”那些坏数据，只专注于好数据，从而在修补、测量和识别任务中变得超级强大。

技术摘要

论文技术总结：Partial Vision Mamba (PVM)

1. 研究背景与问题定义 (Problem)

• 现实挑战：许多现实世界的计算机视觉任务（如深度补全、图像修复、遮挡分类）经常面临输入数据不完整或包含无效区域（arbitrarily shaped regions of missing or invalid data）的情况。这些无效数据通常由传感器误差（如稀疏 LiDAR）、应用限制或隐私遮挡引起，常以占位符（如零值）填充。
• 现有架构的局限性：
  • CNNs：传统的卷积神经网络（CNN）通过**部分卷积（Partial Convolutions, PConvs）**解决了此问题，即仅基于有效像素进行重归一化。
  • Transformer：虽然 Transformer 具有全局上下文能力，但其 $O(N^2)$ 的复杂度限制了高分辨率应用。
  • State Space Models (SSMs/Mamba)：Mamba 及其视觉变体（如 Vision Mamba, VMamba）因其线性复杂度 $O(N)$ 和卓越性能而迅速崛起。然而，现有的 Mamba 架构缺乏处理任意形状无效数据的内在机制。标准 Mamba 将所有数据（包括无效占位符）视为有效输入，导致特征提取被污染，隐藏状态被破坏，最终引发推理失败。
• 核心问题：如何在保持 Mamba 线性复杂度和高性能优势的同时，使其具备像 PConvs 那样的**掩码感知（Mask-aware）**能力，以在推理阶段动态忽略无效数据？

2. 方法论 (Methodology)

作者提出了Partial Vision Mamba (PVM)，这是一种新颖的架构组件，旨在将部分操作（Partial Operations）的原则移植到 Mamba 骨干网络中。

2.1 掩码感知框架 (Mask-aware Framework)
作者定义了一套规则，指导如何设计能够处理无效数据的架构。输入被定义为元组 $(x, m)$，其中 $x$ 是数据张量，$m$ 是布尔有效性掩码。

• 动态掩码更新：掩码 $m$ 不是静态的，而是随数据在网络中流动动态更新。
• 操作规则：
  • 元素级操作：输出有效性是输入掩码的逻辑与（AND）。
  • 特征拼接：通道拼接时，掩码取逻辑与。
  • 感受野操作（卷积/全连接/池化）：标准操作若包含无效输入则输出无效；部分操作（如 PConv）只要有一个有效输入，输出即标记为有效。
  • 序列建模（SSM）：这是关键创新点。标准 SSM 若序列中有一个无效 Token，整个序列会被污染。PVM 要求至少有一个有效 Token 即可产生有效输出。

2.2 Partial Vision Mamba (PVM) 模块设计
PVM 模块旨在解决 Mamba 处理无效数据时的两个主要问题：

1. Patch 内部无效性 (Inner-patch invalidity)：当 Patch 包含部分无效像素时，标准线性投影会产生损坏的 Token。
   • 解决方案：引入部分 Patch 投影（Partial Patch Projection）。将标准的线性层替换为“部分线性层”，即在无效位置先进行均值填充（Mean Padding），再进行线性投影。这使得部分有效的 Patch 能生成有效的 Token。
2. Patch 间无效性 (Inter-patch invalidity)：SSM 在处理序列时，若遇到无效 Token，会污染后续所有状态。
   • 解决方案：引入可学习的掩码 Token (Learned Masked Token)。类似于 BERT 中的 [MASK]，将无效 Token 替换为模型专门学习识别的特殊 Token。这从结构上保证了无效占位符值对有效输出流在数学上是无关的。

2.3 架构集成策略

• 残差连接处理：在残差连接中，PVM 仅更新有效位置的特征，无效位置需保留原始掩码。
• 填充层 (Filling Layer)：对于深度补全等任务，通过迭代使用部分卷积（PConv）作为“填充层”，逐步将稀疏的有效特征图转化为完全有效的特征图。

3. 主要贡献 (Key Contributions)

1. 提出 PVM 组件：设计了第一个基于 Mamba 的、能够处理任意形状无效输入数据的掩码感知架构组件。
2. 形式化掩码感知框架：建立了一套设计原则，定义了如何调整 SSM 架构以支持掩码感知推理，确保无效数据不传播。
3. 跨任务验证：在三个性质迥异的领域（生成式与判别式）证明了方法的通用性：
   • 深度补全 (Depth Completion)：基于稀疏 LiDAR 数据。
   • 图像修复 (Image Inpainting)：填充任意形状的孔洞。
   • 无效数据下的图像分类：在输入被遮挡的情况下进行分类。

4. 实验结果 (Results)

实验在 KITTI-3D（深度补全）、FFHQ（图像修复）和 ImageNet-1k（分类）上进行，对比了标准 Mamba (VM) 和 PVM 变体。

• 深度补全 (KITTI-3D)：

  • 设置：无 RGB 引导，仅使用稀疏深度输入。
  • 结果：PVM-DC 相比标准 VM-DC，RMSE 降低了 23%（从 1.80m 降至 1.38m），MAE 从 0.51m 降至 0.39m。参数增加极小（7.2M vs 7.3M）。
  • 意义：证明了在纯几何数据上，掩码感知机制能显著提升稀疏数据的重建质量。

• 图像修复 (FFHQ)：

  • 结果：所有基于 Mamba 的模型（VM-UNet, PVM-UNet）在感知指标（FID, LPIPS）上均显著优于基于 PConvs 的模型。
  • 对比：PVM-UNet-N（全掩码感知设计）优于 PVM-UNet-1（最小掩码感知修改），FID 从 40.02 降至 37.88。
  • 定性分析：PConvs 产生伪影，标准 VM 模糊，而 PVM 能生成更逼真的细节（如鼻子、头发）。

• 图像分类 (ImageNet-1k)：

  • 设置：输入图像被随机不规则遮挡。
  • 结果：PVM-Cls 的 Top-5 准确率达到 34.93%，相比标准 PlainMamba（25.60%）提升了 36%。
  • 意义：证明了模型在仅依赖部分有效信息时，PVM 架构具有更强的鲁棒性。

• 消融实验：

  • 掩码 Token 填充：验证了“可学习 Token"优于“零填充”或“均值填充”。
  • 掩码鲁棒性：PVM 在不同难度（Easy/Hard/Extreme）的随机掩码下均优于基线，证明其泛化能力不依赖于特定的掩码策略。

5. 意义与结论 (Significance & Conclusion)

• 填补空白：这是首个将掩码感知逻辑成功引入 State Space Models (Mamba) 的工作，解决了 Mamba 在处理稀疏/遮挡数据时的根本缺陷。
• 效率与性能平衡：PVM 在保持 Mamba 线性复杂度 $O(N)$ 优势的同时，实现了比传统 CNN 部分卷积更优的全局上下文建模能力（在修复任务中表现尤为明显）。
• 通用性：提出的设计原则不仅适用于 Mamba，也为其他处理不完整数据的架构提供了理论框架。
• 未来方向：作者计划将此框架扩展到其他深度相关任务（如深度增强），并探索更优的掩码感知扫描策略和位置编码。

总结：该论文通过引入 Partial Vision Mamba (PVM) 和一套严谨的掩码感知设计框架，成功解决了 Mamba 架构在推理阶段无法处理任意形状无效数据的问题，在深度补全、图像修复和分类任务中均取得了显著的性能提升，为高效、鲁棒的视觉推理开辟了新路径。