Task-State Adaptation with Prototype Memory for Multi-Task Dense Prediction

Yangyang Xu; Haobo Yuan; Yuzhu Wang; Duo Su; Xi Ye; Yibo Yang; Jun Zhu

cs.CVAug 28, 2026

Task-State Adaptation with Prototype Memory for Multi-Task Dense Prediction

Authors: Yangyang Xu, Haobo Yuan, Yuzhu Wang, Duo Su, Xi Ye, Yibo Yang, Jun Zhu

Organizations: Tsinghua University · University of California, Merced · Shanghai Jiao Tong University

Abstract

Vision foundation backbones provide strong representations for dense prediction, yet a single shared feature still needs to support tasks with different, image-dependent adaptation requirements. We propose MemMTL, a multi-task dense prediction framework that estimates a compact task state from global visual context and refines it through a learnable task-state prototype memory. The refined state is converted into task-conditioned expert logits and combined with token-level logits before sparse top- $k$ selection over a local expert bank shared by all tasks. A separate task-agnostic residual bank provides a common adaptation path, and both paths are added once to the backbone feature before task-specific prediction. We specify a matched evaluation protocol on NYUD-v2 and PASCAL-Context with SAM 3 and ViT-L backbones to measure predictive quality, computational cost, and the contributions of task-state conditioning, prototype retrieval, and sparse routing. The numerical record in the present working draft predates this canonical implementation and must be regenerated before it can support empirical claims.

Figures & tables

Figure 1: Framework of MemMTL. The backbone extracts multi-scale features shared across tasks. A task-specific task-state MoE models task-dependent residuals, while a shared task-state MoE preserves cross-task structure.

Figure 2: Task-state adaptation router. Pooling feature is mapped to a sample-conditioned coarse task-state through a sparse mixture of global experts, refined by a task-state prototype memory, and then used to modulate downstream local routing.

Model	Backbone	SemSeg	PartSeg	Sal	Normal	Bound	$\Delta_{m}[\%]$ $\uparrow$	FLOPs	Params
Model	Backbone	(mIoU) $\uparrow$	(mIoU) $\uparrow$	(maxF) $\uparrow$	(mErr) $\downarrow$	(odsF) $\uparrow$	$\Delta_{m}[\%]$ $\uparrow$	(G) $\downarrow$	(M) $\downarrow$
ST baseline	HRNet18	62.23	61.66	85.08	13.69	73.06	0.00	-	-
MTI-Net	HRNet18	61.70	60.18	84.78	14.23	70.80	-2.10	161	128
ATRC	HRNet18	57.89	57.33	83.77	13.99	69.74	-4.45	216	96
DeMT	HRNet18	59.23	57.93	83.93	14.02	69.80	-3.79	-	-
ST baseline	ViT-L	81.62	72.21	84.34	13.59	76.79	0.00	-	317

Table 1: Results on PASCAL-Context. ST baseline denotes single task baseline. ‘

\downarrow

’: lower is better. ‘

\uparrow

’: higher is better.

\Delta_{m}

denotes the average per-task performance drop (the higher, the better).

Model	backbone	SemSeg	Depth	Normal	Bound	$\Delta_{m}$
Model	(mIoU) $\uparrow$	(rmse) $\downarrow$	(mErr) $\downarrow$	(odsF) $\uparrow$	[%]
ST baseline	ViT-L	56.77	0.5141	18.56	78.93	0.00
InvPT	ViT-L	53.56	0.5183	19.04	78.10	-2.52
TaskPrompter	ViT-L	55.30	0.5152	18.47	78.20	-0.81
TaskExpert	ViT-L	55.35	0.5157	18.54	78.40	-0.84
MLoRE	ViT-L	55.96	0.5076	18.33	78.43	0.11

Table 2: Results on NYUD-v2.

\uparrow

\downarrow

: higher/lower is better.

Setting	SemSeg	Depth	Normal	Bound
Setting	(mIoU) $\uparrow$	(RMSE) $\downarrow$	(mErr) $\downarrow$	(odsF) $\uparrow$
MT baseline	55.38	0.5013	18.11	77.2
w/ TTM	57.05	0.4916	18.21	78.1
w/ STM	57.42	0.4868	17.74	78.6
w/o TAR	58.96	0.4742	17.38	79.1
MemMTL	60.28	0.4655	17.03	79.8

Table 3: Ablation of the modules on NYUD-v2 with SAM 3. TTM denotes task-specific task-state MoE, STM is shared task-state MoE, and TAR is task-state adaptation router.

Panel of Figure 3 — Figure 3: (a) Routing dynamics. Task-state gate entropy reveals task-specific expert allocation patterns during training. (b) Validation trend. Performance improves with routing evolution, suggesting coupled expert allocation and multi-task optimization. (c) Task-specific expert preference. Mean gate weights show distinct expert preferences across tasks, with some tasks selecting dominant experts and others using more distributed mixtures.

Top- k	SemSeg	Depth	Normal	Bound
Top- k	(mIoU) $\uparrow$	(RMSE) $\downarrow$	(mErr) $\downarrow$	(odsF) $\uparrow$
1	59.34	0.4786	17.42	79.1
2	60.28	0.4655	17.03	79.8
3	59.86	0.4708	17.18	79.5
4	59.57	0.4749	17.31	79.3

Table 4: Ablation study of Top- k in the task-state adaptation router on NYUD-v2 using the SAM 3 backbone. Top- k denotes the number of sparse activations in the global router.

Expert number	SemSeg	Depth	Normal	Bound
Expert number	(mIoU) $\uparrow$	(RMSE) $\downarrow$	(mErr) $\downarrow$	(odsF) $\uparrow$
4	58.95	0.4725	17.29	79.8
6	60.28	0.4655	17.03	79.8
8	60.92	0.4627	17.04	79.8
16	59.46	0.4715	17.32	79.4

Table 5: Ablation on the number of experts in the task-specific task-state MoE on NYUD-v2 using the SAM 3 backbone. The number of experts is set to 6 in all other experiments.

Setting	SemSeg	Depth	Normal	Bound
Setting	(mIoU) $\uparrow$	(RMSE) $\downarrow$	(mErr) $\downarrow$	(odsF) $\uparrow$
MT baseline	55.38	0.5013	18.11	77.2
w/ task-state prototype update	58.72	0.4791	17.46	78.9
w/ task-state prototype retrieval	59.34	0.4746	17.25	79.2
MemMTL w/ TPM	60.28	0.4655	17.03	79.8

Table 6: Ablation of task-state prototype memory (TPM), including task-state-aware update and retrieval for task-state routing.

Figure 4: Qualitative comparison on the five-task PASCAL-Context benchmark with InvPT ( Ye and Xu 2022 ) , TaskPrompter ( Ye and Xu 2023b ) , and MLoRE ( Yang et al. 2024 ) . The columns show the input image and predictions for semantic segmentation, human-part segmentation, surface normals, saliency, and boundaries. Red circles indicate representative errors in competing methods, whereas yellow circles highlight regions where our predictions more closely match the ground truth, particularly in preserving semantic separation and fine-grained human-part structures. Best viewed in color and under magnification.

Multi-Task Learning

Multi-task learning trains models to solve related tasks together by sharing representations, parameters, or supervision.

See all 391 papers

DPNeXt: A Lightweight Multi-Scale Feature Fusion Framework for Efficient ViT-Based Multi-Task Dense Prediction
Jul 17, 2026Jehun Kang, Jungha Wang, Youngjun Hwang +1Multi-Task Learning Multi-Scale Feature Fusion
School of Electrical Engineering, Korea Advanced Institute of Science and Technology (KAIST), Daejeon, South Korea
Task-Instructed Causal Routing of Vision Foundation Models for Multi-Task Learning
Jun 14, 2026Donghyun Han, Yuseok Bae, Jung Uk Kim +1Vision Foundation Model Adaptation Multi-Task Learning
Electronics and Telecommunications Research Institute (ETRI), Korea · Kyung Hee University, Korea · Chonnam National University, Korea
MUSE: Unlocking Timestep as Native Task Steering for One-Step Dense Prediction
Jun 29, 2026Shuo Zhou, Zhaoxin Li, Xiujuan ChaiMonocular Depth Estimation Surface Normal Estimation
Agricultural Information Institute, Chinese Academy of Agricultural Sciences, Beijing, China. · Key Laboratory of Agricultural Big Data, Ministry of Agriculture and Rural Affairs, Beijing, China.

Jul 17, 2026J/K move · Enter open · S save

DPNeXt: A Lightweight Multi-Scale Feature Fusion Framework for Efficient ViT-Based Multi-Task Dense Prediction

Jehun Kang, Jungha Wang, Youngjun Hwang, David Hyunchul Shim

School of Electrical Engineering, Korea Advanced Institute of Science and Technology (KAIST), Daejeon, South Korea

Multi-Task Learning Multi-Scale Feature Fusion Efficient ViTs Robotic Perception

Open paper AskPDF arXiv

Multi-Task Learning (MTL) in robotics perception systems supports comprehensive 3D spatial scene understanding by integrating semantic segmentation and depth estimation. While Vision Foundation Models (VFMs) are increasingly adopted as robust feature encoders, existing decoding strategies present a critical bottleneck. To address this, we propose DPNeXt, a streamlined multi-scale feature fusion decoder and efficient alternative to the standard Dense Prediction Transformer (DPT). DPNeXt uses dual depthwise separable inverted bottlenecks to improve frozen VFM utilization through fusion-centric decoding and independent task modularization. To further mitigate negative inductive transfer between tasks, we introduce the Multi-Task Boundary Guidance (MTBG) strategy. Unlike prior boundary-aware methods that add fusion modules or gating, MTBG applies symmetric boundary-focused supervision to encourage geometric consistency without extra annotation or inference cost. Experiments on Cityscapes show that DPNeXt-S outperforms prior state-of-the-art (SOTA) MTL models, while DPNeXt-B further improves the overall performance and achieves the best results among the compared methods. On NYUv2, DPNeXt-B also achieves the best semantic segmentation and depth estimation results among the compared methods while requiring substantially fewer trainable parameters than prior large-scale MTL models. Compared with the standard DPT, DPNeXt-S reduces trainable parameters by 78.6% and achieves the fastest inference speed among the compared models on resource-constrained laptop hardware. The source code, model checkpoints, and a demo video will be made available at https://github.com/kangjehun/DPNeXt.

Abstract

Multi-Task Learning

Explore similar work

DPNeXt: A Lightweight Multi-Scale Feature Fusion Framework for Efficient ViT-Based Multi-Task Dense Prediction

Task-Instructed Causal Routing of Vision Foundation Models for Multi-Task Learning

MUSE: Unlocking Timestep as Native Task Steering for One-Step Dense Prediction