[論文レビュー] Multi-Compound Transformer for Accurate Biomedical Image Segmentation

Yuanfeng Ji, Ruimao Zhang+5 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network Applications参考文献 23被引用数 15

ひとこと要約

本稿では、自己注意およびクロス注意メカニズムを通じて、複数スケールの特徴依存関係と意味的関係をモデル化することで、生物学的画像分類を向上させる、新しいビジョントランスフォーマー・アーキテクチャ、Multi-Compound Transformer (MCTrans) を提案する。MCTrans は UNet に類似したネットワークと統合され、6つのベンチマークで最先端の性能を達成し、Dice スコアにおいて UNet を最大 4.34% まで上回っている。

ABSTRACT

The recent vision transformer(i.e.for image classification) learns non-local attentive interaction of different patch tokens. However, prior arts miss learning the cross-scale dependencies of different pixels, the semantic correspondence of different labels, and the consistency of the feature representations and semantic embeddings, which are critical for biomedical segmentation. In this paper, we tackle the above issues by proposing a unified transformer network, termed Multi-Compound Transformer (MCTrans), which incorporates rich feature learning and semantic structure mining into a unified framework. Specifically, MCTrans embeds the multi-scale convolutional features as a sequence of tokens and performs intra- and inter-scale self-attention, rather than single-scale attention in previous works. In addition, a learnable proxy embedding is also introduced to model semantic relationship and feature enhancement by using self-attention and cross-attention, respectively. MCTrans can be easily plugged into a UNet-like network and attains a significant improvement over the state-of-the-art methods in biomedical image segmentation in six standard benchmarks. For example, MCTrans outperforms UNet by 3.64%, 3.71%, 4.34%, 2.8%, 1.88%, 1.57% in Pannuke, CVC-Clinic, CVC-Colon, Etis, Kavirs, ISIC2018 dataset, respectively. Code is available at https://github.com/JiYuanFeng/MCTrans.

研究の動機と目的

• 既存のビジョントランスフォーマーが、特にラベル間の複数スケール依存関係と意味的関係をモデル化できないという限界に対処する。
• 医用画像における異なる解剖的構造の間で特徴表現の一貫性と識別性を向上させる。
• 単一スケールの注意を超えて長距離の文脈的モデリングを強化し、サイズが異なる病変をよりよく捉える。
• マルチスケール特徴学習と意味的構造マイニングを統合した包括的なフレームワークを構築し、正確な分類を実現する。
• 提案手法の有効性を、計算コストを最小限に抑えた多様な生物学的画像分類ベンチマークで示す。

提案手法

• MCTrans は、複数の特徴スケールにわたるピクセルレベルの関係をモデル化するためのトランスフォーマー自己注意（TSA）モジュールを採用し、包括的なクロススケール文脈符号化を可能にする。
• 意味的カテゴリを表現するための学習可能なプロキシ埋め込みが導入され、自己注意によりカテゴリ固有の依存関係を捉えるように精練される。
• トランスフォーマークロス注意（TCA）モジュールにより、特徴表現とプロキシ埋め込みの間の相互作用が実現され、特徴の識別性と一貫性が向上する。
• 更新されたプロキシ埋め込みに対して補助損失が適用され、同じカテゴリ内での特徴相関の向上と異なるカテゴリ間の分離性の向上を明示的に促進する。
• MCTrans モジュールは、エンドツーエンドの学習と推論を可能にする、UNet に類似したエンコーダ・デコーダアーキテクチャにスムーズに統合される。
• マルチスケールの畳み込み特徴がトークンとして埋め込まれ、TSA および TCA モジュールを経由して処理された後、アップサンプリングおよびマージされ、分類出力が得られる。

実験結果

リサーチクエスチョン

• RQ1 単一スケールの自己注意と比較して、クロススケール注意メカニズムは生物学的画像分類における特徴表現を改善できるか？
• RQ2 異なるラベル間の意味的関係を学習することは、特徴の識別性と分類精度にどのように影響するか？
• RQ3 補助監視付きの学習可能なプロキシ埋め込みを導入することで、小規模かつ複雑なデータセットでのモデル性能はどの程度向上するか？
• RQ4 標準的な UNet と比較して、提案された MCTrans フレームワークは計算コストのわずかな増加で高い性能を維持できるか？
• RQ5 MCTrans は、病変および解剖的構造分類を含む多様な生物学的分類タスクに一般化可能か？

主な発見

• Pannuke データセットでは、MCTrans が UNet よりも Dice スコアで 3.64% の向上を達成し、CVC-Clinic、CVC-Colon、ETIS、Kavirs、ISIC2018 ではそれぞれ 3.71%、4.34%、2.8%、1.88%、1.57% の向上を示した。
• Pannuke データセットでは、ResNet-34 バックボーンを用いて Dice スコア 68.40% を達成し、パrameter 数が同程度の UNet++（66.08%）および CENet（66.50%）を上回った。
• アブレーションスタディの結果、プロキシ埋め込みの補助損失を削除すると性能が 68.40% から 67.87% に低下し、その損失が特徴相関性と識別性の向上に寄与していることが確認された。
• TSA および TCA モジュールの最適な数は 4 と判明した。これより多くすると性能が低下するため、モデル容量とデータセットサイズのバランスが重要であることが示された。
• FLOPs は 18.065G（UNet は 14.037G）にわずかに増加するのみで、精度の向上と効率性の両方を維持している。
• 図 4 の定性的な結果から、MCTrans は特に小さなまたは不規則な形状の病変に対して、より正確で一貫性のあるセグメンテーションマスクを生成することがわかった。