MTMamba++: Enhancing Multi-Task Dense Scene Understanding via Mamba-Based Decoders

📄 arXiv: 2408.15101v2 📥 PDF

作者: Baijiong Lin, Weisen Jiang, Pengguang Chen, Shu Liu, Ying-Cong Chen

分类: cs.CV, cs.AI

发布日期: 2024-08-27 (更新: 2025-07-26)

备注: Accepted by IEEE Transactions on Pattern Analysis and Machine Intelligence

🔗 代码/项目: GITHUB

💡 一句话要点

MTMamba++：基于Mamba解码器的多任务密集场景理解框架

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱三：空间感知与语义 (Perception & Semantics)

关键词: 多任务学习 密集场景理解 Mamba架构 状态空间模型 跨任务交互 长程依赖 深度学习

📋 核心要点

  1. 多任务密集场景理解面临长程依赖建模和跨任务信息交互的挑战，现有方法难以兼顾性能与效率。
  2. MTMamba++利用Mamba架构，通过自任务Mamba块处理长程依赖，跨任务Mamba块显式建模任务间交互。
  3. 实验表明，MTMamba++在多个数据集上超越了CNN、Transformer和扩散模型，同时保持了较高的计算效率。

📝 摘要（中文）

本文提出MTMamba++，一种用于多任务密集场景理解的新型架构，其特点是基于Mamba的解码器。多任务密集场景理解训练一个模型来处理多个密集预测任务，具有广泛的应用场景。捕获长程依赖关系和增强跨任务交互对于多任务密集预测至关重要。MTMamba++包含两种类型的核心块：自任务Mamba（STM）块和跨任务Mamba（CTM）块。STM利用状态空间模型处理长程依赖关系，而CTM显式地建模任务交互，以促进跨任务的信息交换。我们设计了两种类型的CTM块，即F-CTM和S-CTM，分别从特征和语义的角度增强跨任务交互。在NYUDv2、PASCAL-Context和Cityscapes数据集上的大量实验表明，MTMamba++优于基于CNN、基于Transformer和基于扩散的方法，同时保持了较高的计算效率。代码已在https://github.com/EnVision-Research/MTMamba上发布。

🔬 方法详解

问题定义：多任务密集场景理解旨在同时预测图像中多个像素级别的属性，例如语义分割、深度估计和表面法线预测。现有方法，如基于CNN和Transformer的模型，在捕获长程依赖关系和有效建模跨任务交互方面存在局限性，或者计算复杂度较高，难以满足实际应用的需求。

核心思路：MTMamba++的核心思路是利用Mamba架构的状态空间模型（SSM）来高效地捕获长程依赖关系，并设计专门的跨任务Mamba（CTM）块来显式地建模和增强不同任务之间的信息交互。通过结合自任务Mamba（STM）块和CTM块，MTMamba++能够在保持计算效率的同时，提升多任务场景理解的性能。

技术框架：MTMamba++的整体架构包含一个共享的编码器和一个基于Mamba的解码器。编码器负责提取输入图像的特征表示。解码器由多个STM块和CTM块交替堆叠而成。STM块处理单个任务内的长程依赖，而CTM块则促进不同任务之间的信息交换。最终，解码器的输出被用于预测每个任务的密集预测结果。

关键创新：MTMamba++的关键创新在于提出了基于Mamba的解码器，并设计了两种类型的CTM块：F-CTM和S-CTM。F-CTM从特征层面增强跨任务交互，而S-CTM从语义层面增强跨任务交互。与传统的卷积或Transformer方法相比，Mamba架构能够更高效地处理长程依赖关系，并且CTM块能够显式地建模任务之间的关系，从而提升多任务学习的性能。

关键设计：STM块采用标准的Mamba块结构，利用状态空间模型来建模序列数据。F-CTM块通过注意力机制来融合不同任务的特征表示。S-CTM块则利用语义分割的预测结果来引导跨任务的信息交换。损失函数通常是每个任务的损失函数的加权和，权重可以根据任务的重要性进行调整。具体的参数设置，如Mamba块的维度、CTM块的注意力头数等，需要根据具体的任务和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MTMamba++在NYUDv2、PASCAL-Context和Cityscapes数据集上取得了显著的性能提升。例如，在NYUDv2数据集上，MTMamba++在多个任务上超越了现有的CNN、Transformer和扩散模型。此外，MTMamba++还具有较高的计算效率，使其更适合于实际应用。

🎯 应用场景

MTMamba++在自动驾驶、机器人导航、增强现实等领域具有广泛的应用前景。例如，在自动驾驶中，可以利用MTMamba++同时预测道路分割、交通标志检测和深度估计，从而提高环境感知能力。在机器人导航中，可以利用MTMamba++进行场景理解，帮助机器人更好地理解周围环境，从而实现更安全、更智能的导航。

📄 摘要（原文）

Multi-task dense scene understanding, which trains a model for multiple dense prediction tasks, has a wide range of application scenarios. Capturing long-range dependency and enhancing cross-task interactions are crucial to multi-task dense prediction. In this paper, we propose MTMamba++, a novel architecture for multi-task scene understanding featuring with a Mamba-based decoder. It contains two types of core blocks: self-task Mamba (STM) block and cross-task Mamba (CTM) block. STM handles long-range dependency by leveraging state-space models, while CTM explicitly models task interactions to facilitate information exchange across tasks. We design two types of CTM block, namely F-CTM and S-CTM, to enhance cross-task interaction from feature and semantic perspectives, respectively. Extensive experiments on NYUDv2, PASCAL-Context, and Cityscapes datasets demonstrate the superior performance of MTMamba++ over CNN-based, Transformer-based, and diffusion-based methods while maintaining high computational efficiency. The code is available at https://github.com/EnVision-Research/MTMamba.