[論文レビュー] Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation

Reza Azad, Leon Niggemeier+6 more|arXiv (Cornell University)|Aug 31, 2023

Advanced Neural Network Applications被引用数 7

ひとこと要約

本稿では、医療画像セグメンテーションのためのビジョントランスフォーマーにおける自己注意機構を置き換える、軽量で効率的な注意機構である可変大カーネル注意（D-LKA）を提案する。可変サンプリンググリッドを用いた大規模畳み込みカーネルにより、線形計算量でグローバルなコンテキストを捉えることができ、3Dボリューム特徴の学習を向上させる。得られたD-LKA Netは、パrameter数とFLOPsを減らしながら、Synapse、NIH膵臓、皮膚病変データセットで最先端の性能を達成した。

ABSTRACT

Medical image segmentation has seen significant improvements with transformer models, which excel in grasping far-reaching contexts and global contextual information. However, the increasing computational demands of these models, proportional to the squared token count, limit their depth and resolution capabilities. Most current methods process D volumetric image data slice-by-slice (called pseudo 3D), missing crucial inter-slice information and thus reducing the model's overall performance. To address these challenges, we introduce the concept of extbf{Deformable Large Kernel Attention (D-LKA Attention)}, a streamlined attention mechanism employing large convolution kernels to fully appreciate volumetric context. This mechanism operates within a receptive field akin to self-attention while sidestepping the computational overhead. Additionally, our proposed attention mechanism benefits from deformable convolutions to flexibly warp the sampling grid, enabling the model to adapt appropriately to diverse data patterns. We designed both 2D and 3D adaptations of the D-LKA Attention, with the latter excelling in cross-depth data understanding. Together, these components shape our novel hierarchical Vision Transformer architecture, the extit{D-LKA Net}. Evaluations of our model against leading methods on popular medical segmentation datasets (Synapse, NIH Pancreas, and Skin lesion) demonstrate its superior performance. Our code implementation is publicly available at the: https://github.com/mindflow-institue/deformableLKA

研究の動機と目的

• ビジョントランスフォーマーにおける自己注意機構の高い計算コストを軽減すること。
• 3D医療画像解析における2Dスライス単位処理の制限を克服し、スライス間の特徴統合を可能にすること。
• CNNの局所受容 field の制限を克服し、グローバルコンテキストモデリングを向上させつつ、局所的詳細を保持すること。
• パラメータ効率的で、医療画像セグメンテーションベンチマークで高い性能を維持する軽量な注意機構を開発すること。

提案手法

• 自己注意機構の長距離モデリングを線形計算量で模倣できる大規模畳み込みカーネルを用いる、可変大カーネル注意（D-LKA）を提案する。
• 学習可能で微分可能な可変グリッドを導入し、複雑な解剖的構造に適応するようにサンプリング位置を動的に変形する。
• ボリューム医療画像におけるスライス間依存関係を捉えるために、3D版のD-LKAを設計し、3Dコンテキスト理解を強化する。
• D-LKAブロックとスキップ接続を組み合わせた階層的ビジョントランスフォーマー構造（D-LKA Net）を構築し、特徴の精練を実現する。
• CNNとトランスフォーマーのコンポonentをハイブリッドで組み合わせたエンコーダ・デコーダ構造を採用し、局所的詳細とグローバルコンテキストのバランスを取る。
• スキップ接続によるマルチスケール特徴集約戦略を採用し、空間解像度を保持するとともに、セグメンテーション精度を向上させる。

Figure 1 : Proposed network architecture of the 3D D-LKA model on the left and the 2D D-LKA model on the right.

実験結果

リサーチクエスチョン

• RQ1 大カーネル畳み込み機構が、計算コストを低減させながらも、医療画像セグメンテーションのビジョントランスフォーマーにおける自己注意機構に効果的に置き換え可能か。
• RQ2 大カーネル注意における可変サンプリングが、3D医療画像における不規則な形状の臓器の特徴表現をどのように向上させるか。
• RQ3 2D疑似3Dアプローチと比較して、3D D-LKAはスライス間特徴学習をどの程度向上させるか。
• RQ4 スキップ接続と可変畳み込みが、D-LKA Netアーキテクチャの最終的セグメンテーション性能にどの程度寄与しているか。
• RQ5 軽量な注意機構が、医療画像セグメンテーションベンチマークにおいて、精度とモデル効率の両面で既存のSOTAモデルを上回ることができるか。

主な発見

• Synapseデータセットでは、D-LKA NetがDSC 87.49%を達成し、UNETR++（80.59%）とUNETR（77.42%）を上回り、パラメータ数が37.73Mも少ない。
• NIH膵臓データセットでは、D-LKA NetがDSC 81.22%を達成し、UNETR++より0.63%向上し、HD95が1.04ポイント低下した。
• 可変畳み込みを備えた3D LKAバージョンは、非可変バージョンよりDSCが0.63%高い結果となり、適応的サンプリングの利点を示した。
• アブレーションスタディの結果、スキップ接続が性能向上に顕著に寄与しており、最高レベルのスキップ接続のみを用いた場合でもDSCが0.42%向上した。
• 膵臓データセットでは、SOTA手法の中でも最も少ないパラメータ数（62.07M）を達成し、高い効率性を示した。
• 定性的な結果から、D-LKA NetはUNETRやUNETR++と比較して、膵臓の不規則な形状をよりよく捉えており、セグメンテーションアーチファクトを低減した。

Figure 2 : Architecture of the deformable LKA module.