Parameter Aware Mamba Model for Multi-task Dense Prediction

2025-11-18

AI Summary

Modeling task correlations in multi-task dense prediction remains challenging due to the difficulty of jointly capturing inter-task dependencies and spatial structure. Method: This paper proposes the Parameter-Aware Mamba Model (PAMM), the first to introduce state-space models (specifically S4) into multi-task dense prediction. PAMM features a dual state-space parameter expert mechanism that explicitly incorporates task-specific priors, and employs multi-directional Hilbert scanning to enhance the sequence model’s capacity to capture 2D dense spatial structures. Contribution/Results: PAMM unifies task interaction modeling and global contextual reasoning within an end-to-end trainable framework, enabling joint optimization across tasks. Extensive experiments on NYUD-v2 and PASCAL-Context demonstrate significant improvements over state-of-the-art methods, validating both its effectiveness and generalizability.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multi-instance/Multi-view Learning
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Understanding the inter-relations and interactions between tasks is crucial for multi-task dense prediction. Existing methods predominantly utilize convolutional layers and attention mechanisms to explore task-level interactions. In this work, we introduce a novel decoder-based framework, Parameter Aware Mamba Model (PAMM), specifically designed for dense prediction in multi-task learning setting. Distinct from approaches that employ Transformers to model holistic task relationships, PAMM leverages the rich, scalable parameters of state space models to enhance task interconnectivity. It features dual state space parameter experts that integrate and set task-specific parameter priors, capturing the intrinsic properties of each task. This approach not only facilitates precise multi-task interactions but also allows for the global integration of task priors through the structured state space sequence model (S4). Furthermore, we employ the Multi-Directional Hilbert Scanning method to construct multi-angle feature sequences, thereby enhancing the sequence model's perceptual capabilities for 2D data. Extensive experiments on the NYUD-v2 and PASCAL-Context benchmarks demonstrate the effectiveness of our proposed method. Our code is available at https://github.com/CQC-gogopro/PAMM.

Problem

• Modeling task interactions for multi-task dense prediction
• Enhancing task interconnectivity using state space models
• Improving 2D feature perception with Hilbert scanning sequences

Innovation

• Parameter Aware Mamba Model with dual experts
• State space models enhance task interconnectivity
• Multi-Directional Hilbert Scanning for 2D features