Adaptively Distilled ControlNet: Accelerated Training and Superior Sampling for Medical Image Synthesis

TL;DR

This paper introduces Adaptively Distilled ControlNet, a framework that accelerates training and improves medical image synthesis by using dual-model distillation and adaptive regularization, achieving state-of-the-art results on multiple datasets.

Contribution

The paper presents a novel task-agnostic distillation framework that enhances medical image synthesis efficiency and accuracy, with adaptive regularization based on lesion-background ratios.

Findings

TransUNet improves mDice/mIoU by 2.4%/4.2% on KiTS19

SANet achieves 2.6%/3.5% gains on Polyps datasets

The method enables privacy-preserving medical image generation

Abstract

Medical image annotation is constrained by privacy concerns and labor-intensive labeling, significantly limiting the performance and generalization of segmentation models. While mask-controllable diffusion models excel in synthesis, they struggle with precise lesion-mask alignment. We propose \textbf{Adaptively Distilled ControlNet}, a task-agnostic framework that accelerates training and optimization through dual-model distillation. Specifically, during training, a teacher model, conditioned on mask-image pairs, regularizes a mask-only student model via predicted noise alignment in parameter space, further enhanced by adaptive regularization based on lesion-background ratios. During sampling, only the student model is used, enabling privacy-preserving medical image generation. Comprehensive evaluations on two distinct medical datasets demonstrate state-of-the-art performance: TransUNet improves mDice/mIoU by 2.4%/4.2% on KiTS19, while SANet achieves 2.6%/3.5% gains on Polyps, highlighting its effectiveness and superiority. Code is available at GitHub.