[論文レビュー] GradMask: Reduce Overfitting by Regularizing Saliency

Becks Simpson, Francis Dutil+2 more|arXiv (Cornell University)|Apr 16, 2019

Radiomics and Machine Learning in Medical Imaging参考文献 11被引用数 14

ひとこと要約

GradMask は、勾配に基づく帰属度を用いて、病変領域外の特徴に重要度を割り当てる活性化マップをペナルティ化することで、深層学習における過学習を軽減する新しい正則化手法である。この手法により、小規模な医療画像データセットにおいてテスト AUC が 1–4% 向上し、モデルの注目が正例セグメンテーションと整合的になるように制御される。

ABSTRACT

With too few samples or too many model parameters, overfitting can inhibit the ability to generalise predictions to new data. Within medical imaging, this can occur when features are incorrectly assigned importance such as distinct hospital specific artifacts, leading to poor performance on a new dataset from a different institution without those features, which is undesirable. Most regularization methods do not explicitly penalize the incorrect association of these features to the target class and hence fail to address this issue. We propose a regularization method, GradMask, which penalizes saliency maps inferred from the classifier gradients when they are not consistent with the lesion segmentation. This prevents non-tumor related features to contribute to the classification of unhealthy samples. We demonstrate that this method can improve test accuracy between 1-3% compared to the baseline without GradMask, showing that it has an impact on reducing overfitting.

研究の動機と目的

• スキャナーアーティファクトなどの生物学的でない特徴と誤った相関関係を学習してしまう、小規模な医療画像データセットにおける過学習を是正すること。
• 分類の過程で病変境界外の特徴に依存することを明示的に妨げることで、モデルの一般化性能を向上させること。
• 勾配に基づく活性化マップを活用して、誤った特徴帰属度をペナルティ化する正則化手法を開発すること。
• 正例セグメンテーションと整合的な活性化マップの正則化が、より良いテスト性能をもたらすことを実証すること。

提案手法

• 予測クラススコアの入力画素に関する勾配の L2 ノルムを、バイナリセグメンテーションマスクを用いて非病変領域でマスクすることでペナルティ化する損失項を導入する。
• 活性化マップとして、非健康クラス予測値 $\hat{y}_1$ の入力 $\mathbf{x}$ に関する勾配を用い、病変領域外の勾配を抑制するために $1 - \mathbf{x}_{\text{seg}}$ を適用する。
• 総損失を $\mathcal{L} = \sum_{\mathbf{x} \in D} \mathcal{L}_c + \left\| \frac{\partial \hat{y}_1}{\partial \mathbf{x}} \cdot (1 - \mathbf{x}_{\text{seg}}) \right\|_2^2$ と定式化する。ここで $\mathcal{L}_c$ は標準的な分類損失である。
• 健康と非健康予測の絶対差の勾配を計算することで、特徴の区別に寄与する部分を強調する「コントラスト」版の活性化マップを検討する。
• 医療画像ベンチマーク（MSD と BraTS）から得られる小規模データセットで微調整されたモデルを用い、ハイパーパramータのチューニングに検証 AUC を使用して、訓練中にこの手法を適用する。
• 入力勾配に基づく活性化マップを用いて、病変に関連しない特徴を特定・正則化し、注目が生物学的に関連する領域に局在化したままであるように保証する。

実験結果

リサーチクエスチョン

• RQ1 過学習が一般的に見られる小規模な医療画像データセットにおいて、活性化マップの明示的正則化が一般化性能を向上させるか？
• RQ2 病変境界外の勾配に基づく特徴重要度をペナルティ化することで、標準ベースラインと比較してテスト AUC に顕著な向上が得られるか？
• RQ3 訓練データが限られている（例：64–128 サンプル）状況下で、GradMask は過学習をどれほど軽減できるか？
• RQ4 健康と非健康予測の差に注目するコントラストベースの活性化マップバージョンは、クラス固有の活性化マップ正則化と比較して、より優れた性能を示すか？
• RQ5 Zech ら（2018）の先行研究で示されたように、病院固有の画像アーティファクトのような誤った相関関係への依存度が、GradMask によってどの程度軽減されるか？

主な発見

• GradMask は、4 つの医療画像データセット（肝臓、膵臓、脳、心臓）において、正則化なしのベースラインと比較してテスト AUC を 1–4% 向上させた。
• この手法は、極めて少ないサンプル数でも顕著な向上を示し、最初の正の結果が 64–128 個の訓練サンプルで観測された。
• BraTS 脳腫瘍データセットでは、テスト AUC がベースラインの 0.826 ± 0.026 から 0.798 ± 0.019 に低下したが、他のデータセットでは一貫して向上が確認された。
• コントラストベースの活性化マップバージョンは、複数のシードとデータセットにおいて、より高いテスト AUC を示し、一般化性能が標準バージョンを上回った。
• 検証 AUC を用いたハイパーパramータチューニングにより最適な設定が特定され、3 つの MSD データセットと 1 つの BraTS データセットで 20 個のランダムシードにわたって一貫した性能向上が得られた。
• この手法は、訓練データ内でクラスラベルと相関する非病変領域の特徴に対しても、モデルが重要度を割り当てることを妨げることで、過学習を効果的に軽減した。