Source-linked AI summary
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
TL;DR
Multi-task dense prediction must adapt shared features to differing, image-dependent task requirements without replicating a large decoder per output. MemMTL uses prototype-refined task states and sparse routing, with reported competitive performance and a favorable accuracy–efficiency trade-off, though its numerical evidence requires regeneration.
Problem
Shared vision-backbone features still require task-specific, image-dependent adaptation for heterogeneous dense prediction without duplicating large decoders.
Method
MemMTL refines image-conditioned task states with learnable prototypes, combines global state and token evidence before top-k routing, and adds routed and task-agnostic residuals to shared features.
Results
MemMTL achieves competitive multi-task performance with a favorable accuracy–efficiency trade-off on NYUD-v2 and PASCAL-Context experiments.
Takeaways & Limitations
The modular residual design is compatible with standard task objectives and keeps task-loss balancing independent of routing.
Takeaways & Limitations
The retained numerical tables predate the canonical implementation and must be regenerated before supporting empirical claims.
Abstract
from arXiv · showhide
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.
Introduction
MemMTL addresses image- and task-dependent adaptation in multi-task dense prediction by combining prototype-refined global task states with token-level routing. It uses shared routed and task-agnostic residual banks, with matched experiments specified to evaluate performance, efficiency, and component contributions.
- Motivation: The framework targets adaptation of strong shared features to semantic, geometric, and other pixel-level tasks without replicating a large decoder for every output.The motivating task set includes semantic segmentation, depth, surface normals, and boundaries.
- Method: MemMTL refines an image–task state with learnable prototypes, then combines its expert logits with token-level logits before top-k routing.Both global task context and local token evidence determine dispatched experts.
- Method: A shared routed local expert bank gives expert indices consistent meaning across tasks, while a separate task-agnostic bank supplies a common residual.Both residuals are added once to the backbone feature before task-specific prediction.
- Evaluation: The proposed evaluation covers SAM 3 and ViT-L backbones on NYUD-v2 and PASCAL-Context, measuring multi-task trade-offs, routing sparsity, expert count, and component contributions.The protocol includes component, routing, expert-count, efficiency, and intervention studies.
- Evaluation: The manuscript’s retained numerical tables predate the canonical implementation and are a rerun manifest rather than evidence for empirical claims.The numerical record must be regenerated before supporting the specified claims.
Related Work
Prior multi-task dense prediction work addresses task conflict through optimization, interaction modules, prompts, adapters, and mixture-of-experts. MemMTL focuses on sample-dependent conflict by learning conflict-aware task states while preserving a shared task-invariant path.
- Optimization: Optimization methods balance task losses or modify task gradients to address multi-task training conflicts.Representative approaches include loss balancing, gradient modification, and related optimization strategies.
- Architectures: Architectural methods use task interaction modules, prompts, adapters, and mixture-of-experts to specialize shared models.These approaches provide complementary mechanisms for task interaction and conditional computation.
- Conditional Routing: Task identity alone cannot capture sample-dependent conflicts, so MemMTL learns conflict-aware task states through prototype memory while retaining a shared path for task-invariant structure.The method targets conflicts that vary with the current sample rather than task identity alone.
Method
MemMTL adapts shared foundation features through task-state-conditioned sparse routing that combines global image-level requirements with local token evidence. A prototype memory refines task states, while task-conditioned and shared residual branches preserve specialization and cross-task structure.
- Task-state adaptation: MemMTL addresses heterogeneous dense prediction adaptation by combining task-state and token-level routing signals.The task-state router captures task and image requirements, while the visual router selects experts from local features.
- Task-specific routing: The refined state produces task-conditioned expert preferences that complement token-level top-k routing, separating global task emphasis from fine-grained spatial specialization.Token weights are zero for experts excluded by token-level top-k selection, while task-state features modulate expert outputs.
- Coarse task-state: Global experts transform pooled features into a coarse task state for each image–task pair using top-k sparse normalization.The pooled feature is projected to global expert logits, and only a sparse subset of global experts is activated.
- Prototype memory: A learnable prototype memory retrieves normalized task-state prototypes by temperature-scaled sparse softmax and refines the coarse state through a residual connection and normalization.The memory contains trainable latent task-state modes, and γ controls the retrieved context contribution.
- Shared adaptation: All tasks share the routed local expert bank, while a parallel shared branch uses uniform task-level gating and token-level routing to preserve cross-task structure.Both expert modules return residuals, and the final task feature adds the backbone feature, one routed task residual, and one task-agnostic residual.
- Training: Training uses a standard multi-task objective in which each task loss is weighted by an independent task-balancing coefficient.The routing design remains separate from the task-loss balancing formulation.
Experiments
MemMTL is evaluated on two dense-prediction benchmarks with multiple backbones, task metrics, efficiency measures, and targeted ablations. Results examine overall performance, routing sparsity, expert capacity, prototype memory, and qualitative cross-task behavior.
- Evaluation Protocol: The evaluation covers NYUD-v2 and PASCAL-Context across semantic, geometric, boundary, part-segmentation, and saliency tasks, using task-specific metrics and average relative performance drop ∆m.Metrics include mIoU, RMSE, mErr, odsF, maxF, and ∆m relative to single-task performance.
- Main Results: With ViT-L on PASCAL-Context, MemMTL uses fewer parameters and lower FLOPs than MLoRE while maintaining competitive performance across all five tasks.With SAM 3, it achieves a 0.94% average performance drop from the single-task baseline and reaches 74.68 mIoU on PartSeg and 76.70 on Bound.
- Main Results: With ViT-L on NYUD-v2, MemMTL improves over prior multi-task methods and surpasses the single-task baseline on average, achieving a positive ∆m of 1.52%.With SAM 3, it remains close to the single-task upper bound, with the remaining gap mainly on depth and surface-normal estimation.
- Ablation Study: Adding TTM, STM, and TAR yields complementary gains, with the full model improving the baseline by 4.90 mIoU, reducing depth by 0.0358, reducing normal error by 1.08, and improving boundary by 2.6.Removing TAR causes a clear drop, while STM especially benefits geometry-related tasks.
- Ablation Study: The task-state router performs best at k = 2; k = 1 is inferior, while k = 3 or k = 4 slightly degrades performance by reducing routing selectivity.The ablation uses the SAM 3 backbone on NYUD-v2.
- Ablation Study: Using 6–8 task-specific experts offers the best capacity–specialization trade-off, whereas 4 lacks capacity and 16 over-fragments routing.Six experts improve performance across tasks, eight further benefits semantic segmentation and depth, and sixteen degrades depth and surface-normal estimation.
- Ablation Study: Task-state prototype memory improves all tasks through memory updates, retrieval adds further gains, and the full model shows that both operations are complementary.The prototype-memory ablation compares task-state-aware updates and retrieval for routing.
- Qualitative Analysis: Routing analysis shows task-specific entropy patterns and expert preferences, while qualitative examples report improved region separation, part transitions, surface-normal variation, saliency completeness, and boundary continuity under occlusion.The routing observations connect evolving expert allocation with validation improvement and distinct task preferences.
Conclusion
MemMTL combines task-conditioned sparse routing with a shared task-agnostic residual path for multi-task dense prediction. Experiments on NYUD-v2 and PASCAL-Context report competitive performance, efficiency, and spatially coherent predictions across heterogeneous tasks.
- MemMTL derives an image-conditioned task state and refines it using normalized vectors retrieved from a learnable task-state prototype bank.
- The refined state is combined with token-level evidence before sparse top-k routing, jointly using global task requirements and local visual content for expert dispatch.
- A single routed expert bank is shared across tasks, while an independent task-agnostic bank supplies a common residual added with the routed path through one identity addition.
- Experiments on NYUD-v2 and PASCAL-Context show competitive multi-task performance with a favorable accuracy–efficiency trade-off and spatially coherent predictions.