AMAES: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation

要約

本研究では、3Dセマンティックセグメンテーションモデルの自己教師付き事前学習が、大規模かつ領域固有のデータセットに与える影響を調査する。BRAINS-45Kを導入し、最新のセグメンテーションアーキテクチャを事前学習するための多くの設計選択を、最新の手法を簡略化・最適化し、新しいオーグメンテーション戦略と組み合わせることによって再検討する。その結果、AMAESフレームワークは、マスク画像モデリングと強度ベースのオーグメンテーション反転に基づき、メモリ使用量、実行時間、性能の微調整のバランスをとる。一般的なU-Netと最近のMedNeXtアーキテクチャをバックボーンとして使用し、単一系列、低リソース設定、領域外汎化をカバーする3つの困難な下流タスクに対する事前トレーニングの効果を評価する。その結果、AMAESを用いた提案データセットでの事前学習により、評価されたケースの大部分においてセグメンテーション性能が大幅に向上すること、また、大規模データセットでの事前学習にもかかわらず、オーグメントを用いたモデルの事前学習が有益であることが明らかになった。結果を再現するためのコードとモデルのチェックポイント、およびBRAINS-45Kデータセットは、 \url{https://github.com/asbjrnmunk/amaes} で入手可能である。

要約(オリジナル)

This study investigates the impact of self-supervised pretraining of 3D semantic segmentation models on a large-scale, domain-specific dataset. We introduce BRAINS-45K, a dataset of 44,756 brain MRI volumes from public sources, the largest public dataset available, and revisit a number of design choices for pretraining modern segmentation architectures by simplifying and optimizing state-of-the-art methods, and combining them with a novel augmentation strategy. The resulting AMAES framework is based on masked-image-modeling and intensity-based augmentation reversal and balances memory usage, runtime, and finetuning performance. Using the popular U-Net and the recent MedNeXt architecture as backbones, we evaluate the effect of pretraining on three challenging downstream tasks, covering single-sequence, low-resource settings, and out-of-domain generalization. The results highlight that pretraining on the proposed dataset with AMAES significantly improves segmentation performance in the majority of evaluated cases, and that it is beneficial to pretrain the model with augmentations, despite pretraing on a large-scale dataset. Code and model checkpoints for reproducing results, as well as the BRAINS-45K dataset are available at \url{https://github.com/asbjrnmunk/amaes}.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL