[論文レビュー] TransClaw U-Net: Claw U-Net with Transformers for Medical Image Segmentation

Yao‐Wen Chang, Menghan Hu+2 more|arXiv (Cornell University)|Jul 12, 2021

Advanced Neural Network Applications被引用数 10

ひとこと要約

本論文では、エンコーダーにおいて局所的特徴抽出に畳み込みニューラルネットワーク（CNNs）と、グローバルな文脈モデリングにビジョントランスフォーマー（ViT）ブロックを統合したハイブリッドディーブラーニングアーキテクチャ、TransClaw U-Netを提案する。CNNベースの特徴学習と自己注意機構を組み合わせることで、512×512の入力解像度において、Synapseマルチオルガントンコムスキャンセグメンテーションデータセットで最先端の性能を達成し、平均DSCは80.39%を記録した。

ABSTRACT

In recent years, computer-aided diagnosis has become an increasingly popular topic. Methods based on convolutional neural networks have achieved good performance in medical image segmentation and classification. Due to the limitations of the convolution operation, the long-term spatial features are often not accurately obtained. Hence, we propose a TransClaw U-Net network structure, which combines the convolution operation with the transformer operation in the encoding part. The convolution part is applied for extracting the shallow spatial features to facilitate the recovery of the image resolution after upsampling. The transformer part is used to encode the patches, and the self-attention mechanism is used to obtain global information between sequences. The decoding part retains the bottom upsampling structure for better detail segmentation performance. The experimental results on Synapse Multi-organ Segmentation Datasets show that the performance of TransClaw U-Net is better than other network structures. The ablation experiments also prove the generalization performance of TransClaw U-Net.

研究の動機と目的

• 医用画像セグメンテーションにおけるCNNの長距離空間的依存関係の捉えにくさを是正すること。
• 自己注意機構によるグローバルな文脈モデリングをエンコーダーに統合することで、セグメンテーション精度を向上させること。
• デコーダーにおいて元のClaw U-Netアーキテクチャの詳細特徴回復能力を維持すること。
• 入力解像度、パッチサイズ、スキップ接続数がモデル性能に与える影響を評価すること。
• 腹部CTスキャンにおける複数の臓器にわたる一般化性とロバスト性を示すこと。

提案手法

• エンコーダーは二重スレッドアーキテクチャを採用：一方のスレッドは標準的な畳み込み層で浅い空間的特徴を抽出し、他方のスレッドは特徴マップをパッチとして処理し、自己注意に基づくグローバル表現学習を実施する。
• ビジョントランスフォーマーモジュールは、エンコードされた特徴マップ内の画像パッチ間の長距離依存関係を学習するためにマルチヘッド自己注意を適用する。
• デコーダーは、解像度回復時の細粒度な空間的詳細の保持を目的として、Claw U-Netのトップダウンアップサンプリング構造を踏襲する。
• 対応するエンコーダー・デコーダー層間のスキップ接続を確立することで、特徴の統合を促進し、局在化精度を向上させる。
• モデルは、交差エントロピー損失とDice損失を用い、Synapseデータセット上で最適なパフォーマンスを得られるようにハイパーパrameterを調整して、エンドツーエンドで学習される。
• アブレーションスタディにより、入力サイズ、パッチサイズ、スキップ接続数の影響が体系的に評価される。

Fig. 1 : The architecture of TransClaw U-Net

実験結果

リサーチクエスチョン

• RQ1 U-Netエンコーダーにトランスフォーマーを統合することで、標準的なCNNと比較して長距離特徴学習がどのように向上するか？
• RQ2 セグメンテーション精度と計算コストのバランスを考慮した場合、最適な入力画像解像度は何か？
• RQ3 パッチサイズがビジョントランスフォーマー部のシーケンス長とモデル性能に与える影響は何か？
• RQ4 スキップ接続数を変化させた場合、異なる臓器におけるセグメンテーション精度にどのような影響があるか？
• RQ5 提案されたTransClaw U-Netは、腹部CTスキャンにおける多様な解剖的構造に一般化して適応できるか？

主な発見

• 512×512の入力解像度を使用した場合、Synapseデータセットにおける平均Dice類似係数（DSC）は80.39%に達し、224×224解像度のベースライン（78.09% DSC）を上回った。
• 入力解像度を224×224から512×512に引き上げることでDSCが顕著に向上し、高解像度入力が特徴表現とセグメンテーション精度を向上させることを示した。
• 16×16パッチサイズでDSCが最大の78.09%を記録したが、より大きなパッチサイズ（24×24）では計算コストが増加する一方で性能は77.27%に低下した。
• スキップ接続数に強い正の影響がある：アーティファクトを完全に削除すると、大動脈ではDSCが10.8%低下し、膵臓では10.16%低下し、特徴統合におけるその重要性を確認した。
• アブレーションスタディにより、より大きな入力サイズ、より小さなパッチサイズ、より多くのスキップ接続が性能向上に寄与することが確認されたが、メモリ使用量と学習時間の増加を伴った。
• 最適設定下で、Synapseデータセットの全8臓器にわたって強力な一般化性を示し、肝臓（95.06%）、脾臓（91.16%）、胃（82.82%）において一貫したDSC向上を達成した。

Fig. 2 : Comparison of segmentation results of different models on the Synapse multi-organ CT dataset.