Hybrid diffusion models: combining supervised and generative pretraining for label-efficient fine-tuning of segmentation models

TL;DR

This paper introduces a hybrid diffusion pretraining method combining supervised and self-supervised learning to improve label-efficient segmentation model fine-tuning, demonstrating superior results across multiple datasets.

Contribution

The paper proposes a novel pretext task that fuses denoising and mask prediction, effectively creating a generative model for images and segmentation masks, enhancing domain adaptation.

Findings

Hybrid pretraining outperforms pure supervised or unsupervised methods.

The approach improves segmentation accuracy with limited labeled data.

Empirical results show consistent gains across datasets.

Abstract

We are considering in this paper the task of label-efficient fine-tuning of segmentation models: We assume that a large labeled dataset is available and allows to train an accurate segmentation model in one domain, and that we have to adapt this model on a related domain where only a few samples are available. We observe that this adaptation can be done using two distinct methods: The first method, supervised pretraining, is simply to take the model trained on the first domain using classical supervised learning, and fine-tune it on the second domain with the available labeled samples. The second method is to perform self-supervised pretraining on the first domain using a generic pretext task in order to get high-quality representations which can then be used to train a model on the second domain in a label-efficient way. We propose in this paper to fuse these two approaches by introducing a new pretext task, which is to perform simultaneously image denoising and mask prediction on the first domain. We motivate this choice by showing that in the same way that an image denoiser conditioned on the noise level can be considered as a generative model for the unlabeled image distribution using the theory of diffusion models, a model trained using this new pretext task can be considered as a generative model for the joint distribution of images and segmentation masks under the assumption that the mapping from images to segmentation masks is deterministic. We then empirically show on several datasets that fine-tuning a model pretrained using this approach leads to better results than fine-tuning a similar model trained using either supervised or unsupervised pretraining only.