U-Netは、画像分割、特に生物医学的な文脈のために開発された畳み込みニューラルネットワークアーキテクチャである。これは2015年にフライブルク大学のOlaf Ronneberger、Philipp Fischer、Thomas Broxによって導入され、それ以前の完全畳み込みネットワークを基盤としている。このアーキテクチャは、スキップ接続を備えた対称的なエンコーダ・デコーダ設計によって特徴づけられ、比較的少ない訓練画像から正確な分割マップを生成できる。512×512画像の分割は、U-Netアーキテクチャを使用すると、最新(2015年)のGPUで1秒未満で完了する。
U-Netアーキテクチャは、2014年にEvan Shelhamer、Jonathan Long、Trevor Darrellによって開発された、いわゆる「完全畳み込みネットワーク」(FCN)に由来する。主なアイデアは、通常の収縮ネットワークに連続する層を補足し、プーリング操作をアップサンプリング演算子に置き換えることである。したがって、これらの層は出力の解像度を高める。連続する畳み込み層は、この情報に基づいて正確な出力を組み立てることを学習できる。
ネットワークアーキテクチャ
ネットワークは収縮経路と拡張経路で構成され、これによりU字型のアーキテクチャが形成される。収縮経路は典型的な畳み込みネットワークであり、畳み込みの繰り返し適用で構成され、各畳み込みの後に正規化線形ユニット(ReLU)と最大プーリング操作が続く。収縮中、空間情報は減少する一方で特徴情報は増加する。拡張経路は、一連のアップ畳み込みと、収縮経路からの高解像度特徴との連結を通じて、特徴情報と空間情報を組み合わせる。
U-Netにおける重要な変更点の1つは、アップサンプリング部分に多数の特徴チャネルがあることであり、これによりネットワークはコンテキスト情報を高解像度層に伝播できる。その結果、拡張経路は収縮部分とほぼ対称になり、U字型のアーキテクチャが得られる。ネットワークは、完全に接続された層を一切使用せず、各畳み込みの有効部分のみを使用する。したがって、ネットワークは混合の入力解像度で訓練および実行できる。
画像の境界領域のピクセルを予測するために、画像境界を超えた欠落コンテキストは、画像境界で入力画像をミラーリングすることによって外挿される。元のU-Net論文はさらにタイル戦略を提案しており、大きな画像は重複するタイルに切り分けられ、各タイルは独立して処理される。これにより、利用可能なGPUメモリを超える高解像度画像の処理が可能になる。最近では、医用画像分割のための受容野ベースのU-Netモデルへの関心も高まっている。
訓練と損失関数
U-Netは、ピクセル単位の損失関数(通常はクロスエントロピーまたはその重み付き変種)を用いた確率的勾配降下法(SGD)を使用してエンドツーエンドで訓練される。エネルギー関数は、最終特徴マップ上のピクセル単位のソフトマックスとクロスエントロピー損失の組み合わせによって計算される。著者らは、接触するオブジェクト間の境界近くのピクセルに高い重要度を与える重みマップを導入し、これによりネットワークが生物医学画像内の個々のインスタンスを分離することを学習できる。この重みマップは各訓練画像に対して事前計算され、境界領域の誤差をより重く罰するために使用される。
ネットワークはデータ拡張を使用して実効的な訓練セットサイズを増加させ、弾性変形、回転、強度変動を含む。これは、注釈付きデータが乏しい生物医学画像において特に重要である。元の実装では、安定した収束を確保するために高いモーメンタム(0.99)と低い学習率(0.00001)を使用した。
生物医学画像における応用
U-Netの生物医学画像分割における応用は多数あり、コンピュータ断層撮影(CT)および磁気共鳴画像法(MRI)スキャン内のさまざまな臓器や身体系の分割が含まれる。具体的な例としては、脳画像分割(BRATS)や肝臓画像分割(siliver07)、タンパク質結合部位予測が挙げられる。U-Net実装は物理科学でも使用されており、例えば材料の顕微鏡写真の解析などがある。
U-Netの変種は医用画像再構成にも適用されている。以下にU-Netのいくつかの変種と応用を示す。
- ピクセル単位回帰を用いたU-Netとパンシャープニングへの応用。
- 3D U-Net:スパース注釈からの密な体積分割の学習。
- TernausNet:画像分割のためのImageNetで事前訓練されたVGG11エンコーダを備えたU-Net。
- 蛍光染色を推定するための画像間変換。
- タンパク質構造の結合部位予測。
拡散モデルでの使用
U-Netアーキテクチャは、反復画像ノイズ除去のための拡散モデルでも採用されている。この技術は、DALL-E、Midjourney、Stable Diffusionなど、多くの現代の画像生成モデルの基盤となっている。これらのモデルでは、U-Netは拡散プロセスの各ステップで画像に追加されたノイズを予測するように訓練され、モデルがランダムノイズ入力を徐々に一貫した画像に洗練できるようにする。U-Netのスキップ接続は、ノイズ除去プロセス中に高周波詳細を保持するため、ここで特に有益である。
言語モデルへの探求
U-Netは言語モデルでも探求されている。トークン化は別個のステップではなく、モデルが綴りをより容易に理解し、同時に高次概念をベクトル化・トークン化できる。これは、U-Netのマルチスケール特徴捕捉能力を活用する新興研究分野であり、自然言語処理における支配的なTransformer (architecture)アーキテクチャの代替を提供する可能性がある。
歴史と遺産
U-Netは2015年にOlaf Ronneberger、Philipp Fischer、Thomas Broxによって作成され、論文「U-Net: Convolutional Networks for Biomedical Image Segmentation」で報告された。これはFCNの改良および発展であり、Evan Shelhamer、Jonathan Long、Trevor Darrell(2014年)の「Fully convolutional networks for semantic segmentation」に基づく。この論文は深層学習分野で最も引用されたものの1つとなり、アーキテクチャは生物医学画像を超えた無数のタスクに適応されてきた。
U-Netの成功は多くの変種を生み出し、異なるエンコーダ(VGGやResNetなど)、注意機構、マルチスケール処理を含む。その影響は機械学習の他の分野にも及び、Generative AIやコンピュータビジョンを含む。このアーキテクチャは医用画像チャレンジの一般的なベースラインでもあり、TensorFlowやPyTorchなどのフレームワークで広く実装されている。
実装
U-Netのオープンソース実装がいくつか利用可能である。J Akeret(2017年)によるTensorflow Unetは人気のある実装である。元のU-Netソースコードは、ドイツのフライブルク大学コンピュータサイエンス学科のパターン認識および画像処理グループから入手可能である。これらの実装により、U-Netは研究と産業の両方で広く採用されることが促進された。