[論文レビュー] Disruptive Autoencoders: Leveraging Low-level features for 3D Medical Image Pre-training

Jeya Maria Jose Valanarasu, Yucheng Tang+9 more|arXiv (Cornell University)|Jul 31, 2023

Radiomics and Machine Learning in Medical Imaging被引用数 5

ひとこと要約

本稿では、3D医療画像の自己教師付き事前学習フレームワークとして、局所的マスキング（トークンではなくチャネル埋め込みをマスキング）と低レベルの摂動（ノイズやダウンサンプリング）を組み合わせることで、低レベル特徴の学習を強化する、新しい自己符号化器であるDisruptive Autoencoders（DAE）を提案する。この手法は、BTCVマルチオルガニズムセグメンテーションチャレンジの公開テストリーダーボードで最良の性能を達成し、複数の3D医療画像セグメンテーションベンチマークで最先端の性能を発揮する。

ABSTRACT

Harnessing the power of pre-training on large-scale datasets like ImageNet forms a fundamental building block for the progress of representation learning-driven solutions in computer vision. Medical images are inherently different from natural images as they are acquired in the form of many modalities (CT, MR, PET, Ultrasound etc.) and contain granulated information like tissue, lesion, organs etc. These characteristics of medical images require special attention towards learning features representative of local context. In this work, we focus on designing an effective pre-training framework for 3D radiology images. First, we propose a new masking strategy called local masking where the masking is performed across channel embeddings instead of tokens to improve the learning of local feature representations. We combine this with classical low-level perturbations like adding noise and downsampling to further enable low-level representation learning. To this end, we introduce Disruptive Autoencoders, a pre-training framework that attempts to reconstruct the original image from disruptions created by a combination of local masking and low-level perturbations. Additionally, we also devise a cross-modal contrastive loss (CMCL) to accommodate the pre-training of multiple modalities in a single framework. We curate a large-scale dataset to enable pre-training of 3D medical radiology images (MRI and CT). The proposed pre-training framework is tested across multiple downstream tasks and achieves state-of-the-art performance. Notably, our proposed method tops the public test leaderboard of BTCV multi-organ segmentation challenge.

研究の動機と目的

• 3D医療画像において、臓器や病変のセグメンテーションなどの細分化タスクに不可欠な、頑健な低レベル表現を学ぶ課題に対処すること。
• 従来のマスク付き自己符号化器（MAE）が3Dレントゲン画像における細かな解剖的構造の再構成に限界を示す問題を克服すること。
• 複数の医療画像モodal（例：MRI、CT）を統一的かつ対照学習に配慮した形で効果的に活用できる事前学習フレームワークを設計すること。
• 大規模かつ多モodalな3D医療画像データセットを収集・活用し、効果的な事前学習を可能にすること。
• より強力な低レベル特徴の学習が、医療画像セグメンテーションにおける下流タスクの性能向上に寄与することを実証すること。

提案手法

• 空間的トークンではなくチャネル埋め込みをマスキングする、新しい局所的マスキング戦略を提案し、局所的で細分化された特徴の学習能力を向上させる。
• 事前処理段階で適用する破壊的摂動（加法的ノイズとダウンサンプリング）を導入し、さらに低レベル表現学習を刺激する。
• これらの摂動をトランスフォーマーに基づくエンコーダーと畳み込みデコーダーと組み合わせ、損傷を加えた潜在表現から元の3D画像を再構築する。
• 異なる画像モダリティ（例：MRI、CT）間の特徴を統一的にアライメントするため、クロスモダリティ対照損失（CMCL）を導入し、モダリティに配慮した表現学習を可能にする。
• 多様な解剖的構造と画像取得条件の変動をカバーする、大規模かつ収集済みの3D MRIおよびCTボリュームデータセットを用いてモデルを事前学習する。
• 自己教師付き事前学習に続き、下流のセグメンテーションタスクで微調整を行い、BTCV や PROMISE12 といった標準ベンチマークで評価する。

Figure 1: Disruptive Autoencoders: Here, we disrupt the 3D medical image with a combination of low-level perturbations - noise and downsampling, followed by tokenization and local masking. These disrupted tokens are then passed through a transformer encoder and convolutional decoder to learn to reco

実験結果

リサーチクエスチョン

• RQ1 局所的摂動に焦点を当てた事前学習フレームワークが、3D医療画像の表現学習を向上させることができるか？
• RQ2 空間的トークンではなくチャネル埋め込みをマスキングする局所的マスキングは、標準的なトークンマスキングと比較して、より優れた再構築性能と下流タスク性能を達成できるか？
• RQ3 ノイズとダウンサンプリングを含む低レベル摂動と局所的マスキングの組み合わせは、個別の摂動と比較して、事前学習効果に優れているか？
• RQ4 クロスモダリティ対照学習は、3D医療画像における多モダリティ事前学習をどの程度向上させるか？
• RQ5 より強力な低レベル特徴の学習は、特にデータが少ない状況下でも、より良い一般化性能をもたらすか？

主な発見

• Disruptive Autoencodersは、BTCVマルチオルガニズムセグメンテーションチャレンジの公開テストリーダーボードで、MAE や Swin-UNETR といった先行手法を上回る最先端の性能を達成した。
• BTCVデータセットでは、Diceスコアが0.8472に達し、Swin-UNETRの3つの事前学習タスクを組み合わせたものでさえも上回った。
• PROMISE12データセットでは、20%および50%の学習データにおいて、ランダム初期化とMAEの両方を上回る性能を示し、特にデータ効率に優れた成果を上げた。
• アブレーションスタディの結果、局所的マスキングと低レベル摂動の組み合わせが最良の性能を発揮し、局所的マスキング単体でもノイズやダウンサンプリング単体よりも優れた性能を示した。
• CKA解析により、事前学習から微調整までの段階で、ネットワークの初期層がより多くの低レベル特徴を保持していることが確認され、低レベル表現がより安定的かつ下流タスクに重要であることが裏付けられた。
• クロスモダリティ対照損失（CMCL）の導入により、微調整性能が向上した。これは、モダリティに配慮した事前学習が、画像モダリティ間での一般化性能を向上させることを示している。

Figure 2: Comparison of reconstruction quality. It can be observed that masked image modelling produces a coarse reconstruction for radiology without local context while the proposed disruptions in this work obtain sharper reconstructions recovering meaningful fine details.