FCT: 医療画像セグメンテーション論文読影用の完全畳み込みトランスフォーマー

医療画像セグメンテーション用の完全畳み込みトランスフォーマー

论文：医用画像セグメンテーション用の完全畳み込みトランスフォーマー (thecvf.com)

コード: Thanos-DB/FullyConvolutionalTransformer (github.com)

ジャーナル/カンファレンス: WACV 2023

まとめ

我々は、異なるモダリティの医療画像をセグメント化できる新しい変換器を提案します。医用画像解析のきめ細かい性質によってもたらされる課題は、トランスフォーマーの解析への適応がまだ初期段階にあることを意味します。UNet の圧倒的な成功は、セグメンテーションタスクのきめ細かな性質を認識できる能力にあります。この能力は、既存のトランスフォーマーベースのモデルには現在備わっていません。この欠点に対処するために、私たちは、効果的な画像表現を学習する畳み込みニューラルネットワークの機能に基づいて構築され、入力の特徴を効果的にキャプチャするトランスフォーマーの機能と組み合わせた完全畳み込みトランスフォーマー (FCT) を提案します。。FCT は、医療画像文献における最初の完全な畳み込み Transformer モデルです。入力は 2 段階で処理されます。まず、入力画像から長期的な意味の依存関係を抽出する方法を学習し、次に特徴から階層的なグローバルプロパティを取得する方法を学習します。FCT はコンパクト、正確、堅牢です。私たちの結果は、事前トレーニングを必要とせずに、異なるデータモダリティを備えた複数の医用画像セグメンテーションデータセットに対して、既存のすべての変換アーキテクチャよりも大幅に優れたパフォーマンスを発揮することを示しています。FCT は、ACDC データセットで 1.3%、Synapse データセットで 4.4%、Spleen データセットで 1.2%、ISIC 2017 データセットで 1.1%、およびダイスメトリクスで直接モデルを上回っています。パラメーターは最大 5 分の 1 です。ACDC Post-2017MICCAI-Challenge オンラインテストセットでは、当社のモデルは未確認の MRI テストケースで新しい最先端を設定し、大規模アンサンブルモデルやパラメータの少ない nnUNet を上回ります。

1 はじめに

医療画像のセグメンテーションは、コンピューター支援診断の重要なツールです。画像内の病変の境界を検出して位置を特定するのに役立ち、腫瘍や癌領域の潜在的な存在を迅速に特定するのに役立ちます。これにより、診断が迅速化され、腫瘍を検出する可能性が向上し、臨床医が時間をより効率的に使えるようになり、患者の転帰に利益がもたらされる可能性があります [15]。従来、最新の医療画像セグメンテーションアルゴリズムは、最初に入力画像を潜在空間に圧縮 (エンコード) し、次に画像内の関心領域の位置をデコードする方法を学習する、対称的なトップダウンのエンコーダー/デコーダー構造として構築されています。この垂直情報フローに中間信号の水平伝播 (スキップ接続) を追加すると、UNet アーキテクチャが得られます。これは、おそらく最近のセグメンテーションアルゴリズムで最も影響力のある進歩です。現在のほとんどの最新のセグメンテーションシステムには、UNet またはその亜種が含まれています。UNet の成功の鍵は、その完全な畳み込みの性質にあります。UNet は、その構造内で非畳み込み学習可能なパラメーターを推定しません。

畳み込みニューラルネットワーク (CNN) に基づく UNet モデルは、医療画像セグメンテーションタスクの精度とパフォーマンスの点で大きな成功を収めています。しかし、臨床医による疾患の早期診断を真に支援するには、さらにさらなる改善が必要です。畳み込み演算子の固有のローカルな性質は、入力画像からの長距離の意味論的依存関係を悪用できないため、CNN にとって重要な問題です。CNN にグローバルコンテキストを追加するさまざまな方法が提案されており、最も注目すべき方法は、アテンションメカニズムの導入と、カーネルの受容野を高めるための畳み込みカーネルの拡大です。ただし、これらの方法には独自の欠点があります。トランスフォーマーは、非常に広範囲のシーケンス依存関係を効率的に処理できるため、言語学習タスクで大きな成功を収めています。これが、最近のさまざまな視覚タスクへの適応につながりました [7、18、21、22]。ViT [7] などの最近提案されたアーキテクチャは、ベンチマーク画像処理タスクで CNN のパフォーマンスを上回っていますが、CvT [36]、CCT [10]、Swin Transformer [25] などの ViT に対する最近の多くの改良により、トランスフォーマーのパフォーマンスが向上していることが示されています。大量のデータを消費するモデルは必要なく、少量のデータでも処理できるため、CNN のパフォーマンスを上回ります。通常、ViT スタイルのモデルは、まず画像から重複しない個別のパッチ (NLP ではトークンと呼ばれます) を抽出します。次に、位置エンコーディングを介してこれらのパッチに空間的ローカリゼーションを挿入し、この表現を標準のトランスフォーマー層に渡して、データ内の長期的な意味論的な依存関係をモデル化します。

CNN と Transformer の明白な利点を考慮すると、医用画像セグメンテーションの次のステップは、医用画像の長期的な意味論的依存関係を効果的に利用できる完全畳み込みエンコーダ/デコーダ深層学習モデルであると考えられます。この目標を達成するために、医療画像セグメンテーション用の最初の完全畳み込み Transformer を提案します。私たちの新しい完全畳み込み Transformer レイヤーは、モデルの主要な構成要素を形成します。これには、畳み込みアテンションモジュールと完全畳み込みワイドフォーカスモジュールという 2 つの主要なコンポーネントが含まれています (セクション 3 を参照)。私たちの貢献を次のように正式に表現します。

我々は、医療画像セグメンテーション用の最初の完全畳み込み Transformer を提案します。これは、複数のバイナリ分類およびセマンティックセグメンテーションデータセットに対する既存のすべての畳み込みおよびトランスフォーマーベースの医用画像セグメンテーションアーキテクチャのパフォーマンスを上回ります。
我々は、畳み込みアテンションモジュールを使用して長距離のセマンティックコンテキストを学習し、次にワイドフォーカスモジュールを介して多重解像度アトラス畳み込みを使用して階層的なローカルからグローバルコンテキストを作成する、新しい完全畳み込みトランスフォーマー層を提案します。
広範なアブレーション研究を通じて、モデルのパフォーマンスに対する影響という観点から、モデルのさまざまな構成要素の影響を実証します。

2.関連作品

初期の CNN およびアテンションモデル: UNet [29] は、医療画像セグメンテーション用の最初の CNN モデルでした。アテンションモデルを医療画像セグメンテーションに導入する最初の研究の 1 つは、UNet のエンコーダからデコーダへの特徴伝播にゲート関数を適用することでした [26]。FocusNet [17] などのメソッドは、デュアルエンコーダ/デコーダ構造を採用しており、アテンションゲート学習により、関連する特徴が 1 つの UNet のデコーダから次の UNet のエンコーダに伝播されます。FocesNet++ [19] は、グループ化された畳み込みのさまざまなフィルターバンクにアテンションメカニズムを統合した最初の作品の 1 つです。特徴抽出を強化するためにさまざまな残差ブロックを使用する UNet のバリアントも多数あります [32、28、33、20、16]。UNet++ [43] は、エンコーダとデコーダの間に入れ子になった階層的な高密度スキップ接続パスを作成し、エンコーダとデコーダ間の学習された特徴の意味論的なギャップを削減します。最近最も影響力のある UNet の亜種である nnUNet [14] は、人間の介入なしにデータを前処理し、タスクに最適なネットワークアーキテクチャを選択するように自動的に適応します。

Transformer モデル: オリジナルの Transformer アーキテクチャ [31] は、自然言語処理タスクに革命をもたらし、すぐに視覚理解タスクのモデルになりました [7]。トランスフォーマーは、長距離の視覚コンテキストを作成できるため、視覚には適していますが、CNN などの画像の空間コンテキストを利用できないという固有の欠点があります。最近の研究では、この欠点を克服するための可能な解決策が検討されています。CvT [36]、CCT [10]、および Swin Transformer [25] はすべて、変圧器に十分な空間環境を統合する試みです。医療画像セグメンテーションにおける既存の研究のほとんどは、特徴処理のための Transformer と CNN のハイブリッドモデルの作成に焦点を当てています。アテンション UNet [26] と同様に、UNet Transformer [27] はスキップ接続内にマルチヘッドアテンションを追加することで CNN を強化します。TransUNet [5] は、医療画像セグメンテーション用に最初に提案された Transformer-CNN ハイブリッドモデルの 1 つで、カスケード畳み込みデコーダに供給される Transformer エンコーダを使用します。TransUNet と同様に、UNETR [12] と Swin UNETR [11] は、エンコーダー上でトランスフォーマーと畳み込みデコーダーを使用してセグメンテーションマップを構築します。Transfuse [40] は、畳み込み層を備えたエンコーダと変換層を備えたもう 1 つのデュアルブランチエンコーダを実行し、その機能を新しい BiFusion モジュールと組み合わせています。ただし、このモデルのデコーダは畳み込み型です。

現在の作業: 最近、ハイブリッド Transformer-CNN モデルの作成から、医療画像の微妙な違いを処理するために Transformer ブロック自体を改善することに移行しています。Swin UNet [3] は、医療画像を処理するための純粋な変換アーキテクチャを最初に提案しました。ここでの純粋とは、事前にトレーニングされたバックボーンネットワークアーキテクチャを必要とせず、トランスフォーマー層のみによって抽出および処理される画像特徴を指します。DS-TransUNet[24] では、グローバルな依存関係をより適切に表現するために、Transformer Interactive Fusion モジュールが導入されています。両方のモデルの計算コアは Swin Transformer ブロックです。nnFormer [42] や DFormer [37] などの同時研究では、医療画像内のローカルおよびグローバルコンテキストを利用して、特別に作成されたマルチヘッドセルフアテンションブロックを通じてこのタスクを達成しようとしています。これらのモデルの主な欠点は、注意の投影と特徴処理の固有の線形的な性質であり、FCT はこれを軽減することを目指しています。

既存の医用画像セグメンテーションモデルには現在、次の 3 つの制限のうち少なくとも 1 つがあります。これらは CNN バックボーンネットワークに基づいているか、畳み込み層を使用して作成されているため、受容野を超えて画像の意味論的コンテキストを取得する能力が制限されています(初期の CNN 手法)。彼らは、長期的なセマンティックコンテキストを作成する機能を活用するために、Transformer を特徴処理パイプラインに統合しようとしましたが、その結果、モデルが大きくなり、計算が複雑になってしまいました (Hybrid Transformer-CNN)。彼らは、低レベルの特徴抽出段階 (同時作業) でローカルの空間コンテキストをモデル化しようとせずに、セグメンテーション用の純粋な Transformer モデルを作成することで、この計算負荷を軽減しようとしています。既存の方法とは異なり、当社の完全畳み込み Transformer にはこれらの欠点がなく、依然として純粋な Transformer ベースの医療画像セグメンテーションアーキテクチャです。補足資料の表 4 には、既存モデルと比較した FCT の主な違いがさらにまとめられています。

3.完全畳み込みトランスフォーマー

データセット $\{\mathbf{X}, \mathbf{Y}\} が与えられたとします。$ 、ここで、 $\mathbf{X}$ はモデル $\mathbf{Y}$ は、対応するセマンティックマップまたはバイナリセグメンテーションマップです。各画像について、 $\mathbf{x}_i∈\mathbb{R}^{H×W ×C}$ 、ここで $H$ 和 $W$ は画像の空間解像度です。 $=\{3,\ldots, N\}$ は入力チャネルの数であり、モデルは出力セグメンテーションマップ $\mathbf{y}_i∈\mathbb{R}^{H×W ×K} を生成します。$ 、ここで $K∈\{1,\ldots,D\}$ . FCT への入力は、入力 3D 画像の各スライスからサンプリングされた 2D パッチです。私たちのモデルは、基本的な構成要素として FCT 層を備えた、おなじみの UNet の形状に従っています。既存の手法とは異なり、私たちのモデルは CNN と Transformer のハイブリッドでも、Transformer と UNet 構造でもありません。既製のトランスフォーマー層を使用して入力特徴をエンコードまたはリファインします。まず画像から重複するパッチを抽出し、次にパッチベースのスキャン埋め込みを作成し、これらのパッチにマルチヘッドセルフアテンションを適用して特徴表現を構築します。指定された画像の出力投影は、ワイドフォーカスモジュールを通じて処理され、投影から詳細な情報が抽出されます。図 1 は、ネットワークアーキテクチャの概要を示しています。

3.1 FCT層

各 FCT レイヤーはLayerNormalization-Conv-Conv-Maxpoolオペレーションから始まります。3×3これらの逐次畳み込みをカーネルサイズが小さいパッチに順次適用すると、最初に画像のパッチごとの投影を直接作成する場合と比較して、画像情報をより適切にエンコードできることが経験的にわかります。各畳み込み層の後にはGelu活性化関数が続きます。FCT ブロックが他のモデルブロックと異なる最初の例は、医療画像用の畳み込み注意アプリケーションにあります。

MaxPoolの出力は変換関数 $\mathbf{T}(·)に入力されます。$ 、変換関数 $\mathbf{T}(·)$ を使用して、新しいトークンマッピングに変換します。 $\mathbf{T}(·)$ を選択します。 $て (\cdot ）$ ですDepthwise-Convolution operator。小さいカーネルサイズ3×3、 strides×s、および効率的なパディングを選択して、(1) ほとんどの既存の作業とは異なり、抽出されたパッチが重複すること、および (2) 畳み込み演算によって常に出力サイズが変更されないことを保証します。次にLayerNormalization操作です。得られたトークンマッピング $p_{i+1}∈\mathbb{R}^{W_t×H_t×C_t}$ $W_tH_t ×C_t$ に平坦化されます $W H \times C$ 、パッチ埋め込み入力を作成します。次の例は、FCT レイヤーが、注意の投影によって既存のトランスベースの医用画像アプリケーション手法とは異なることです。既存のすべてのモデルは、マルチヘッドセルフアテンション (MHSA) の計算に線形ポイントバイポイント線形マッピングを採用しています。これにより、Transformer モデルは画像処理アプリケーションにとって非常に重要な空間情報を失います。既存の方法は、畳み込み強調を画像処理タスクに適応させることでこの問題を軽減しようとしています。ただし、これにより、提案されたモデルに追加の計算コストが追加されます。[36] で提案された方法に触発され、MHSA ブロックの点単位の線形マッピングを置き換えて、Depthwise-Convolution計算コストを削減し、画像からより優れた空間コンテキスト情報を活用します。パッチ埋め込みと畳み込みアテンション投影は、畳み込みアテンションのコンポーネントを形成します。[36] とは異なり、LayerNormalization置換の使用がBatchNormalizationパフォーマンスの向上に役立つことに注意してください。さらに、削除により、Point-wise Convolutionパフォーマンスが低下することなく、より単純なモデルが得られます。Depthwise-Convolution提供される空間コンテキストにより、入力に空間情報を挿入し、各パッチの位置を順次追跡するために使用される位置エンコーディングの必要性がさらに排除され、アーキテクチャ設計がさらに簡素化されます。

一般的なトランスフォーマー層は線形層の MHSA ブロックに続くため、画像内のすべての空間コンテキストが失われます。これらの線形層を畳み込みに直接置き換えることは、この問題を緩和してパフォーマンスを向上させる比較的簡単な方法です。ただし、医療画像にはきめの細かい情報処理が必要です。これを念頭に置いて、マルチブランチ畳み込み層を採用します。この層では、1 つの層が空間畳み込みを MHSA 出力に適用し、他の層が非トラウス畳み込みを適用して、より良い空間コンテキストを得るために受容野を増加させます。次に、これらの特徴を加算によって融合し、特徴集約レイヤーに渡します。この特徴の集約は、別の空間畳み込み演算子によって実現されます。このモジュールをワイドフォーカスと呼びます。残留接続は、層全体にわたる特徴の伝播を強化するために使用されます。最終的な特徴は再形成され、さらに次の FCT 層に伝播されます。図 1 (上) は FCT 層を示しています。

3.2 エンコーダ

私たちのモデルのエンコーダーには、特徴の抽出と伝播を担当する 4 つの FCT レイヤーが含まれています。のために $l$ 変換層、畳み込みアテンションモジュールの出力は、 $\mathbf{z}_l' = \mathbf{MHSA}(z_{l－ 1) })+ \mathbf{z}^{q/k/v}_{l−1}$ 、特に、 $\mathbf{z}^{q/k/v}_{l−1} = \ mathbf{Flatten}(\mathbf{DepthConv}(\mathbf{Reshape}(z_{l−1})))$ MHSA $\mathbf{MHSA}(\mathbf{z}_{l−1})= ソフトマックス(QK^T/\sqrt{d})V$ 。次に、 $z_l' は$ , $\mathbf{z}_l = \mathbf{WF}(\mathbf{z}_l) + \mathbf{z}_l'$ 。さらに、さまざまなカテゴリーやさまざまなスケールでのより小さな ROI 特徴を強調表示することを目的として、ピラミッドスタイルの画像入力をエンコーダーに挿入します。特に、私たちのモデルは、そのようなマルチスケール画像ピラミッド入力がなくても最先端の結果を得ることができます。(ボトルネックとなる) データの潜在的なエンコーディングは、別の FCT レイヤーを使用して作成されます。

3.3 デコーダ

デコーダはボトルネック表現を入力として受け取り、この情報からバイナリ分類またはセマンティックセグメンテーションマップをリサンプリングする方法を学習します。デコーダ層でより適切なコンテキスト相関を作成するために、エンコーダからデコーダへのスキップ接続も使用されます。この場合、エンコーダ層からの同じ解像度の特徴マップがデコーダ層に接続されます。デコーダの形状はエンコーダと対称です。デコーダーのレイヤーはエンコーダーの画像ピラミッドレイヤーに対応し、追加の監視を提供してモデルの予測力を向上させる中間セグメンテーションマップを出力します。コンテキスト依存関係は、最初に特徴量をアップサンプリングし、次にそれらを FCT レイヤーに渡して可能な限り最良の表現を学習することによって作成されます。FCT の最低スケールでは深い規制を採用していないため、私たちのモデルは「完全に深い規制」ではありません。これは、入力画像スキャンの関心領域 (ROI) が小さすぎて28 × 28最小スケール ( ) でセグメント化できない場合があり、その結果、モデルのパフォーマンスが低下することが観察されたためです。この低スケールの出力は、一部の出力 ROI をバックグラウンドクラスとして予測するための強いバイアスをモデルに追加します。

4. 実験

データ集：(MRI) 自動心臓診断チャレンジ (ACDC) [2]、(CT) シナプス多臓器セグメンテーションチャレンジ 1、(CT) 脾臓セグメンテーションデータセット [1]、および (ダーモスコピー) ISIC 2017 [6] 皮膚がんセグメンテーションチャレンジ。

実験の詳細: モデル入力には 2 つのサイズがあります224 x 224。384 x 384Adam オプティマイザー、学習率 1e-3。

5. 結果

6. まとめ

既存のモデルよりも少ないパラメータでバイナリ分類とセマンティックセグメンテーションタスクを正確に実行できる完全な畳み込み変換器を提案します。FCT は、パラメータ数の点で nnFormer の 5 分の 1 以上、TransUNet および LeViT-UNet の 3 分の 1 以上です。FCT レイヤーは、畳み込みアテンションとワイドフォーカスという 2 つの主要なコンポーネントで構成されます。畳み込みアテンションは、深さ方向に分離可能な畳み込みを使用してモデルの重複するパッチを作成することにより、パッチ作成段階での位置エンコーディングの必要性を排除します。当社の深さ方向に分離可能な畳み込みベースの MHSA ブロックは、空間情報を統合して、医療画像のコンテキストにおける長距離の意味論的依存性を初めて推定します。私たちのアブレーション実験から、ワイドフォーカスは医療画像に存在するきめ細かい特徴情報を活用するのに役立ち、トランスブロックのパフォーマンスを向上させる重要な要素であることがわかります。私たちは、さまざまなモダリティと次元の競争力の高い複数のセグメンテーションデータセットに対する最先端の結果を通じて、モデルの機能を実証します。当社の FCT ブロックは、医用画像アプリケーション向けに提案された最初の完全な畳み込み変換ブロックであり、医用画像の他の分野やアプリケーションに簡単に拡張できます。私たちは、このモデルが将来のセグメンテーションタスクの効果的なバックボーンネットワークとして機能し、トランスベースの医療画像処理における革新への道を開くことができると信じています。