[論文レビュー] ConvFormer: Plug-and-Play CNN-Style Transformers for Improving Medical Image Segmentation

Xian Lin, Zengqiang Yan+3 more|arXiv (Cornell University)|Sep 9, 2023

Advanced Neural Network ApplicationsComputer Science被引用数 3

ひとこと要約

本稿では、注意の崩壊を軽減するため、プーリング、CNNスタイルの自己注意（CSA）、畳み込み型フィードフォワードネットワークを用いて、標準的なビジョントランスフォーマー部品を置き換える、プラグアンドプレイ式のCNNスタイルのトランスフォーマー・モジュールであるConvFormerを提案する。2次元特徴マップ上で直接処理を行い、適応的かつカーネルスケーラブルな注意を可能にすることで、長距離依存性の学習を向上させ、ACDC、ISIC、ICHのデータセットにおいて、複数の最先端のトランスフォーマー基盤モデルで一貫した性能向上を達成する。

ABSTRACT

Transformers have been extensively studied in medical image segmentation to build pairwise long-range dependence. Yet, relatively limited well-annotated medical image data makes transformers struggle to extract diverse global features, resulting in attention collapse where attention maps become similar or even identical. Comparatively, convolutional neural networks (CNNs) have better convergence properties on small-scale training data but suffer from limited receptive fields. Existing works are dedicated to exploring the combinations of CNN and transformers while ignoring attention collapse, leaving the potential of transformers under-explored. In this paper, we propose to build CNN-style Transformers (ConvFormer) to promote better attention convergence and thus better segmentation performance. Specifically, ConvFormer consists of pooling, CNN-style self-attention (CSA), and convolutional feed-forward network (CFFN) corresponding to tokenization, self-attention, and feed-forward network in vanilla vision transformers. In contrast to positional embedding and tokenization, ConvFormer adopts 2D convolution and max-pooling for both position information preservation and feature size reduction. In this way, CSA takes 2D feature maps as inputs and establishes long-range dependency by constructing self-attention matrices as convolution kernels with adaptive sizes. Following CSA, 2D convolution is utilized for feature refinement through CFFN. Experimental results on multiple datasets demonstrate the effectiveness of ConvFormer working as a plug-and-play module for consistent performance improvement of transformer-based frameworks. Code is available at https://github.com/xianlin7/ConvFormer.

研究の動機と目的

• 小規模な医療画像データセットに適用した際のビジョントランスフォーマーにおける注意の崩壊を解消すること。
• トランスフォーマー基盤の医療画像分類モデルにおける自己注意マップの収束性と多様性を向上させること。
• アーキテクチャの大幅な見直しを伴わずに、純粋なトランスフォーマーおよびCNN-トランスフォーマーハイブリッドアーキテクチャの両方を強化できる、プラグアンドプレイモジュールの設計。
• 自己注意を学習可能な畳み込みカーネルとしてモデル化することで、スケーラブルかつ適応的な長距離特徴学習を可能にすること。
• 多様なデータセットおよびバックボーンアーキテクチャにおいて、一貫した性能向上を示すこと。

提案手法

• 空間局所性を保持し解像度を低減するために、ViTの1次元トークン化を2次元畳み込みとマックスプーリングに置き換える。
• 自己注意行列を適応的かつスケーラブルな畳み込みカーネルとして構築するCNNスタイルの自己注意（CSA）を導入し、ピクセル単位の長距離依存性を処理する。
• CSAの後に畳み込み型フィードフォワードネットワーク（CFFN）を用いて特徴を精緻化し、2次元空間構造を維持する。
• CSAにおける受容 field のサイズを制御するための学習可能なスケーリング要因 α を使用し、局所的およびグローバルなコンテキストの間で動的適応を可能にする。
• 標準的なViTが使用する1次元シーケンス変換を避けて、入力画像を2次元そのままで処理する。
• プラグアンドプレイモジュールとして設計されており、既存のトランスフォーマー基盤の分類フレームワークに容易に統合可能である。

Figure 1: Visualization of attention maps from the selected layers of the first head in different transformer frameworks. The darker the color, the closer the dependency.

実験結果

リサーチクエスチョン

• RQ1 CNNスタイルのトランスフォーマー設計は、低データ医療画像分類における注意の崩壊を緩和できるか？
• RQ2 1次元シーケンスではなく2次元特徴マップ上で処理することで、注意の多様性とモデルの収束性が向上するか？
• RQ3 適応的かつカーネルスケーラブルな注意に基づくプラグアンドプレイモジュールは、多様なトランスフォーマー基盤分類アーキテクチャにおいて一貫して性能向上をもたらすか？
• RQ4 CSAにおける受容 field のスケーラビリティが、医療画像ベンチマークにおける分類精度に与える影響は何か？
• RQ5 提案されたモジュールは、医療画像分類における異なるデータセットおよびバックボーンモデルに一般化可能か？

主な発見

• Pure-transformerバックボーンであるSETRに統合した場合、ACDCデータセットで平均Diceスコアが3.86%向上した。
• ISICデータセットでは、SETRを用いた場合にDiceスコアが1.38%向上し、Patcherを用いた場合に1.07%向上した。
• ICHデータセットでは、SETRを用いた場合にDiceスコアが1.39%向上し、Patcherを用いた場合に1.15%向上し、データセット全体で一貫した向上を示した。
• TransUNet や TransFuse といったCNN-トランスフォーマーハイブリッドモデルにおいても、CNNの支配的影響があるにもかかわらず、ConvFormerは安定した性能向上を示した。
• アブレーションスタディの結果、α = 0.4 が最良のパフォーマンス（ACDCで91.00%のDice）を達成したことが確認された。
• 可視化結果から、ConvFormerはベースラインと比較してより多様で構造的な注意マップを生成しており、注意の崩壊を効果的に緩和していることが示された。

Figure 2: Comparison between vanilla vision transformer and ConvFormer. CBR is short for the combination of convolution, batch normalization, and Relu. Multiple heads are omitted for simplicity.