Enhancing Mamba Decoder with Bidirectional Interaction in Multi-Task Dense Prediction

TL;DR

This paper introduces a Bidirectional Interaction Mamba (BIM) model that enhances multi-task dense prediction by efficiently capturing cross-task interactions through novel scanning mechanisms, balancing accuracy and computational efficiency.

Contribution

The work proposes a new BIM framework with BI-Scan and MS-Scan mechanisms, enabling effective multi-granularity scene modeling with linear complexity.

Findings

Outperforms state-of-the-art on NYUD-V2 and PASCAL-Context benchmarks.

Efficiently balances interaction completeness and computational cost.

Demonstrates superior multi-task dense prediction performance.

Abstract

Sufficient cross-task interaction is crucial for success in multi-task dense prediction. However, sufficient interaction often results in high computational complexity, forcing existing methods to face the trade-off between interaction completeness and computational efficiency. To address this limitation, this work proposes a Bidirectional Interaction Mamba (BIM), which incorporates novel scanning mechanisms to adapt the Mamba modeling approach for multi-task dense prediction. On the one hand, we introduce a novel Bidirectional Interaction Scan (BI-Scan) mechanism, which constructs task-specific representations as bidirectional sequences during interaction. By integrating task-first and position-first scanning modes within a unified linear complexity architecture, BI-Scan efficiently preserves critical cross-task information. On the other hand, we employ a Multi-Scale Scan~(MS-Scan) mechanism to achieve multi-granularity scene modeling. This design not only meets the diverse granularity requirements of various tasks but also enhances nuanced cross-task feature interactions. Extensive experiments on two challenging benchmarks, \emph{i.e.}, NYUD-V2 and PASCAL-Context, show the superiority of our BIM vs its state-of-the-art competitors.

Tables

Table 1. Table 1 : Quantitative comparison with other state-of-the-art methods on NYUD-v2 dataset.

Model	Semseg	Depth	Normal	Boundary
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$
CNN based
Cross-Stitch[22]	36.34	0.6290	20.88	76.38
PAP[42]	36.72	0.6178	20.82	76.42
PSD[40]	36.69	0.6246	20.87	76.42
PAD-Net[31]	36.61	0.6270	20.85	76.38
MTI-Net[29]	45.94	0.5365	20.27	77.86
ATRC[2]	46.33	0.5363	20.18	77.94
Transformer based
InvPT[35]	53.66	0.5183	19.04	78.10
TaskPrompter[36]	55.30	0.5152	18.47	78.20
InvPT++[37]	53.85	0.5096	18.67	78.10
MQTransformer[32]	54.84	0.5325	19.67	78.20
Mamba based
MTMamba[16]	55.82	0.5066	18.63	78.70
BIM(Ours)	57.40	0.4733	18.55	78.72

Table 2. Table 2 : Quantitative comparison with other state-of-the-art methods on Pascal-Context dataset.

Model	Semseg	Parsing	Saliency	Normal	Boundary
	mIoU $\uparrow$	mIoU $\uparrow$	maxF$\uparrow$	mErr $\downarrow$	odsF $\uparrow$
CNN based
PAD-Net[31]	53.60	59.60	65.80	15.30	72.50
ASTMT[20]	68.00	61.10	65.70	14.70	72.40
MTI-Net[29]	61.70	60.18	84.78	14.23	70.80
ATRC[2]	62.69	59.42	84.70	14.20	70.96
ATRC-ASPP[2]	63.60	60.23	83.91	14.30	70.86
ATRC-BMTAS[2]	67.67	62.93	82.29	14.24	72.42
Transformer based
InvPT[35]	79.03	67.61	84.81	14.15	73.00
TaskPrompter[36]	80.89	68.89	84.83	13.72	73.50
InvPT++[37]	80.22	69.12	84.74	13.73	74.20
MQTransformer[32]	78.93	67.41	83.58	14.21	73.90
Mamba based
MTmamba[16]	81.11	72.62	84.14	14.14	78.80
BIM(Ours)	81.25	73.15	84.14	14.13	79.05

Table 3. Table 3 : Complexity comparisons on NYUD-v2. †denotes the result is from  [ 19 ] .

Model	Semseg	Depth	Normal	Boundary	FLOPs	# Params
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$	(G)$\downarrow$	(M)$\downarrow$
InvPT[35]	53.66	0.5183	19.04	78.10	598	402
TaskPrompter[36]	55.30	0.5152	18.47	78.20	470	392
InvPT++†[37]	53.85	0.5096	18.67	78.10	$-$	$\sim$402
MTMamba[16]	55.82	0.5066	18.63	78.70	541	308
DBIM (Ours)	56.41	0.4789	18.66	78.80	495	290
BIM (Ours)	57.40	0.4733	18.55	78.72	547	388

Table 4. Table 4 : Effectiveness of different components on NYUD-v2.

Model	Semseg	Depth	Normal	Boundary	MTL Gain	FLOPs	# Params
	mIoU $\uparrow$	RMSE $\downarrow$	mErr $\downarrow$	odsF $\uparrow$	${\Delta }_m\uparrow$	(G) $\downarrow$	(M) $\downarrow$
STL	54.32	0.5166	19.21	77.30	+0.00	1075	889
MTL	53.72	0.5239	19.97	76.50	-1.87	466	303
Baseline	55.82	0.5066	18.63	78.70	+2.38	541	308
BI-Scan	57.11	0.4856	18.66	78.90	+4.02	547	290
MS-Scan	57.40	0.4829	18.79	78.60	+4.00	541	397
BIM	57.40	0.4733	18.55	78.72	+4.83	547	388

Table 5. Table 5 : Ablation on task order on NYUD-v2.

Order	Semseg	Depth	Normal	Boundary	MTL Gain
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$	${\Delta }_m\uparrow$
S-D-N-B	57.11	0.4856	18.66	78.90	+4.02
N-S-D-B	56.15	0.4832	18.66	79.00	+3.72
B-S-D-N	56.97	0.4836	18.73	78.90	+3.96
S-N-D-B	57.52	0.4843	18.71	78.90	+4.20
Random	55.41	0.4861	18.69	78.90	+3.17

Table 6. Table 6 : Different scanning scales of MS-Scan on NYUD-v2.

Scan Scale	Semseg	Depth	Normal	Boundary	MTL Gain
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$	${\Delta }_m\uparrow$
{1,2}	57.39	0.4756	18.60	78.80	+4.68
{1,3}	57.53	0.4739	18.68	78.50	+4.62
{1,4}	57.40	0.4733	18.55	78.72	+4.83

Table 7. Table 7 : Different scanning number of MS-Scan on NYUD-v2.

Scan Number	Semseg	Depth	Normal	Boundary	MTL Gain
	mIoU $\uparrow$	RMSE $\downarrow$	mErr $\downarrow$	odsF $\uparrow$	${\Delta }_m\uparrow$
Type 1	57.11	0.4856	18.66	78.90	+4.02
Type 2	57.40	0.4733	18.55	78.72	+4.83
Type 3	56.76	0.4756	18.59	78.70	+4.35

Table 8. Table 8 : Performance on different encoders.

Model	Semseg	Depth	Normal	Boundary
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$
MTMamba-Base	53.62	0.5126	19.28	77.70
BIM-Base	54.32	0.4915	18.89	77.74
MTMamba-Large	55.82	0.5066	18.63	78.70
BIM-Large	57.40	0.4733	18.55	78.72

Table 9. Table 9: Effect of scan mode order in BI-Scan.

Setting	Semseg	Depth	Normal	Boundary	FLOPs	# Params
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$	(G) $\downarrow$	(M) $\downarrow$
TF	56.36	0.4905	18.67	78.70	510	289
PF	56.96	0.4883	18.68	78.70	510	289
TF $\to$ PF	57.11	0.4856	18.66	78.80	547	290
PF $\to$ TF	56.55	0.4806	18.71	78.70	547	290

Table 10. Table 10: Model efficiency with varying task quantities.

Methods	Semseg	Depth	Normal	Boundary	FLOPs	# Params
	mIoU $\uparrow$	RMSE$\downarrow$	mErr $\downarrow$	odsF $\uparrow$	(G) $\downarrow$	(M) $\downarrow$
2 tasks	57.20	0.4715	-	-	375	304
3 tasks	55.91	0.4891	18.63	-	461(+86)	346(+42)
4 tasks	57.40	0.4733	18.55	78.72	547(+86)	388(+42)

Equations

$$\{F_1,F_2,\dots ,F_T\}=\Theta (\Phi (\text{I})),$$

$${\hat{\text{Y}}}_t=H_t(F_t)={\mathcal{P}}_t(\Psi (F_t)),$$

$$\text{G}=\{{\text{G}}^1,{\text{G}}^2,{\text{G}}^3,{\text{G}}^4\}=\Phi (\text{I}),$$

$${\text{F}}_t^s={\mathcal{F}}_t^s+{\mathcal{G}}_t^s\times {\mathcal{F}}_{sh}^s[t]+(1-{\mathcal{G}}_t^s)\times {\mathcal{F}}_{t,N}^s,$$

$${\mathcal{F}}^s={\mathcal{F}}_{in}^s+\mathcal{P}({\mathcal{F}}_h^s\times {\mathcal{G}}^s).$$

$$S_1={\mathcal{O}}_1({\mathcal{U}}_1({\mathcal{F}}_t^s,{\mathcal{F}}_2^s,\dots ,{\mathcal{F}}_T^s))\in {\mathbb{R}}^{(T\times (H\times W))\times \frac{C}{2}}.$$

$$S_2={\mathcal{U}}_2({\mathcal{O}}_2({\bar{\mathcal{F}}}_1^s,{\bar{\mathcal{F}}}_2^s,\dots ,{\bar{\mathcal{F}}}_T^s))\in {\mathbb{R}}^{((H\times W)\times T)\times \frac{C}{2}}.$$

$$\{{\bar{x}}_1,{\bar{x}}_2,\dots ,{\bar{x}}_N\}=\{{\mathcal{W}}_1,{\mathcal{W}}_2,\dots ,{\mathcal{W}}_N\}(\mathcal{S}(x)).$$

$${\bar{y}}_i=\text{SS2D}({\bar{x}}_i)\in {\mathbb{R}}^{C_i\times H_i\times W_i},i\in \{1,2,\dots ,N\}.$$

$$y={\mathcal{S}}^{-1}({\mathcal{W}}_1^{-1}({\bar{y}}_1),{\mathcal{W}}_2^{-1}({\bar{y}}_2),\dots ,{\mathcal{W}}_N^{-1}({\bar{y}}_N)),$$

$$L=\sum _{t\in \mathcal{T}}{\lambda }_t{\mathcal{L}}_t({\hat{\text{Y}}}_t,{\text{Y}}_t),$$

Full text

1 Introduction

Multi-task dense prediction is a critical visual task designed to simultaneously predict outputs for various pixel-level tasks, such as semantic segmentation, depth prediction, surface normal estimation, and saliency detection. Existing approaches made mainly efforts [30, 5] to refine task-general representations, obtained from the encoder, into task-specific representations. In these approaches, cross-task interaction has been proven to be significant in multi-task dense prediction [31, 39, 35, 26], as it can enhance the robustness and effectiveness of task representations.

Initially, multi-task dense prediction methods primarily utilized Convolutional Neural Networks (CNNs) with pair-wise interactions, as illustrated in Fig. 1 (a). These networks [31, 8, 28] were carefully designed with specialized branches for each task, augmented by interaction modules that facilitated information exchanges between each pair of tasks. These pair-wise interaction approaches maintain the integrity of task information during interaction, promoting fine-grained cross-task interaction. While, the inherently limited receptive fields of CNN architectures often led to insufficient spatial interactions. In response to these challenges, transformer-style networks [35] demonstrate remarkable capability in modeling long-range dependencies. This capability substantially improves the interaction effectiveness of models in handling multi-task scenarios, leading to enhanced performance. However, these sufficient cross-task interactions usually introduce high interaction complexity, which typically increases quadratically with the number of tasks. The high interaction complexity limits their applicability in certain powerful models, particularly transformers.

To enhance the efficiency of cross-task interactions, a crowd of methods [24, 1, 39, 16] with linear interaction complexity has been proposed, as depicted in Fig. 1 (b). Typically, these methods strategically select a representative task [24, 1] or extract representative information from all tasks [39, 16] to construct a compact bridge representation that aims to preserve essential interaction information. However, in these methods, the necessity of maintaining linear interaction complexity inherently limits the representational capacity of these bridge representations. This information compression inevitably leads to the degradation of task features, potentially discarding critical interaction details and ultimately resulting in suboptimal performance. Therefore, how to enhance sufficient interaction capability while maintaining linear interaction complexity remains an unresolved issue.

In response to these challenges, based on Mamba architecture, we propose a simple yet effective Bidirectional Interaction Mamba (BIM) method that achieves comprehensive preservation of task information during interaction while maintaining linear interaction complexity. Specifically, we introduce a Bidirectional Interaction Scan (BI-Scan) mechanism to achieve sufficient cross-task interaction and a Multi-Scale Scan (MS-Scan) mechanism to deliver a comprehensive visual representation. The BI-Scan mechanism constructs each task representation into sequences and deploys a bidirectional scanning strategy for modeling cross-task interactions, efficiently mitigating information loss while preserving linear interaction complexity. Furthermore, BI-Scan enables multi-level cross-task interactions by integrating task-first and position-first scanning modes, facilitating more nuanced modeling of task relationships. In the MS-Scan, features are initially partitioned into multiple spaces, where the scene structure of images is modeled at various scales. Specific tasks subsequently integrate multi-scale scene structural information to enhance task-specific representations. During cross-task interaction, different task branches adaptively integrate interaction representations from multiple scales, addressing the varying granularity demands of each task for scene structural information.

The main contributions are summarized as follows:

• •

We propose a simple yet effective Bidricational Interaction Mamba (BIM) for multi-task dense prediction, enabling comprehensive cross-task interaction while preserving linear interaction complexity.

• •

We introduce a novel Bidirectional Interaction Scan (BI-Scan) mechanism to maintain the integrity of task information during interactions. The BI-Scan incorporates two scanning modes: task-first mode and position-first mode, facilitating multi-level cross-task interactions.

• •

We design a Multi-Scale Scan (MS-Scan) mechanism to capture scene structural information at multiple scales, thereby addressing the diverse granularity requirements of various tasks and enhancing more detailed cross-task interactions in multi-task dense prediction.

2 Related Work

**Multi-Task Learning. **Most existing multi-task learning works focused on training optimization and network design. Optimization approaches balance tasks via gradient manipulation [38, 23, 13, 34] or loss-balancing [4, 14] to coordinate resource allocation for various tasks during training. The approaches of network design seek to enhance task representation learning by developing diverse mechanisms. Some CNN-based methods [31, 29] manually design interaction mechanisms to extract useful information across tasks. With the advancement of Transformer, current methods have gained improved global task interaction capabilities to enhance task representation efficiency. They enable superior pairwise [1, 24] or global [35, 32, 36, 15] task interactions. To mitigate computational complexity, MTMamba [16] introduce Mamba [9] to multi-task learning, which showcases effective long sequence modeling capabilities and achieving satisfactory performance. However, achieving both efficient and sufficient cross-task interaction remains a challenge.

State Space Models. In efficient long-range dependency modeling methods, state space models [11, 27] has become a striking alternative to Transformers. [10] proposed a Structured State Space Sequence (S4) model based on a new parameterization, which alleviates the computational and memory efficiency issues faced by SSM. Subsequently, numerous efforts [7, 21] are dedicated to bridging the performance disparity between SSMs and Transformers. Mamba [9] introduced an input-based parameterization method and hardware-aware algorithm, achieving performance on par with Transformers in natural language processing. This success inspired various efforts [43] towards Mamba’s adaptation for visual tasks.

The preservation of comprehensive image structural information poses a critical challenge in Mamba’s sequential processing model, which has attracted considerable attention and effort [43, 17, 33, 12, 41]. VMamba [17] proposes the 2D Selective Scan (SS2D), a four-way scanning mechanism tailored for spatial domain traversal, aimed at enhancing Mamba’s image modeling capabilities. Subsequent research studies have explored various scan patterns and combinations tailored to different tasks or scenarios [33, 12, 41]. However, by utilizing fixed token sizes, these methods overlook the importance of hierarchical spatial structural information in visual tasks.

3 Method

We first outline the overall architecture of BIM for multi-task dense prediction in Section 3.1, Subsequently, Section 3.2 focuses on the BIM Encoder, while Section 3.3 examines the BIM Decoder and its two scanning mechanisms: BI-Scan and MS-Scan. Finally, Section 3.4 discusses the optimization objectives.

3.1 Framework Overview

As shown Fig. 2 (a), the proposed BIM framework consists of three main components: a task-shared encoder $\Phi$ for extracting task-generic representations, a Mamba decoder $\Theta$ for refining task features and T task-specific heads $\{H_{t}\}_{t=1}^{T}$ generating predictions $\{\hat{\textbf{Y}}_{t}\}_{t=1}^{T}$ for individual tasks. This can be formulated as:

[TABLE]

where $\textnormal{I}\in\mathbb{R}^{3\times H\times W}$ denotes the RGB input image, $F_{t}$ denotes the features generated by the decoder, corresponding to the distinct task $t$. $\hat{\textbf{Y}}_{t}$ represents the prediction for task $t$, which has the same height $H$ and width $W$ as I. $\Psi$ denotes a module designed to expand the task feature resolution to $H\times W$, which consists of a linear projection and a reshape operation. $\mathcal{P}_{t}$ represents a linear layer that projects the feature channels to a specific number required by the corresponding task.

3.2 BIM Encoder

The encoder shares similarities with other methods [16]. We utilized a pretrained Swin Transformer [18] to extract task-generic features, which begins by dividing the input image $\textnormal{I}\in\mathbb{R}^{H\times W\times 3}$ into $\frac{H}{w}\times\frac{W}{w}$ tokens of dimension $C$ through patch partitioning and linear layers, where $w$ denotes the partition size. These tokens are then processed through multiple stages involving alternating patch merging and Swin Transformer block processing, ultimately yielding hierarchical image representations:

[TABLE]

where G represents the task-generic features extracted from the encoder $\Phi$. In our practical implementation, we utilize a partition size of 4, resulting in the following shapes for G: $C\times\frac{H}{4}\times\frac{W}{4},2C\times\frac{H}{8}\times\frac{W}{8},4C\times\frac{H}{16}\times\frac{W}{16}$ and $8C\times\frac{H}{32}\times\frac{W}{32}$, respectively.

3.3 BIM Decoder

The BIM Decoder consists of three Mamba Feature Refinement (MFR) blocks, as depicted in Fig. 2 (a). These MFRs are designed to bridge the gap between task-generic and task-specific representations. They generate task representations $\{\mathbf{F}_{1}^{i},\mathbf{F}_{2}^{i},\ldots,\mathbf{F}_{T}^{i}\}_{i=1}^{3}$ corresponding to the first three encoder stages, with dimensions of $4C\times\frac{H}{16}\times\frac{W}{16},2C\times\frac{H}{8}\times\frac{W}{8}$, and $C\times\frac{H}{4}\times\frac{W}{4}$, respectively. Finally, the last refined features $\{\textbf{F}_{t}^{3}\}_{t=1}^{T}$ are input into the task heads as $\{F_{t}\}_{t=1}^{T}$ to generate the final predictions $\{\hat{\textbf{Y}}_{t}\}_{t=1}^{T}$.

We introduce two specialized Mamba blocks, Bidirectional Cross-Task Feature Refine (BCFR) and Multi-Scale Single-task (MSST) in each MFR block, as illustrated in Fig. 2 (b). The BCFR Block is designed to facilitate effective cross-task interaction while maintaining the integrity of task information. It begins by concatenating features $\{\mathcal{F}^{s}_{t}\}_{t=1}^{T}$ from different task branches to construct task-shared features $\mathcal{F}_{sh}^{s}$ using the BI-Scan mechanism. Subsequently, each task branch adaptively merges $\mathcal{F}_{sh}^{s}$ by a selection value $\mathcal{G}^{s}$ to obtain fine-tuned task representations $\textbf{F}^{s}=\{\textbf{F}^{s}_{t}\}_{t=1}^{T}$:

[TABLE]

where $\mathcal{F}^{s}_{t,N}$ is is derived from $\mathcal{F}^{s}_{t}$ through layer normalization, and $\mathcal{G}^{s}_{t}$ is obtained from $\mathcal{F}^{s}_{t,N}$ via a linear layer followed by a Sigmoid function. $\mathcal{F}^{s}_{sh}[t]$ is the $t$-th component of $\mathcal{F}^{s}_{sh}$ corresponding to the $t$-th task. The MSST primarily aims to construct comprehensive representations through task-internal interactions. Its architecture is illustrated in Fig. 2 (b) and comprises two main branches: the scan branch and the gating branch. In the scan branch, we integrate a hierarchical global scene structure representation $\mathcal{F}_{h}^{s}$ using a novel MS-Scan mechanism. Simultaneously, a gating signal $\mathcal{G}^{s}$ from the gate branch regulates the flow of information. Subsequently, we adjust the channel dimensions of the multi-scale scene representation by applying a linear projection $\mathcal{P}$ and establish a residual connection between $\mathcal{F}_{h}^{s}$ and the input $\mathcal{F}_{in}^{s}$. This process is repeated twice to produce the final output $\mathcal{F}^{s}$:

[TABLE]

Bidirectional Interaction Scan.

To maintain task information integrity during interaction with linear complexity, we propose a novel BI-Scan that serves as a cornerstone of BIM, as shown in Fig. 3. This module takes the representations of different tasks $\{\mathcal{F}_{t,N}^{s}\}_{t=1}^{T}$ as input, which is initially partitioned along the channel dimension into dual-branch representations, each with the dimensions $H\times W\times\frac{C}{2}$. The former undergoes processing through our Forward Scan module, while the latter is reversed along the task dimension before the Backward Scan processing. Specifically, the Forward Scan architecture implements hierarchical cross-task interaction through two scanning modes: Task-First mode and Position-First mode.

Task-First Mode: Following a similar configuration to SS2D in Vmamba [17], we implement multiple distinct Scan Patterns for each scanning mode to facilitate enhanced task interaction. Without loss of generality, we use the first scan pattern as an example. As shown in Fig. 3, in the task-first mode, the construction of global task information contains two processes: (1) Pattern-guided feature serialization ($\mathcal{U}_{1}$): each task feature is serialized according to the scan pattern, resulting in $T$ subsequences, each of length $H\times W$. (2) Task order-guided subsequence aggregation ($\mathcal{O}_{1}$): these subsequences are concatenated along the length dimension according to the specified task order, producing a global sequence $S$ of length $T\times(H\times W)$. These processes can be expressed as:

[TABLE]

This sequence undergoes global modeling through SSM [9], yielding task-level interaction information.

Position-First Mode: In the position-first mode, the output of Task-First mode is first restored from the sequence to the input task feature size, recorded as $\{\bar{F}_{t}^{s}\}_{t=1}^{T}\in\mathbb{R}^{H\times W\times\frac{C}{2}}$. Subsequently, the construction of interaction representations also involves two steps: (1) Task order-guided feature serialization ($\mathcal{O}_{2}$): Spatially aligned task representation tokens are constructed into $H\times W$subsequences, each of length T, according to the specified task order. (2) Pattern-guided subsequence aggregation ($\mathcal{U}_{2}$): These subsequences are concatenated along the length dimension according to the scan pattern, resulting in a sequence of length $(H\times W)\times T$. The above processes can be expressed as:

[TABLE]

This sequence is also processed through SSM, allowing for the derivation of task-level interaction insights.

Subsequently, the outputs corresponding to all patterns are fused via element-wise summation, yielding the output of Forward Scan, denoted as $\mathcal{F}_{for}^{s}$. Finally, the output of the Backward Scan, $\mathcal{F}_{back}^{s}$, is reversed along the task dimension and concatenated with $\mathcal{F}_{for}^{s}$ to yield the refined task representations $\mathcal{F}_{sh}^{s}$. In practical application, when employing BI-Scan independently, the selection of Scan Patterns aligns with those used in SS2D [17], specifically, the four directions shown in Fig. 3. When MS-Scan is introduced within the BI-Scan module, we incorporate Scan Patterns of varying scales without increasing computational complexity.

Multi-Scale Scan.

Mamba processes features by converting them into sequences of tokens. While a variety of tokenization strategies exist for text, image tokenization requires careful consideration of multiple approaches. Previous research has demonstrated that multi-scale processing is particularly effective for image data. To leverage this advantage, We propose the MS-Scan mechanism that requires minimal additional computational cost compared to single-scale scanning, as illustrated in Fig. 4 (a). In this framework, we employ multiple scanning scales, denoted as $\{s_{i}\}_{i=1}^{N}$, in multiple branches $\{\mathcal{B}_{i}\}_{i=1}^{N}$, and transform the input image feature $x\in\mathbb{R}^{C\times H\times W}$ into token sequences of varying dimensions for Mamba modeling. For instance, the initial image feature can be tokenized into a sequence with a total length of $H\times W$, where each token has a dimension of $C$. When applying a scanning scale of $s_{i}=2$, the image is divided into non-overlapping patches, resulting in a tokenized feature sequence with a dimensionality of $4C\times\left(\frac{H}{2}\times\frac{W}{2}\right)$. Specifically, MS-Scan comprises three key components: input handling, multi-scale scanning, and multi-scale fusion.

Input Handling: To construct inputs for $N$ different scanning branches $\mathcal{B}=\{\mathcal{B}_{i}\}_{i=1}^{N}$, we perform two main operations. (1) Channel Split ($\mathcal{S}$): We begin by splitting the input representation $x$ into $N$ sub-features $\{x_{i}\}_{i=1}^{N}\in\mathbb{R}^{m\times H\times W}$ along the channel dimension, where $m=\frac{C}{N}$. (2) Window Tokenization ($\mathcal{W}_{i}$): For the $i$-th branch $\mathcal{B}_{i}$ with scanning scale $s_{i}$, we first divide $x_{i}$ into $\frac{H}{s_{i}}\times\frac{W}{s_{i}}$ non-overlapping patches, each of size $m\times s_{i}\times s_{i}$. Subsequently, we concatenate the pixel feature values in each patch along the channel dimension, resulting in the scan input $\bar{x}_{i}$ with a shape of $(m\times s_{i}\times s_{i})\times\frac{H}{s_{i}}\times\frac{W}{s_{i}}$. Ultimately, we obtain the input for all branches:

[TABLE]

Multi-Scale Scanning: Following the input handling process, we employ distinct scanning scales to construct a multi-scale scene representation in each branch. For all branches $\mathcal{B}$, we utilize the four-way scanning method (SS2D) from VMamba [17] to generate scene features at the specified scale. This method creates four token sequences, each shaped as $C_{i}\times(H_{i}\times W_{i})$, by scanning the input features $\bar{x}_{i}\in\mathbb{R}^{C_{i}\times H_{i}\times W_{i}}$ in four directions. The resulting sequences are then processed by SSM [9] and combined to produce the output feature $\bar{y}=\{\bar{y}_{i}\}_{i=1}^{N}$:

[TABLE]

Multi-Scale Fusion: In our approach to multi-scale feature fusion, we adopt a methodology that reverses the input handling process, consisting of two key steps. (1) Token Windowing ($\mathcal{W}_{i}^{-1}$): For each branch $\mathcal{B}_{i}$, we split each pixel feature into $s_{i}\times s_{i}$ segments along the channel dimension, resulting in a set of feature segments $\{\bar{y}_{i,j}\}_{j=1}^{s_{i}\times s_{i}}\in\mathbb{R}^{m\times 1\times 1}$. These segments are then concatenated along the spatial dimensions (height and width) to form patches, which are subsequently combined to produce the output for $\mathcal{B}_{i}$. (2) Channel Concatenation ($\mathcal{S}^{-1}$): We concatenate the features from all branches along the channel dimension, yielding the final output feature $y\in\mathbb{R}^{C\times H\times W}$:

[TABLE]

where $\mathcal{S}^{-1}$ and $\mathcal{W}_{i}^{-1}$ refer to the inverse operation of $\mathcal{S}$ and $\mathcal{W}_{i}$ respectively.

3.4 Optimization Objective

We jointly train all tasks to optimize BIM encoder $\Phi$, BIM decoder $\Theta$, and task heads $\{H\}_{i=1}^{T}$ simultaneously. To maintain consistency with previous approaches, we use $L1$ loss for depth estimation and surface normal estimation tasks and the cross-entropy loss for other tasks, therefore, the optimization objective can be expressed as follows:

[TABLE]

where $\mathcal{T}$ is the set of all tasks, $\lambda_{t}$, $\mathcal{L}_{t}$ and $\textbf{Y}_{t}$ are the loss weight, loss function, and task label for image I in task $t$, respectively.

4 Experiment

4.1 Experimental Setup

Datasets. We performed experiments using the benchmark datasets NYUD-v2 [25] and PASCAL-Context [3]. NYUD-v2 primarily focuses on indoor scenes, with 795 and 654 RGB images for training and testing purposes. Tasks in NYUD-v2 include semantic segmentation, monocular depth estimation, surface normal estimation, and object boundary detection. PASCAL-Context encompasses indoor and outdoor scenes, offering pixel-level labels for tasks like semantic segmentation, human parsing, object boundary detection, surface normal estimation, and saliency detection tasks. The PASCAL-Context dataset contains 4,998 training images and 5,105 test images.

Implementation Details. We employ a pretrained Swin-Large Transformer [18] on ImageNet-22K [6] as our encoder. Our models are trained on the NYUD-v2 dataset for 50,000 iterations with a batch size of 4, and on the Pascal Context dataset for 75,000 iterations with a batch size of 6. Across all datasets, we use the Adam optimizer with a learning rate of $5\times 10^{-5}$ and a weight decay rate of $1\times 10^{-5}$, alongside a polynomial learning rate scheduler. The preliminary decoder has an output channel number of 768. We follow common practice [35, 16] in resizing the input images and applying data augmentation techniques. Specifically, we resize the input images of NYUD-v2 and PASCAL-Context to $448\times 576$ and $512\times 512$, respectively, and apply random color jittering, random cropping, random scaling, and random horizontal flipping.

4.2 Comparisons with State-of-the-art Methods

Main Results. Tab. 1 and Tab. 2 report a comparison of the proposed BIM against previous state-of-the-art methods, including MTmamba [16], MQTransformer [32], TaskPrompter [36],InvPT++ [37],InvPT [35], ATRC [2], MTI-Net [29], PAD-Net [31], PSD [40], PAP [42], Cross-Stitch [22] and ASTMT [20] on NYUD-V2 and PASCAL-Context dataset respectively. Notably, the previous best method, i.e., MTmamba, and our BIM are built upon the Transformer-encoder Mamba-decoder architecture with the same backbone. On NYUD-v2, the performance of Semseg is clearly boosted from the previous best, i.e., 55.82 to 57.40 (+1.52). On Pascal-Context, we achieved superior performance on all tasks compared to MTMamba.

Comparisons on Model Size and Computational Cost. In Tab. 3, we compare the FLOPs and parameter counts with other State-of-the-art methods. With the comparable computational cost, our BIM shows significant outperforms. our light-vision BIM (DBIM), which implements dilated sampling in MS-Scan, shows outperformance with reduced computational cost and parameter numbers, showcasing the advantage of these designs. Specifically, in MS-Scan, we use dilated sampling in each scan branch $\mathcal{B}$ to lighten the operation of MS-Scan. Specifically: (1) We only sample one token in a window when generating multi-scale sequences from image features, instead of all tokens. (2) When restoring the sequence to image features, we perform linear interpolation between these selected tokens. These operations do not introduce any parameters, and the computational burden is reduced because only a subset of tokens are used for modeling.

Qualitative Comparisons. We qualitatively compared our proposed BIM with the previous best-performing method, as shown in Fig. 5. Our method shows clear improvements in detail, as highlighted in the circled regions.

4.3 Ablation Study

Effectiveness of Different Components. We performed ablation experiments on BI-Scan and MS-Scan using the NYUD-V2 dataset. These experiments all used Swin-Large Transformer as encoder. We use the same STL and MTL base model as in MTMamba. (1) The term “STL” indicates that each task is equipped with a separate network, utilizing two Swin Transformer blocks in each decoder stage, and (2) the “MTL” indicates that all tasks share a common encoder. (3) “Baseline” refers to MTMamba. (4) “BI-Scan” indicates that the MS-Scan is replaced with SS2D throughout the entire model. (5) “MS-Scan” refers to the implementation of the same cross-task interaction strategy as that used in the baseline. Lastly, (6) “BIM” represents the default method. As depicted in Tab. 4, The BI-Scan improved model performance with fewer parameters, emphasizing the necessity of retaining task-specific details in cross-task interactions. Meanwhile, the MS-Scan enabled multi-granularity scene modeling, further enhancing performance.

Effectiveness of Bidirectional Scan in BI-Scan. As shown in Fig. 6 (left), we compared the performance of unidirectional (Uni) and bidirectional (Bi) cross-task scanning. The results show that the introduction of bidirectional scanning significantly enhances the cross-task modeling capability, thereby improving the overall performance of the model with only a negligible increase in computational cost.

Effectiveness of Two Scan Modes in BI-Scan. As shown in Fig. 6 (right), we conducted ablation experiments on the Task-First Scan mode and Position-First Scan mode within BI-Scan. The results indicate that both Task-First (Task) and Position-First (Position) interaction modes achieve superior performance compared to the baseline (Base) at a lower computational cost. Furthermore, employing both modes together (All) results in further performance improvements. Ultimately, we incorporate both Task-First mode and Position-First mode as the final configuration of BI-Scan for BIM.

Ablation on Task Order in BI-Scan. We experiment on NYUD-v2 to study the effect of the task order during scan interaction. In addition to fixed task orders, we also implemented a random task order modeling approach (Random), in which the order of tasks in BI-Scan is randomly shuffled during both training and testing. For clarity, we refer to the tasks Semseg, Depth, Normal, and Boundary as S, D, N, and B, respectively. As shown in Tab. 5, the overall performance of various fixed order models is relatively similar, while the Random setting exhibits a slight decrease in overall performance, which remains within an acceptable range. This indicates that our proposed BI-Scan demonstrates robustness to the task order. Ultimately, we adopt the first task order setting as the final configuration of BI-Scan for BIM.

Ablation on Scan Scale and Number of MS-Scan. We experimented with the impact of varying scanning scales and numbers on model performance, as shown in Tab. 6 and Tab. 7. Experiments were conducted on the NYUD-v2 dataset with Swin-Large Transformer as encoder. We experimented with three different settings for scale ablation: all MFRs use {1,2}, {1,3}, and {1,4} two scanning scales, respectively. For number ablation, We compared three different settings, Type 1: all three MFRs use {1} scale; Type 2: all three MFRs use {1,4} two scales; Type 3: three MFRs use {1,2}, {1,2,4} and {1,2,4,6}, respectively. The results showed that in multi-scale scanning, employing appropriate scanning scale partitions can effectively enhance overall performance. Significantly, all settings yield notable performance enhancements. Ultimately, we adopt $\{1,4\}$ scale setting for all MFRs as the final configuration.

Performance on Different Encoders. We evaluate the effect of model size on experimental performance, as presented in Tab. 8. All experiments were conducted on the NYUD-v2 dataset. We compare our method with the previous best-performing model, MTMamba, using two different encoders: Swin-Base Transformer (denoted as ‘-Base’) and Swin-Large Transformer (denoted as ‘-Large’). The results suggest that models with greater capacities typically exhibit superior performance. Furthermore, our approach has demonstrated superior performance across all encoder.

5 Conclusion

We introduce the Bidirectional Interaction Mamba (BIM) framework, a simple yet effective solution designed to address the inherent conflict between adequacy and complexity in cross-task interactions for multi-task dense prediction. We first propose a Bidirectional Interaction Scan (BI-Scan) mechanism, which constructs task information as bidirectional sequences to model cross-task interactions efficiently. By integrating task-first and position-first modes within linear complexity framework, BI-Scan preserves the integrity of task-specific information during interaction while maintaining computational efficiency. We further develop a Multi-Scale Scan (MS-Scan) mechanism, which captures scene structure information at multiple scales, alleviating the difficulty of feature learning and enhancing nuanced cross-task feature interactions. Extensive experiments show our method significantly enhances model performance.

6 Acknowledgement

This work was supported in part by National Natural Science Foundation of China under Grants 62088102, U24A20325 and 12326608, Key Research and Development Plan of Shaanxi Province under Grant 2024PT-ZCK-80, and the Public Welfare Research Program of Ningbo City under Grant 2024S063.

Appendix A Details of DBIM

To further validate the effectiveness of our method, we present Dilated BIM (DBIM), a lightweight version of BIM, which achieves superior performance with reduced computational complexity and parameter count compared to MTMamba. In DBIM, we replace MS-Scan with a more lightweight variant, DMS-Scan, which conducts sparse scanning within each scanning branch $\mathcal{B}$. Specifically, as shown in Fig. 7, we perform dilated sampling in generating multi-scale sequences from image features instead of using all tokens. When restoring sequences to image features, we perform linear interpolation. These operations do not introduce any parameters and exhibit a reduced computational burden due to sampling a subset of tokens for modeling.

Appendix B More Ablation Studies

Effect of scan mode order in BI-Scan. We performed an ablation study to assess the impact of the scanning order in BI-Scan, containing two sequences: TF $\rightarrow$ PF (Task-First mode then Position-First mode) and PF $\rightarrow$ TF (the reverse). The results in Tab. 9 show that the TF $\rightarrow$ PF setting achieves better performance. Importantly, both combined strategies surpass the performance of using either TF or PF individually, which verifies the inherent complementarity between the two scanning modes.

Model efficiency with varying task quantities. To further validate the linear complexity of our method, we conducted an experiment on the NYUD dataset with a varying number of tasks. Specifically, we benchmarked our model’s performance using 2, 3, and 4 tasks. The results, presented in Tab. 10, show that as the number of tasks increases, the incremental computational cost (GFLOPs) and the number of additional parameters both remain constant for each new task. This observation empirically verifies that our model’s complexity scales linearly with the number of tasks.

Appendix C More Visual Comparison Results

Effect of BI-Scan and MS-Scan on Task Attention Patterns. To systematically validate the efficacy of our proposed scanning mechanisms in task representation enhancement, we conduct a quantitative visual analysis of attention patterns in the final MFR block, as illustrated in Figure 8. Our comparative investigation examines three distinct feature configurations: (a) Task-specific features (SP), (b) Features enhanced by MS-Scan in MSST block (W/ MS), and (c) Features refined via the BI-Scan in BCFR block (W/ BI). The results reveal two critical insights: First, the MS-Scan mechanism substantially enriches task-specific feature details through multi-scale contextual integration. Second, the BI-Scan induces task-aligned attention redistribution, effectively suppressing irrelevant spatial responses while amplifying task-critical regions. This synergistic effect is demonstrated by the progressive attention focusing observed in white-circled areas of Fig. 8.

Effect of Bidirectional Scan on Task Attention Patterns. To validate the bidirectional scan efficacy in the BI-Scan mechanism, we conduct a controlled comparative study of task attention patterns in the final MFR block under two experimental configurations: (1) Unidirectional modeling (w/o BD) with forward scanning only, and (2) Bidirectional modeling (w/ BD), while maintaining identical feature channel dimensions for fair comparison. As shown by the white-circled regions in Fig. 9, in the upper panel exemplars, bidirectional scanning exhibits enhanced attention localization in task-critical regions while preserving the structural integrity of target areas. The lower panel reveals the model without bidirectional scanning, which erroneously groups architecturally distinct elements (doors/walls) due to a flawed understanding of the scene structure. whereas bidirectional implementation achieves precise separation through detailed and comprehensive cross-task interactions.

Qualitative Comparison with state-of-the-art method. We present more qualitative results compared with the SOTA methods. In Figs. 10, 11 and 12, the results indicate that our method generates more detailed multi-task predictions, as highlighted in the circled regions.