DiNO-Diffusion. Scaling Medical Diffusion via Self-Supervised Pre-Training

Guillermo Jimenez-Perez,Pedro Osorio,Josef Cersovsky,Javier Montalt-Tordera,Jens Hooge,Steffen Vogler,Sadegh Mohammadi

2024-07-16

Abstract:Diffusion models (DMs) have emerged as powerful foundation models for a variety of tasks, with a large focus in synthetic image generation. However, their requirement of large annotated datasets for training limits their applicability in medical imaging, where datasets are typically smaller and sparsely annotated. We introduce DiNO-Diffusion, a self-supervised method for training latent diffusion models (LDMs) that conditions the generation process on image embeddings extracted from DiNO. By eliminating the reliance on annotations, our training leverages over 868k unlabelled images from public chest X-Ray (CXR) datasets. Despite being self-supervised, DiNO-Diffusion shows comprehensive manifold coverage, with FID scores as low as 4.7, and emerging properties when evaluated in downstream tasks. It can be used to generate semantically-diverse synthetic datasets even from small data pools, demonstrating up to 20% AUC increase in classification performance when used for data augmentation. Images were generated with different sampling strategies over the DiNO embedding manifold and using real images as a starting point. Results suggest, DiNO-Diffusion could facilitate the creation of large datasets for flexible training of downstream AI models from limited amount of real data, while also holding potential for privacy preservation. Additionally, DiNO-Diffusion demonstrates zero-shot segmentation performance of up to 84.4% Dice score when evaluating lung lobe segmentation. This evidences good CXR image-anatomy alignment, akin to segmenting using textual descriptors on vanilla DMs. Finally, DiNO-Diffusion can be easily adapted to other medical imaging modalities or state-of-the-art diffusion models, opening the door for large-scale, multi-domain image generation pipelines for medical imaging.

Computer Vision and Pattern Recognition,Artificial Intelligence,Machine Learning

What problem does this paper attempt to address?

The problem that this paper attempts to solve is:

In the field of medical imaging, although Diffusion Models (DMs) have shown strong capabilities in image generation, enhancement, reconstruction, and segmentation, they rely on a large number of annotated datasets for training. However, medical imaging datasets are usually small and sparsely annotated, which limits the application of diffusion models in this field. To solve this problem, the paper proposes DiNO - Diffusion, which is a self - supervised method for training Latent Diffusion Models (LDMs). It trains by eliminating the dependence on annotations and using more than 868,000 unlabeled images extracted from the public Chest X - Ray (CXR) dataset.

Specifically, DiNO - Diffusion solves the above problems in the following ways:
1. **Self - supervised training**: Use the frozen DiNO model to extract image embeddings as conditions for the generation process, thus not relying on label data.
2. **Utilization of large - scale data**: It can use a large amount of unlabeled medical imaging data for training, improving the generalization ability and generation quality of the model.
3. **Multi - task applications**: The generated synthetic data can be used for data augmentation to improve classification performance; at the same time, the model also shows the ability of zero - shot segmentation and can generate high - quality segmentation masks without additional training.
4. **Privacy protection**: By generating synthetic data, downstream AI models can be trained without using real data, which helps protect patient privacy.

The paper verifies the effectiveness of DiNO - Diffusion through multiple experiments, including image quality assessment, data augmentation, fully synthetic data training, and zero - shot segmentation, showing its potential in medical image generation and processing.