[論文レビュー] More than Encoder: Introducing Transformer Decoder to Upsample

Yijiang Li, Wentian Cai+3 more|arXiv (Cornell University)|Jun 20, 2021

Radiomics and Machine Learning in Medical Imaging参考文献 52被引用数 6

ひとこと要約

本稿では、医用画像セグメンテーションにおける特徴再構築を向上させるために、変換器のデコーダーを活用する新しいトランスフォーマーに基づくアップサンプリング手法であるWindow Attention Upsample (WAU) を提案する。ダウンサンプリングパスからの特徴に局所的なウィンドウ付き自己注意を適用することで、WAUは空間的詳細と意味的整合性を向上させ、2Dおよび3DのU-Net、FCN、DeepLabV3を含む複数のアーキテクチャにおいて、Synapse、MSD Brain、ACDCデータセットで最先端の性能を達成した。

ABSTRACT

Medical image segmentation methods downsample images for feature extraction and then upsample them to restore resolution for pixel-level predictions. In such a schema, upsample technique is vital in restoring information for better performance. However, existing upsample techniques leverage little information from downsampling paths. The local and detailed feature from the shallower layer such as boundary and tissue texture is particularly more important in medical segmentation compared with natural image segmentation. To this end, we propose a novel upsample approach for medical image segmentation, Window Attention Upsample (WAU), which upsamples features conditioned on local and detailed features from downsampling path in local windows by introducing attention decoders of Transformer. WAU could serve as a general upsample method and be incorporated into any segmentation model that possesses lateral connections. We first propose the Attention Upsample which consists of Attention Decoder (AD) and bilinear upsample. AD leverages pixel-level attention to model long-range dependency and global information for a better upsample. Bilinear upsample is introduced as the residual connection to complement the upsampled features. Moreover, considering the extensive memory and computation cost of pixel-level attention, we further design a window attention scheme to restrict attention computation in local windows instead of the global range. We evaluate our method (WAU) on classic U-Net structure with lateral connections and achieve state-of-the-art performance on Synapse multi-organ segmentation, Medical Segmentation Decathlon (MSD) Brain, and Automatic Cardiac Diagnosis Challenge (ACDC) datasets. We also validate the effectiveness of our method on multiple classic architectures and achieve consistent improvement.

研究の動機と目的

• 既存のアップサンプリング手法が、ダウンサンプリングパスからの局所的で詳細な特徴を十分に活用していないという限界を是正すること。
• 変換器デコーダーの特徴拡大への潜在的応用を検討し、通常のシーケンス生成用途を超えて活用すること。
• 微細な詳細を保持する上で重要な医用画像において、メモリおよび計算コストに優れたアップサンプリング機構を開発すること。
• スキップ接続を用いる任意のセグメンテーションモデルと互換性のある汎用的アップサンプリングモジュールを構築すること。
• 2Dおよび3DのU-Net、FCN、DeepLabV3を含む多様なアーキテクチャにおいて、一貫した性能向上を実証すること。

提案手法

• 変換器ベースの注意デコーダー（AD）と双線形アップサンプリングを残差接続として組み合わせた二ストリームモジュール「Attention Upsample（AU）」を提案する。
• 自己注意計算を局所的な空間的ウィンドウに制限することで、計算コストを低減しながらも長距離モデリングを維持する「ウィンドウ付き注意デコーダー（WAD）」を設計する。
• 各ウィンドウ内での効率的な特徴投影を実現するため、深度可分畳み込みを採用し、小さな医用データセットでも性能向上を図る。
• 局所的特徴ウィンドウからクエリ、キー、バリュー行列を計算するためのウィンドウ付き自己注意を適用し、文脈に応じたアップサンプリングを可能にする。
• U-Netに類似したアーキテクチャのスキップ接続パスにWAUを統合し、標準のセグメンテーションヘッドと併せてエンドツーエンド学習を可能にする。
• 双線形アップサンプリングが注意ベースの最適化を補完する残差学習スキームを採用することで、学習の安定性を高め、特徴品質を向上させる。

実験結果

リサーチクエスチョン

• RQ1 変換器のデコーダーは、医用画像セグメンテーションにおける特徴拡大に効果的に再利用可能か？
• RQ2 ダウンサンプリングパスからの局所的で詳細な特徴を組み込むことで、従来の転置畳み込みや双線形アップサンプリングと比較して、アップサンプリング品質はどのように向上するか？
• RQ3 ウィンドウ付き注意機構において、局所的詳細と長距離コンテキストモデリングのバランスを最適化するための最適なウィンドウサイズは何か？
• RQ4 提案手法WAUは、2Dおよび3Dモデルを含むさまざまなセグメンテーションアーキテクチャにどのように一般化可能か？
• RQ5 グローバルな自己注意と局所的なウィンドウ付き自己注意を比較した場合、アップサンプリングにおける計算コストと性能向上のトレードオフはどのようなものか？

主な発見

• WAUは、多臓器セグメンテーションデータセットであるSynapseで最先端の性能を達成し、既存手法を上回った。
• Medical Segmentation Decathlon（MSD）Brainデータセットでは、ベースラインのResUNetに対して平均DSCが1.47ポイント向上した。
• 3D U-Netでは0.36 DSCの向上を達成し、3D医用ボリュームセグメンテーションへの強力な一般化能力を示した。
• WAUの最適なウィンドウサイズは4であり、それより大きいまたは小さいウィンドウでは、焦点の喪失またはコンテキスト不足により性能が低下した。
• 畳み込みタイプの中で、深度可分畳み込みが最も優れた性能を示し、特にパrameter数が限られた小さな医用データセットにおいて顕著だった。
• 可視化結果から、WAUの注意重みが関連する解剖的構造に空間的に集中していることが確認され、特徴の区別能と病変の可視性が向上した。