TransformerがCNNに取って代わる!Transformerによる医用画像のセグメンテーション!
3つの要点
✔️ 3D医療画像のセグメンテーションのための初のCNN free モデル
✔️ 3つの異なるデータセットにおいて、CNNよりも優れた、またはCNNと競合するセグメンテーション精度を実現
✔️ CNNよりもはるかに優れた伝達学習能力を発揮
Convolution-Free Medical Image Segmentation using Transformers
written by Davood Karimi, Serge Vasylechko, Ali Gholipour
(Submitted on 26 Feb 2021)
Comments: Accepted to arXiv.
Subjects: Image and Video Processing (eess.IV); Computer Vision and Pattern Recognition (cs.CV)
code:
はじめに
深層ニューラルネットワークは、医療分野で大きな影響を与えています。悪性画像の分類、画像内の悪性領域の検出とセグメント化など、さまざまなタスクを効率的に実行できます。人手による作業が最も信頼できる方法とされていますが、DNNモデルはより高速で、拡張性があり、多少の初期費用はかかりますが安価です。医療用画像処理タスクにおけるDNNの成功の裏には、畳み込みニューラルネットワーク(CNN)があります。CNNは画像に対する強い帰納的バイアスを持っており、さまざまな視覚タスクで成功を収めています。アーキテクチャや損失関数、DNNモデルの学習方法などは大きく変わったが、基本的な構造要素である畳み込み層は変わりません。
トランスフォーマーは、いくつかのNLPタスクで非常に優れた性能を発揮しています。ビジュアルトランスフォーマー(ViT)の導入により、コンピュータビジョンのタスクにおいても自己学習が効率的であることが証明されました。このことから、トランスフォーマーが医療用ビジョンタスクの現状を改善するのに役立つのではないかと考えています。
本論文では、3D医用画像のセグメンテーションのために、CNNを使用しない変換ベースのモデルを紹介します。このモデルは、CNNベースのモデルと同等以上の性能を持ち、ラベル付けされた画像が20~200枚しかないデータセットでも微調整が可能です。さらに、このモデルは最先端のCNNよりも優れた伝達学習能力を持っていることもわかります。
アーキテクチャ
上の図は、CNNフリーネットワークのモデルアーキテクチャを示しています。まず、3D画像からブロックを抽出し、n3個のパッチに分割します。ブロック B ∈ RW×W×W×c の点をn3個の重ならないパッチ{pi ∈ Rw×w×w×c}に分割するとすると、w =W/n、cは画像の次元数は3です。次に各パッチをw3c次元のベクトルに平坦化します。これらのw3c次元のベクトルは、学習した線形マッピングを用いてD次元の空間に変換されます。このようにして得られたベクトルに、位置エンコーディングを加えます。D: X0 = [Ep1; ...; EpN ] + Epos このようにして、それぞれがD次元のn3個のブロックの列を形成します。他の多くのタスクとは異なり、ここでの位置エンコーディングは学習可能なパラメータであることに注意してください。
変圧器エンコーダーにはKステージがあり、それぞれに1つのマルチヘッド自己注意レイヤー(MSA)と、それに続くレイヤー正規化を使用する2つのフィードフォワードネットワーク(FFN)があります。 エンコーダは標準のトランスと非常によく似ています。 query(Q)、key(K)、values(V)が計算され、自己注意は次の式を使用して計算されます。ここで、Dhは非表示の次元に等しい倍率です。
最後のFNNの後、シーケンスは全結合層を使用して次元nclassの空間に投影されます。ここで、nclassはクラスの数です(バイナリセグメンテーションの場合は2)。 この行列は、IRn×n×n×nclassクラスのセグメンテーションマスクY に再形成されます(ブロックはセグメント化されており、個々のピクセルではありません)。
これは、ブロックの中央パッチ専用のセグメンテーションマスクです。 このプロセスは、3D画像ブロック全体に対して繰り返す必要があります。
実験と評価
このモデルは、DICE係数(DSC)に基づく医療画像セグメンテーションのためのCNNベースの最先端モデルである3D UNet ++と比較します。
Pre-Training
ラベル付きのトレーニングインスタンスの数が少ない場合のモデルの精度をさらに向上させるために、ノイズ除去とインペインティング(画像再構成)のタスクのための大規模な非ラベルのデータセットでモデルを学習しています。ノイズ除去では、SNR=10dBのガウスノイズを画像ブロックの中央のパッチに加えます。またインペインティングでは、画像ブロックのセンターパッチの値を0に設定して、画像を再構成します。どちらのタスクも、実画像と構築された画像の間のL2距離を最小化するように学習されます。モデルはソフトマックス層なしで事前学習されており、微調整のためにセグメンテーションマスクを予測するソフトマックス層がモデルに追加されています。また、ネットワーク全体を微調整する方が、最後の層だけを微調整するよりも効果的であることがわかりました。
評価
このモデルは、脳の皮質板、海馬、膵臓の画像の3つの異なるデータセットでベンチマークを行いました。
上の図は、異なるデータセットでの実験結果を示しています。ご覧のように、ほとんどすべてのケースでモデルは異なるメトリクスにおいて、UNet++モデルよりも優れた性能を示しています。下の図は、皮質板(左)と膵臓(右)について、極めて少ないインスタンス(5,10,15)で、提案モデルとUNet++の性能を示しています。どちらのケースでも、提案モデルの方が適応性が高く、またインペインティングを用いた事前学習の方が、デノイジングを用いた事前学習よりも効果的であることがわかります。また学習可能な位置エンコーディングを用いた学習は、固定された位置エンコーディングを用いた学習よりもうまくいくことがわかりました。下の図は、いくつかのアブレーション研究の結果を示しています。
予測結果画像
膵臓のセグメンテーションは難しいタスクとされています。そんな中でもある程度のセグメンテーションマスクの予測ができていると思われます。
まとめ
脳の皮質板を手動で分割するような作業は非常に複雑で、専門家でも数時間かかるものもあります。このような場合には、今回のようなモデルが有利であることは間違いありません。医療用コンピュータビジョンのタスクでは、真の陽性の学習が少ないことを考えると、より少ないインスタンスから学習できる今回のようなモデルは必然的に必要になります。最後に、「3次元点変換モデル」が医用画像のセグメンテーションタスクでどのように機能するかも興味深いところです。このモデルは医療画像ではテストされていませんが、他の3Dセグメンテーションタスクでは非常に素晴らしい結果を示しています。
この記事に関するカテゴリー