[論文レビュー] Adapting Pre-trained Vision Transformers from 2D to 3D through Weight Inflation Improves Medical Image Segmentation

Yuhui Zhang, Shih-Cheng Huang+3 more|arXiv (Cornell University)|Feb 8, 2023

Advanced Neural Network Applications被引用数 4

ひとこと要約

本論文は、事前学習された2次元ビジョントランスフォーマーを3次元医療画像セグメンテーションに適応させるための単純ながらも効果的な重み拡張戦略を提案する。この戦略により、深さ情報の保持を伴いながら転移学習が可能となり、計算コストの増加はたった0.75%にとどまる。12種類の多様な3次元医療画像データセットにおいて、最先端の性能を達成した。

ABSTRACT

Given the prevalence of 3D medical imaging technologies such as MRI and CT that are widely used in diagnosing and treating diverse diseases, 3D segmentation is one of the fundamental tasks of medical image analysis. Recently, Transformer-based models have started to achieve state-of-the-art performances across many vision tasks, through pre-training on large-scale natural image benchmark datasets. While works on medical image analysis have also begun to explore Transformer-based models, there is currently no optimal strategy to effectively leverage pre-trained Transformers, primarily due to the difference in dimensionality between 2D natural images and 3D medical images. Existing solutions either split 3D images into 2D slices and predict each slice independently, thereby losing crucial depth-wise information, or modify the Transformer architecture to support 3D inputs without leveraging pre-trained weights. In this work, we use a simple yet effective weight inflation strategy to adapt pre-trained Transformers from 2D to 3D, retaining the benefit of both transfer learning and depth information. We further investigate the effectiveness of transfer from different pre-training sources and objectives. Our approach achieves state-of-the-art performances across a broad range of 3D medical image datasets, and can become a standard strategy easily utilized by all work on Transformer-based models for 3D medical images, to maximize performance.

研究の動機と目的

• 事前学習された2次元ビジョントランスフォーマーを3次元医療画像セグメンテーションに適応させる課題に取り組み、転移学習の利点と深さ情報を保持する。
• 重み拡張が、2次元自然画像と3次元医療スキャンの間のドメインギャップを効果的に埋め合わせられるかを調査する。
• 3次元医療ビジョンタスクにおいて、事前学習された2次元モデルを活用するための汎用的で標準化可能な戦略を確立する。
• 異なる事前学習ソースと目的が3次元セグメンテーション性能に与える影響を評価する。
• 最小限のハイパーパrameterチューニングで、多様な3次元医療画像データセットにおいて一貫した性能向上を示す。

提案手法

• 中心化戦略を用いて、2次元重みを3次元に拡張することで、事前学習された2次元ビジョントランスフォーマーを3次元に適応させる。具体的には、中心スライスの重みをコピーし、他のスライスはゼロ初期化する。
• 3次元入力ボリュームを、数個の隣接スライスを含む小さな3次元ウィンドウに分割し、空間的文脈を保持する。
• 各ウィンドウ内で中心スライスのセグメンテーション予測のみを訓練対象とすることで、リーマン学習に類似させ、段階的かつ深さ方向の統合を可能にする。
• 標準的な3D U-Netスタイルのアーキテクチャを採用し、拡張された2次元トランスフォーマーをエンコーダーとして用いて、エンドツーエンドのセグメンテーションを実現する。
• アブレーションスタディを実施し、予測ターゲット、重み初期化、ストライド、入力スライス数などの異なる拡張戦略を比較する。
• 異なる解剖的領域と画像モodalities（CT/MRI）を有する11の追加3次元医療画像データセットを用いて、汎用性を評価する。

Figure 1: Approach overview. Large-scale pre-trained Transformers are used as the encoder in the segmentation model for transfer learning, in which weights are adapted using the inflation strategy to support 3D inputs. Each 3D image is split into windows, which contain a small number of neighbor sli

実験結果

リサーチクエスチョン

• RQ1 アーキテクチャの変更なしに、重み拡張戦略が事前学習された2次元ビジョントランスフォーマーを3次元医療画像セグメンテーションに効果的に適応できるか？
• RQ2 自己教師あり学習と教師あり学習のどちらの事前学習ソースが、3次元医療セグメンテーションにおける転移性能に良い影響を与えるか？
• RQ3 重み拡張を用いた3次元適応において、最適なウィンドウサイズと予測戦略（例：中心スライスのみ）は何か？
• RQ4 本手法は、異なる解剖的構造と画像プロトコルを持つ多様な3次元医療画像データセットに一般化可能か？
• RQ5 標準的な2次元または3次元のランダム初期化と比較して、重み拡張は性能をどの程度向上させるか？

主な発見

• 提案された重み拡張戦略により、強力なベースラインより平均DSCが1.95%向上し、計算コストの増加はたった0.75%にとどまる。
• 事前学習済み重みで初期化されたモデルは、マルチオルガンセグメンテーションデータセットでランダム初期化された3次元モデルよりも11.18%高い性能を示し、転移学習の重要性を示している。
• 中心化拡張戦略（中心スライスに2次元重みをコピーし、他のスライスはゼロ初期化）が最良の性能を達成しており、段階的かつ深さ方向の統合が有効であることが示唆されている。
• 隣接スライスを数枚（例：3〜5枚）含めるだけで十分な性能が得られ、遠く離れたスライスよりも隣接スライスがより関連性の高い文脈を提供することがわかった。
• 12のデータセットにおいて最先端の性能を達成し、主な評価を除く11の追加データセットのうち10つでも同様の優れた性能を示しており、強い汎用性を示している。
• 教師あり事前学習ソースが自己教師あり学習ソースよりも優れた転移性能を示しており、ラベル豊富な事前学習が3次元医療セグメンテーションにおいてより有益であることが示唆されている。

Figure 2: Transfer learning reduces noise in local predictions, and depth information further improves prediction consistency across slices. We show five consecutive input images, ground-truth labels, predictions from Ours , Ours w/o D , Ours w/o T&D .