Advancing Image Classification with Discrete Diffusion Classification Modeling

2025-11-25

AI Summary

To address the insufficient robustness of image classification under high uncertainty—such as image corruptions and data scarcity—this paper proposes DiDiCM, a Discrete Diffusion Classification Modeling framework. DiDiCM is the first to introduce discrete diffusion models into classification tasks, directly modeling the posterior distribution of class labels given input images and enabling efficient iterative inference over the discrete label space. Its tunable diffusion step count allows flexible trade-offs between computational cost and memory usage; even with few steps, it substantially outperforms conventional classifiers, with gains magnified under higher uncertainty. Experiments on ImageNet demonstrate significant improvements in classification accuracy across diverse image degradation scenarios and few-shot settings, empirically validating the efficacy of explicit posterior distribution modeling for robust classification.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Multi-class/Multi-label Learning & Extreme Classification
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Humans versus LLMs for data annotation and labeling

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 .

Problem

• Improving image classification under high-uncertainty and corrupted input conditions
• Modeling posterior class distributions using diffusion processes for better accuracy
• Enhancing classification performance on challenging datasets like ImageNet

Innovation

• Leverages diffusion modeling for class posterior distribution
• Supports diffusion on class probabilities or discrete labels
• Achieves higher accuracy with few diffusion iterations