[論文レビュー] Transfer Learning for Microstructure Segmentation with CS-UNet: A Hybrid Algorithm with Transformer and CNN Encoders

Khaled Alrfou, Tian Zhao+1 more|arXiv (Cornell University)|Aug 26, 2023

Mineral Processing and Grinding被引用数 5

ひとこと要約

本論文は、顕微鏡画像におけるマイクロ構造セグメンテーションのため、CNNとSwin Transformerエンコーダーを組み合わせたハイブリッドディーブラーニングモデルであるCS-UNetを提案する。独自の顕微鏡データセット（MicroLite）で事前学習することで、特に少データ（few-shot）および分布外（out-of-distribution）データに対して、セグメンテーション性能が顕著に向上する。これは、ドメイン固有のデータで事前学習されたCNNとTransformerエンコーダーの統合が、CNNオンリーモデルを上回ることを示している。

ABSTRACT

Transfer learning improves the performance of deep learning models by initializing them with parameters pre-trained on larger datasets. Intuitively, transfer learning is more effective when pre-training is on the in-domain datasets. A recent study by NASA has demonstrated that the microstructure segmentation with encoder-decoder algorithms benefits more from CNN encoders pre-trained on microscopy images than from those pre-trained on natural images. However, CNN models only capture the local spatial relations in images. In recent years, attention networks such as Transformers are increasingly used in image analysis to capture the long-range relations between pixels. In this study, we compare the segmentation performance of Transformer and CNN models pre-trained on microscopy images with those pre-trained on natural images. Our result partially confirms the NASA study that the segmentation performance of out-of-distribution images (taken under different imaging and sample conditions) is significantly improved when pre-training on microscopy images. However, the performance gain for one-shot and few-shot learning is more modest with Transformers. We also find that for image segmentation, the combination of pre-trained Transformers and CNN encoders are consistently better than pre-trained CNN encoders alone. Our dataset (of about 50,000 images) combines the public portion of the NASA dataset with additional images we collected. Even with much less training data, our pre-trained models have significantly better performance for image segmentation. This result suggests that Transformers and CNN complement each other and when pre-trained on microscopy images, they are more beneficial to the downstream tasks.

研究の動機と目的

• ハイブリッドCNN-Transformerエンコーダーを用いたトランスファー学習により、顕微鏡画像におけるマイクロ構造セグメンテーションを向上させること。
• ImageNetでの事前学習と比較して、顕微鏡特化データで事前学習することで、セグメンテーション性能が向上するかどうかを評価すること。
• 統一されたアーキテクチャ（CS-UNet）におけるCNNとTransformerエンコーダーの統合が、材料マイクロ構造解析において有効であるかを調査すること。
• ドメイン特化の顕微鏡データで事前学習されたハイブリッドモデルが、低データ環境（ワンショットおよび少データ学習）で優れた性能を示すことを実証すること。
• 材料科学コミュニティが下流の分析タスクを加速できるよう、公開可能な事前学習済みモデルを提供すること。

提案手法

• CNN（例：EfficientNet、ResNeXt）とSwin Transformer（Swin-T）の2つのエンコーダーを備えた、U-Netベースのアーキテクチャ（CS-UNet）を提案する。
• CNNエンコーダーをMicroNet（11万枚の顕微鏡画像）で事前学習した重みで初期化し、Swin-TエンコーダーをImageNetで事前学習した重みで初期化した後、MicroLiteデータセットで微調整する。
• CNNとSwin-Tエンコーダーの特徴量を統合し、その統合表現をデコーダーに渡してピクセル単位のセグメンテーションを実行する。
• ニッケル基超合金および環境バリアコーティングをカバーする74クラスの顕微鏡画像約5万枚からなるカスタムデータセット（MicroLite）でモデルを学習する。
• ワンショットおよび少データ学習設定下で、7つのテストサブセット（Super-1 から Super-4、EBC-1 から EBC-3）における交差率（IoU）を用いて性能を評価する。
• ImageNetとMicroLiteで事前学習したモデルを比較し、個々のエンコーダーとハイブリッドエンコーダー構成の両方を評価する。

実験結果

リサーチクエスチョン

• RQ1 顕微鏡特化データで事前学習することで、顕微鏡画像のセグメンテーション性能がImageNetでの事前学習よりも向上するか？
• RQ2 少データおよびワンショット学習状況下で、ハイブリッドCNN-Transformerエンコーダーの性能はCNNオンリーエンコーダーと比べてどうか？
• RQ3 顕微鏡データで事前学習した場合、CNNとSwin Transformerエンコーダーの組み合わせが、単体のエンコーダーと比較してより高いセグメンテーション精度を達成できるか？
• RQ4 ドメイン特化の事前学習が、撮影条件や試料状態が異なる分布外画像に与える影響は何か？
• RQ5 限られた顕微鏡データセットで微調整した場合、Swin-TベースのモデルとCNNベースのモデルの性能はどのように比較されるか？

主な発見

• MicroLiteデータセット（5万枚の顕微鏡画像）で事前学習することで、特にワンショットおよび少データ学習設定下で、分布外画像のセグメンテーション性能が顕著に向上する。
• CNNとSwin-Tエンコーダーを併用したハイブリッドCS-UNetモデルは、CNNオンリーモデルよりも高いIoUスコアを達成し、両エンコーダーが顕微鏡データで事前学習された場合に最も高い性能を示す。
• EBC-1データセットでは、MicroLiteで事前学習した中間Swin-Tエンコーダーを搭載したTransDeepLabV3+が96.59%の最高IoUスコアを記録し、ImageNetオンリーモデルを上回った。
• Super-3データセットでは、MicroLiteで事前学習することで、TransDeepLabV3+アーキテクチャのIoUがImageNetの70.74%からMicroLiteの92.23%に向上した。
• Transformerモデルに対しては顕微鏡データでの事前学習による性能向上がやや限定的であったが、CNNモデルに対しては低データ環境下で顕著な恩恵が得られた。
• 約5万枚の限定的なトレーニングデータでさえも、事前学習済みモデルが最先端の性能を達成した。これは、材料マイクロ構造セグメンテーションにおける効果的なトランスファー学習のため、ドメイン特化の事前学習が不可欠であることを示している。