[論文レビュー] SegTransVAE: Hybrid CNN -- Transformer with Regularization for medical image segmentation

Quan-Dung Pham, Hai Nguyen-Truong+2 more|arXiv (Cornell University)|Jan 21, 2022

AI in cancer detection参考文献 12被引用数 5

ひとこと要約

SegTransVAE は、医用画像分類のためのハイブリッド 3D CNN-Transformer アーキテクチャを提案し、変分自己オートエンコーダ（VAE）ブランチを備える。3D CNN による局所的特徴抽出、トランスフォーマーによるグローバルな文脈モデリング、VAE による自己教師あり再構成を統合することで、一般化性能を向上させる。BraTS および KiTS19 データセットにおいて、3D U-Net や UNETR、SegresnetVAE よりもDiceスコアおよび95％ヒューベルト距離で優れた性能を達成し、軽量CNNと同等の推論速度を維持する。

ABSTRACT

Current research on deep learning for medical image segmentation exposes their limitations in learning either global semantic information or local contextual information. To tackle these issues, a novel network named SegTransVAE is proposed in this paper. SegTransVAE is built upon encoder-decoder architecture, exploiting transformer with the variational autoencoder (VAE) branch to the network to reconstruct the input images jointly with segmentation. To the best of our knowledge, this is the first method combining the success of CNN, transformer, and VAE. Evaluation on various recently introduced datasets shows that SegTransVAE outperforms previous methods in Dice Score and $95\%$-Haudorff Distance while having comparable inference time to a simple CNN-based architecture network. The source code is available at: https://github.com/itruonghai/SegTransVAE.

研究の動機と目的

• 医用画像分類におけるCNNの長距離空間的依存関係の捉えにくさ、およびトランスフォーマーのインダクティブバイアスの欠如と低データ環境下での一般化の難しさを解決すること。
• VAEブランチを介した分類と自己再構成の共同学習により、過学習を低減し、特徴学習を強化することで、分類精度を向上させること。
• 3D CNN（局所的不変性）、トランスフォーマー（グローバルアテンション）、VAE（正則化）の長所を統合したエンコーダ-デコーダフレームワークに統合すること。
• 挑戦的な医用画像ベンチマークで最先端の性能を達成するとともに、競争力ある推論速度を実現すること。

提案手法

• エンコーダは、インスタンス正則化とLeaky ReLUを用いた変更済みResNetブロックをスタックし、入力ボリュームからの階層的3D局所特徴を抽出する。
• 3D特徴マップをフラット化し、学習可能な線形投影によりd次元の埋め込み空間に変換する。空間構造を保持するために、学習可能な位置埋め込みを追加する。
• マルチヘッドアテンションとフィードフォワードネットワークのスタックにより、トランスフォーマーのエンコーダを構築し、すべてのパッチにわたる自己アテンションによりグローバル特徴モデリングを実現する。
• 特徴マッピングモジュールにより、トランスフォーマー出力を再び3D特徴マップに変換し、畳み込みブロックによるチャネル次元の低減を経てデコーダの入力とする。
• 最終的なトランスフォーマーレイヤーにVAEブランチを接続し、入力画像の再構成を実行することで、正則化を提供し、限られたデータ環境下での一般化を向上させる。
• デコーダは、エンコーダからのスキップ接続を伴う3Dトランスポジット畳み込みを用い、高解像度のセグメンテーションマスクを再構成する。

実験結果

リサーチクエスチョン

• RQ1 VAEベースの自己教師あり再構成を備えたハイブリッドCNN-Transformerアーキテクチャは、医用画像分類タスクにおける分類性能を向上させることができるか？
• RQ2 VAE正則化の統合により、過学習が低減され、特に小規模な医用データセットにおいて一般化性能が向上するか？
• RQ3 局所的（CNN）およびグローバル的（トランスフォーマー）特徴学習の組み合わせは、純粋なCNNまたは純粋なトランスフォーマーのモデルを上回る分類精度を達成できるか？
• RQ4 既存の最先端モデルと比較して、提案アーキテクチャの推論速度およびパラメータ効率はどの程度か？

主な発見

• BraTS 2021 データセットにおいて、SegTransVAEは強化腫瘍（ET）で85.48％、全腫瘍（WT）で90.52％、腫瘍コア（TC）で92.60％のDiceスコアを達成し、3D U-Net、UNETR、SegresnetVAEを上回った。
• 95％ヒューベルト距離は、ETで2.89 mm、WTで3.57 mm、TCで5.84 mmであり、ベースライン手法よりも顕著な改善を示した。
• KiTS19 データセットにおいて、SegTransVAEは腎臓で95.28％、腫瘍で66.31％のDiceスコアを達成し、3D U-Net、UNETR、SegresnetVAEの両方の指標を上回った。
• 腎臓セグメンテーションの95％ヒューブエルト距離は3.28 mmであり、3D U-Net（6.32 mm）や UNETR（8.84 mm）よりも顕著に低かった。
• SegTransVAEは4470万パラメータを有し、推論時間は0.45秒であり、3D U-Net（560万パラメータ、0.45秒）と同等で、SegresnetVAE（750万パラメータ、0.55秒）より高速であった。
• 3D U-Net よりもパラメータ数が多いものの、GFLOPsは607.5と3D U-Net（1000以上）を下回っており、より高い計算効率を示した。