Steerable Conditional Diffusion for Out-of-Distribution Adaptation in Medical Image Reconstruction

TL;DR

This work tackles the challenge of out-of-distribution degradation in diffusion-based medical image reconstruction by proposing Steerable Conditional Diffusion (SCD), a method that adapts the pretrained diffusion model during reverse sampling using a single measurement. By injecting a low-rank residual pathway via LoRA and updating only these added parameters at each step, SCD enforces data consistency without expensive full fine-tuning, preserving the original prior. Across diverse datasets and imaging modalities, SCD yields substantial improvements over baselines in OOD scenarios, including sparse-view CT, accelerated MRI, and super-resolution, while maintaining parameter efficiency and flexibility. The approach offers a practical path toward robust deployment of diffusion priors in real-world medical imaging, with potential extensions to broader measurement regimes and multi-measurement settings.

Abstract

Denoising diffusion models have emerged as the go-to generative framework for solving inverse problems in imaging. A critical concern regarding these models is their performance on out-of-distribution tasks, which remains an under-explored challenge. Using a diffusion model on an out-of-distribution dataset, realistic reconstructions can be generated, but with hallucinating image features that are uniquely present in the training dataset. To address this discrepancy during train-test time and improve reconstruction accuracy, we introduce a novel sampling framework called Steerable Conditional Diffusion. Specifically, this framework adapts the diffusion model, concurrently with image reconstruction, based solely on the information provided by the available measurement. Utilising our proposed method, we achieve substantial enhancements in out-of-distribution performance across diverse imaging modalities, advancing the robust deployment of denoising diffusion models in real-world applications.

Figures (12)

• Figure 1: Conditional sampling with diffusion models for out-of-distribution data for sparse-view computed tomography with $60$ angles. Left: Ground truth image. Middle: Sample with diffusion model trained on synthetic ellipses. Right: Sample with diffusion model trained on CT images. For the conditional sampling we made use of DDS chung2024decomposed, more details in Section \ref{['sec:results']}. The artefacts in the middle image are due to the mismatch of the ground truth and the training dataset.
• Figure 2: An illustration of the Steerable Conditional Diffusion (SCD) sampling process. In addition to the measurement consistency steps (green), SCD includes an adaptation step (red) to fine-tune the diffusion model on the provided data.
• Figure 3: DPS, Red-diff, DDS, DIP+TV and SCD (ours) are compared to reconstruct real-measured $\mu$CT data of a walnut from $60$ angles and $128$ detector pixels. The diffusion model was trained on Ellipses. The non-adaptation methods clearly show ellipsoid artefacts in this OOD scenario. Top: full image. Bottom: zoomed-in part.
• Figure 4: Varying the number of angles for $\mu$CT reconstruction for SCD and DDS from $30$ to $120$ angles. For all experiments, SCD is able to outperform DDS on this OOD task. Both SCD and DDS were tuned for $60$ angles and then applied to the different sparse-view settings.
• Figure 5: Left: Results for sparse view CT for training on Ellipses and testing on AAPM. Right: Results for sparse view CT for training on AAPM to testing on Ellipses. We compare SCD against DDS and the filtered-back projection.