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

arXiv: https://arxiv.org/abs/2603.06173 · HTML: https://mlsanity.com/paper/2603.06173/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 14 of its 21 references that are in the map.
- First posted on arXiv 6 Mar 2026.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)
- Deep reinforcement learning methods / [RL for clinical treatment and surgery](https://mlsanity.com/deeprl/rl-for-clinical-treatment-and-surgery/)

## Builds on (most cited first, 14 of 14)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [DiffusionNFT: Online Diffusion Reinforcement with Forward Process](https://mlsanity.com/paper/2509.16117/) (2025)
- [Learning Spatio-Temporal Features with 3D Residual Networks for Action Recognition](https://mlsanity.com/paper/1708.07632/) (2017)
- [Conditional Diffusion Models for Semantic 3D Brain MRI Synthesis](https://mlsanity.com/paper/2305.18453/) (2023)
- [Generation of 3D Brain MRI Using Auto-Encoding Generative Adversarial Networks](https://mlsanity.com/paper/1908.02498/) (2019)
- [Generating 3D Brain Tumor Regions in MRI using Vector-Quantization Generative Adversarial Networks](https://mlsanity.com/paper/2310.01251/) (2023)
- [Enhancing Fetal Plane Classification Accuracy with Data Augmentation Using Diffusion Models](https://mlsanity.com/paper/2501.15248/) (2025)
- [RL4Med-DDPO: Reinforcement Learning for Controlled Guidance Towards Diverse Medical Image Generation using Vision-Language Foundation Models](https://mlsanity.com/paper/2503.15784/) (2025)
- [Fine-Tuning Diffusion Models via Intermediate Distribution Shaping](https://mlsanity.com/paper/2510.02692/) (2025)
- [Improving 2D Diffusion Models for 3D Medical Imaging with Inter-Slice Consistent Stochasticity](https://mlsanity.com/paper/2602.04162/) (2026)

## Built on by (most cited first, 1 of 1)

- [Beyond Representation Learning: A Systematic Study of Joint-Embedding Predictive Generation for 3D Brain MRI](https://mlsanity.com/paper/2608.28787/) (2026)

## Abstract (arXiv)

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.