3D Wavelet-Based Structural Priors for Controlled Diffusion in Whole-Body Low-Dose PET Denoising

The paper proposes WCC-Net, a 3D wavelet-conditioned ControlNet framework that integrates explicit frequency-domain structural priors into a diffusion model to achieve superior anatomical consistency and denoising performance in whole-body low-dose PET imaging compared to existing CNN, GAN, and diffusion-based methods.

The Gist

The Big Problem: The "Blurry Flashlight"

Imagine you are trying to take a photo of a dark room using a very weak flashlight. To protect your eyes (or in this case, the patient's body) from too much light, you use a dimmer setting. The result? You get a picture, but it's incredibly grainy and full of "static" (noise).

In medicine, this is Low-Dose PET Scanning. Doctors use a radioactive tracer to see inside the body (like finding a tumor). Usually, they need a strong dose to get a clear picture, but that exposes the patient to more radiation. If they lower the dose to keep the patient safe, the image becomes so noisy that it's hard to see small details, like a tiny tumor or a thin blood vessel. It's like trying to read a map through a foggy, scratched-up window.

The Old Solutions: The "Over-zealous Editor"

For years, scientists tried to fix this using computer programs (AI).

• The "Smoothing" Approach: Some programs just tried to blur the noise away. Think of this like taking a photo and applying a "soft focus" filter. It removes the grain, but it also blurs the edges of the buildings. The picture looks smooth, but you can't tell where the roof ends and the sky begins.
• The "Guessing" Approach: Newer AI (like GANs) tries to "hallucinate" or guess what the missing details should look like. While this makes the picture look sharp, sometimes the AI gets creative in the wrong way, inventing fake details that aren't actually there.

The New Solution: WCC-Net (The "Architect's Blueprint")

The authors of this paper propose a new method called WCC-Net. Instead of just guessing or smoothing, they use a clever trick involving Wavelets and ControlNet.

Here is how it works, broken down into three simple steps:

1. The Wavelet "Sieve" (Separating the Wheat from the Chaff)

Imagine you have a bucket of mixed sand and gold nuggets. The sand is the noise (grainy static), and the gold nuggets are the real anatomy (the shape of the organs).

• Standard AI tries to pick out the gold while looking at the whole messy bucket.
• WCC-Net uses a Wavelet Transform, which is like a magical sieve. It separates the bucket into two piles:
  • Pile A (Low Frequency): The big, solid gold nuggets (the general shape of the liver, heart, etc.). This pile is very clean and stable.
  • Pile B (High Frequency): The fine sand and dust (the noise and tiny, jagged edges).
• The AI looks at Pile A to understand the structure of the body. It knows, "Okay, the liver is roughly this big and in this shape."

2. The Frozen Backbone (The "Master Painter")

The researchers use a pre-trained AI model (a Diffusion Model) that is already an expert at painting realistic medical images. Think of this as a Master Painter who knows exactly how a human body should look.

• Usually, if you give this painter a blurry, noisy photo, they might get confused and paint the noise as if it were real.
• In WCC-Net, the Master Painter is "Frozen." This means we don't let them change their style or memory. We trust their knowledge of anatomy.

3. The ControlNet "Guide" (The "Architect's Blueprint")

This is the secret sauce. We take Pile A (the clean structural blueprint from the Wavelet sieve) and hand it to the Master Painter as a Guide.

• We tell the Painter: "You know how to paint a realistic body. But here is a blueprint of the exact shape we need. Please paint your masterpiece, but make sure you follow this blueprint strictly."
• The Painter uses their artistic skill to fill in the details and remove the noise, but they are forced to keep the structural integrity of the blueprint. They can't invent fake tumors, and they can't blur the edges of the organs.

Why is this a Game Changer?

The paper tested this on "Ultra-Low-Dose" scans (where the image is almost unusable).

• The Result: WCC-Net produced images that were sharper, clearer, and more accurate than any previous method.
• The Analogy: If the old methods were like trying to clean a muddy window by wiping it with a wet cloth (smearing the dirt), WCC-Net is like having a laser-guided cleaning robot that knows exactly where the glass is and only removes the mud without scratching the surface.

The Bottom Line

This new method allows doctors to give patients much less radiation (making scans safer) while still getting crystal-clear images (making diagnoses more accurate). It does this by teaching the AI to look at the "skeleton" of the image (the structure) separately from the "dust" (the noise), ensuring the final picture is both beautiful and medically truthful.

Technical Summary

1. Problem Statement

Low-dose Positron Emission Tomography (PET) is critical for reducing patient radiation exposure but suffers from severe noise degradation, which compromises image quality and diagnostic reliability.

• The Core Challenge: While Diffusion Models (DMs) have shown strong denoising capabilities, their stochastic nature makes it difficult to enforce anatomically consistent structures, especially in low signal-to-noise ratio (SNR) regimes.
• The Limitation of Existing Methods: Current DM-based approaches typically condition on raw low-dose spatial-domain images. In these images, anatomical structure is entangled with noise. Consequently, the model faces a trade-off: suppressing pervasive high-frequency noise often leads to the loss of fine anatomical edges (blurring) or the generation of hallucinated structures, which is critical for clinical diagnosis.
• The Gap: There is a need for a method that can explicitly separate robust structural information from noise to guide the diffusion process without compromising the generative expressiveness of the model.

2. Methodology: Wavelet-Conditioned ControlNet (WCC-Net)

The authors propose WCC-Net, a fully 3D diffusion-based framework that integrates explicit frequency-domain structural priors via Discrete Wavelet Transform (DWT) into a frozen diffusion backbone using a ControlNet architecture.

Key Architectural Components:

1. Frozen Diffusion Backbone:

   • The core is a pre-trained 3D Denoising Diffusion Probabilistic Model (DDPM) (specifically a 3D U-Net) trained to map low-dose PET to normal-dose PET.
   • The backbone parameters ($\theta$) remain frozen during the training of WCC-Net to preserve the model's pre-trained generative capacity and prevent catastrophic forgetting.

2. Wavelet-Based Conditioning Prior:

   • Instead of conditioning on the raw noisy image, the input low-dose PET image ($y$) is decomposed using a 3D Discrete Wavelet Transform (DWT) (Haar wavelet).
   • This decomposition separates the image into frequency subbands:
     • Low-frequency (LLL): Encodes global coarse anatomical structure and intensity distribution.
     • High-frequency (H): Encodes fine details, edges, and noise.
   • The method specifically extracts the low-frequency subband ($y_{LLL}$) as the structural prior ($c_{wav}$), as it is robust to noise and represents the stable anatomical framework.

3. ControlNet Integration:

   • A lightweight, trainable Control Branch (a copy of the encoder blocks of the backbone) processes the wavelet prior.
   • Zero-Initialized Convolutions (ZeroConv): The features from the Control Branch are injected into the frozen backbone's skip connections via zero-initialized $1\times1\times1$ convolutions.
   • Mechanism: At initialization, the ZeroConv weights are zero, meaning the model behaves exactly like the pre-trained baseline. During training, the weights gradually learn to inject the wavelet-based structural guidance.
   • Decoupling: This design effectively decouples anatomical structure (guided by the wavelet prior) from stochastic noise removal (handled by the frozen diffusion backbone).

3. Key Contributions

1. Novel Framework: Proposal of the first fully 3D wavelet-conditioned diffusion framework specifically designed for whole-body low-dose PET denoising, ensuring volumetric anatomical continuity.
2. Decoupled Conditioning Strategy: Introduction of a strategy that injects frequency-domain structural priors via a trainable ControlNet branch into a frozen backbone. This allows the model to leverage explicit structural guidance without disrupting the pre-trained generative distribution.
3. Robust Generalization: Demonstration that the method generalizes effectively to unseen dose levels (both lower and higher than training), maintaining superior structural fidelity compared to state-of-the-art baselines.

4. Experimental Results

The method was evaluated on the Ultra-Low-Dose PET (UDPET) dataset (Siemens Biograph Vision Quadra), using 1/20 dose for training and 1/50 and 1/4 doses for external testing.

Quantitative Performance:

• Internal Test (1/20 Dose): WCC-Net outperformed all baselines (CNN, GAN, and standard Diffusion).
  • PSNR: Improved by +1.21 dB over the strongest diffusion baseline (3D DDPM).
  • SSIM: Improved by +0.008.
  • Structural Distortion (GMSD) & Intensity Error (NMAE): Significantly reduced, indicating better edge preservation and quantitative accuracy.
• External Test (Unseen Doses):
  • 1/50 Dose (Severe Noise): Achieved +1.24 dB PSNR gain over 3D DDPM.
  • 1/4 Dose (Moderate Noise): Achieved the highest SSIM (0.992) and lowest GMSD (0.007) among all methods.

Qualitative Findings:

• Visual Quality: Unlike CNNs/GANs which often produce inconsistent textures or artifacts, and standard DMs which tend to blur fine details, WCC-Net preserved thin cortical boundaries and small lesions while effectively suppressing noise.
• Error Analysis: Signed error maps showed that WCC-Net had lower spatial bias and reduced over/under-estimation near organ boundaries compared to competitors.
• Ablation Study: Confirmed that conditioning on low-frequency wavelet components (LLL) yields the best results. Conditioning on high-frequency components alone or mixing all bands resulted in inferior performance, validating the hypothesis that low-frequency priors provide the most stable structural guidance.

5. Significance and Impact

• Clinical Relevance: By enabling high-quality image reconstruction from ultra-low-dose acquisitions, WCC-Net has the potential to significantly reduce patient radiation exposure in clinical PET scans without sacrificing diagnostic accuracy.
• Methodological Advancement: The paper addresses a fundamental limitation in diffusion models for medical imaging: the entanglement of noise and structure. By moving conditioning from the spatial domain to the frequency domain, it offers a principled way to enforce anatomical constraints.
• Generalizability: The ability to generalize to unseen dose levels suggests the model learns robust anatomical features rather than simply memorizing noise patterns, making it suitable for diverse clinical protocols.

Limitations & Future Work:
The current implementation uses a fixed single-level Haar wavelet, which may limit multi-scale modeling. Future work aims to explore adaptive wavelet bases, multi-level decompositions, and validation across multiple scanner vendors and radiotracers.