MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model

In plain words

Segmentation means drawing the outline of an organ or lesion in a medical image. Most methods predict that outline in one pass. This paper instead starts from a random scribble and cleans it up over many small steps, the way image generators turn noise into a picture. Each step looks at the original scan and at the current, still rough outline, and makes the outline a bit better.

Two changes make this work for medical images. First, the scan is looked at again at every step, and the current outline is fed back into how the scan is read. This helps when the target is faint and hard to tell from its background. Second, the model learns a filter in the frequency domain, which is a way of describing an image by its fine and coarse patterns. The filter weakens the fine, noisy patterns that the rough outline brings in.

The authors test the method on three kinds of scans: eye photographs, brain MRI and ultrasound. They report better overlap with the expert outlines than the methods they compare against. These results come from their own experiments.

TL;DR

MedSegDiff trains a diffusion probabilistic model whose reverse chain turns noise into a segmentation mask, conditioned on the input image §2 eq. 1 eq. 2. Two additions carry the method: dynamic conditional encoding, which fuses the current noisy mask into the image encoder §2.1 eq. 3, and FF-Parser, a learned Fourier-domain filter on skip features §2.2 eq. 4 eq. 5 eq. 6. Author-reported Dice for the largest variant, MedSegDiff++, is 87.5 (optic cup), 90.5 (brain tumour) and 86.6 (thyroid nodule), above the compared methods on all three tasks tab. 1. Inference fuses 25 sampled masks with STAPLE §3.2, so each image needs many network evaluations.

Problem

Medical targets such as optic cups, tumours and nodules are often low contrast and hard to separate from their background §2.1. The authors argue that a diffusion model with only the raw image as its condition gives one static signal at every step, which is too weak for such targets §2.1. The method targets general medical segmentation across modalities, and the authors say it is the first diffusion-based model aimed at that task §1.

The evaluation covers three tasks: optic-cup segmentation in fundus images, brain-tumour segmentation in MRI and thyroid-nodule segmentation in ultrasound §1 §3.1.

Method

The paper uses the standard diffusion probabilistic model (DPM) setup of Ho et al. The mask is noised in a forward process. A network learns the reverse process, and starting from Gaussian noise it recovers a clean mask §2 eq. 1. The reverse chain is the product of one-step transitions eq. 1.

Step 1: conditioning §2 eq. 2 fig. 1. The noise estimate is conditioned on the image prior I. The image embedding E^I and the current mask embedding E^x are added and passed to a UNet decoder D. The step index t is embedded with a shared learned look-up table and added to the decoder features §2 eq. 2.

Step 2: dynamic conditional encoding §2.1 eq. 3. In the image encoder, the image feature map m_I^k at every scale k is merged with the mask feature map m_x^k, which has the same shape. Both maps are layer-normalised and multiplied element-wise to form an affinity map, which then multiplies m_I^k eq. 3. The fusion is applied in the two middle stages of the encoder, and each stage follows the ResNet34 layout §2.1. This step is applied at every sampling step, so the conditioning changes over the chain §2.1.

Step 3: FF-Parser §2.2 eq. 4 eq. 5 eq. 6 fig. 2. FF-Parser is placed on each skip path of the feature integration §2.2. For a feature map m of height H, width W and C channels, it takes the 2D FFT over the spatial dimensions eq. 4, multiplies the spectrum element-wise by a learned complex map A of the same shape eq. 5, and returns to space with the inverse FFT eq. 6. The paper describes this as a learnable frequency filter §2.2.

Step 4: training §2.3 eq. 7. The loss is the standard noise-prediction mean-squared error. Each iteration draws a random pair of image I_i and segmentation label S_i, a random step t from a uniform distribution, and Gaussian noise §2.3 eq. 7.

Step 5: architecture §2.3 fig. 1. The backbone is a modified ResUNet. The image I and the noisy mask x_t are encoded by two separate encoders with three convolution stages, each with residual blocks following the ResNet34 block counts. Blocks use group normalisation, SiLU and a convolution, and receive the time embedding through two linear layers §2.3. The two encoder outputs are summed and sent to the last encoding stage, and a standard convolutional decoder predicts the final result §2.3.

Step 6: inference §3.2. The paper uses 100 diffusion steps at inference, draws several samples, and fuses them with the STAPLE algorithm. The paper states 25 samples for every variant §3.2.

# Training step  [§2.3] [eq. 7] [code: guided_diffusion/train_util.py#L251]
def training_step(image, label):
    # Binary target: any non-zero label is foreground  [§3.1]
    # [code: guided_diffusion/gaussian_diffusion.py#L992-L993]
    target_mask = label > 0
    # Draw a random step and Gaussian noise  [§2.3] [eq. 7]
    t = sample_uniform_step()
    noise = gaussian_like(target_mask)
    # Noise the mask: sqrt(a_t) * x0 + sqrt(1 - a_t) * eps  [eq. 7]
    noisy_mask = q_sample(target_mask, t, noise)
    # Input is the image channels plus the noisy mask channel  [§2.3]
    x_t = concat(image, noisy_mask)                  # (B, 4, 256, 256)
    # Predict the noise and an image-only foreground map  [§2.1]
    # [code: guided_diffusion/unet.py#L786]
    eps_hat, aux_mask = model(x_t, t)
    loss_diff = mse(noise, eps_hat)                  # [eq. 7]
    # Auxiliary loss, weight 10; not in eq. 7
    # [code: guided_diffusion/train_util.py#L251]
    loss_aux = mse(target_mask, aux_mask)
    return loss_diff + 10 * loss_aux

# Dynamic conditioning on one skip path  [§2.1] [eq. 3]
# Paper: LayerNorm affinity gate (eq. 3). Code: see Paper vs code.
def dynamic_skip(image_skip, noisy_feature, level):
    # Project the noisy-path feature to the gate width  [code: unet.py#L2483]
    projected = conv1x1_gate(noisy_feature)
    # Learned Fourier-domain filter  [§2.2] [code: unet.py#L2484]
    filtered = ff_parser(projected, level)
    # Spatial gate, one channel  [code: unet.py#L2485]
    spatial_gate = conv1x1_spatial(filtered)
    # Per-channel gate: spatial mean of the filtered map  [code: unet.py#L2486]
    channel_gate = mean(filtered, over_height_and_width)
    return image_skip * spatial_gate * channel_gate

# FF-Parser  [§2.2] [eq. 4] [eq. 5] [eq. 6]
# Code uses a half-spectrum real FFT with orthonormal scaling
# [code: guided_diffusion/unet.py#L476-L480]
def ff_parser(feature, level):
    # feature: (B, C, H, W); learned complex weight: (C, H, W // 2 + 1)
    # 2D real FFT over the spatial dimensions  [eq. 4]
    spectrum = rfft2(feature, norm="ortho")
    # Element-wise modulation by the learned complex map  [eq. 5]
    spectrum = spectrum * learned_weight[level]
    # Back to the spatial domain  [eq. 6]
    return irfft2(spectrum, size=(H, W), norm="ortho")

# Sampling and ensemble fusion  [§3.2]
# [code: scripts/segmentation_sample.py#L101-L103]
def segment(image, num_samples, steps):
    # Start the mask channel from Gaussian noise  [§2]
    # [code: scripts/segmentation_sample.py#L102-L103]
    noise = randn_like(image[:, :1])
    x_T = concat(image, noise)
    candidates = []
    for each sample in range(num_samples):
        # Reverse chain over the given number of steps  [§2] [eq. 1] [§3.2]
        x_0, aux_map = reverse_chain(model, x_T, steps)
        # Blend the two maps; the blend depends on their Dice agreement
        # [code: guided_diffusion/gaussian_diffusion.py#L589-L593]
        if dice(x_0, aux_map) < 0.65:
            final = clip(aux_map + 0.25 * x_0, 0, 1)
        else:
            final = clip(0.5 * aux_map + 0.5 * x_0, 0, 1)
        candidates.append(final)
    # Fuse the ensemble with STAPLE  [§3.2]
    # [code: guided_diffusion/utils.py#L45]
    return staple(stack(candidates))

Hyperparameters

Values from the paper are cited to the paper. Values taken from the code or its README are marked as such. Where the paper is silent, the table says so.

Results

Author-reported. Tasks and datasets: optic-cup segmentation on REFUGE-2 (fundus images), brain-tumour segmentation on BraTS-2021 (MRI), and thyroid-nodule segmentation on DDTI (ultrasound) §3.1. Metrics: Dice and IoU §3.3. How the scores are averaged over images or slices is not stated. The baseline numbers are as printed in the paper; whether they were re-run or copied is not stated, except EnsemDiff, which the authors say they reproduced under the same setting §3.2.

The paper prints the CENet optic-cup IoU as "69,4" with a comma; this card reads it as 69.4. The paper marks the methods proposed for each task with a grey background; the table above does not reproduce that.

Ablation, Dice (%), on the three tasks tab. 2. The model size used for the ablation is not stated.

Headline claims in the text: MedSegDiff outperforms the compared methods on all three tasks §3.3. Against EnsemDiff on brain tumour, the text gives a gain of 2.3% Dice and 2.4% IoU §3.3. Table 1 gives 88.7 → 90.5 Dice and 80.9 → 82.8 IoU for MedSegDiff++, i.e. 1.8 and 1.9 points. Dynamic conditioning gives 2.1, 1.6 and 1.8 points on the three tasks §3.4 tab. 2. The text says FF-Parser adds "near 1%" §3.4; Table 2 gives 0.8, 0.7 and 0.7 points.

Paper vs code

The official repository is linked in the abstract §abstract, and its owner (WuJunde) matches the first author. Its HEAD commit is 28b343f. The code includes a second model class, labelled "new" and used by default, which the README describes as the newer version README.md. The paper's model is, by inference from the code, the other class, UNetModel_v1preview, selected when the version argument is anything other than new script_util.py segmentation_sample.py. The rest of this section compares the paper with that class.

• Attention, eq. 3. The paper fuses the features with the LayerNorm affinity map of eq. 3 §2.1. The v1 forward pass never calls enhance, the method that implements it unet.py. The code instead gates the image skip features multiplicatively with a spatial map and a per-channel mean unet.py.
• Where the fusion happens. The paper applies the fusion in the two middle stages §2.1. The code applies it at four skip levels, the first four of the encoder's skip levels, each one computed from an image-only encoder unet.py.
• What the FFT acts on. Eq. 4 applies the FFT to a decoder feature map m §2.2. The code applies it to a 1 × 1 convolution of the noisy-path features that are fed into the gate, not to a decoder feature map unet.py. The code also uses a half-spectrum real FFT with orthonormal scaling unet.py, and initialises the learned weight with a scaled random normal (std 0.02) unet.py. The paper states neither.
• Two encoders. The paper describes two separate encoders for I and x_t, summed before the last encoding stage §2.3. The code runs one UNet over the concatenated input and adds the bottleneck output of an image-only encoder to its bottleneck, just before the middle block unet.py.
• Auxiliary loss. The paper's objective is eq. 7 alone §2.3. The code adds an MSE term on the image-only foreground map with weight 10 train_util.py gaussian_diffusion.py.
• Learned variance. The README's training commands pass --learn_sigma True README.md README.md. With that flag, the code adds a variational term to the loss gaussian_diffusion.py. The paper does not mention a learned variance.
• Final map. The paper says the samples are fused with STAPLE §3.2. For the v1 path, each sample is first blended with the image-only map, the blend depending on their Dice agreement gaussian_diffusion.py. The default path uses the sample's mask channel directly segmentation_sample.py. The paper describes neither.
• Ensemble size. The paper states 25 samples §3.2. The sampling script defaults to 5 segmentation_sample.py, and the README uses 5 for MedSegDiff-B.
• Inference steps. The paper states 100 steps §3.2. The code defaults to 1000 script_util.py, and the README suggests 50 with DPM-Solver README.md.
• Learning rate and batch size for B. The paper gives 1e-4 and 32 §3.2. The README gives 5e-5 and 8 for MedSegDiff-B README.md.
• Hard-coded sizes. The highway branch assumes 128 base channels and a 256-pixel input. Its channel counts are int(d/2 + 1) * 128 unet.py, and each FF-Parser's spatial size is computed from 256 unet.py. Other image sizes or base widths need code changes. The paper does not mention this.
• Dataset text. The paper's introduction says "brain CT" §1, while the abstract and §3.1 say MRI. §3.1 names "glaucoma, thyroid cancer and melanoma diagnosis", but the evaluated tasks are segmentation tasks.

Pitfalls

• Ensemble cost. At the paper's settings, 25 samples of 100 steps each give 2,500 denoiser calls per image §3.2. This is arithmetic from the stated numbers, not a reported runtime.
• Undocumented output rule. Inference blends maps with a Dice-based rule gaussian_diffusion.py. A reimplementation that uses only the sampled mask will not match the reported numbers.
• Auxiliary branch. The image-only head is trained with a weight-10 MSE term train_util.py. Dropping it changes training, and the paper does not mention it.
• Threshold choice. The 0.65 Dice threshold and the 0.25 and 0.5 blend weights are not in the paper gaussian_diffusion.py.
• Variance learning. --learn_sigma True changes the loss, and the paper does not mention it gaussian_diffusion.py.
• Input layout. The code takes the first channels as the image and the last channel as the noisy mask; the image-only path drops the last channel unet.py segmentation_sample.py. Any other channel order needs a matching change in the code.
• FFT numerics. The FFT runs on float32 after an explicit cast unet.py. The paper does not mention precision.
• Unstated evaluation. How Dice and IoU are averaged over slices and images is not stated. The sampling script averages per image through its own Dice function segmentation_env.py.
• Ablation settings. The model size used for the ablation is not stated tab. 2.
• Text versus table. The 2.3% and 2.4% gains over EnsemDiff and the "near 1%" FF-Parser gain do not match Table 1 and Table 2 (see Results).

Lineage

• Builds on: the diffusion probabilistic model of Ho et al. §2 eq. 1 and the UNet detail of Nichol et al. (2021) §2.3. The paper's frequency filter is framed as a learnable version of classical frequency filters §2.2.
• Compared against: ResUNet, BEAL, TransBTS, EnsemDiff, MTSeg, UltraUNet, CENet, MRNet, SegNet, nnUNet and TransUNet §3.3 tab. 1.
• Superseded by: not stated in the paper. The code README points to a second version, MedSegDiff-V2, at arXiv:2301.11798 README.md. That paper has no card here.