[論文レビュー] 3D U$^2$-Net: A 3D Universal U-Net for Multi-Domain Medical Image Segmentation

Chao Huang, Hu Han+3 more|arXiv (Cornell University)|Sep 4, 2019

Advanced Neural Network Applications参考文献 18被引用数 9

ひとこと要約

本論文は、分離可能畳み込みを用いてドメイン間の特徴を共有しながら、チャネル単位の畳み込みによってドメイン固有のパターンを学習する、多分野の医用画像セグメンテーションに適応可能な3次元U²-Netを提案する。このモデルは、タスク固有のネットワークと同等のセグメンテーション精度を達成しながら、パラメータのわずか1%しか使用せず、新しい臓器に対してもたった0.3%の追加パラメータで効果的に適応可能である。

ABSTRACT

Fully convolutional neural networks like U-Net have been the state-of-the-art methods in medical image segmentation. Practically, a network is highly specialized and trained separately for each segmentation task. Instead of a collection of multiple models, it is highly desirable to learn a universal data representation for different tasks, ideally a single model with the addition of a minimal number of parameters steered to each task. Inspired by the recent success of multi-domain learning in image classification, for the first time we explore a promising universal architecture that handles multiple medical segmentation tasks and is extendable for new tasks, regardless of different organs and imaging modalities. Our 3D Universal U-Net (3D U$^2$-Net) is built upon separable convolution, assuming that {\it images from different domains have domain-specific spatial correlations which can be probed with channel-wise convolution while also share cross-channel correlations which can be modeled with pointwise convolution}. We evaluate the 3D U$^2$-Net on five organ segmentation datasets. Experimental results show that this universal network is capable of competing with traditional models in terms of segmentation accuracy, while requiring only about $1\%$ of the parameters. Additionally, we observe that the architecture can be easily and effectively adapted to a new domain without sacrificing performance in the domains used to learn the shared parameterization of the universal network. We put the code of 3D U$^2$-Net into public domain. \url{https://github.com/huangmozhilv/u2net_torch/}

研究の動機と目的

• 各医用画像セグメンテーションタスクに対して個別に深層学習モデルを訓練するという課題に対処する。これは非効率的で、拡張が困難である。
• 再訓練なしに、さまざまな解剖的構造および画像モality（例：CT、MRI）をカバーする1つのニューラルネットワークの汎用性を実現する。
• 複数のドメインにわたってモデルサイズを最小限に抑えつつ高いセグメンテーション精度を維持する、パラメータ効率の良いアーキテクチャを開発する。
• 既存のタスクでのパフォーマンス低下を伴わずに、新しい臓器（例：脾臓）への転移学習を最小限の追加パラメータで可能にする。
• 画像分類からの成功を、密度予測タスクにまで拡張することで、医用画像セグメンテーションにおける多ドメイン学習のギャップを埋める。

提案手法

• ドメイン固有の空間的パターンを捉えるためにドメインごとのチャネル単位畳み込みを、共通の相関関係をモデル化するために共有されたポイントワイド畳み込みを用いた、3次元U-Netに基づく汎用アーキテクチャを設計する。
• 新しいドメインに適応する際、共有コンponent（例：ポイントワイド畳み込み）を固定し、ドメイン固有のパラメータ（例：入力層、チャネル単位畳み込み）を導入する。
• 最適化中の各ドメインの寄与をバランスさせるために、全ドメインをラウンドロビン方式でサンプリングして、汎用モデルを訓練する。
• 学習率のスケジューリングに、訓練損失の指数移動平均を用い、損失が停滞した際に5倍の割合で減少させる。
• 5つの臓器セグメンテーションデータセットにおけるエンドツーエンド訓練に、ADAM最適化法を用い、重み減衰と初期学習率3×10⁻⁴を設定する。
• 共有レイヤーを固定し、ドメイン固有モジュール（例：新しい入力層、チャネル単位畳み込み）のみを微調整することで、ゼロショット適応を可能にする。

実験結果

リサーチクエスチョン

• RQ1 異なる臓器と画像モダリティ（例：心臓、肝臓、前立腺）を持つ複数の医用画像ドメインにおいて、1つの深層ニューラルネットワークが競争力のあるセグメンテーション性能を達成できるか？
• RQ2 共有部とドメイン固有部を有する汎用アーキテクチャが、独立モデルおよび共有のみのモデルに比べて、セグメンテーション精度とパラメータ効率の両面で優れているか？
• RQ3 新しい未観測臓器（例：脾臓）に対して、最小限の追加パラメータでどれほど汎用モデルが一般化できるか？また、既存タスクのパフォーマンスに劣化は生じないか？
• RQ4 チャネル単位畳み込みによるドメイン固有特徴とポイントワイド畳み込みによる共有特徴の組み合わせにより、分離可能畳み込みが多ドメインセグメンテーションにおけるモデル効率とパフォーマンスをどのように向上させるか？
• RQ5 新しい医用画像タスクへの拡張において、モデルサイズ、パラメータ共有、セグメンテーション精度の間のトレードオフはどのようなものか？

主な発見

• 3D U²-Netは、5つのベースドメイン（心臓、肝臓、海馬、前立腺、膵臓）で平均Diceスコア83.10%を達成し、独立モデルの84.28%に非常に近い性能を示したが、パラメータは1.7Mと126.7Mのわずか1%にとどまった。
• 汎用モデルは、すべてのベースドメインで共有モデル（82.56%平均Dice）を上回り、特に前立腺の周辺部（PZ）と遷移部（TZ）のセグメンテーションで顕著な向上（68.50% vs. 68.38%）を示した。
• 新しいドメイン（脾臓）に適応した際、汎用モデルは91.60%のDiceスコアを達成し、共有モデル（90.67%）を上回った。また、完全な独立モデルのパラメータのわずか0.3%で実現した。
• 汎用モデルのパラメータ数は、独立モデルの約1%（1.7M vs. 126.7M）であり、極めて高いパラメータ効率を示した。
• 新しい臓器への適応後も、すべての元のドメインで優れたパフォーマンスを維持しており、共有表現が安定的かつ汎用的であることを確認した。
• 共有コンponentを固定し、ドメイン固有モジュールのみを訓練することで、ゼロショット適応が効果的に可能となり、パフォーマンス低下が最小限で、高いスケーラビリティを有する。