[論文レビュー] 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.

研究の動機と目的

• レアだが深刻な症例（例：がん）を誤分類すると臨床的コストが高いため、特に医療画像分野におけるクラス不均衡の深刻な課題に対処する。
• 高再現率（TPR）における誤検出率（FPR）の低減に焦点を当てることで、マイノリティ（深刻）クラスのモデル性能を向上させる。
• 既存の損失関数と統合可能な原理的で微分可能な制約を提案し、高TPRにおけるFPR低減を優先する非対称なAUC最適化を実現する。
• より広範な適用性を実現するため、二値分類用制約を二つの提案された定式化を通じて多クラス設定に拡張する。
• 医療用MRI、CIFAR10、CIFAR100を含む多様なデータセットにおいて、さまざまなクラス不均衡比の下で一貫した性能向上を示す。

提案手法

• AUCの最大化を目的とし、高TPRにおけるFPR低減に非対称的重みを付ける制約付き最適化問題として、二値DNN学習を定式化する。
• Mann-Whitney U統計量を用いてAUCの目的関数を定義し、最適化フレームワーク内に制約として埋め込む。
• 増強ラグランジュ法（ALM）を用いて制約付き問題を、確率的勾配降下法で解ける双対非制約問題に変換する。
• ALMベースの制約を、BCE、ファーコス損失、クラスバランスBCE、LDAMを含む9種類のベースライン損失関数と統合する。
• 多クラス問題への拡張として、2つの拡張形（ALM m,1 および ALM m,2）を提案し、複数の深刻クラスを扱えるように制約を適応させる。
• 標準的なディープラーニングのトレーニングパイプラインに従い、既存のアーキテクチャおよびトレーニング手順と互換性を保ったまま、得られた損失関数を最適化する。

実験結果

リサーチクエスチョン

• RQ1 制約付き最適化フレームワークは、クラス不均衡データセットにおけるマイノリティで深刻なクラスのDNN性能向上に寄与するか？
• RQ2 高TPRにおけるFPR低減を優先する非対称的AUC最大化は、二値分類タスクにおいて臨床的有用性を向上させるか？
• RQ3 提案されたALMベースの制約は、FPRとTPRのトレードオフの観点から、既存のクラス不均衡損失関数と比較して優れているか？
• RQ4 二値分類用制約は、性能向上を維持したまま多クラス設定に効果的に拡張可能か？
• RQ5 本手法は、医療画像および標準ビジョンベンチマークを含む多様なデータセットにおいて一貫した性能向上を示すか？

主な発見

• 自社MRIデータセットでは、ALM + BCEがFNR 0%でFPR 54.0%を達成し、ベースライン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低減の観点から、深刻クラスの優先的処理の有効性が確認された。