Advancing Image Classification with Discrete Diffusion Classification Modeling

TL;DR

This paper introduces Discrete Diffusion Classification Modeling (DiDiCM), a diffusion-based framework that improves image classification accuracy, especially under high-uncertainty conditions, by modeling class posterior distributions with flexible prediction options.

Contribution

The paper presents a novel diffusion-based approach for image classification that enhances performance in challenging scenarios, outperforming standard classifiers on ImageNet.

Findings

DiDiCM achieves higher accuracy with fewer diffusion steps.

Performance gains increase under more challenging conditions.

Code is publicly available for reproducibility.

Abstract

Image classification is a well-studied task in computer vision, and yet it remains challenging under high-uncertainty conditions, such as when input images are corrupted or training data are limited. Conventional classification approaches typically train models to directly predict class labels from input images, but this might lead to suboptimal performance in such scenarios. To address this issue, we propose Discrete Diffusion Classification Modeling (DiDiCM), a novel framework that leverages a diffusion-based procedure to model the posterior distribution of class labels conditioned on the input image. DiDiCM supports diffusion-based predictions either on class probabilities or on discrete class labels, providing flexibility in computation and memory trade-offs. We conduct a comprehensive empirical study demonstrating the superior performance of DiDiCM over standard classifiers, showing that a few diffusion iterations achieve higher classification accuracy on the ImageNet dataset compared to baselines, with accuracy gains increasing as the task becomes more challenging. We release our code at https://github.com/omerb01/didicm .