[論文レビュー] Boundary loss for highly unbalanced segmentation

Hoel Kervadec, Jihene Bouchtiba+4 more|Espace ÉTS (ETS)|Dec 17, 2018

Advanced Neural Network Applications被引用数 6

ひとこと要約

本稿では、領域内画素和ではなく輪郭距離に基づく境界損失を提案し、極度に不均衡な医療画像セグメンテーションの訓練安定性と性能を向上させる。非対称L2距離をソフトマックス出力による領域的積分として表現することで、既存のネットワークとの容易な統合が可能となり、WMH や ISLES のような不均衡データセットにおいて、Dice スコアと安定性が顕著に向上する。

ABSTRACT

Widely used loss functions for CNN segmentation, e.g., Dice or cross-entropy, are based on integrals over the segmentation regions. Unfortunately, for highly unbalanced segmentations, such regional summations have values that differ by several orders of magnitude across classes, which affects training performance and stability. We propose a boundary loss, which takes the form of a distance metric on the space of contours, not regions. This can mitigate the difficulties of highly unbalanced problems because it uses integrals over the interface between regions instead of unbalanced integrals over the regions. Furthermore, a boundary loss complements regional information. Inspired by graph-based optimization techniques for computing active-contour flows, we express a non-symmetric $L_2$ distance on the space of contours as a regional integral, which avoids completely local differential computations involving contour points. This yields a boundary loss expressed with the regional softmax probability outputs of the network, which can be easily combined with standard regional losses and implemented with any existing deep network architecture for N-D segmentation. We report comprehensive evaluations and comparisons on different unbalanced problems, showing that our boundary loss can yield significant increases in performances while improving training stability. Our code is publicly available: https://github.com/LIVIAETS/surface-loss .

研究の動機と目的

• 白質高信号病変（WMH）など、小規模な病変検出を含む極度に不均衡なセグメンテーションタスクにおいて、標準的な領域ベース損失関数（例：交差エントロピー、Dice）が示す不安定さと低性能を是正すること。
• クラスの不均衡により、背景領域が大きな影響を及ぼすため、最適化が多数クラスに偏る問題を克服すること。
• セグメンテーション誤差が最も深刻な境界領域に注目するため、輪郭レベルの距離メトリクスを活用した損失関数を開発すること。
• 任意の既存のディープラーニングアーキテクチャおよびN次元セグメンテーションタスクと互換性を持たせ、アーキテクチャの変更なしに利用可能であること。
• 特に白質高信号病変のような小規模で稀な構造に対して、セグメンテーション精度と訓練安定性の両方を向上させること。

提案手法

• 輪郭空間における非対称L2距離に基づく境界損失を提案し、輪郭点への直接計算を回避する。
• ネットワークのソフトマックス確率出力に基づいて、この輪郭ベースの距離を領域的積分として表現することで、標準的なバックプロパゲーションによる微分可能訓練が可能になる。
• 予測済みおよび真値の確率マップにのみ依存する領域損失として境界損失を定式化し、任意の既存のCNNアーキテクチャと互換性を確保する。
• 一般化Dice損失などの標準的な領域損失と、重み付き和により組み合わせ、学習可能またはスケジューリング可能なハイパーパrameter α を用いる。
• 訓練中にαの再バランス戦略を用い、境界損失の寄与度を動的に調整することで、汎化性能の向上を図る。
• 特別なレイヤー型や複雑な後処理を必要とせず、標準的なディープラーニングフレームワークを用いて効率的に実装可能である。

実験結果

リサーチクエスチョン

• RQ1 標準的な領域損失と比較して、輪郭幾何に基づく損失関数が、極度に不均衡な医療画像データセットにおけるセグメンテーション性能を向上させることができるか？
• RQ2 境界までの距離情報に基づく境界損失は、小構造セグメンテーションにおける訓練の安定性と収束特性にどのように影響を与えるか？
• RQ3 境界損失と領域損失（例：一般化Dice）を組み合わせる最適な方法は何か？（ハイパーパrameterのスケジューリングと重み付けの観点から）
• RQ4 提案された境界損失は、滑らかで解剖学的に妥当なセグメンテーション境界をもたらす空間的正則化効果を有するか？
• RQ5 提案された境界損失は、二値セグメンテーションにとどまらず、多クラスまたは多構造セグメンテーション問題に対しても効果的に適用可能か？

主な発見

• WMH データセットでは、再バランスα戦略を用いてGDLと組み合わせた場合、DSCがGDL単独時の0.210から0.392に向上し、顕著な性能向上を示した。
• ISLES データセットでは、再バランスαを用いた境界損失がDSC 0.644（±0.026）を達成し、GDL単独（0.511）およびすべての定数α設定を上回った。
• 訓練DSCが類似している場合でも、検証DSCが顕著に向上したため、より良い一般化性能と過学習の低減が示された。
• 境界損失の訓練時間は標準損失と同等であり、平均バッチ時間0.210秒であり、計算オーバーヘッドは最小限に抑えられた。
• 境界損失は、特に小構造に対して滑らかで正確な輪郭を生成し、図1の可視化で確認されたように、見逃された病変が回復された。
• 再バランスαスケジューリング戦略は、α=1 や α=2 などのすべての定数α値を上回り、広範なハイパーパrameterチューニングなしに一貫した改善をもたらした。