Image Classification Using a Diffusion Model as a Pre-Training Model
Kosuke Ukita, Ye Xiaolong, Tsuyoshi Okita
TL;DR
This work addresses the challenge of learning robust image representations from unlabeled data by leveraging diffusion models as pre-training engines. It introduces the Representation-Conditioned Latent Diffusion Transformer (RC-LDT), which injects ViT-derived representations into the denoising process to produce more discriminative features for downstream tasks. The framework operates in two stages—self-supervised pre-training on unlabeled data and downstream adaptation for classification—achieving state-of-the-art zero-shot hematoma detection in brain CT scans, with accuracy improved by +6.15% and F1 by +13.60% compared to DINOv2. The results demonstrate that conditioning diffusion models on representations enhances both image generation quality and discriminative power, offering a scalable alternative to contrastive learning for medical imaging and beyond. The practical impact lies in enabling high-performance classification with limited labeled data, reducing annotation costs while maintaining高 fidelity in generated samples.
Abstract
In this paper, we propose a diffusion model that integrates a representation-conditioning mechanism, where the representations derived from a Vision Transformer (ViT) are used to condition the internal process of a Transformer-based diffusion model. This approach enables representation-conditioned data generation, addressing the challenge of requiring large-scale labeled datasets by leveraging self-supervised learning on unlabeled data. We evaluate our method through a zero-shot classification task for hematoma detection in brain imaging. Compared to the strong contrastive learning baseline, DINOv2, our method achieves a notable improvement of +6.15% in accuracy and +13.60% in F1-score, demonstrating its effectiveness in image classification.