Stochastic Differential Equation Approach as Uncertainty-Aware Feature Recalibration Module in Image Classification

Abstract

Despite significant advancements in image classification, deep learning models struggle to accurately discern fine details in images, producing overly confident and imbalanced predictions for certain classes. These models typically employ feature recalibration techniques but do not account for the underlying uncertainty in predictions—particularly in complex sequential tasks like image classification. These uncertainties can significantly impact the reliability of subsequent analyses, potentially compromising accuracy across various applications. To address these limitations, we introduce the Stochastic Differential Equation Recalibration Module (SDERM), a novel approach designed to dynamically adjust the channel-wise feature responses in convolutional neural networks. It integrates a stochastic differential equation (SDE) framework into a feature recalibration module to capture the inherent uncertainties in the data and its model predictions. To the best of our knowledge, our study is the first to explore the integration of SDE-based feature recalibration modules in image classification. We build SDERM based on two interconnected networks—drift and diffusion network. The drift network serves as a deterministic component that approximates the predictive function of the model that systematically influences recalibrations of the predictions without considering the randomness. Concurrently, the diffusion network uses the Wiener process that captures the inherent uncertainties within the data and the network's predictions. We tested the classification accuracy of SDERM in ResNet50, ResNet101, and ResNet152 against other recalibration modules, including Squeeze-Excitation (SE), Convolutional Block Attention Module (CBAM), Gather and Excite (GE), and Position-Aware Recalibration Module (PARM), as well as the original Bottleneck architecture. Public image classification datasets were used, including CIFAR-10, SVHN, FashionMNIST, and HAM10000, and their classification accuracies were evaluated using the F1 score. The proposed ResNetSDE architecture achieved state-of-the-art F1 scores across four of five benchmark datasets. On Fashion-MNIST, ResNetSDE attained an F1 score of 0.937 (CI: 0.932–0.941), outperforming all baseline recalibration methods by margins of 0.9%–1.3%. For CIFAR-10 and CIFAR-100, ResNetSDE achieved 0.886 (CI: 0.879–0.892) and 0.962 (CI: 0.958–0.965), respectively, surpassing ResNet-GE and ResNet-CBAM by 3.5% and 1.3%, respectively. ResNetSDE dominated SVHN with an F1 of 0.956 (CI: 0.953–0.958), a significant improvement over ResNet-CBAM's 0.948 (CI: 0.945–0.951). While ResNet-CBAM led on the class-imbalanced HAM10000 (0.770, CI: 0.758–0.782), ResNetSDE remained competitive (0.768, CI: 0.749–0.786) since its consistent superiority—evidenced by narrow confidence intervals—validates its efficacy as a feature recalibration framework. Our experiments demonstrate that SDERM can outperform existing feature recalibration modules in image classification. The integration of SDERM to ResNet enables leveraging adaptability toward the stochasticity of each dataset at various depths of the architecture in image classification where uncertainty plays a fundamental role.