[論文レビュー] TEC-Net: Vision Transformer Embrace Convolutional Neural Networks for Medical Image Segmentation

Rui Sun, Tao Lei+4 more|arXiv (Cornell University)|Jun 7, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

TEC-Netは、動的可変畳み込み（DDConv）をCNNブランチに統合し、局所的特徴抽出を適応的に行うとともに、変則的（シフトされた）ウィンドウに依存する自己注意モジュール（S）W-ACAMをTransformerブランチに導入し、多次元の長距離依存関係をモデル化するハイブリッド医療画像セグメンテーションネットワークを提案する。本手法は、事前学習を必要とせず、パラメータ数と推論コストを削減しながら、最先端の性能を達成する。

ABSTRACT

The hybrid architecture of convolution neural networks (CNN) and Transformer has been the most popular method for medical image segmentation. However, the existing networks based on the hybrid architecture suffer from two problems. First, although the CNN branch can capture image local features by using convolution operation, the vanilla convolution is unable to achieve adaptive extraction of image features. Second, although the Transformer branch can model the global information of images, the conventional self-attention only focuses on the spatial self-attention of images and ignores the channel and cross-dimensional self-attention leading to low segmentation accuracy for medical images with complex backgrounds. To solve these problems, we propose vision Transformer embrace convolutional neural networks for medical image segmentation (TEC-Net). Our network has two advantages. First, dynamic deformable convolution (DDConv) is designed in the CNN branch, which not only overcomes the difficulty of adaptive feature extraction using fixed-size convolution kernels, but also solves the defect that different inputs share the same convolution kernel parameters, effectively improving the feature expression ability of CNN branch. Second, in the Transformer branch, a (shifted)-window adaptive complementary attention module ((S)W-ACAM) and compact convolutional projection are designed to enable the network to fully learn the cross-dimensional long-range dependency of medical images with few parameters and calculations. Experimental results show that the proposed TEC-Net provides better medical image segmentation results than SOTA methods including CNN and Transformer networks. In addition, our TEC-Net requires fewer parameters and computational costs and does not rely on pre-training. The code is publicly available at https://github.com/SR0920/TEC-Net.

研究の動機と目的

• 標準的な畳み込みネットワークにおける固定受容 field とパラメータ共有の制限を克服すること。
• 医療画像のTransformerで用いられる従来の自己注意機構が、チャネルおよび多次元自己注意を軽視している問題を解消すること。
• モデルの複雑さを低減し、小規模な医療データセットにおける高いセグメンテーション精度を維持しながら、事前学習への依存を排除すること。
• CNNとTransformerブランチを効果的に統合し、医療画像セグメンテーションにおける補完的な局所的およびグローバル特徴学習を実現すること。

提案手法

• タスクに応じた特徴抽出を可能にするために、空間的に変化するカーネル重みとオフセットを学習する動的可変畳み込み（DDConv）をCNNブランチに導入する。
• 空間的、チャネル的、および多次元自己注意を統合的にモデル化する（S）W-ACAMモジュールをTransformerブランチに提案し、長距離依存関係のモデル化を向上させる。
• Transformerエンコーダーにおける計算コストとパラメータ数を削減するために、コンactな畳み込み投影層を採用する。
• Transformerの前向き伝搬ネットワークにおける標準的なMLPの代わりに、学習可能なパrameterモジュール（LPM）を採用し、モデルサイズを縮小するとともに、事前学習への依存を排除する。
• CNNとTransformerブランチの特徴を統合するためのスイッチバック接続付きの二重ブランチ型U字型エンコーダー・デコーダー構造を設計し、局所的局所化とセグメンテーションの精度を向上させる。
• マルチスケール特徴統合と深層ネットワーク層における段階的精錬を適用し、輪郭やオブジェクト境界の正確性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 標準畳み込みと比較して、動的可変畳み込みは医療画像セグメンテーションにおける局所的特徴表現を改善できるか？
• RQ2 Transformerで空間的・チャネル的（多次元）自己注意をモデル化することで、コントラストが低くノイズが多い複雑な医療画像のセグメンテーションが向上するか？
• RQ3 ハイブリッドCNN-Transformerアーキテクチャは、パラメータ数を減らし、計算コストを低減しながらも最先端の性能を達成できるか？
• RQ4 標準的なMLPを学習可能なパrameterモジュール（LPM）に置き換えることで、医療画像セグメンテーションにおける事前学習の必要性がどの程度排除できるか？
• RQ5 CNNとTransformerブランチの統合は、医療画像セグメンテーションにおいて局所的ディテールの保持とグローバルコンテキストモデリングの両方をどのように向上させるか？

主な発見

• TEC-NetはISIC2018データセットで90.72%のDiceスコアを達成し、Swin-Unet や U-Net の変種を含む既存の最先端手法を上回る性能を示した。
• アブレーションスタディの結果、DDConv、（S）W-ACAM、LPMを併用した場合に最良の性能が得られ、すべてのモジュールを組み合わせた際のDiceスコアは90.88%に達した。
• TEC-Netはモデルパラメータ数を11.58Mにまで削減し、U-Net + Swin-Unetベースラインの46.92Mと比較して顕著に低いパラメータ効率を示した。
• 事前学習を必要とせず高い精度を達成しており、LPM単体で88.43%、（S）W-ACAMとLPMを併用で89.46%のDiceスコアを示した。これは、小規模データセットにおいても高い頑健性を示している。
• 可視化結果から、デコーダーにおける特徴統合を通じて、特に境界や輪郭の詳細に注目しながら、段階的にセグメンテーションが精錬されていることが確認された。
• （S）W-ACAMモジュールは、病変境界や臓器構造に注目する改善された注意マップを示し、多次元依存関係を効果的に捉えていることが裏付けられた。