[論文レビュー] U-Net vs Transformer: Is U-Net Outdated in Medical Image Registration?

Xi Jia, Joseph W. Bartlett+4 more|arXiv (Cornell University)|Aug 7, 2022

Brain Tumor Detection and Classification被引用数 5

ひとこと要約

本論文では、受容 field を拡大するために大口径畳み込みブロックを組み込んだ変更版 U-Net である LKU-Net を提案する。3D 医療画像登録において、TransMorph などの最新のトランスフォーマー基盤手法を上回ることを示しており、パラメータが TransMorph の 1.12%、乗加算演算量が 10.8% のみで、アトラスベースおよび被験者間登録の両タスクで優れた性能を達成している。これは、U-Net がこの分野で陳腐化しているという考えを覆すものである。

ABSTRACT

Due to their extreme long-range modeling capability, vision transformer-based networks have become increasingly popular in deformable image registration. We believe, however, that the receptive field of a 5-layer convolutional U-Net is sufficient to capture accurate deformations without needing long-range dependencies. The purpose of this study is therefore to investigate whether U-Net-based methods are outdated compared to modern transformer-based approaches when applied to medical image registration. For this, we propose a large kernel U-Net (LKU-Net) by embedding a parallel convolutional block to a vanilla U-Net in order to enhance the effective receptive field. On the public 3D IXI brain dataset for atlas-based registration, we show that the performance of the vanilla U-Net is already comparable with that of state-of-the-art transformer-based networks (such as TransMorph), and that the proposed LKU-Net outperforms TransMorph by using only 1.12% of its parameters and 10.8% of its mult-adds operations. We further evaluate LKU-Net on a MICCAI Learn2Reg 2021 challenge dataset for inter-subject registration, our LKU-Net also outperforms TransMorph on this dataset and ranks first on the public leaderboard as of the submission of this work. With only modest modifications to the vanilla U-Net, we show that U-Net can outperform transformer-based architectures on inter-subject and atlas-based 3D medical image registration. Code is available at https://github.com/xi-jia/LKU-Net.

研究の動機と目的

• U-Net 基盤の手法が、最新のトランスフォーマー基盤アーキテクチャと比較して、医療画像登録分野で陳腐化しているかどうかを調査すること。
• トランスフォーマーがモデル化する長距離依存性が、正確な 3D 医療画像登録に不可欠かどうかを評価すること。
• 自己注意機構に依存せずに受容 field を拡大するように変更された U-Net アーキテクチャを設計すること。
• 軽量で効率的なアーキテクチャを用いて、公的 3D 医療画像登録ベンチマークで最先端の性能を達成すること。
• 単一の洗練された U-Net 変種が、複雑でパラメータが多いトランスフォーマー モデルを上回ることを実証すること。

提案手法

• 標準 U-Net のエンコーダーおよびデコーダー パathに並列で大口径畳み込みブロックを統合した大口径畳み込み U-Net である LKU-Net を提案する。
• 大口径畳み込み（例：15×15×15 のカーネルサイズ）を大口径ブロックに用い、深さを増さずに広い領域の空間的文脈を捉える。
• 特徴の融合と勾配の流れを保つために、標準の U-Net スキップ接続アーキテクチャを維持する。
• エンド・トゥ・エンドで訓練可能な、ワーピング層、データ適合性項、正則化を含む標準の非教師あり損失を採用する。
• IXI および Learn2Reg 2021 データセットで、公平な比較のため、TransMorph と同一のデータ前処理、訓練プロトコル、評価指標を適用する。
• ヤコビアン行列式正則化を有無にかかわらず、同じネットワークアーキテクチャで微分可能および非微分可能な登録を評価する。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマー基盤のネットワークが医療画像登録で性能向上を達成するのは、主にその長距離モデリング能力のおかげなのか？
• RQ2 自己注意機構を使用しない変更版 U-Net アーキテクチャが、最先端の性能を達成できるか？
• RQ3 3D 脳画像における平均変位が 2.1–2.3 ピxls にとどまることは、登録に長距離依存性が不要であることを示唆するのか？
• RQ4 軽量な U-Net 変種が、パラメータ数が多く複雑なトランスフォーマー基盤モデル（例：TransMorph）を、正確性および効率性の両面で上回れるか？
• RQ5 U-Net アーキテクチャは現代の 3D 医療画像登録分野で陳腐化しているのか、それともアーキテクチャ的改良によって競争力を持続できるのか？

主な発見

• 3D IXI 脳データセットにおいて、LKU-Net は 0.757 の Dice スコアを達成し、TransMorph（0.746）を上回り、パラメータはその 1.12% のみを消費した。
• LKU-Net は計算コストを TransMorph の 10.8% まで削減しながら、登録の正確性を維持または向上させた。
• MICCAI Learn2Reg 2021 タスク3 データセットにおいて、LKU-Net は 0.8861 の Dice スコアを達成し、公的検証リーダーボードで第1位となり、TransMorph-LC や優勝した LapIRN を上回った。
• 8 チャネルおよびカーネルサイズ 15 を有する LKU-Net 変種は、アトラスから被験者への登録で 0.757 の Dice スコアを達成し、TransMorph や nnFormer を含む他のすべての手法を上回った。
• 微分可能登録において、LKU-Net-diff 8,5 は 0.753 の Dice スコアを達成し、TransMorph-diff（0.599）を顕著に上回り、TransMorph-bspl（0.752）と同等の性能を示した。
• 可視化比較では、LKU-Net は特に複雑な解剖的構造を有する領域で滑らかな変形とより現実的な変形画像を生成した。