[論文レビュー] RaBiT: An Efficient Transformer using Bidirectional Feature Pyramid Network with Reverse Attention for Colon Polyp Segmentation

Nguyen Hoang Thuan, Thi-Oanh Nguyen+3 more|arXiv (Cornell University)|Jul 12, 2023

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

ひとこと要約

RaBiTは、逆注意（RaBiFPN）を備えた双方向特徴ピラミッドとボトルネック最適化逆注意モジュールを用いて、マルチスケール特徴の統合と境界の微調整を向上させる、軽量なTransformerベースのエンコーダデコーダネットワークを提案する。高解像度ETIS-Laribで4.5%のmIoU向上を達成し、ベースラインモデルと比較してパラメータを37%削減する。

ABSTRACT

Automatic and accurate segmentation of colon polyps is essential for early diagnosis of colorectal cancer. Advanced deep learning models have shown promising results in polyp segmentation. However, they still have limitations in representing multi-scale features and generalization capability. To address these issues, this paper introduces RaBiT, an encoder-decoder model that incorporates a lightweight Transformer-based architecture in the encoder to model multiple-level global semantic relationships. The decoder consists of several bidirectional feature pyramid layers with reverse attention modules to better fuse feature maps at various levels and incrementally refine polyp boundaries. We also propose ideas to lighten the reverse attention module and make it more suitable for multi-class segmentation. Extensive experiments on several benchmark datasets show that our method outperforms existing methods across all datasets while maintaining low computational complexity. Moreover, our method demonstrates high generalization capability in cross-dataset experiments, even when the training and test sets have different characteristics.

研究の動機と目的

• 既存のモデルが多スケール特徴を捉えることや、多様な内視鏡画像データセットに一般化することに課題を抱えるのを是正すること。
• 新規の逆注意機構を用いた反復的微調整により、ポリープ境界のセグメンテーションを向上させること。
• 高い精度を維持しながら計算コストを低減すること。
• エンコーダデコーダの各レベル間での特徴統合を、双方向特徴ピラミッド構造を用いて強化すること。
• 多クラス対応の逆注意モジュールを導入することで、効果的な多クラスセグメンテーションを実現すること。

提案手法

• エンコーダは、複数スケールにわたる階層的グローバルな意味的関係を捉えるために、軽量なMiT-B3ベースのTransformerを採用する。
• デコーダは、重み付き連結とスカイプ接続を用いて、異なるエンコーダレベルからの特徴を統合する、複数の双方向特徴ピラミッド（RaBiFPN）モジュールのスタックを採用する。
• パラメータと計算コストを低減しつつ、ポリープ境界の繰り返し微調整を可能にする、軽量なボトルネックベースの逆注意（RA）モジュールを導入する。
• 非新生物性および新生物性ポリープのより良いセグメンテーションを実現するため、多クラスRAモジュールを設計する。
• RAモジュールではシグモイドまたはソフトマックス活性化関数を用い、実験により小規模オブジェクトのセグメンテーションにおいてソフトマックスが優れた性能を示した。
• アーキテクチャは、交差エントロピー損失とDice損失を用いてエンドツーエンドで訓練され、アブレーションスタディにより各構成要素の有効性が検証された。

実験結果

リサーチクエスチョン

• RQ1 軽量なTransformerエンコーダは、コロニー・ポリープセグメンテーションにおいてグローバルコンテキストと多スケール特徴を効果的にモデル化できるか？
• RQ2 逆注意を備えた双方向特徴ピラミッドは、標準的なFPNやBiFPNと比較して、特徴統合と境界微調整を向上させるか？
• RQ3 ボトルネック最適化された逆注意モジュールは、計算コストを低減しつつ、セグメンテーション精度を維持または向上させられるか？
• RQ4 本手法は、解像度やポリープ特性が異なる多様な内視鏡画像データセットに一般化できるか？
• RQ5 多クラス対応の逆注意モジュールは、新生物性および非新生物性ポリープの両方のセグメンテーション性能を向上させるか？

主な発見

• 標準的なBiFPNと比較して、高解像度ETIS-Laribデータセットで4.5%のmIoU向上を達成し、大規模画像において顕著な効果を示した。
• 逆注意モジュールにソフトマックス活性化関数を用いることで、Neo-SmallデータセットにおけるDice非メトリックが6.43%向上し、小規模オブジェクトのセグメンテーションが改善された。
• 提案されたボトルネックモジュールにより、パラメータ数が81.79Mから51.72Mに、GFLOPsが41.37から23.12に削減され、計算コストが顕著に低減した。
• 4つのRaBiFPNモジュールを用いることで、RaBiTはSegFormer-B3や他の最先端モデルを上回る、全ベンチマークデータセットで最高の総合的性能を達成した。
• クロスデータセット一般化において、トレーニングセットとテストセットの分布が異なる場合でも、RaBiTは高い性能を維持し、高いロバスト性を示した。
• アブレーションスタディにより、RaBiFPNとボトルネックモジュールがそれぞれ独立してmIoUを向上させていることが確認され、特に高解像度データセットで最大の向上が得られた。