Paper deep dive
SNR-Adaptive Unified Diffusion for Multi-Task Medical Image Segmentation
Jiahao Liu, Hang Wei, Shuai Wu
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 7/7/2026, 6:28:17 AM
Summary
The paper introduces UniT-Diff, a unified diffusion-based framework for multi-task medical image segmentation that consolidates semi-supervised learning, unsupervised domain adaptation, and domain generalization into a single model. It addresses the semantic collision and gradient imbalance inherent in naive joint training through three mechanisms: an 11-channel task-specific output space to prevent gradient sign reversal, SNR-Adaptive Task Conditioning (SATC) to modulate task guidance based on diffusion noise levels, and Task-Type-Aware Conditional Dropout (TTACD) to enforce vendor-agnostic inference for domain generalization. Evaluated on LA, MMWHS, and MNMS cardiac datasets, UniT-Diff outperforms independent task-specific baselines across all benchmarks simultaneously.
Entities (10)
Relation Signals (9)
MMWHS → usedfor → Unsupervised domain adaptation
confidence 95% · the Multi-Modality Whole Heart Segmentation dataset (MMWHS, unsupervised domain adaptation)
LA → usedfor → Semi-supervised learning
confidence 95% · the Left Atrium dataset (LA, semi-supervised learning)
MNMS → usedfor → Domain generalisation
confidence 95% · the Multi-Centre Cardiac dataset (MNMS, domain generalisation)
UniT-Diff → uses → SNR-Adaptive Task Conditioning (SATC)
confidence 95% · SNR-Adaptive Task Conditioning (SATC) scales the task token by the log signal-to-noise ratio of the current diffusion timestep
UniT-Diff → uses → Task-Type-Aware Conditional Dropout (TTACD)
confidence 95% · Task-Type-Aware Conditional Dropout (TTACD) permanently removes the task token for domain-generalisation inputs
UniT-Diff → builton → DiffVNet
confidence 90% · Building on DiffVNet [4] as the backbone, this work investigates how to fuse three heterogeneous cardiac datasets
Naive joint training → causes → Semantic collision
confidence 90% · Semantic collision arises because the left ventricle is foreground in MMWHS but background in LA
Naive joint training → causes → Gradient imbalance
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Clinical cardiac imaging pipelines currently deploy separate models for each dataset and modality, incurring redundant training costs and precluding knowledge sharing across anatomically related tasks. Consolidating semi-supervised learning, unsupervised domain adaptation, and domain generalisation into one model is therefore a practical necessity, yet naive joint training exposes a fundamental barrier: conflicting label semantics between datasets collapse LA Dice from 90.31\% to 83.38\%, while gradient imbalance across tasks of unequal complexity suppresses the weaker tasks throughout training. We present UniT-Diff, a unified diffusion segmentation framework that resolves these conflicts through three targeted mechanisms. An 11-channel task-specific output space physically partitions label categories, eliminating cross-task gradient sign reversal by construction. SNR-Adaptive Task Conditioning (SATC) scales the task token by the log signal-to-noise ratio of the current diffusion timestep, suppressing domain-specific bias during coarse denoising and restoring full task guidance as the signal clears. Task-Type-Aware Conditional Dropout (TTACD) permanently removes the task token for domain-generalisation inputs, routing them through a shared neutral pathway that draws on cross-dataset cardiac anatomy rather than source-vendor statistics. Under a single parameter set, UniT-Diff surpasses independently trained task-specific baselines on all three benchmarks simultaneously: +0.87\% on LA, +1.77\% on MMWHS, and +0.88\% on MNMS.
Tags
Links
- Source: https://arxiv.org/abs/2607.03103v1
- Canonical: https://arxiv.org/abs/2607.03103v1
Trouble viewing inline? Open PDF directly →
Full Text
34,521 characters extracted from source content.
Expand or collapse full text
SNR-Adaptive Unified Diffusion for Multi-Task Medical Image Segmentation Jiahao Liu1, Hang Wei2∗, Shuai Wu2∗ 1School of Telecommunications Engineering, Xidian University 2School of Computer Science and Technology, Xidian University 23012100032@stu.xidian.edu.cn, weihang, wushuai@xidian.edu.cn Abstract Clinical cardiac imaging pipelines currently deploy separate models for each dataset and modality, incurring redundant training costs and precluding knowledge sharing across anatomically related tasks. Consolidating semi-supervised learning, unsupervised domain adaptation, and domain generalisation into one model is therefore a practical necessity, yet naive joint training exposes a fundamental barrier: conflicting label semantics between datasets collapse LA Dice from 90.31% to 83.38%, while gradient imbalance across tasks of unequal complexity suppresses the weaker tasks throughout training. We present UniT-Diff, a unified diffusion segmentation framework that resolves these conflicts through three targeted mechanisms. An 11-channel task-specific output space physically partitions label categories, eliminating cross-task gradient sign reversal by construction. SNR-Adaptive Task Conditioning (SATC) scales the task token by the log signal-to-noise ratio of the current diffusion timestep, suppressing domain-specific bias during coarse denoising and restoring full task guidance as the signal clears. Task-Type-Aware Conditional Dropout (TTACD) permanently removes the task token for domain-generalisation inputs, routing them through a shared neutral pathway that draws on cross-dataset cardiac anatomy rather than source-vendor statistics. Under a single parameter set, UniT-Diff surpasses independently trained task-specific baselines on all three benchmarks simultaneously: +0.87% on LA, +1.77% on MMWHS, and +0.88% on MNMS. Index Terms—Medical Image Segmentation, Diffusion Models, Multi-task Learning, Domain Generalization. I Introduction Medical image segmentation underpins quantitative analysis, surgical planning, and longitudinal monitoring across modern clinical workflows [1]. In practice, a hospital’s imaging pipeline must contend with MRI acquired from multiple vendors, CT collected under varying contrast protocols, and multi-centre cohorts whose annotation conventions differ organ by organ. The prevailing response is to maintain a separate model for each dataset and modality—a “one-model-per-task” paradigm whose storage overhead, version management burden, and inability to share anatomical knowledge across related structures scale poorly as data diversity grows. Consolidating semi-supervised learning (SSL), unsupervised domain adaptation (UDA), and domain generalisation (DG) into a single deployable parameter set is therefore not merely an academic exercise but a practical clinical necessity. Diffusion probabilistic models have demonstrated strong anatomical priors in single-task medical segmentation [2, 3], and the generic semi-supervised diffusion framework DiffVNet [4] has established state-of-the-art results on each of the three tasks in isolation. The central question of this work is whether these single-task strengths can be preserved—and ideally reinforced—when all three datasets are trained jointly within DiffVNet’s diffusion prior. A naive answer is to merge all datasets under a shared 9-channel output head. The result is immediate and severe: LA Dice collapses from 90.31%90.31\% to 83.38%83.38\%. Post-hoc gradient analysis identifies two coupled failure modes. Semantic collision arises because the left ventricle is foreground in MMWHS but background in LA; any shared output channel receives opposing gradients from the two tasks simultaneously. Gradient imbalance arises because the five-class MMWHS task produces systematically larger gradient norms than the two-class LA task, biasing shared encoder updates toward the more complex dataset. Fig. 1(a) illustrates how the three tasks span increasing levels of distribution diversity, which together make naive joint optimisation unstable. (a) SSL, UDA, and SemiDG represent increasing levels of distributional diversity between training and test data. (b) Training paradigm comparison: (i) Independent single-task models, (i) GenericSSL baseline, (i) Proposed UniT-Diff. Figure 1: Overview of the problem setting and proposed paradigm. (a) The three tasks span increasing distributional diversity, making naive joint optimisation unstable. (b) UniT-Diff consolidates all three tasks into a single parameter set, enabling positive knowledge transfer across SSL, UDA, and SemiDG. Beyond these structural conflicts, task conditioning itself introduces a subtler tension. A fixed task token injected uniformly throughout the denoising chain biases the model toward source-domain statistics at every noise level, even when the latent representation is too coarse to benefit from task-specific guidance. Moreover, for domain-generalisation inputs—where the test scanner is never seen during training—the task token associates samples with source-vendor statistics rather than generalisable anatomy, actively degrading out-of-distribution accuracy. Task conditioning is therefore simultaneously a necessity for cross-modal alignment and a liability for cross-centre generalisation, depending on which task is being processed. Building on DiffVNet [4] as the backbone, this work investigates how to fuse three heterogeneous cardiac datasets into a single joint training pipeline without negative transfer. Fig. 1(b) summarises the resulting framework, UniT-Diff, relative to prior paradigms. The contributions are threefold. • 11-channel unified label space for joint optimisation. For the first time in diffusion-based medical image segmentation, we demonstrate that SSL, UDA, and DG can be optimised jointly within a single parameter set. By assigning each task a non-overlapping output interval, we physically eliminate cross-task gradient sign reversal, enabling truly simultaneous training without task-specific network branches or sequential task switching. • SNR-Adaptive Task Conditioning (SATC). The task embedding is gated by the log signal-to-noise ratio of the current diffusion timestep via a per-task learnable temperature, suppressing domain-specific bias during coarse denoising and restoring full task guidance as the signal-to-noise ratio increases. • Task-Type-Aware Conditional Dropout (TTACD). Token dropout probabilities are assigned according to each task’s learning objective—20%20\% for SSL (LA) combined with loss re-weighting to compensate for limited labeled data, 0%0\% for UDA (MMWHS) to preserve cross-modal alignment, and 100%100\% for DG (MNMS) to enforce vendor-agnostic inference through the shared neutral pathway. Evaluated on LA, MMWHS, and MNMS, UniT-Diff surpasses the independently trained DiffVNet baselines on all three benchmarks simultaneously under a single parameter set: +0.87%+0.87\% on LA, +1.77%+1.77\% on MMWHS, and +0.88%+0.88\% on MNMS. I Related Work I-A Semi-supervised Medical Image Segmentation Reducing annotation dependence in medical segmentation has two dominant research lines. Consistency-based and pseudo-labeling methods [6, 18, 15] regularize unlabeled predictions via perturbation invariance and label propagation, achieving strong single-task performance under low-label regimes. However, all these methods rely on a fixed, homogeneous label space, and cannot resolve annotation conflicts across heterogeneous multi-dataset settings. Diffusion models have recently emerged as a powerful semi-supervised segmentation framework. Our backbone DiffVNet [4] frames segmentation as iterative mask denoising, achieving state-of-the-art results via its strong anatomical prior. Yet it assumes a unified label space throughout training, which becomes a critical barrier when combining datasets with conflicting class definitions. I-B Domain Adaptation and Generalisation Cross-domain cardiac segmentation follows two complementary paradigms. Unsupervised domain adaptation (UDA) aligns labeled source and unlabeled target modalities, with recent diffusion-based methods [29, 30] leveraging DDPM priors for cross-modality alignment. Domain generalisation (DG) learns vendor-agnostic representations without access to target data at training time [31, 26]. A shared limitation across both families is that they are designed for a single organ or modality pair. No existing mechanism can selectively enable or suppress domain-specific signals based on the task objective, and their training pipelines cannot accommodate additional tasks with conflicting label conventions. I-C Multi-task Learning and Gradient Conflict Multi-task learning (MTL) trains a shared encoder across multiple objectives to improve generalisation, but suffers from conflicting gradients when tasks have opposing optimisation targets. Gradient surgery methods such as PCGrad [10] and GradNorm [11] resolve this via gradient projection or adaptive loss reweighting, and are widely used in multi-organ medical segmentation. A core, unchallenged assumption of existing MTL methods is that explicit task identity signals are uniformly beneficial. Our work directly challenges this: we empirically show that task conditioning aids cross-modal alignment in UDA, but degrades out-of-distribution generalisation in DG by encoding source-domain statistics absent at test time. Neither gradient surgery nor uniform task conditioning addresses this fundamental asymmetry, which our framework resolves via task-aware conditioning policies. I Methodology I-A Problem Formulation Let =(k,k)k=1KD=\(D_k,C_k)\_k=1^K with K=3K=3 denote three heterogeneous cardiac benchmarks: the Left Atrium dataset (LA, semi-supervised learning), the Multi-Modality Whole Heart Segmentation dataset (MMWHS, unsupervised domain adaptation), and the Multi-Centre Cardiac dataset (MNMS, domain generalisation). The class counts |k||C_k| differ across tasks: |1|=2|C_1|=2 (LA), |2|=5|C_2|=5 (MMWHS), |3|=4|C_3|=4 (MNMS), so a naive shared head of size ∑k|k|−(K−1)=9 _k|C_k|-(K-1)=9 would conflate background classes across datasets.Each kD_k contains a small labeled subset kLD_k^L and a substantially larger unlabeled subset kUD_k^U. Rather than instantiating separate parameters θk _k per task, a single model fθf_θ is trained to maximise maxθ∑k=1K(,)∼kL[Dice(fθ()[sk:ek],)], _θ\; _k=1^KE_(x,y) _k^L [Dice\! (f_θ(x)[\,s_k:e_k],\;y ) ], (1) where [sk,ek)[s_k,e_k) denotes the task-specific output channel interval. The reverse diffusion process pθ(xt−1|xt,)p_θ(x_t-1|x_t,x) reconstructs the segmentation mask conditioned on the input volume x, following the formulation of DiffVNet [4]. Fig. 1 contrasts our joint training strategy against single-task and sequential baselines. Figure 2: Architecture of the proposed UniT-Diff framework. Three heterogeneous cardiac datasets feed a shared DiffVNet backbone whose 11-channel output is partitioned into non-overlapping task intervals to eliminate semantic collision. SNR-Adaptive Task Conditioning (SATC) injects a learnable task token scaled by a per-task temperature and the instantaneous log-SNR, providing strong guidance at low noise levels while suppressing domain-specific bias during coarse denoising. Task-Type-Aware Conditional Dropout (TTACD) permanently removes the task token for the domain-generalisation stream (MNMS), routing those samples through the shared neutral pathway at inference. I-B Unified 11-Channel Label Space The overall architecture is illustrated in Fig. 2. A fundamental obstacle in joint training over heterogeneous datasets is semantic collision: the same anatomical region carries contradictory labels across tasks. The left ventricle, annotated as foreground in MMWHS, belongs to the background class in LA. Any shared output channel therefore receives opposing gradients ∇ℒLA _LA and ∇ℒMMWHS _MMWHS simultaneously, causing destructive interference that degrades both tasks. A 9-channel design that merges the background class across all tasks confirms this effect empirically: LA Dice drops from 90.31%90.31\% to 83.38%83.38\%, a degradation of 6.936.93 percentage points. To eliminate this interference by construction, the output space is expanded to 11 non-overlapping channels with exclusive task intervals: [0,2)[0,2) for LA, [2,7)[2,7) for MMWHS, and [7,11)[7,11) for MNMS. For a voxel labelled c at position (h,w,d)(h,w,d) in task k, the one-hot encoding populates only the interval [sk,ek)[s_k,e_k): yoh[b,τk(c),h,w,d]=1⇔label[b,h,w,d]=c,y_oh[\,b,\, _k(c),\,h,w,d\,]=1\;\; \;\;label[\,b,h,w,d\,]=c, (2) and the diffusion prior is scaled to xstart=yoh×2−1∈−1,+111x_start=y_oh× 2-1∈\-1,+1\^11. Because each task’s softmax operates exclusively on its own channels, gradient sign reversal is structurally prevented. One consequence of Eq. (2) is that channels outside [sk,ek)[s_k,e_k) are padded with −1-1 in xstartx_start, encoding them as “background” rather than “unlabelled”—a prior mismatch whose effect on generative capacity is discussed in Sect. IV-F. Two implementation details are necessary to realise the benefit of this design in practice. First, three independent loss instances must be maintained per task; sharing a single loss object across batches of varying channel dimensionality corrupts internally cached class-weight statistics, producing erratic gradient magnitudes that destabilise convergence. Second, each task retains an independent Difficulty monitor [4] that tracks per-class Dice progress and dynamically adjusts the supervised loss weight, preserving the adaptive difficulty-aware training of the single-task baseline within the unified framework. I-C SNR-Adaptive Task Conditioning Channel decoupling resolves label-space conflicts but leaves the shared encoder without any cue as to which input modality it is processing. Adding a fixed task token to the timestep embedding is the natural remedy, yet uniform injection across all noise levels is theoretically unsound. At high noise levels the latent state retains only coarse, low-frequency anatomy; a strong task signal at this stage biases the denoising trajectory toward training-domain statistics before any task-relevant fine structure is recoverable. We quantify the noise level via SNR(t)=α¯t1−α¯t,SNR(t)= α_t1- α_t, (3) where α¯t α_t is the cumulative noise schedule [2]. The proposed SNR-Adaptive Task Conditioning (SATC) gates the task embedding by the log-SNR: hcond=htime+σ(τk⋅logSNR(t))⏟λk(t)⋅htask,h_cond=h_time+ σ\! ( _k· (t) )_ _k(t)· h_task, (4) where htimeh_time is the sinusoidal timestep embedding, htask∈ℝdh_task ^d is a learnable task token, and τk _k is a per-task learnable temperature. The gate satisfies λk(t)→0 _k(t)\!→\!0 as t→Tt\!→\!T (shared anatomy dominates under high noise) and λk(t)→1 _k(t)\!→\!1 as t→0t\!→\!0 (task-specific detail is fully injected near the clean signal). The per-task parameterisation of τk _k is central to this design. MMWHS requires aggressive cross-modality alignment—CT and MRI intensity distributions differ substantially—and converges to a steeper gate. LA, operating entirely within single-centre MRI, benefits from a shallower slope that avoids over-specialising the shared encoder. A shared scalar temperature would allow the larger gradient of the five-class MMWHS task to dictate the conditioning profile of the two-class LA task, a coupling that consistently harms LA performance in ablation. This formulation draws on classifier-free guidance [5], where conditioning strength is itself a trainable quantity, and extends it to the multi-task regime by making that strength task-dependent and noise-level-aware simultaneously. I-D Task-Type-Aware Conditional Dropout SATC regulates the temporal conditioning profile, but a second, qualitatively distinct failure mode persists for domain generalisation. At MNMS inference, any task embedding associates inputs with the source vendors (B, C, D) seen during training; since the target vendor (A) is never observed, this vendor-specific association provides no useful inductive bias and actively degrades out-of-distribution accuracy. The core tension is that task conditioning is simultaneously necessary for cross-modal alignment and harmful for cross-centre generalisation, depending on which task is being processed. Task-Type-Aware Conditional Dropout (TTACD) resolves this tension by assigning each task a deterministic token-dropout probability grounded in its learning objective. Unlike the uniform dropout of classifier-free guidance [5], which treats all conditions symmetrically, TTACD differentiates by task type: TABLE I: Task-Type-Aware Conditional Dropout (TTACD) probabilities. Dataset Type pdropp_drop Rationale LA SSL 20% Reinforces neutral path; mild regulariser for the 8-scan labeled set MMWHS UDA 0% Full conditioning for CT↔ alignment MNMS DG 100% Enforces vendor-agnostic inference at test time For MMWHS, any suppression of the task token eliminates the cross-modality alignment signal central to UDA performance. For MNMS, permanent suppression—applied identically during training and inference—forces segmentation through the neutral pathway, the shared cardiac representation accumulated from both LA and MMWHS, without access to source-vendor statistics. The 20%20\% rate for LA serves a dual purpose: it periodically reinforces the neutral pathway with single-centre MRI samples, improving the quality of that path for MNMS inference, and acts as a mild regulariser on the limited labeled set. Task loss re-weighting. A complementary challenge arises from gradient imbalance across tasks. LA provides only 8 labeled volumes against MMWHS’s 20; the effective gradient contribution of LA—further reduced by 20% dropout episodes—risks being overwhelmed by the more data-rich MMWHS. A per-task scalar weight wkw_k is therefore applied to the composite loss before back-propagation: ℒtotalk=wk(ℒdenok+ℒdiffk+μ(e)ℒuk),L^k_total=w_k (L^k_deno+L^k_diff+μ(e)\,L^k_u ), (5) where μ(e)μ(e) is a sigmoid ramp-up function [4] and wLA=1.5w_LA=1.5, wMMWHS=wMNMS=1.0w_MMWHS=w_MNMS=1.0. The weight 1.51.5 is selected by grid search on the validation set over wLA∈1.0,1.5,2.0,3.0w_LA∈\1.0,1.5,2.0,3.0\; values beyond 1.51.5 recover LA accuracy at the cost of MNMS generalisation, reflecting the fundamental trade-off between task-specific performance and shared representation quality. Crucially, re-weighting rescales loss magnitudes before differentiation, preserving the natural gradient directions that SATC and TTACD are calibrated to operate on—a property that would be violated by gradient-projection methods such as PCGrad [10]. Together, the four components—11-channel physical decoupling, SNR-adaptive temporal gating (SATC), task-type-aware token dropout (TTACD), and gradient-balanced loss re-weighting—constitute an integrated optimisation strategy that addresses semantic, representational, and gradient-level conflicts within a single parameter set, without introducing task-specific subnetworks or architectural branching. IV Experiments IV-A Datasets and Evaluation Protocol Three public cardiac benchmarks cover distinct clinical learning paradigms. LA [12]: 100 gadolinium-enhanced MR scans (80/20 train-test split); the standard 10%-labeled SSL protocol (8 labeled, 72 unlabeled) is adopted. MMWHS [13]: 20 labeled MR and 20 unlabeled CT volumes for MR→ UDA; targets are ascending aorta (A), left atrium cavity (LAC), left ventricle cavity (LVC), and myocardium (MYO). MNMS [14]: four scanner vendors (Domains A–D); Domains B, C, D provide 5% labeled training data and Domain A is the unseen test target. Dice (%) and Jaccard (%) measure volumetric overlap; 95HD (m) and ASD (m) measure boundary accuracy. IV-B Implementation Details UniT-Diff is built on DiffVNet [4] with an 11-channel output head and trained for 300 epochs on a single NVIDIA RTX 5090 (24 GB) using SGD (momentum 0.9, weight decay 3×10−53×10^-5, lr 0.01, poly schedule) with AMP. A multi-dataloader samples one mini-batch per task per iteration; dataset-length imbalance is corrected by sample repetition. Patch sizes are 112×112×80112×112×80 (LA), 128×128×128128×128×128 (MMWHS), 32×128×12832×128×128 (MNMS). Hyper-parameters τk _k and wLA=1.5w_LA=1.5 follow the ablation in Sect. I-D. IV-C Comparison with State-of-the-Art IV-C1 Semi-Supervised Learning on LA TABLE I: Results on LA (SSL, 10% labeled). Best in bold. Method Dice↑ Jaccard↑ 95HD↓ ASD↓ MC-Net (2021) [15] 87.62 78.25 10.03 1.82 S-Net (2022) [16] 88.55 79.62 7.49 1.90 Simcvd (2022) [17] 89.03 80.34 8.34 2.59 BCP (2023) [18] 89.62 81.31 6.81 1.76 MLRPL (2024) [19] 89.86 81.68 6.91 1.85 UGPL (2025) [28] 89.95 81.81 5.73 1.79 GenericSSL (2024) [4] 90.31 82.40 5.55 1.64 UniDeg (2025) [27] 91.07 83.67 4.96 1.65 UniT-Diff (Ours) 91.18 83.85 4.82 1.53 UniT-Diff reaches 91.18%91.18\% Dice (Table I), surpassing the single-task baseline GenericSSL by 0.87p0.87\,p and the unified competitor UniDeg by 0.11p0.11\,p. The 0.71m0.71\,m reduction in 95HD suggests that cross-task cardiac priors from MMWHS and MNMS sharpen boundary localisation beyond single-task capacity. IV-C2 Unsupervised Domain Adaptation on MMWHS TABLE I: Results on MMWHS (UDA, MR→ ). Method Dice (%)↑ Avg↑ ASD↓ A LAC LVC MYO SIFA (2020) [33] 81.3 79.5 73.8 61.6 74.1 7.0 DSFN (2020) [20] 84.7 76.9 79.1 62.4 75.8 N/A DSAN (2021) [21] 79.9 84.8 82.8 66.5 78.5 5.9 LMISA-3D (2022) [22] 84.5 82.8 88.6 70.1 81.5 2.3 Diffuse-UDA (2024) [29] – – – – 88.30 1.60 ReTri (2026) [30] – – 91.99 77.53 84.76 0.856 GenericSSL (2024) [4] 93.2 89.5 91.7 86.2 90.10 1.7 UniDeg (2025) [27] 85.4 92.9 91.0 95.1 91.10 1.60 UniT-Diff (Ours) 89.8 92.1 92.0 93.6 91.87 1.60 UniT-Diff achieves 91.87%91.87\% average Dice (Table I), a 1.77p1.77\,p gain over GenericSSL and 0.77p0.77\,p over UniDeg. MYO rises from 86.2%86.2\% to 93.6%93.6\%, consistent with the thin myocardial wall benefiting from the shared denoising prior across all three datasets. A declines from 93.2%93.2\% to 89.8%89.8\%: the aortic arch borders the pulmonary trunk, a region unannotated in LA and MNMS, leaving no corrective gradient signal in the unified setting (see Sect. IV-F). IV-C3 Domain Generalisation on MNMS TABLE IV: Results on MNMS (DG, 5% labeled, Dice %). Method A B C D Avg nnUNet (2021) [1] 65.30 79.73 78.06 81.25 76.09 SDNet+Aug (2019) [23] 71.21 77.31 81.40 79.95 77.47 LDDG (2020) [24] 66.22 69.49 73.40 75.66 71.29 SAML (2020) [25] 67.11 76.35 77.43 78.64 74.88 BCP (2023) [18] 73.66 79.04 77.01 78.49 77.05 DGNet (2021) [31] 72.40 80.30 82.51 83.77 79.75 TEGDA (2025) [32] – 83.78 79.34 82.04 81.72 vMFNet (2022) [26] 77.06 82.29 84.01 85.13 82.12 GenericSSL (2024) [4] 81.71 85.44 82.18 83.90 83.31 UniT-Diff (Ours) 82.56 85.26 84.16 84.78 84.19 On the unseen Domain A scanner (Table IV), UniT-Diff reaches 84.19%84.19\% average Dice, exceeding GenericSSL by 0.88p0.88\,p and vMFNet—a dedicated DG method—by 2.07p2.07\,p. Gains hold across three of four domains (+0.85+0.85, +1.98+1.98, +0.88p+0.88\,p on A, C, D); Domain B trails GenericSSL by a marginal 0.18p0.18\,p. IV-D Ablation Study TABLE V: Component ablation (Avg Dice %). Baseline: naive 9-channel joint training. 11-cls SATC TTACD wLAw_LA LA MMWHS MNMS Naive 9-cls joint 83.38 87.52 76.25 ✓ 90.29 88.94 82.34 ✓ ✓ 90.72 91.32 83.44 ✓ ✓ 90.84 91.34 83.50 ✓ ✓ ✓ 90.93 91.67 83.86 ✓ ✓ ✓ ✓ 91.18 91.87 84.19 The 9-channel baseline collapses LA to 83.38%83.38\%; 11-channel expansion alone recovers most of this loss by eliminating gradient sign reversal. SATC and TTACD act on the same task-embedding pathway and yield comparable single-component gains (+0.55+0.55/+0.64p+0.64\,p on LA, +2.38+2.38/+2.40p+2.40\,p on MMWHS), meaning either can partially compensate for the absence of the other. Their combination is synergistic rather than additive: SATC gates temporally while TTACD enforces structural pathway separation, two orthogonal effects. Loss re-weighting (wLA=1.5w_LA=1.5) closes the residual LA gap without affecting MNMS, identifying gradient imbalance between the 8-scan LA set and the larger MMWHS as the final bottleneck. IV-E Qualitative Analysis Task 2D GT 2D Pred 3D GT 3D Pred LA MMWHS MNMS Figure 3: Qualitative results (2D GT / Pred, 3D GT / Pred renders). LA: no false positives in adjacent ventricular tissue, confirming semantic collision is eliminated. MMWHS: all four substructures correctly separated; thin myocardium recovered with high fidelity (+7.4p+7.4\,p MYO vs. GenericSSL). MNMS: biventricular shape preserved on unseen Domain A, with boundary roughness only in high-gradient slices. Fig. 3 complements the quantitative results. On LA, no false positives extend into adjacent ventricular tissue, confirming 11-channel decoupling suppresses semantic collision. On MMWHS, the thin myocardial wall is reconstructed with high surface fidelity despite the MR→ gap; the ascending aorta is slightly under-extended at the arch boundary, consistent with the A regression. On MNMS, the neutral pathway produces anatomically plausible reconstructions on Domain A without test-domain exposure. IV-F Discussion Ablation results follow a consistent ordering: MMWHS gains are large and stable, MNMS gains are moderate, and LA is most sensitive to design choices. MMWHS benefits most because CT/MRI dissimilarity makes task conditioning directly effective; MNMS profits indirectly through the cross-dataset cardiac anatomy; LA sits between the two, where single-modality data limits the benefit of explicit task guidance while MMWHS gradients subtly reshape the shared encoder. The A regression (93.2%→89.8%93.2\%→89.8\% vs. GenericSSL) illustrates an inherent limitation of unified optimisation under annotation heterogeneity: the aortic arch borders the pulmonary trunk, a region neither LA nor MNMS annotates, so no corrective gradient exists in the unified setting. A second limitation is the false background prior in xstartx_start: inactive channels are padded with −1-1 rather than masked, constraining denoising capacity on 7–9 of 11 channels. Dynamic channel routing is a natural extension that would remove this constraint without restructuring the overall framework. V Conclusion UniT-Diff consolidates SSL, UDA, and DG cardiac segmentation into a single parameter set by resolving task conflicts at three levels: label space, timestep conditioning, and token dropout policy. Two findings generalise beyond the evaluated benchmarks. Task conditioning in diffusion models is not uniformly beneficial: its value depends on both the noise level and the learning objective, and withholding the token from DG inputs provides implicit regularisation through cross-dataset cardiac priors that single-task training cannot replicate. The per-task temperature in SATC converges to distinct values across tasks, offering a diagnostic measure of how much each task benefits from explicit manifold guidance. The primary limitation—inactive channels padded with −1-1 rather than masked—introduces a false background prior that constrains denoising capacity; dynamic channel routing is the natural remedy. Extending the framework to broader anatomical domains and larger task repertoires remains future work. References [1] F. Isensee, et al., “nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation,” Nature Methods, vol. 18, p. 203–211, 2021. [2] J. Ho, A. Jain, P. Abbeel, “Denoising diffusion probabilistic models,” NeurIPS, vol. 33, p. 6840–6851, 2020. [3] J. Wu, et al., “MedSegDiff: Medical image segmentation with diffusion probabilistic model,” arXiv:2211.00611, 2022. [4] L. Wang, Y. Li, “Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation,” NeurIPS, 2023. [5] J. Ho, T. Salimans, “Classifier-free diffusion guidance,” arXiv:2207.12598, 2022. [6] L. Yu, Y. Yang, S. Zhang, Q. Dou, P. Heng, “Uncertainty-aware self-ensembling model for semi-supervised 3D left atrium segmentation,” in Proc. International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2019, p. 605–613. [7] C. Chen, Q. Dou, H. Chen, J. Qin, and P. Heng, “Unsupervised bidirectional cross-modality adaptation via deeply synergistic image and feature alignment for medical image segmentation,” IEEE Transactions on Medical Imaging, vol. 39, no. 7, p. 2494–2505, 2020. [8] H. Zheng, Q. Dou, and P. Heng, “Cross-domain medical image segmentation via label-efficient domain adaptation,” IEEE Transactions on Medical Imaging, vol. 40, no. 10, p. 2724–2735, 2021. [9] Q. Liu, Q. Dou, F. Feng, and P. Heng, “FedDG: Federated domain generalization on medical image segmentation via continuous frequency space adaptation,” in Proc. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021, p. 1011–1020. [10] T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn, “Gradient surgery for multi-task learning,” Advances in Neural Information Processing Systems (NeurIPS), vol. 33, p. 5824–5836, 2020. [11] Z. Chen, V. Badrinarayanan, C.-Y. Lee, and A. Rabinovich, “Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks,” International Conference on Machine Learning (ICML), p. 794–803, 2018. [12] Z. Xiong, et al., “A global benchmark of algorithms for segmenting the left atrium from late gadolinium-enhanced cardiac magnetic resonance imaging,” Medical Image Analysis, vol. 67, p. 101832, 2021. [13] X. Zhuang and J. Shen, “Multi-scale patch and multi-modality atlases for whole heart segmentation of MRI,” Medical Image Analysis, vol. 31, p. 77–87, 2016. [14] V. M. Campello, et al., “Multi-centre, multi-vendor and multi-disease cardiac segmentation: the M&Ms challenge,” IEEE Transactions on Medical Imaging, vol. 40, no. 12, p. 3543–3554, 2021. [15] Y. Wu, M. Xu, Z. Ge, J. Cai, and L. Zhang, “Semi-supervised left atrium segmentation with mutual consistency training,” in Proc. MICCAI, 2021, p. 297–306. [16] Y. Wu, Z. Wu, Q. Wu, Z. Ge, and J. Cai, “Exploring smoothness and class-separation for semi-supervised medical image segmentation,” in Proc. MICCAI, 2022, p. 34–43. [17] C. You, R. Zhao, L. H. Staib, and J. S. Duncan, “Momentum contrastive voxel-wise representation learning for semi-supervised volumetric medical image segmentation,” in Proc. MICCAI, 2022, p. 639–652. [18] J. Bai, et al., “Bidirectional Copy-Paste for Semi-Supervised Medical Image Segmentation,” in Proc. CVPR, 2023, p. 11514–11524. [19] H. Su, et al., “MLRPL: Multi-label representation learning for semi-supervised medical image segmentation,” arXiv preprint arXiv:2401.05540, 2024. [20] D. Zou, Q. Zhu, and P. Yan, “Unsupervised domain adaptation with dual-scheme fusion network for medical image segmentation,” in Proc. IJCAI, 2020, p. 3291–3298. [21] S. Lu, et al., “Deep symmetric adaptation network for cross-modality medical image segmentation,” IEEE Transactions on Medical Imaging, vol. 41, no. 1, p. 121–132, 2021. [22] M. Jafari, et al., “LMISA: A lightweight multi-modality image segmentation network via domain adaptation,” Medical Image Analysis, vol. 81, p. 102536, 2022. [23] A. Chartsias, et al., “Disentangled representation learning for domain generalisation in cardiac segmentation,” in Proc. STACOM (MICCAI Workshop), 2019. [24] H. Li, et al., “LDDG: Learning domain disentanglement for domain generalization in medical image segmentation,” in Proc. MICCAI, 2020, p. 518–528. [25] Q. Liu, Q. Dou, and P.-A. Heng, “Shape-aware meta-learning for generalizing prostate MRI segmentation to unseen domains,” in Proc. MICCAI, 2020, p. 475–485. [26] X. Liu, et al., “vMFNet: Compositionality meets domain-generalised segmentation,” in Proc. MICCAI, 2022, p. 704–714. [27] S. Kumari and P. Singh, “A unified degradation-robust approach to SSL and UDA for 3D medical images,” in Proc. AAAI, vol. 39, no. 4, p. 4428–4436, 2025. [28] J. Li, Y. Zhang, Z. Wang, et al., “Uncertainty-guided progressive learning for semi-supervised medical image segmentation,” IEEE Trans. Med. Imaging, vol. 44, no. 11, p. 3872–3883, 2025. [29] H. Gong, Y. Wang, Y. Wang, J. Xiao, X. Wan, and H. Li, “Diffuse-UDA: Addressing unsupervised domain adaptation in medical image segmentation with appearance and structure aligned diffusion models,” arXiv preprint arXiv:2408.05985, 2024. [30] Y. Gao, X. Liu, L. Chen, et al., “ReTri: Recursive translation and consistency learning for cross-modality cardiac segmentation,” Neural Networks, vol. 169, p. 1068–1082, 2026. [31] X. Liu, S. Thermos, A. O’Neil, and S. A. Tsaftaris, “Semi-supervised meta-learning with disentanglement for domain-generalised medical image segmentation,” in Proc. Int. Conf. Med. Image Comput. Comput.-Assist. Intervention (MICCAI), p. 307–317, 2021. [32] Y. Zhou, J. Wu, W. Liao, S. Zhang, S. Zhang, and G. Wang, “TEGDA: Test-time evaluation-guided dynamic adaptation for medical image segmentation,” in Proc. Int. Conf. Med. Image Comput. Comput.-Assist. Intervention (MICCAI), p. 145–155, 2025. [33] X. Chen, Q. Dou, H. Chen, Y. Jin, and P. A. Heng, “SIFA: Unsupervised domain adaptation for cross-modality medical image segmentation via style-image feature alignment,” in Proc. Int. Conf. Med. Image Comput. Comput.-Assist. Intervention (MICCAI), p. 536–546, 2020.