[論文レビュー] Optimized U-Net for Brain Tumor Segmentation

Michał Futrega, Alexandre Milesi+2 more|arXiv (Cornell University)|Oct 7, 2021

Brain Tumor Detection and Classification被引用数 8

ひとこと要約

この論文は、BraTS21チャレンジにおける脳腫瘍セグメンテーションの最適化されたU-Netアーキテクチャを提示しており、検証フェーズで1位、テストフェーズで3位を達成した。手法は、ディープサブスクリプション、チャネル数を増加させた深層エンコーダー、フォアグラウンドボクセルのワンホットエンコーディング、およびDiceスコアの向上を図る特化したポストプロセッシング戦略を組み合わせたものである。

ABSTRACT

We propose an optimized U-Net architecture for a brain tumor segmentation task in the BraTS21 challenge. To find the optimal model architecture and the learning schedule, we have run an extensive ablation study to test: deep supervision loss, Focal loss, decoder attention, drop block, and residual connections. Additionally, we have searched for the optimal depth of the U-Net encoder, number of convolutional channels and post-processing strategy. Our method won the validation phase and took third place in the test phase. We have open-sourced the code to reproduce our BraTS21 submission at the NVIDIA Deep Learning Examples GitHub Repository.

研究の動機と目的

• マルチモodal MRIスキャンにおける自動脳腫瘍セグメンテーションのための高性能なU-Netベースモデルの開発。
• 分類精度の向上を図るためのアーキテクチャ的要素とトレーニング戦略の最適な組み合わせの特定。
• 限られた医療トレーニングデータと腫瘍の外観のばらつきに直面し、体系的なアブレーションスタディを通じて解決策を模索。
• アーキテクチャ的最適化とトレーニングスケジュール最適化を用いて、BraTS21ベンチマークで最先端のパフォーマンスを達成すること。

提案手法

• 勾配の流れとトレーニングの安定性を向上させるために、ディープサブスクリプションを用いたU-Netアーキテクチャを採用。
• エンコーダーの深さを6レベルから7レベルに増加させ、畳み込みフィルタ数をスケーリングすることで表現力の向上を図った。
• フォアグラウンド腫瘍ボクセルのワンホットエンコーディングチャンネルを入力に追加し、空間的認識を強化。
• 小さな低信頼度の強化腫瘍（ET）領域を除去する多段階ポストプロセッシング戦略を適用し、誤検出を低減。
• 腫瘍セグメンテーションにおけるクラス不均衡に対処するため、Focal損失とDice損失を組み合わせたハイブリッド損失関数を採用。
• 複数のU-Netバリアントと要素を対象に、5分割交差検証を用いた広範なアブレーションスタディを実施。

実験結果

リサーチクエスチョン

• RQ1 基本U-Net、アテンションU-Net、リーマンスU-Net、SegResNetVAE、またはUNETRの中から、BraTS21で最も優れた性能を示すU-Netバリアントはどれか？
• RQ2 ディープサブスクリプション、ドロップブロック、リーマンス接続といったアーキテクチャ的要素が、セグメンテーションパフォーマンスに与える影響は何か？
• RQ3 エンコーダーの深さ、チャネル数、入力前処理の最適な組み合わせは何か？ これは最大のDiceスコアを達成するためのものである。
• RQ4 ポストプロセッシングは、特に小さな低信頼度ET領域に対して、最終的なセグメンテーションパフォーマンスにどのように影響を与えるか？
• RQ5 BraTS21データセットにおける一般化性能を最大化するための最適なトレーニングスケジュールと損失関数の設定は何か？

主な発見

• ディープサブスクリプションを備えたU-Netが、全テストバリアントの中で最高の平均Diceスコア0.9149を達成し、ベースラインU-Net（0.9130）を上回った。
• より深いエンコーダー（深さ7）、チャネル数の増加（64–512）、フォアグラウンドワンホット入力の組み合わせにより、Diceスコアは0.9156に向上した。
• すべての最適化（より深いエンコーダー、チャネル数の変更、ワンホット入力、ポストプロセッシング）を適用した最終モデルは、検証フェーズで平均Diceスコア0.9163を達成した。
• ポストプロセッシングにより、低信頼度の小さなET成分が除去され、誤検出が低減され、スコアは0.9156から0.9163に向上した。
• モデルはBraTS21の検証フェーズで優勝し、最終テストフェーズで3位を獲得した。
• Focal損失とドロップブロックは、ベースラインに対して顕著な性能向上をもたらさなかったことから、この設定では限定的な利点しか得られなかった。