Parameter-Aware Mamba Model for Multitask Dense Prediction

Understanding the inter-relations and interactions between tasks is crucial for multitask 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 multitask learning (MTL) setting. Distinct from approaches that employ Transformers to model holistic task relationships, PAMM leverages the rich, scalable parameters of state-space models (SSMs) to enhance task interconnectivity. It features dual state-space parameter experts (PEs) that integrate and set task-specific parameter priors (PPs), capturing the intrinsic properties of each task. This approach not only facilitates precise multitask interactions but also allows for the global integration of task priors through the structured state-space sequence (S4) model. Furthermore, we employ the multidirectional Hilbert scanning (MDHS) method to construct multiangle feature sequences, thereby enhancing the sequence model's perceptual capabilities for 2-D 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.