Application of a diffusion model with improved adaptive class-balancing criteria for imbalanced datasets

Abstract

This article addresses the problem of class imbalance, which remains one of the key limitations in building image recognition systems in computer vision applications. The imbalance between classes, in which some categories are represented by thousands of samples while others contain only a few dozen, leads to a systematic bias of models towards dominant classes and to poor recognition of rare categories, which is unacceptable in a clinical context. Generative models have become a common instrument for addressing this problem, and diffusion models in particular now offer a more stable and higher-quality alternative to generative-adversarial networks. However, one component of such pipelines has received considerably less attention than the generative part itself: the criteria that decide how a limited number of synthetic images should be distributed among classes. Existing approaches derive this decision from the accuracy of the classifier measured on a single epoch, which is unreliable for classes represented by a few dozen samples and carries no information about which classes the errors are directed to. The article proposes improved adaptive class-balancing criteria that address both shortcomings. The first component is an exponentially smoothed estimate of per-class accuracy, which preserves sensitivity to persistent weakness of a class while suppressing fluctuations caused by measurement on very few samples. The second component is an error-concentration term derived from the confusion matrix, which distinguishes a class whose errors are directed almost entirely at one neighbouring category from a class whose errors are spread across many, and adjusts its share of generated images accordingly. The proposed criteria were evaluated on the HAM10000 dataset, which contains over 10 000 dermatoscopic images of seven classes with an imbalance ratio of 58:1, using a class-conditional diffusion model as the generator and a single shared pool of synthetic images, so that the criteria could be isolated from every other factor. The results demonstrate that the proposed criteria improve balanced accuracy, macro-averaged F1-score and the Matthews correlation coefficient compared with the accuracy-based criterion, with the improvement concentrated on the smallest classes and the recognition of the dominant class left unchanged. The variability of the per-class allocation across update points is reduced by 63% on average and by 95% for the vascular lesion class, which confirms directly the instability of accuracy-based criteria. The results obtained show that the way synthetic data is distributed among classes affects the outcome more than the fact of its addition, and that careful design of the allocation criteria can compensate for a substantial part of the computational budget normally spent on generation.