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

TL;DR

This paper introduces a reinforcement learning approach with multi-scale feedback to improve 3D diffusion models for medical imaging, leading to higher quality synthetic MRI data and better downstream task performance.

Contribution

It proposes a novel RL-based fine-tuning method for 3D diffusion models using multi-scale rewards, bridging the gap between image quality and clinical relevance.

Findings

Improved FID scores on BraTS 2019 and OASIS-1 datasets.

Enhanced utility of synthetic data in classification tasks.

Effective multi-scale reward system guiding model optimization.

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\'echet Inception Distance (FID) and, crucially, the synthetic data demonstrates enhanced utility in downstream tumor and disease classification tasks compared to non-optimized baselines.