Flat-Aware Cross-Stage Distilled Framework for Imbalanced Medical Image Classification

Medical data often follow imbalanced distributions, which poses a long-standing challenge for computer-aided diagnosis systems built upon medical image classification. Most existing efforts are conducted by applying re-balancing methods for the collected training samples, which improves the predictive performance for the minority class but at the cost of decreasing the performance for the majority. To address this paradox, we adopt a flat-aware cross-stage distilled framework (FCD), where we first search for flat local minima of the base training objective function on the original imbalanced dataset, and then continuously finetune this classifier within the flat region on the re-balanced one. To further prevent the performance decreasing for the majority, we propose a cross-stage distillation regularizing term to promote the optimized features to remain in the common optimal subspace. Extensive experiments on two imbalanced medical image datasets demonstrate the effectiveness of our proposed framework and its generality in improving the performance of existing imbalanced methods. The code of this work will be released publicly.