Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Anisotropic Diffusion Probabilistic Model for Imbalanced Image Classification

Jingyu Kong, Yuan Guo, Yu Wang, Yuping Duan

TL;DR

This work tackles long-tailed medical image classification by tailoring diffusion probabilistic models to per-class data distributions. It introduces Anisotropic Diffusion Probabilistic Model (ADPM), which assigns class-specific diffusion speeds through a noise level parameter $ ext{\lambda}_j$ guided by an imbalance-sensitive bound, enabling tail classes to learn more robust decision boundaries. The model also integrates global/local priors in the forward process and a feature-conditioned reverse prior via cross-attention, training with a multi-task objective that includes MMD regularization to align noise samples with estimated Gaussian noise. Empirically, ADPM improves tail-class accuracy across four medical datasets (PAD-UFES, HAM10000, SCIN, Hyper-Kvasir) while maintaining head-class performance, achieving up to ~4% gains in F1 on PAD-UFES and ~3% on HAM10000. These results demonstrate the practical impact of principled anisotropic diffusion and prior-informed learning for imbalanced medical imaging tasks, with potential for broader applications in long-tailed visual recognition.

Abstract

Real-world data often has a long-tailed distribution, where the scarcity of tail samples significantly limits the model's generalization ability. Denoising Diffusion Probabilistic Models (DDPM) are generative models based on stochastic differential equation theory and have demonstrated impressive performance in image classification tasks. However, existing diffusion probabilistic models do not perform satisfactorily in classifying tail classes. In this work, we propose the Anisotropic Diffusion Probabilistic Model (ADPM) for imbalanced image classification problems. We utilize the data distribution to control the diffusion speed of different class samples during the forward process, effectively improving the classification accuracy of the denoiser in the reverse process. Specifically, we provide a theoretical strategy for selecting noise levels for different categories in the diffusion process based on error analysis theory to address the imbalanced classification problem. Furthermore, we integrate global and local image prior in the forward process to enhance the model's discriminative ability in the spatial dimension, while incorporate semantic-level contextual information in the reverse process to boost the model's discriminative power and robustness. Through comparisons with state-of-the-art methods on four medical benchmark datasets, we validate the effectiveness of the proposed method in handling long-tail data. Our results confirm that the anisotropic diffusion model significantly improves the classification accuracy of rare classes while maintaining the accuracy of head classes. On the skin lesion datasets, PAD-UFES and HAM10000, the F1-scores of our method improved by 4% and 3%, respectively compared to the original diffusion probabilistic model.

Anisotropic Diffusion Probabilistic Model for Imbalanced Image Classification

TL;DR

This work tackles long-tailed medical image classification by tailoring diffusion probabilistic models to per-class data distributions. It introduces Anisotropic Diffusion Probabilistic Model (ADPM), which assigns class-specific diffusion speeds through a noise level parameter guided by an imbalance-sensitive bound, enabling tail classes to learn more robust decision boundaries. The model also integrates global/local priors in the forward process and a feature-conditioned reverse prior via cross-attention, training with a multi-task objective that includes MMD regularization to align noise samples with estimated Gaussian noise. Empirically, ADPM improves tail-class accuracy across four medical datasets (PAD-UFES, HAM10000, SCIN, Hyper-Kvasir) while maintaining head-class performance, achieving up to ~4% gains in F1 on PAD-UFES and ~3% on HAM10000. These results demonstrate the practical impact of principled anisotropic diffusion and prior-informed learning for imbalanced medical imaging tasks, with potential for broader applications in long-tailed visual recognition.

Abstract

Real-world data often has a long-tailed distribution, where the scarcity of tail samples significantly limits the model's generalization ability. Denoising Diffusion Probabilistic Models (DDPM) are generative models based on stochastic differential equation theory and have demonstrated impressive performance in image classification tasks. However, existing diffusion probabilistic models do not perform satisfactorily in classifying tail classes. In this work, we propose the Anisotropic Diffusion Probabilistic Model (ADPM) for imbalanced image classification problems. We utilize the data distribution to control the diffusion speed of different class samples during the forward process, effectively improving the classification accuracy of the denoiser in the reverse process. Specifically, we provide a theoretical strategy for selecting noise levels for different categories in the diffusion process based on error analysis theory to address the imbalanced classification problem. Furthermore, we integrate global and local image prior in the forward process to enhance the model's discriminative ability in the spatial dimension, while incorporate semantic-level contextual information in the reverse process to boost the model's discriminative power and robustness. Through comparisons with state-of-the-art methods on four medical benchmark datasets, we validate the effectiveness of the proposed method in handling long-tail data. Our results confirm that the anisotropic diffusion model significantly improves the classification accuracy of rare classes while maintaining the accuracy of head classes. On the skin lesion datasets, PAD-UFES and HAM10000, the F1-scores of our method improved by 4% and 3%, respectively compared to the original diffusion probabilistic model.
Paper Structure (18 sections, 3 theorems, 34 equations, 7 figures, 9 tables, 1 algorithm)

This paper contains 18 sections, 3 theorems, 34 equations, 7 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Anisotropic diffusion probabilistic model for classification
  3. Notations and preliminaries
  4. Denoising Diffusion Probabilistic Model
  5. Anisotropic Diffusion Probabilistic Model
  6. Imbalance-sensitive noise distribution
  7. Priors for Diffusion Probabilistic Model
  8. Forward prior in diffusion model
  9. Reverse prior in diffusion probabilistic model
  10. Training objective and algorithm
  11. Numerical experiments
  12. Implementation details
  13. Datasets
  14. Evaluation on imbalance-sensitive noise
  15. Ablation study
  16. ...and 3 more sections

Key Result

Lemma 2.1

For $f\in C(\mathcal{F})$ of the fixed size $m$, let the size of dataset $\mathcal{S}_1$ and $\mathcal{S}_2$ be $n_1$ and $n_2$ with $n_1<n_2$. Then the expected of loss $L(f|m,n_1) \geq L(f|m,n_2)$.

Figures (7)

  • Figure 1: The generalization error of the classification model, where the tail classes contribute more to the generalization error.
  • Figure 1: Overview of our anisotropic diffusion probabilistic model, where (a) the training phase (forward process) and (b) the inference phase (reverse process).
  • Figure 1: The t-SNE obtained from the final feature embedding by the diffusion reverse process on PAD-UFES dataset, where $c$ is chosen as $c=5$ and $\alpha$ is chosen from $[0,1]$. \newlabelt-sne0
  • Figure 2: The motivation behind our framework is to apply anisotropic noise to different classes during the diffusion process, effectively balancing the diffusion speed between head and tail classes.
  • Figure 2: The t-SNE obtained from the denoised feature embedding by the diffusion reverse process during inference on PAD-UFES and HAM10000 dataset, where $t$ is the current diffusion time step for inference. As the time step encoding progresses, the noise is gradually removed, thereby obtaining a clear distribution of classes. \newlabelconvergence0
  • ...and 2 more figures

Theorems & Definitions (6)

  • Lemma 2.1
  • Proof 1
  • Proposition 2.2
  • Proof 2
  • Theorem 2.3
  • Proof 3