[论文解读] Constrained Optimization to Train Neural Networks on Critical and Under-Represented Classes

Sara Sangalli, Ertunç Erdil+3 more|arXiv (Cornell University)|Feb 21, 2021

Artificial Intelligence in Healthcare被引用 10

一句话总结

本文提出一种基于增广拉格朗日法（ALM）的约束优化框架，用于在类别极度不平衡的数据集中，提升对关键但样本稀少类别的模型性能。通过在高真正率（TPR）下不对称地优先降低假正率（FPR），该方法在保持高敏感性的同时，提升了AUC并减少了对少数类（关键类）的误分类，在医学影像和视觉基准测试中优于现有损失函数。

ABSTRACT

Deep neural networks (DNNs) are notorious for making more mistakes for the classes that have substantially fewer samples than the others during training. Such class imbalance is ubiquitous in clinical applications and very crucial to handle because the classes with fewer samples most often correspond to critical cases (e.g., cancer) where misclassifications can have severe consequences. Not to miss such cases, binary classifiers need to be operated at high True Positive Rates (TPRs) by setting a higher threshold, but this comes at the cost of very high False Positive Rates (FPRs) for problems with class imbalance. Existing methods for learning under class imbalance most often do not take this into account. We argue that prediction accuracy should be improved by emphasizing reducing FPRs at high TPRs for problems where misclassification of the positive, i.e. critical, class samples are associated with higher cost. To this end, we pose the training of a DNN for binary classification as a constrained optimization problem and introduce a novel constraint that can be used with existing loss functions to enforce maximal area under the ROC curve (AUC) through prioritizing FPR reduction at high TPR. We solve the resulting constrained optimization problem using an Augmented Lagrangian method (ALM). Going beyond binary, we also propose two possible extensions of the proposed constraint for multi-class classification problems. We present experimental results for image-based binary and multi-class classification applications using an in-house medical imaging dataset, CIFAR10, and CIFAR100. Our results demonstrate that the proposed method improves the baselines in majority of the cases by attaining higher accuracy on critical classes while reducing the misclassification rate for the non-critical class samples.

研究动机与目标

• 解决深度学习中类别不平衡的严峻挑战，特别是在医学影像中，将罕见但严重的病例（如癌症）误分类具有高昂的临床代价。
• 通过聚焦于在高真正率工作点上降低假正率，提升少数（关键）类别的模型性能。
• 提出一种系统化且可微分的约束，可与现有损失函数集成，以不对称方式优化AUC，优先降低高TPR下的FPR。
• 通过两种提出的公式将二分类约束扩展至多分类场景，以增强广泛适用性。
• 在包括自研医学MRI、CIFAR10和CIFAR100在内的多种数据集上，展示在不同类别不平衡比例下的一致性能提升。

提出的方法

• 将二分类DNN训练建模为约束优化问题，以在高TPR下不对称地强调降低FPR，从而最大化AUC。
• 使用Mann-Whitney U统计量定义AUC目标，并将其作为约束嵌入优化框架。
• 应用增广拉格朗日法（ALM）将约束问题转化为可通过随机梯度下降求解的对偶无约束问题。
• 将基于ALM的约束与9种不同的基线损失函数集成，包括二元交叉熵（BCE）、焦点损失、类别平衡BCE和LDAM。
• 提出两种用于多分类问题的扩展：ALM m,1 和 ALM m,2，以适应处理多个关键类别。
• 使用标准深度学习训练流程优化最终损失，同时保持与现有架构和训练流程的兼容性。

实验结果

研究问题

• RQ1 约束优化框架是否能提升深度神经网络在类别极度不平衡数据集中对样本稀少关键类别的性能？
• RQ2 不对称AUC最大化——在高TPR下优先降低FPR——是否能在二分类任务中提升临床实用性？
• RQ3 所提出的基于ALM的约束与现有类别不平衡损失函数相比，在FPR与TPR权衡方面表现如何？
• RQ4 该二分类约束能否有效扩展至多分类设置，同时保持对关键类别的性能增益？
• RQ5 该方法是否在包括医学影像和标准视觉基准在内的多种数据集中持续提升性能？

主要发现

• 在自研MRI数据集中，ALM + BCE在FNR为0%时实现54.0%的FPR，显著优于基线BCE（80.0%）和CB-BCE（100.0%）。
• ALM + BCE实现平均AUC为76.8%，集成AUC达85.4%，超过最佳基线（MBAUC）的82.4%集成AUC。
• 在长尾CIFAR10（不平衡比1:100）中，ALM m,2 + LDAM在80% TPR下实现24.76%错误率，90% TPR下为31.02%，优于所有基线，包括LDAM和CB-BCE。
• ALM m,2 + LDAM在80% TPR下实现73.32%总体准确率，90% TPR下为67.41%，表明在各项指标上均保持一致增益。
• 在大多数实验中，ALM m,2扩展相比ALM m,1进一步提升了总体准确率，尤其在与LDAM和CB-BCE结合时效果更显著。
• 在所有数据集和不平衡比例下，所提方法在多数情况下优于基线模型，特别是在高TPR下降低FPR方面表现突出，证实其在关键类别优先处理中的有效性。