DPL: Spatial-Conditioned Diffusion Prototype Enhancement for One-Shot Medical Segmentation

One-shot medical image segmentation faces fundamental challenges in prototype representation due to limited annotated data and significant anatomical variability across patients. Traditional prototype-based methods rely on deterministic averaging of support features, creating brittle representations that fail to capture intra-class diversity essential for robust generalization. This work introduces Diffusion Prototype Learning (DPL), a novel framework that reformulates prototype construction through diffusion-based feature space exploration. DPL models one-shot prototypes as learnable probability distributions, enabling controlled generation of diverse yet semantically coherent prototype variants from minimal labeled data. The framework operates through three core innovations: (1) a diffusion-based prototype enhancement module that transforms single support prototypes into diverse variant sets via forward-reverse diffusion processes, (2) a spatial-aware conditioning mechanism that leverages geometric properties derived from prototype feature statistics, and (3) a conservative fusion strategy that preserves prototype fidelity while maximizing representational diversity. DPL ensures training-inference consistency by using the same diffusion enhancement and fusion pipeline in both phases. This process generates enhanced prototypes that serve as the final representations for similarity calculations, while the diffusion process itself acts as a regularizer. Extensive experiments on abdominal MRI and CT datasets demonstrate significant improvements respectively, establishing new state-of-the-art performance in one-shot medical image segmentation.

💡 Research Summary

One‑shot medical image segmentation suffers from two intertwined problems: the scarcity of annotated data and the large anatomical variability across patients. Traditional prototype‑based approaches address the first issue by collapsing the support features into a single deterministic prototype, typically via masked averaging. While computationally cheap, this deterministic representation cannot capture the intra‑class diversity that is crucial for robust generalization in clinical settings.

The paper introduces Diffusion Prototype Learning (DPL), a framework that reconceptualizes prototype construction as a probabilistic process driven by diffusion models. Instead of treating the prototype as a fixed point, DPL models it as a learnable distribution from which diverse, semantically coherent variants can be sampled. The method comprises three core components:

1. Diffusion‑based Prototype Enhancement – Starting from the conventional masked‑average prototype (p_0), a forward diffusion process injects Gaussian noise across a short schedule (T = 20 steps, cosine‑beta). This creates a set of noisy prototypes (p_t) that explore the neighbourhood of the original feature in a principled manner.

2. Spatial‑aware Conditioning – To preserve anatomical plausibility, DPL extracts four geometric descriptors from the prototype statistics: normalized centroid coordinates (x, y), compactness, and elongation. These are encoded by a lightweight network into a spatial condition vector (c_{\text{spatial}}). During reverse diffusion, the predicted noise (\epsilon_\theta(p_t, t)) is augmented with a time‑dependent injection of (c_{\text{spatial}}) (scaled by (\alpha_t)), providing strong spatial guidance in early denoising steps and gradually fading it as the process converges.

3. Conservative Fusion – The original prototype (p_0) and the diffusion‑refined prototype (p_{\text{diff}}) are combined through learnable weights (w_{\text{fidelity}}) and (w_{\text{diversity}}) (sigmoid‑activated). The final enhanced prototype is a weighted average that balances fidelity to the original, reliable representation with the diversity introduced by diffusion.

Training optimizes a composite loss:

• Segmentation loss ((L_{\text{seg}})) – cross‑entropy computed from pixel‑wise similarity between query features and the enhanced prototype.
• Alignment loss ((L_{\text{align}})) – L2 distance between support‑enhanced and query‑enhanced prototypes, encouraging consistency across domains.
• Diffusion loss ((L_{\text{diffusion}})) – standard DDPM denoising objective that forces the noise predictor to accurately recover the clean prototype from noisy versions, conditioned on the spatial vector.

The total loss is (L_{\text{total}} = L_{\text{seg}} + L_{\text{align}} + \beta L_{\text{diffusion}}) with (\beta = 0.02), a value determined by ablation to provide regularization without overwhelming the primary segmentation objective. Crucially, the same diffusion‑enhancement and fusion pipeline is employed during both training and inference, guaranteeing consistency between the two phases.

Experiments were conducted on two widely used abdominal datasets: a 3‑D MRI collection (20 scans, four organs) and a 3‑D CT collection (30 scans, same organ set). Following a strict few‑shot protocol where test classes are completely unseen during training, the authors trained DPL for three epochs (≈30 k iterations) on a single RTX 4090 GPU. Performance was measured with Dice Similarity Coefficient (DSC).

Results show that DPL outperforms state‑of‑the‑art one‑shot methods such as PANet, ALPNet, Q‑Net, and SE‑Net by 4–7 percentage points in average DSC. The gains are especially pronounced for organs with high shape variability (kidneys, liver, spleen). Ablation studies reveal: (i) removing spatial conditioning reduces DSC by ~4 pp, confirming the importance of anatomical guidance; (ii) replacing the conservative fusion with a simple average degrades performance, highlighting the need to preserve the original prototype’s reliability; (iii) increasing the diffusion regularization weight ((\beta)) leads to excessive noise and a drop in segmentation quality. Visualizations of prototype heatmaps across diffusion timesteps illustrate how the forward process introduces controlled variance while the reverse process, guided by spatial cues, restores smooth, anatomically plausible features.

The authors discuss limitations: the current implementation operates on 2‑D feature maps extracted from 3‑D volumes, incurring additional memory and compute overhead for full‑volume diffusion; the geometric descriptors are simple statistics and may not capture complex pathological deformations; and the method has been evaluated only in a strict 1‑shot setting. Future work is suggested in extending DPL to multi‑shot and multi‑institution scenarios, integrating more sophisticated shape priors, and designing lightweight diffusion modules for real‑time clinical deployment.

In summary, DPL demonstrates that diffusion models can be repurposed from data generation to feature‑space exploration, providing a principled mechanism to endow one‑shot medical segmentation prototypes with the diversity needed to handle anatomical variability while maintaining spatial coherence. The approach sets a new performance benchmark for one‑shot medical image segmentation and opens a promising research direction at the intersection of diffusion modeling and prototype‑based few‑shot learning.