17日前

MIST:畳み込み注意メカニズム混合(CAM)デコーダを備えた医療画像セグメンテーションTransformer

Md Motiur Rahman, Shiva Shokouhmand, Smriti Bhatt, Miad Faezipour
MIST:畳み込み注意メカニズム混合(CAM)デコーダを備えた医療画像セグメンテーションTransformer
要約

医療画像セグメンテーションに用いられる代表的かつ有望なディープラーニング手法の一つとして、自己注意(self-attention)を活用することで画素間の長距離依存関係を捉えることができるトランスフォーマーが挙げられる。医療画像セグメンテーションにおいて高い成果を上げている一方で、トランスフォーマーは多モーダルな次元における画素の局所的文脈を捉える点で限界がある。本研究では、この課題を解決するために、新たな畳み込み型アテンションミキサー(Convolutional Attention Mixing, CAM)デコーダを組み込んだ医療画像セグメンテーション用トランスフォーマー(Medical Image Segmentation Transformer, MIST)を提案する。MISTは2つの構成要素からなる:まず、事前学習済みの多軸視覚トランスフォーマー(multi-axis vision transformer, MaxViT)をエンコーダとして用い、次に得られた特徴表現をCAMデコーダに通して画像のセグメンテーションを実行する。CAMデコーダでは、マルチヘッド自己注意、空間注意、およびシーザー・アテンション(squeeze and excitation)モジュールを統合したアテンションミキサーを導入し、すべての空間次元において長距離依存関係を効果的に捉える。さらに、空間情報の強化を図るため、深層畳み込み(deep convolution)を特徴抽出に、浅層畳み込み(shallow convolution)を受容 field の拡大にそれぞれ活用する。異なるネットワーク段階から得られる低レベルと高レベルの特徴をスキップ接続により統合することで、MISTは不要な情報を効果的に抑制できる。実験の結果、ACDCおよびSynapseデータセットにおいて、MISTは医療画像セグメンテーションに特化して設計された最先端モデルを上回る性能を示した。また、階層的トランスフォーマーにCAMデコーダを追加することで、セグメンテーション性能が顕著に向上することが明らかになった。本研究のモデルおよびコードは、GitHub上で公開されている。