Diffusion with Adversarial Fine-Tuning for Improving Rare Retinal Disease Diagnosis

Abstract

As machine-aided disease diagnosis becomes more common, there is a rising need for high volumes of quality data, which might be unavailable for rare diseases. Generative methods offer a solution, allowing for synthesising realistic-looking data that can improve diagnosis accuracy. We investigate the applications of diffusion to a small, imbalanced dataset of Optical Coherence Tomography (OCT) images. We propose modifying the basic Denoising Diffusion Probabilistic Model with attention mechanisms, a class-aware training strategy, and the addition of adversarial fine-tuning. We demonstrate that this model is capable of synthesising realistic-looking images with class-specific features even for diseases with as little as 22 samples. We achieve values of FID at 62.58, and CLIP Similarity at 0.96. We show that the addition of generated data in the training dataset improves the overall and class-specific performance of a ResNet18 classifier on the OCT data, offering an improvement for downstream tasks such as rare retinal disease diagnosis