U-NetとTransformerを合体！医用画像セグメンテーションのための新しいネットワークSwin Unetを紹介｡

medical 2022年05月20日

3つの要点
✔️ 医用画像解析においてもCNNは画期的な成果をあげているが、畳込み演算のせいで大域的な情報を学習できない。
✔️ Transformerは大域的な情報を学習できるため、本論文ではTransformerベースのU-NetであるSwin-Unetを提案する。
✔️ 多臓器セグメンテーションタスクで検証した結果、Swin-UnetはCNNベースのU-Netよりも優れていることを実証した。

Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation
written by Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, Manning Wang
(Submitted on 12 May 2021)
Comments: Published on arxiv.
Subjects: Image and Video Processing (eess.IV); Computer Vision and Pattern Recognition (cs.CV)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

概要

ここ数年、畳み込みニューラルネットワーク（convolutional neural network, CNN）は医用画像解析に画期的な成果をもたらしてきました。特にセグメンテーションタスクにおいてU-Netは、ディープラーニングの医療分野において広く用いられています。

しかしCNNは、畳込み演算の局所性のため大域的な特徴をうまく学習することができないと指摘されるようになりました。これに比較してTransformerは大域的な特徴を学習するため、医用分野でも利用されることになりました。

本研究では、TransformerベースのU-NetであるSwin-Unetを提案し、多臓器のセグメンテーションを行いました。その結果Swin-Unetは、CNNやTransformer+CNNの手法よりも優れていることが実証されました。

背景

近年、自然言語分野においてTransformerと呼ばれるネットワークが報告され、その高い性能から急激に人気が出ています。Transformerの核となるアイデアはAttention（注意）機構と呼ばれており、翻訳タスクにおける入力のどの単語と出力のどの単語が結びついているのかを明示し、さらにその単語が文章の中でのどれくらい重要なのかを計算する手法です。文全体に対するその単語の重要性を考慮し出力が決定されるため、Transformerは大域的な情報を学習することに長けています。以下にTransformerのアーキテクチャを示します｡

これに対して、画像分野におけるCNNの核となるアイデアは畳み込みです。畳み込みとは、画像を構成する画素のまとまりから情報を集約することです。まるで画像を畳んで小さくするような処理をするため畳み込みと呼ばれています。つまり形の変化、色の変化といった情報を集約し、積み上げていくことによって出力が決定されます。そのためCNNは局所的な情報を集約することに長けており、逆に言えば、画像上で離れた部位同士の情報は集約できません。

このようにCNNとTransformerのアプローチは全く異なるものでしたが、のちにTransformerのAttentionが画像分野にも応用されました。これをVision Transformer（ViT）といい、ViTは猛烈な勢いで流行することとなり、医用分野でも素晴らしい成果を達成しています。ViTの基本的なアイデアは、画像を細分化して、それを自然言語分野における単語とみなした点にあります。つまり画像を文章のように処理したわけですね。その結果、画像上の大域的な情報を学習することに成功しました。

そして医用画像のセグメンテーションタスクではU-Netが成功を収めてきました。そして3D U-Net、Res-Unet、U-Net+など様々な改良型も報告されています。しかしU-Netの構造はCNNベースであるために、本質的に大域的な除法を集約できません。そこで本研究ではU-NetにViTを適用し、Swin-Unetと呼ばれる新しいネットワークを提案しました。

Swin-Unetにおいて、入力された画像はTransformerベースのエンコーダーに供給され、空間的に広い特徴を学習します。本研究では、多臓器セグメンテーションと心臓セグメンテーションで検証した結果、提案手法は優れたセグメンテーション精度と頑健な汎化性を有することが示されました。

実験方法

Swin-Unetのアーキテクチャを以下に示します｡

入力画像はまず4x4ピクセルの正方形で区切られ､パッチ化されます｡これがSwin Transformerブロックに入力されます｡Patch Merging層ではダウンサンプリングが施され､Patch Expanding層はその逆です｡ダウンサンプリングによる空間情報の損失を補完するためにスキップコネクションを通じてエンコーダからアップサンプリングされた特徴にマージされます｡

Swin-Unetの基本単位であるSwin Transformerを以下に示します｡

Swin Transformerでは､従来のMulti-head Self Attentionモジュールがそのまま使われてはいません｡Swin Transformerは､LayerNorm層（LN）とウィンドウベースのMulti-head Self Attention（W-MHA）層からなります｡

実験

30症例（註：CTなので1症例あたり約3800枚あります）が含まれるSynapse multi-organセグメンテーションデータセット（以後､シナプス）を用いて､セグメンテーションタスクを行いました｡先に結果を示します｡

実際に使用したのは腹部8臓器（大動脈､胆嚢､脾臓､左右腎臓､肝臓､膵臓､胃）で､評価には平均ダイス係数（DSC）と平均ハウスルドルフ距離（HD）を用いています｡太字が最大スコアを達成しているモデルであり､例えば大動脈（Aorta）ではAtt-UNetが最も良い成績を出しています｡提案されたSwin Unetでは､左腎､肝臓､脾臓､胃のスコアで最高値を出しており､全体的な平均では既存手法を上回っています｡

セグメンテーションの例を上に示します｡左から順に､正解ラベル､Swin Unet､TransUNet､AttUnet､UNetと並んでいます｡黄色の膵臓は扁平なので検出が難しい臓器です｡下段の行では膵臓のセグメンテーションに大きな差があることがわかります｡

結論

本論文では､医用画像セグメンテーションのためにTransformerベースのU-Netを提案しました｡Transformerによる大域的な情報の学習のため､Swin Transformer Blockを用いました｡その結果､多臓器セグメンテーションタスクにおいてSwin Unetが優れた性能を発揮することがわかりました｡