DiffusionMTL

Learning Multi-Task Denoising Diffusion Model from Partially Annotated Data
CVPR 2024
Hanrong Ye,
Dan Xu
Department of Computer Science and Engineering, HKUST
arXiv Github
logo

DiffusionMTL is a multi-task denoising diffusion framework for improving prediction quality under a partially annotated multi-task learning setting, where the initial multi-task predictions are noisy due to the severe lack of ground-truth supervision.

Abstract

Recently, there has been an increased interest in the practical problem of learning multiple dense scene understanding tasks from partially annotated data, where each training sample is only labeled for a subset of the tasks. The missing of task labels in training leads to low-quality and noisy predictions, as can be observed from state-of-the-art methods. To tackle this issue, we reformulate the partially-labeled multi-task dense prediction as a pixel-level denoising problem, and propose a novel multi-task denoising diffusion framework coined as DiffusionMTL. It designs a joint diffusion and denoising paradigm to model a potential noisy distribution in the task prediction or feature maps and generate rectified outputs for different tasks. To exploit multi-task consistency in denoising, we further introduce a Multi-Task Conditioning strategy, which can implicitly utilize the complementary nature of the tasks to help learn the unlabeled tasks, leading to an improvement in the denoising performance of the different tasks. Extensive quantitative and qualitative experiments demonstrate that the proposed multi-task denoising diffusion model can significantly improve multi-task prediction maps, and outperform the state-of-the-art methods on three challenging multi-task benchmarks, under two different partial-labeling evaluation settings.

Partially Annotated Multi-Task Learning Setting

Each training sample contains labels for a subset of the studied tasks, rather than all tasks. As there is a lack of multi-task labels for each training sample, the partially supervised multi-task learning problem is more challenging compared to the fully supervised multi-task learning problem.

Framework

Illustration of the proposed DiffusionMTL (Prediction Diffusion). DiffusionMTL first uses an initial backbone model for producing starter prediction maps for all tasks. Then, to denoise the initial prediction maps and generate rectified maps, we propose a Multi-Task Denoising Diffusion Network (MTDNet). MTDNet involves a diffusion process and a denoising process. During training, the initial prediction map of the labeled target task is gradually degraded by applying noise, resulting in the noisy prediction map. Then, we utilize a Multi-Task Conditioned Denoiser (referred to as the “Denoiser”) to denoise iteratively over S steps, resulting in a clean prediction map that is supervised by the ground-truth label. For better learning of unlabeled tasks, we propose a Multi-Task Conditioning mechanism in the denoising process to stimulate information sharing across different tasks. During inference, the diffusion and denoising processes are applied to all tasks to produce denoised multi-task prediction maps.

Illustration of the proposed DiffusionMTL (Feature Diffusion). Different from Prediction Diffusion, Feature Diffusion conducts noise decay and denoising on initial feature maps. The denoised feature maps are projected to the final prediction map with a task head after the denoising.

Experiments

DiffusionMTL effectively denoises the noisy prediction maps of both depth estimation and semantic segmentation maps on Cityscapes dataset.

Our DiffusionMTL is able to rectify noisy input and generate clean prediction maps. The model used in this comparison is trained on the Cityscapes dataset under the one-label MTPSL setting.

Quantitative Results

logo
logo
logo

BibTeX

@InProceedings{diffusionmtl,
  title={DiffusionMTL: Learning Multi-Task Denoising Diffusion Model from Partially Annotated Data},
  author={Ye, Hanrong and Xu, Dan},
  booktitle={CVPR},
  year={2024}
}
The website template was adapted from SegGen.