Synthetic Augmentation with Large-scale Unconditional Pre-training

TL;DR

HistoDiffusion is a pre-trained latent diffusion model that generates realistic, category-specific images from unlabeled data, significantly improving medical image classification accuracy with minimal labeled data.

Contribution

The paper introduces HistoDiffusion, a novel synthetic augmentation method pre-trained on unlabeled data and fine-tuned with classifier guidance, reducing dependence on labeled datasets.

Findings

Pre-training on unlabeled datasets enhances image synthesis quality.

Synthetic augmentation improves classification accuracy by 6.4%.

Method effectively generalizes to unseen histopathology datasets.

Abstract

Deep learning based medical image recognition systems often require a substantial amount of training data with expert annotations, which can be expensive and time-consuming to obtain. Recently, synthetic augmentation techniques have been proposed to mitigate the issue by generating realistic images conditioned on class labels. However, the effectiveness of these methods heavily depends on the representation capability of the trained generative model, which cannot be guaranteed without sufficient labeled training data. To further reduce the dependency on annotated data, we propose a synthetic augmentation method called HistoDiffusion, which can be pre-trained on large-scale unlabeled datasets and later applied to a small-scale labeled dataset for augmented training. In particular, we train a latent diffusion model (LDM) on diverse unlabeled datasets to learn common features and generate realistic images without conditional inputs. Then, we fine-tune the model with classifier guidance in latent space on an unseen labeled dataset so that the model can synthesize images of specific categories. Additionally, we adopt a selective mechanism to only add synthetic samples with high confidence of matching to target labels. We evaluate our proposed method by pre-training on three histopathology datasets and testing on a histopathology dataset of colorectal cancer (CRC) excluded from the pre-training datasets. With HistoDiffusion augmentation, the classification accuracy of a backbone classifier is remarkably improved by 6.4% using a small set of the original labels. Our code is available at https://github.com/karenyyy/HistoDiffAug.