Multi-Channel Uncertainty-Weighted Score Matching for Conditional Diffusion in Medical UDA
Abstract
Robust medical image segmentation across modalities remains challenging due to severe domain shifts and the lack of targetdomain labels. While diffusion models have been explored for crossdomain generation and augmentation, target-domain conditional diffusion training typically relies on highly noisy pseudo masks; naively conditioning on a single Arg-Max pseudo-label can corrupt diffusion training and downstream segmentation. We propose UPDiff-UDA, a unified UDA framework whose core is an uncertainty-guided training objective for target-domain conditional diffusion. Given an imperfect sourcetrained segmenter, we use its per-pixel softmax distribution to form ranked pseudo-label maps (Arg-Max, Arg-2nd, Arg-3rd, ...). Each map yields a conditional score estimate, and we aggregate them via pixel-wise confidence weighting to obtain an uncertainty-reweighted score for score matching, improving robustness to pseudo-label noise while leveraging alternative plausible labels in uncertain regions. We further provide a theoretical justification showing that confidence-weighted aggregation follows a minimum-MSE convex-combination principle under the segmenterinduced surrogate label distribution. To improve pseudo-condition quality, we also introduce a feature-guided, low-degree-of-freedom Bézier curve adaptation to reduce appearance gaps. Experiments on multiple public datasets and modality shifts show that UPDiff-UDA generates highfidelity labeled target-style samples for augmentation and consistently outperforms strong UDA baselines. The code for this project is available at: https://github.com/superlc1995/Multi-Channel-UncertaintyDiffusion-UDA