# StableMTL: Repurposing Latent Diffusion Models for Multi-Task Learning from Partially Annotated Synthetic Datasets

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 66 of its 87 references that are in the map.
- First posted on arXiv 9 Jun 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 66)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Scalability in Perception for Autonomous Driving: Waymo Open Dataset](https://mlsanity.com/paper/1912.04838/) (2019)
- [Towards Robust Monocular Depth Estimation: Mixing Datasets for Zero-shot Cross-dataset Transfer](https://mlsanity.com/paper/1907.01341/) (2019)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Deep High-Resolution Representation Learning for Human Pose Estimation](https://mlsanity.com/paper/1902.09212/) (2019)
- [Depth Map Prediction from a Single Image using a Multi-Scale Deep Network](https://mlsanity.com/paper/1406.2283/) (2014)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [A Large Dataset to Train Convolutional Networks for Disparity, Optical Flow, and Scene Flow Estimation](https://mlsanity.com/paper/1512.02134/) (2015)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Argoverse 2: Next Generation Datasets for Self-Driving Perception and Forecasting](https://mlsanity.com/paper/2301.00493/) (2023)
- [Taskonomy: Disentangling Task Transfer Learning](https://mlsanity.com/paper/1804.08328/) (2018)
- [Hypersim: A Photorealistic Synthetic Dataset for Holistic Indoor Scene Understanding](https://mlsanity.com/paper/2011.02523/) (2020)
- [Cross-stitch Networks for Multi-task Learning](https://mlsanity.com/paper/1604.03539/) (2016)
- [3D Packing for Self-Supervised Monocular Depth Estimation](https://mlsanity.com/paper/1905.02693/) (2019)
- [Semi-Supervised Semantic Segmentation with Cross-Consistency Training](https://mlsanity.com/paper/2003.09005/) (2020)
- [Which Tasks Should Be Learned Together in Multi-task Learning?](https://mlsanity.com/paper/1905.07553/) (2019)
- [Virtual KITTI 2](https://mlsanity.com/paper/2001.10773/) (2020)
- [Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation](https://mlsanity.com/paper/2312.02145/) (2023)
- [GeoWizard: Unleashing the Diffusion Priors for 3D Geometry Estimation from a Single Image](https://mlsanity.com/paper/2403.12013/) (2024)
- [DIODE: A Dense Indoor and Outdoor DEpth Dataset](https://mlsanity.com/paper/1908.00463/) (2019)
- [YouTube-VOS: Sequence-to-Sequence Video Object Segmentation](https://mlsanity.com/paper/1809.00461/) (2018)
- [Pareto Multi-Task Learning](https://mlsanity.com/paper/1912.12854/) (2019)
- [Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout](https://mlsanity.com/paper/2010.06808/) (2020)
- [PandaSet: Advanced Sensor Suite Dataset for Autonomous Driving](https://mlsanity.com/paper/2112.12610/) (2021)
- [Lotus: Diffusion-based Visual Foundation Model for High-quality Dense Prediction](https://mlsanity.com/paper/2409.18124/) (2024)
- [RGBleftrightarrowX: Image decomposition and synthesis using material- and lighting-aware diffusion models](https://mlsanity.com/paper/2405.00666/) (2024)

## Built on by (most cited first, 2 of 2)

- [NexusFlow: Unifying Disparate Tasks under Partial Supervision via Invertible Flow Networks](https://mlsanity.com/paper/2512.06251/) (2025)
- [Unified Video Dense Prediction from Disjoint Data](https://mlsanity.com/paper/2607.21592/) (2026)

## Abstract (arXiv)

Multi-task learning for dense prediction is limited by the need for extensive annotation for every task, though recent works have explored training with partial task labels. Leveraging the generalization power of diffusion models, we extend the partial learning setup to a zero-shot setting, training a multi-task model on multiple synthetic datasets, each labeled for only a subset of tasks. Our method, StableMTL, repurposes image generators for latent regression. Adapting a denoising framework with task encoding, per-task conditioning and a tailored training scheme. Instead of per-task losses requiring careful balancing, a unified latent loss is adopted, enabling seamless scaling to more tasks. To encourage inter-task synergy, we introduce a multi-stream model with a task-attention mechanism that converts N-to-N task interactions into efficient 1-to-N attention, promoting effective cross-task sharing. StableMTL outperforms baselines on 7 tasks across 8 benchmarks.