Improved Generation of Synthetic Imaging Data Using Feature-Aligned Diffusion

TL;DR

This work aligns intermediate features of the diffusion model to the output features of an expert, and preliminary findings show an improvement of 9% in generation accuracy and ~0.12 in SSIM diversity.

Abstract

Synthetic data generation is an important application of machine learning in the field of medical imaging. While existing approaches have successfully applied fine-tuned diffusion models for synthesizing medical images, we explore potential improvements to this pipeline through feature-aligned diffusion. Our approach aligns intermediate features of the diffusion model to the output features of an expert, and our preliminary findings show an improvement of 9% in generation accuracy and ~0.12 in SSIM diversity. Our approach is also synergistic with existing methods, and easily integrated into diffusion training pipelines for improvements. We make our code available at \url{https://github.com/lnairGT/Feature-Aligned-Diffusion}.

Figures (8)

• Figure 1: Overview of feature-aligned training of diffusion models. Example shows synthesis of Adipose tissue image.
• Figure 2: Fine-tuning validation accuracy of expert models -- ResNet50 (left) with 93% and ViT with 87%.
• Figure 3: Loss during fine-tuning with feature alignment.
• Figure 4: Generation accuracy of feature-aligned vs. baseline diffusion for two fine-tuning pipelines: typical and DreamBooth
• Figure 5: (a) Confusion matrix of expert on feature-aligned diffusion generations, (b) Confusion matrix of expert on baseline diffusion generations (c) Confusion matrix of expert on original dataset. The expert does not mis-classify tumor, stroma, and lympho as debris within the original dataset.