Source-linked AI summary

Unsupervised Anatomical Feature Learning via Diffusion Models: Enhanced Medical Image Segmentation with Denoising Diffusion Probabilistic Models

Akshat G, Divyansh Gupta, Shaleen Bhatnagar, Shilpa Ankalaki, Tusar Kanti Mishra

arXiv:2608.25693v1cs.CVcs.AIcs.LG

TL;DR

Medical segmentation needs costly expert annotations, and conventional U-Nets can lack global anatomical structure. The paper pretrains DDPM encoders on unlabeled abdominal CT images and transfers them to segmentation, reporting improved accuracy, boundary precision, variance, and label efficiency, with scope limited by organ and single-institution evaluation.

  • Problem

    Medical image segmentation requires substantial expert-labeled data, while conventional models learn from random initialization and need large labeled datasets.

  • Method

    The study uses unsupervised DDPM denoising to learn anatomical representations and transfers the resulting encoder weights to downstream segmentation.

  • Results

    Diffusion pretraining improved segmentation across overlap and boundary metrics, reduced variance by 47–68%, and frozen encoders retained over 80% of fine-tuned liver performance.

  • Takeaways & Limitations

    Unlabeled diffusion pretraining can embed anatomical features before supervised optimization and support data-efficient, anatomy-aware segmentation.

  • Takeaways & Limitations

    Evaluation focused on large, well-contrasted liver and kidney structures in BTCV data from a single institution, leaving broader generalization unvalidated.

Abstract

from arXiv · show

Acquiring pixel-level annotations for medical image segmentation is a severe bottleneck. Traditional U-Net architectures, while effective, learn local texture patterns and lack awareness of global anatomical structures, leading to boundary delineation failures in low-data regimes. This research paper proposes utilizing unsupervised Denoising Diffusion Probabilistic Models (DDPMs) to extract anatomical features. We train a DDPM on 21 unlabeled abdominal CT scans to learn structural representations, transferring the encoder weights to a downstream segmentation task evaluated on the BTCV multi-organ dataset. Diffusion pretraining significantly improved liver segmentation: Dice increased from $0.75\pm0.36$ to $0.93\pm0.16$ ($p < 5.33\times10^{-26}$, 0.529 Cohen's d), Average Surface Distance (ASD) decreased by 66%, and 95th-percentile Hausdorff Distance (HD95) reduced by 45%. For kidney segmentation, Dice improved from $0.90\pm0.19$ to $0.95\pm0.10$ ($p < 4.01\times10^{-11}$). Multi-organ pooled performance showed a 68% variance reduction and a 74% improvement in boundary precision (Dice $0.95\pm0.07$). Crucially, frozen encoder models retained > 80% of fine-tuned performance without exposure to segmentation labels, proving the existence of learned anatomical priors. In low-data scenarios, diffusion-pretrained models maintained robust performance with only 50% (Dice: 0.92 liver, 0.94 kidney), 25%, and even 10% (Dice: 0.89 liver, 0.71 kidney) of labeled data. Using unlabeled images for diffusion-based pretraining successfully embeds robust anatomical features prior to human supervision, transforming U-Nets into anatomy-aware systems.

I. INTRODUCTION

Medical image segmentation is constrained by costly, variable expert annotations, while standard U-Nets primarily learn local texture patterns and lack global anatomical context. The study proposes DDPM pretraining to learn structural representations before supervised segmentation.

  • Motivation: 30–60 minutes and over $100 per volume can be required to annotate a 100–300-slice CT scan.Annotation burden is especially acute because trained radiologists must trace structures across three-dimensional volumes.
  • Motivation: Expert boundary disagreement, particularly for diffuse-edged organs such as the liver, introduces inconsistent supervision and limits reliability.Inter-labeler variability propagates into the training data.
  • Existing limitations: Randomly initialized U-Nets learn local textures, edge orientations, intensity gradients, and co-occurrence statistics through pixel-wise classification loss.These learned features do not provide an anatomical prior during initialization.
  • Existing limitations: High HD95 can reveal missing lobes, incorrect attachments, or disconnected regions despite reasonable overall segmentation.Boundary errors may therefore remain hidden when evaluating overlap alone.
  • Existing limitations: Convolutional locality and random initialization leave U-Nets without prior knowledge of anatomical topology, paired organs, or predictable size ranges.This limits global spatial context for segmentation.
  • Proposed direction: DDPM pretraining uses annotation-free denoising to learn abdominal structural patterns, then transfers encoder weights to downstream segmentation.The approach targets anatomy-aware boundary prediction in limited-label settings rather than image generation at deployment.
  • Study contributions: Diffusion pretraining is reported to improve segmentation accuracy and boundary precision, reduce prediction variance by 47–68%, and preserve over 80% of fine-tuned liver performance with frozen encoders.These contributions are presented as evidence for learned anatomical priors.

B. SELF-SUPERVISED REPRESENTATION LEARNING

Self-supervised methods learn from unlabeled images through reconstruction or invariance objectives, but some may discard spatial or boundary information needed for dense prediction. This study instead transfers encoder features learned through diffusion denoising, using domain-specific abdominal CT preprocessing and patient-level splitting.

  • Self-supervised learning: Self-supervised learning defines pretext tasks that learn representations from unlabeled data without manual annotation.Examples include contrastive learning, autoencoding, and masked reconstruction.
  • Self-supervised learning: Contrastive learning may discard fine-grained spatial information, standard autoencoders may reproduce texture and intensity, and MAE masking may miss subtle boundaries.These limitations are relevant to dense medical image prediction.
  • Diffusion pretraining: Diffusion denoising learns hierarchical representations from fine to global scales while encouraging spatial relationships and structural coherence.This objective is presented as advantageous for segmentation pretraining.
  • Diffusion pretraining: DDPMs gradually corrupt images with noise and train a network to reverse that process across timesteps.Here, the reverse process is used to learn encoder features rather than generate images during segmentation inference.
  • Study design: The proposed method extracts diffusion-trained encoder weights for conventional segmentation deployment instead of using computationally expensive diffusion inference.The study emphasizes limited data, boundary precision, and anatomically grounded features.
  • Evaluation data: BTCV provides 30 abdominal CT scans with expert annotations for 13 anatomical structures, including liver, kidneys, spleen, stomach, gallbladder, pancreas, and vessels.The dataset includes clinical variation in slice thickness and in-plane resolution.
  • Preprocessing and evaluation: Patient-level splitting assigns each patient exclusively to training, validation, or test data to prevent leakage from adjacent slices.The reported split is 70% training, 15% validation, and 15% test, with the test set withheld from model selection.

B. SELF-SUPERVISION USING DIFFUSION

DDPMs learn to reverse progressively added noise, using unlabeled images to acquire structural representations that can transfer to segmentation. The formulation combines a forward corruption process, a learned reverse process, and a simplified noise-prediction objective designed to capture global anatomy.

  • Diffusion formulation: DDPMs use a forward Markov process to corrupt images with Gaussian noise and a learned reverse process to reconstruct them.The forward process is controlled by a variance schedule β1, …, βT, while the reverse process approximates computationally intractable reversal steps.
  • Data preparation: The preprocessing pipeline extracts target-containing axial slices, applies Hounsfield Unit windowing, normalizes and resizes images, maps liver and kidney labels, and splits patients strictly.Patient-level splitting assigns each patient’s slices exclusively to training, validation, or test sets.
  • Diffusion formulation: The forward process permits direct sampling of a noisy image at any timestep from the original image using reparameterization.This avoids iterating through all preceding timesteps when constructing training examples.
  • Training objective: The reverse network predicts the noise added at timestep t, replacing direct mean prediction in the simplified training objective.The objective removes variance-weighting factors and trains εθ(xt, t) to estimate the injected noise.
  • Anatomical representation learning: The unweighted objective is intended to improve training stability and force the encoder to learn global anatomical structures across noise scales.These structures include organ shapes, positions, relative arrangements, boundaries, and spatial relationships needed for denoising.

2) Architecture and Implementation

The DDPM uses a U-Net backbone configured so its encoder can transfer to segmentation. Timestep conditioning is injected at the bottleneck while preserving compatibility with standard U-Net representations.

  • Architecture: The DDPM architecture employs a U-Net backbone with an encoder-decoder structure and skip connections for downstream decoder integration.The design preserves a standard U-Net organization while supporting diffusion pretraining.
  • Conditioning: Timestep t is encoded with sinusoidal position embeddings and injected exclusively at the bottleneck.This provides temporal conditioning while maintaining −1 to 1 structural compatibility with standard U-Nets.

C. SUPERVISED FINE-TUNING AND TRANSFER LEARNING

After unsupervised DDPM pretraining, encoder weights are transferred to segmentation models using supervised baselines, fully fine-tuned encoders, or frozen encoders. These strategies separate adaptation benefits from the value of learned anatomical features.

  • Transfer strategies: DDPM-pretrained encoder weights are transferred to downstream segmentation tasks through three distinct transfer strategies.The study compares random initialization, fine-tuned pretrained encoders, and frozen pretrained encoders.
  • Supervised baseline: Randomly initialized U-Nets are trained end-to-end using labeled segmentation data as the performance reference.Models M1–M3 use He initialization and provide the baseline for measuring pretraining benefits.
  • Fine-tuned transfer: Models M4–M6 initialize the encoder with DDPM-pretrained weights and fine-tune the encoder, decoder, and segmentation head using segmentation loss.This strategy adapts pretrained anatomical features to pixel-wise segmentation while retaining structural knowledge.

3) Transfer Study 3: Frozen DDPM Encoder

The frozen-encoder study tests whether DDPM pretraining supplies task-useful anatomical features without segmentation-specific encoder adaptation. It evaluates liver, kidney, and multi-organ variants using overlap, boundary, and statistical metrics, including restricted-label experiments.

  • Transfer Study 3: Frozen DDPM Encoder: Frozen DDPM encoders keep encoder parameters fixed while only the decoder and segmentation head learn from labeled data.Strong performance in this condition is intended to test for useful anatomical priors without task-specific adaptation.
  • Experimental tasks: The transfer study trains liver, kidney, and three-class multi-organ segmentation variants for each transfer strategy.The models are evaluated across single-organ and joint multi-organ tasks.
  • Low-data simulation: Restricted-label experiments keep DDPM pretraining on the complete dataset while reducing only the labeled fine-tuning set.This design evaluates robustness to supervised data scarcity.
  • Evaluation: Dice measures volumetric overlap, IoU penalizes false positives and negatives more heavily, and HD95 and ASD quantify boundary accuracy.Together, the metrics capture both region overlap and boundary quality.
  • Statistical analysis: Wilcoxon signed-rank tests with Bonferroni correction assess paired comparisons, while Cohen’s d quantifies practical effect sizes.The study interprets d values from 0.2 to below 0.5 as small, 0.5 to below 0.8 as medium, and at least 0.8 as large.

IV. RESULTS

Diffusion-pretrained models improved volumetric overlap, boundary precision, and prediction consistency across segmentation comparisons. Frozen encoders retained substantial liver performance, supporting learned anatomical structure without segmentation labels.

  • All primary statistical comparisons remained significant after Bonferroni correction, with effect sizes ranging from small to large.
  • 23.6%: Liver Dice increased from 0.751±0.363 to 0.928±0.160 after diffusion pretraining.Variance decreased by over 55%.
  • 68%: Multi-organ Dice variance decreased, with Dice increasing from 0.783±0.225 to 0.950±0.071.
  • 44.7%: Liver HD95 decreased from 11.76 to 6.50 mm, while ASD decreased by 65.5% from 4.38 to 1.51 mm.
  • Diffusion pretraining changed error distributions toward spatially consistent predictions with plausible anatomy, beyond improving pixelwise accuracy.
  • 86.9%: The frozen liver encoder retained this fraction of fine-tuned performance and exceeded the random baseline, whereas frozen kidney performance declined below baseline.

4) Low-Data Performance

Diffusion pretraining improved segmentation data efficiency, but its benefits differed sharply by organ. Liver performance degraded gradually under label scarcity, whereas kidney performance deteriorated substantially below 50% of the labeled data.

  • 99.4%: Liver segmentation retained this fraction of full performance with 50% of the labeled data.Performance remained strong with approximately 10% of labels, or 3–4 annotated patients.
  • 85.7%: Kidney segmentation retained this fraction of performance at 25% of the labeled data, declining to 75.0% at 10%.Lower-data kidney models also showed increased variance.
  • 50%: Using half the labeled data yielded liver Dice above 0.92 while reducing estimated annotation cost and time by roughly half.
  • Low-data models were presented across representative cases alongside baseline, fine-tuned diffusion, and frozen diffusion configurations for visual comparison.

V. DISCUSSION

The study interprets diffusion pretraining as learning anatomical information before supervised segmentation. Its benefits are strongest for large, distinctive organs, while complex smaller-organ morphology remains a scope boundary.

  • V. DISCUSSION: Diffusion pretraining embeds anatomical information in encoders before supervised optimization through image-reconstruction learning.
  • V. DISCUSSION: Diffusion pretraining improved Dice, IoU, HD95, and ASD by supplying broader anatomical context than pixelwise supervision alone.
  • V. DISCUSSION: The denoising objective learns representations from local textures through complete anatomical configurations across hierarchical scales.
  • V. DISCUSSION: 86.9%: Frozen liver encoders retained this fraction of fine-tuned performance without segmentation labels, supporting transfer of anatomical structure rather than generic image features.
  • V. DISCUSSION: Frozen-encoder performance differed by organ, with positive liver results and negative kidney results associated with differences in structural complexity and organ size.
  • V. DISCUSSION: The discussion proposes organ-specific or hybrid pretraining approaches because smaller or anatomically complex organs may require supervised information.

C. ANALYSIS OF VARIANCE REDUCTION

Diffusion-pretrained models reduced performance variance and preserved useful behavior under label scarcity. The analysis attributes this pattern to anatomical priors that constrain outliers and reduce local overfitting, with organ-dependent degradation.

  • C. ANALYSIS OF VARIANCE REDUCTION: 47–68%: Diffusion pretraining reduced variance across liver, kidney, and multi-organ analyses, with consistency emphasized as clinically relevant.
  • C. ANALYSIS OF VARIANCE REDUCTION: Diffusion-pretrained models were described as encoding anatomical invariants robust to patient-specific variation.
  • C. ANALYSIS OF VARIANCE REDUCTION: Distributional priors and global structural constraints were proposed as mechanisms that suppress outliers and prevent local overfitting.
  • C. ANALYSIS OF VARIANCE REDUCTION: 99.4%: Liver performance was retained with 50% of labels, while 95.5% was retained with only 10%, indicating robustness under label scarcity.
  • C. ANALYSIS OF VARIANCE REDUCTION: Liver degradation was graceful, whereas kidney degradation was catastrophic as labeled data decreased, consistent with organ-specific anatomical demands.

VI. LIMITATIONS

The study identifies scope, generalization, dimensionality, and interpretability limitations while reporting convergent quantitative, qualitative, mechanistic, and practical evidence for diffusion-based anatomical feature learning.

  • Scope and generalization: Evaluation focused on liver and kidney, so generalization to smaller, pathological, or lower-contrast structures requires validation.The frozen kidney encoder’s failure suggests that organ-specific characteristics influence diffusion pretraining effectiveness.
  • Scope and generalization: All experiments used BTCV data from a single institution, leaving cross-dataset generalization across scanners, protocols, and demographics unresolved.Reduced variance may partly reflect dataset-specific memorization rather than purely anatomical learning.
  • Dimensionality: The slice-based approach compromises volumetric context because it cannot enforce three-dimensional anatomical constraints.The authors propose 3D diffusion models to capture volumetric relationships and improve robustness and 3D boundary metrics.
  • Interpretability: Grad-CAM attention maps are approximate, and bounded patterns may reflect genuine anatomical priors or learned shortcuts.This ambiguity limits the certainty of mechanistic interpretation.
  • Supporting evidence: Diffusion pretraining improved all reported metrics by 5.7%–23.6%, reduced boundary errors by 38%–74%, and reduced variance by 47%–68%.Visual assessments also reported smooth, anatomically plausible boundaries with global structural coherence.
  • Supporting evidence: Frozen diffusion encoders developed organ-specific attention patterns, while 10% of labels yielded 95.5% of fine-tuned liver performance.The denoising objective is described as learning structural priors that regularize downstream segmentation.

A. AUTHOR CONTRIBUTIONS STATEMENT

The statement assigns distinct responsibilities across study design, implementation, experimentation, supervision, analysis, manuscript preparation, and supporting administrative disclosures.

  • Author contributions: A. G. conceived the study, designed the methodology, implemented DDPM and transfer learning, conducted experiments, analyzed results, prepared figures, and drafted the manuscript.
  • Author contributions: D. G. assisted with result tabulation, result analysis, figure and table review, and the initial manuscript draft.
  • Author contributions: S. B., S. A., and T. K. M. supervised the research, provided conceptual guidance, and contributed to interpreting results.
  • Additional information: The authors report no human or animal studies, no competing financial interests or personal relationships, and publicly accessible BTCV data.
  • Author affiliations: S. B., S. A., and T. K. M. are identified as academics at Manipal Institute of Technology, while A. G. and D. G. are identified as computer science students there.
Loading 2608.25693v1…