Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Optimizing 3D Diffusion Models for Medical Imaging via Multi-Scale Reward Learning

Yueying Tian, Xudong Han, Meng Zhou, Rodrigo Aviles-Espinosa, Rupert Young, Philip Birch

TL;DR

This paper presents a method to enhance 3D diffusion models using Reinforcement Learning (RL) with multi-scale feedback, and indicates that incorporating RL feedback effectively steers the generation process toward higher quality distributions.

Abstract

Diffusion models have emerged as powerful tools for 3D medical image generation, yet bridging the gap between standard training objectives and clinical relevance remains a challenge. This paper presents a method to enhance 3D diffusion models using Reinforcement Learning (RL) with multi-scale feedback. We first pretrain a 3D diffusion model on MRI volumes to establish a robust generative prior. Subsequently, we fine-tune the model using Proximal Policy Optimization (PPO), guided by a novel reward system that integrates both 2D slice-wise assessments and 3D volumetric analysis. This combination allows the model to simultaneously optimize for local texture details and global structural coherence. We validate our framework on the BraTS 2019 and OASIS-1 datasets. Our results indicate that incorporating RL feedback effectively steers the generation process toward higher quality distributions. Quantitative analysis reveals significant improvements in Fréchet Inception Distance (FID) and, crucially, the synthetic data demonstrates enhanced utility in downstream tumor and disease classification tasks compared to non-optimized baselines.

Optimizing 3D Diffusion Models for Medical Imaging via Multi-Scale Reward Learning

TL;DR

This paper presents a method to enhance 3D diffusion models using Reinforcement Learning (RL) with multi-scale feedback, and indicates that incorporating RL feedback effectively steers the generation process toward higher quality distributions.

Abstract

Diffusion models have emerged as powerful tools for 3D medical image generation, yet bridging the gap between standard training objectives and clinical relevance remains a challenge. This paper presents a method to enhance 3D diffusion models using Reinforcement Learning (RL) with multi-scale feedback. We first pretrain a 3D diffusion model on MRI volumes to establish a robust generative prior. Subsequently, we fine-tune the model using Proximal Policy Optimization (PPO), guided by a novel reward system that integrates both 2D slice-wise assessments and 3D volumetric analysis. This combination allows the model to simultaneously optimize for local texture details and global structural coherence. We validate our framework on the BraTS 2019 and OASIS-1 datasets. Our results indicate that incorporating RL feedback effectively steers the generation process toward higher quality distributions. Quantitative analysis reveals significant improvements in Fréchet Inception Distance (FID) and, crucially, the synthetic data demonstrates enhanced utility in downstream tumor and disease classification tasks compared to non-optimized baselines.
Paper Structure (22 sections, 5 equations, 1 figure, 5 tables)

This paper contains 22 sections, 5 equations, 1 figure, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Latent Models in Medical Imaging
  4. Fine-tuning Diffusion Models
  5. Methods
  6. Latent 3D Diffusion and the Fidelity Gap
  7. Multi-Scale Reward Learning
  8. Synthetic Trajectories.
  9. Noised-Reconstruction Trajectories.
  10. Reward Objective and Intuition
  11. Multi-Scale Feedback
  12. 3D Volumetric Reward ($R_{3D}$).
  13. 2D Slice-wise Reward ($R_{2D}$).
  14. RL Fine-tuning via PPO
  15. Experiments
  16. ...and 7 more sections

Figures (1)

  • Figure 1: Illustration for the proposed framework. Stage I involves pretraining a 3D-VQGAN and a latent diffusion model. Stage II (Middle) shows the self-supervised reward generation strategy. We utilize synthetic trajectories (FID $\approx$ 50) and noised-reconstruction trajectories (FID $\approx$ 25) to fill the fidelity gap. Reward values are computed using the objective $R = \exp(-(FID - 25) / 15)$. Stage III (Bottom) illustrates the RL fine-tuning phase where the policy $\pi_\theta$ is optimized via PPO based on 3D volumetric and 2D slice-wise feedback.