[論文レビュー] CiT-Net: Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation

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

Advanced Neural Network Applications参考文献 38被引用数 5

ひとこと要約

本稿では、CNNブランチに動的可変畳み込みを統合し、Transformerブランチにシフトドウインドウ適応補完的アテンションモジュールを組み込んだハイブリッドCNN-Transformerアーキテクチャ、CiT-Netを提案する。モデルは事前学習を必要としない状態でも、SOTA手法よりも少ないパラメータ数とFLOPsで最先端の性能を達成し、ISIC2018およびLiTS-LiverデータセットにおいてSwin-Unet や TransUNet よりも優れた性能を示した。

ABSTRACT

The hybrid architecture of convolutional neural networks (CNNs) and Transformer are very popular for medical image segmentation. However, it suffers from two challenges. First, although a CNNs branch can capture the local image features using vanilla convolution, it cannot achieve adaptive feature learning. Second, although a Transformer branch can capture the global features, it ignores the channel and cross-dimensional self-attention, resulting in a low segmentation accuracy on complex-content images. To address these challenges, we propose a novel hybrid architecture of convolutional neural networks hand in hand with vision Transformers (CiT-Net) for medical image segmentation. Our network has two advantages. First, we design a dynamic deformable convolution and apply it to the CNNs branch, which overcomes the weak feature extraction ability due to fixed-size convolution kernels and the stiff design of sharing kernel parameters among different inputs. Second, we design a shifted-window adaptive complementary attention module and a compact convolutional projection. We apply them to the Transformer branch to learn the cross-dimensional long-term dependency for medical images. Experimental results show that our CiT-Net provides better medical image segmentation results than popular SOTA methods. Besides, our CiT-Net requires lower parameters and less computational costs and does not rely on pre-training. The code is publicly available at https://github.com/SR0920/CiT-Net.

研究の動機と目的

• 標準CNNにおける固定受容 field とパラメータ共有の制限を是正すること。
• 標準Transformerにおけるチャネルおよびクロス次元アテンションの欠如が、複雑な医療画像における性能低下を引き起こすのを克服すること。
• MLPを軽量パーセプトロンモジュールに置き換えることで、モデルの複雑さを低減し、事前学習への依存を排除すること。
• 既存のハイブリッドCNN-Transformerモデルよりも、計算コストとパラメータ数を低く抑えつつ、高いセグメンテーション精度を達成すること。
• 局所的およびグローバルな特徴を保持する、堅牢で効率的かつ汎用性の高い医療画像セグメンテーションアーキテクチャの開発

提案手法

• 固定カーネル制約を回避し、カーネルオフセットを学習することで、空間的特徴抽出を適応的に行える動的可変畳み込み（DDConv）をCNNブランチに導入する。
• 空間的特徴とチャネル特徴の間のクロス次元的依存関係を適応アテンションによってモデル化する、シフトドウインドウ適応補完的アテンションモジュール（(S)W-ACAM）をTransformerブランチに提案する。
• Transformerにおける標準MLPの代わりに軽量パーセプトロンモジュール（LPM）を採用し、モデルパラメータと計算コストを顕著に削減する。
• CNNとTransformerブランチが並列に特徴を処理するデュアルブランチアーキテクチャを設計し、局所的およびグローバルな表現を効果的に統合する。
• Transformerブランチにおける長距離依存関係のモデル化を可能にしつつ、効率性を維持するためのコンact畳み込み投影を採用する。
• 事前学習を一切行わず、エンドツーエンドでモデルを訓練することで、限られた医療データセットにおける過学習リスクを低減し、一般化性能を向上させる。

実験結果

リサーチクエスチョン

• RQ1 標準CNNと比較して、動的可変畳み込みは医療画像セグメンテーションにおける局所的特徴表現を向上させることができるか？
• RQ2 Transformerブランチにクロス次元アテンションを組み込むことで、複雑な医療画像におけるセグメンテーション精度が向上するか？
• RQ3 軽量パーセプトロンモジュールは、性能を損なわずに標準MLPを効果的に置き換えることができるか？
• RQ4 本稿で提案するハイブリッドアーキテクチャは、事前学習なしで、SOTA手法と比較してセグメンテーション精度、パラメータ数、FLOPsの観点でどのように差をつけるか？
• RQ5 適応アテンションと動的畳み込みの統合により、医療画像における局所的およびグローバルな特徴の保持はどの程度向上するか？

主な発見

• CiT-Net-BはISIC2018データセットで90.88%のDiceスコアを達成し、Swin-Unet、TransUNet、CvTよりもそれぞれ1.20%、1.03%、1.01%高い性能を示したが、事前学習を一切使用していない。
• CiT-Net-Tはわずか11.58Mパラメータと4.53 GFLOPsで、ISIC2018で最高のDiceスコア90.72%を記録し、最も効率的なSOTAモデルとなった。
• LiTS-Liverデータセットでは、CiT-Net-BがDI、VOE、RVD、ASD、RMSDの全5つの指標で優れた性能を示し、強力な汎化性とロバスト性を示した。
• アブレーションスタディの結果、DDConvと(S)W-ACAMを併用することで最高の性能が得られ、両方を適用した場合、ベースラインから1.2%のDiceスコア向上が確認された。
• LPMにより、最良の設定でモデルパラメータを9.67Mにまで削減し、計算コストを顕著に低減しながらも、高い精度を維持した。
• ISIC2018において、CiT-NetはBAT、CvT、CrossFormよりもそれぞれ1.02%、3.00%、3.79%高いDiceスコアを達成したが、パラメータ数とFLOPsが少ないにもかかわらずである。