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

arXiv: https://arxiv.org/abs/2608.28078 · HTML: https://mlsanity.com/paper/2608.28078/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 43 references that are in the map.
- First posted on arXiv 28 Aug 2026.

## Builds on (most cited first, 26 of 26)

- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [Detect What You Can: Detecting and Representing Objects using Holistic Models and Body Parts](https://mlsanity.com/paper/1406.2031/) (2014)
- [PAD-Net: Multi-Tasks Guided Prediction-and-Distillation Network for Simultaneous Depth Estimation and Scene Parsing](https://mlsanity.com/paper/1805.04409/) (2018)
- [Channel-wise Knowledge Distillation for Dense Prediction](https://mlsanity.com/paper/2011.13256/) (2020)
- [MTI-Net: Multi-Scale Task Interaction Networks for Multi-Task Learning](https://mlsanity.com/paper/2001.06902/) (2020)
- [M3ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design](https://mlsanity.com/paper/2210.14793/) (2022)
- [InvPT: Inverted Pyramid Multi-task Transformer for Dense Scene Understanding](https://mlsanity.com/paper/2203.07997/) (2022)
- [Merging Multi-Task Models via Weight-Ensembling Mixture of Experts](https://mlsanity.com/paper/2402.00433/) (2024)
- [FaPN: Feature-aligned Pyramid Network for Dense Image Prediction](https://mlsanity.com/paper/2108.07058/) (2021)
- [Exploring Relational Context for Multi-Task Dense Prediction](https://mlsanity.com/paper/2104.13874/) (2021)
- [VMT-Adapter: Parameter-Efficient Transfer Learning for Multi-Task Dense Scene Understanding](https://mlsanity.com/paper/2312.08733/) (2023)
- [Multi-Task Dense Prediction via Mixture of Low-Rank Experts](https://mlsanity.com/paper/2403.17749/) (2024)
- [Polyhistor: Parameter-Efficient Multi-Task Adaptation for Dense Vision Tasks](https://mlsanity.com/paper/2210.03265/) (2022)
- [TaskExpert: Dynamically Assembling Multi-Task Representations with Memorial Mixture-of-Experts](https://mlsanity.com/paper/2307.15324/) (2023)
- [DeMT: Deformable Mixer Transformer for Multi-Task Learning of Dense Prediction](https://mlsanity.com/paper/2301.03461/) (2023)
- [Vision Transformer Adapters for Generalizable Multitask Learning](https://mlsanity.com/paper/2308.12372/) (2023)
- [Prompt Guided Transformer for Multi-Task Dense Prediction](https://mlsanity.com/paper/2307.15362/) (2023)
- [NTKMTL: Mitigating Task Imbalance in Multi-Task Learning from Neural Tangent Kernel Perspective](https://mlsanity.com/paper/2510.18258/) (2025)

## Abstract (arXiv)

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.
