[論文レビュー] UNesT: Local Spatial Representation Learning with Hierarchical Transformer for Efficient Medical Segmentation

Xin Yu, Qi Yang+13 more|arXiv (Cornell University)|Sep 28, 2022

Advanced Neural Network Applications被引用数 4

ひとこと要約

UNesTは、ブロック単位の特徴量集約を通じて局所的な空間表現を強化する階層的3次元トランスフォーマーに基づく医療画像セグメンテーションモデルを提案する。このモデルは、133クラスの全脳および腎臓亜構造セグメンテーションという挑戦的なタスクで最先端の性能を達成した。コリンデータセットでは0.8505の平均DSCを達成し、27ネットワークアンサンブルを上回る性能を示した。また、少量のデータでも高いデータ効率と頑健性を示した。

ABSTRACT

Transformer-based models, capable of learning better global dependencies, have recently demonstrated exceptional representation learning capabilities in computer vision and medical image analysis. Transformer reformats the image into separate patches and realizes global communication via the self-attention mechanism. However, positional information between patches is hard to preserve in such 1D sequences, and loss of it can lead to sub-optimal performance when dealing with large amounts of heterogeneous tissues of various sizes in 3D medical image segmentation. Additionally, current methods are not robust and efficient for heavy-duty medical segmentation tasks such as predicting a large number of tissue classes or modeling globally inter-connected tissue structures. To address such challenges and inspired by the nested hierarchical structures in vision transformer, we proposed a novel 3D medical image segmentation method (UNesT), employing a simplified and faster-converging transformer encoder design that achieves local communication among spatially adjacent patch sequences by aggregating them hierarchically. We extensively validate our method on multiple challenging datasets, consisting of multiple modalities, anatomies, and a wide range of tissue classes, including 133 structures in the brain, 14 organs in the abdomen, 4 hierarchical components in the kidneys, inter-connected kidney tumors and brain tumors. We show that UNesT consistently achieves state-of-the-art performance and evaluate its generalizability and data efficiency. Particularly, the model achieves whole brain segmentation task complete ROI with 133 tissue classes in a single network, outperforming prior state-of-the-art method SLANT27 ensembled with 27 networks.

研究の動機と目的

• ビジョントランスフォーマーが局所的な空間情報を保持できないことや、3次元医療画像における大規模で多様な組織構造を処理できないという限界を解決すること。
• 医療画像で一般的な少量データ環境におけるデータ効率とモデルの頑健性を向上させること。
• 脳や腎臓の亜構造のような、数100もの組織クラスや階層的に接続された構造を正確にセグメンテーションできること。
• アンサンブルモデルを上回る性能と効率を実現する単一モデルソリューションを開発すること。
• 放射線科医によるアノテーションを実施した、独自に構築した腎臓亜構造CTデータセットを整備し、検証および一般化性の評価を可能にすること。

提案手法

• UNesTは、空間的に隣接するパッチをブロックにグループ化する階層的3次元ブロック集約モジュールを採用しており、局所的な通信を可能にしつつ、階層レベル間でグローバル自己注意を維持する。
• U-Netに類似したアーキテクチャを採用し、ネストされたトランスフォーマー畳み込みエンコーダーと畳み込みデコーダーを、複数の解像度でスキップ接続によって接続する。
• 局所的空間表現を強化するために、トランスフォーマー・エンコーダーを階層的にスタックし、グローバルな文脈を維持しながら、局所構造に対する誘導的バイアスを向上させる。
• 計算コストを低減しつつ、高解像度3次元ボリュームでも性能を維持できる、簡素化され、収束が速いトランスフォーマー設計を統合する。
• エンコーダーとデコーダー間のスキップ接続により、マルチスケール特徴量の統合が可能となり、小さなまたは複雑な構造の局所化精度が向上する。
• T1強調MRIやCTを含むマルチモodalデータセット上で、広範なデータ拡張と正則化を用いて、エンドツーエンドで訓練される。

実験結果

リサーチクエスチョン

• RQ1 階層的トランスフォーマー・アーキテクチャは、3次元医療画像においてグローバルな文脈を保持しつつ、局所的な空間表現を効果的に学習できるか？
• RQ2 UNesTは、27ネットワークアンサンブルを含む既存のモデルと比較して、133組織クラスを有する全脳セグメンテーションタスクで優れた性能を示せるか？
• RQ3 特に希少または複雑な亜構造（例：腎臓部品）に対して、UNesTは少量データ環境でも優れた性能を示せるか？
• RQ4 モデルは、異なる解剖的領域、モodal、組織の複雑さレベルに対してどの程度一般化できるか？
• RQ5 豊富な手動アノテーションを必要とせずに、UNesTは腎臓亜構造セグメンテーションにおいて放射線科医レベルの再現性を達成できるか？

主な発見

• コリンデータセットを用いた全脳セグメンテーションタスクにおいて、UNesTは0.8505の平均DSCを達成し、従来の最先端技術（0.7444）を著しく上回った。
• CANDIデータセットでは、平均DSCを0.6968から0.7025に向上させ、挑戦的で公開済みのベンチマークで高い性能を示した。
• 27ネットワークアンサンブルであるSLANT27を上回り、単一モデルで高い精度を達成した。これにより、推論コストと複雑さが低減された。
• 腎臓亜構造セグメンテーションタスクにおいて、UNesTは0.7906のDSCを達成し、複雑で階層的に接続された腎臓部品においても優れた性能を示した。
• モデルは高い再現性を示し、放射線科医との間で体積測定の平均差が0.03にとどまり、間接的評価の一致度（0.29）を上回った。
• 外れ値分析では、90％のセグメンテーションが信頼できる結果を示したが、残りの10％は誤って過小または過大にセグメンテーションされていた。これらは視覚的検査で検出可能であり、後処理により補正可能である可能性がある。