[論文レビュー] A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities

Simon Köhl, Bernardino Romera‐Paredes+6 more|arXiv (Cornell University)|May 30, 2019

AI in cancer detection参考文献 42被引用数 27

ひとこと要約

本稿では、粗いスケールから細かいスケールへの空間的に構造化された潜在変数マップの階層を用いて、医療画像および自然画像のセグメンテーションにおけるマルチスケールのあいまいさをモデル化する、U-Netアーキテクチャと統合された条件付き変分オートエンコーダである階層的確率的U-Net（HPU-Net）を提案する。本手法は、複雑で構造的な分布をスケール間で学習しつつ、高精度な再構成と詳細なセグメンテーションのサンプリングを可能にし、インスタンスセグメンテーションおよび不確実性を考慮した予測タスクにおいて、先行モデルを上回る性能を発揮する。

ABSTRACT

Medical imaging only indirectly measures the molecular identity of the tissue within each voxel, which often produces only ambiguous image evidence for target measures of interest, like semantic segmentation. This diversity and the variations of plausible interpretations are often specific to given image regions and may thus manifest on various scales, spanning all the way from the pixel to the image level. In order to learn a flexible distribution that can account for multiple scales of variations, we propose the Hierarchical Probabilistic U-Net, a segmentation network with a conditional variational auto-encoder (cVAE) that uses a hierarchical latent space decomposition. We show that this model formulation enables sampling and reconstruction of segmenations with high fidelity, i.e. with finely resolved detail, while providing the flexibility to learn complex structured distributions across scales. We demonstrate these abilities on the task of segmenting ambiguous medical scans as well as on instance segmentation of neurobiological and natural images. Our model automatically separates independent factors across scales, an inductive bias that we deem beneficial in structured output prediction tasks beyond segmentation.

研究の動機と目的

• 画像の証拠が単一の正しい解釈を決定するのに不十分な場合に、医療画像セグメンテーションにおける複雑でマルチスケールのあいまいさをモデル化する課題に対処すること。
• 階層的な潜在空間の分解を可能にすることで、細部にまで及ぶ構造的詳細を保持した、生成および再構成されたセグメンテーションマップの再現性を向上させること。
• 空間スケールおよび位置ごとに独立して変化する、構造的かつ分離可能なセグメンテーションの分布を学習すること。
• 曇りや隠蔽がある領域においても、不確実性を考慮したセグメンテーションとインスタンスレベルのラベリングを可能にすること。
• 複数の妥当なセグメンテーション仮説を生成することで、診断予測などの下流タスクを支援すること。

提案手法

• U-Netエンコーダ・デコーダと、複数の解像度における空間的に構造化された潜在マップを用いる階層的条件付き変分オートエンコーダ（cVAE）を統合する。
• 各スケールにおける潜在変数がU-Netのデコーダパスに組み込まれており、不確実性および構造のスケール別モデリングが可能になる。
• 階層的な事前分布および事後分布が、空間的に配置された潜在マップ上で定義されており、粗いスケールから細かいスケールへの段階的生成が可能になる。
• インスタンスレベルのラベリングのために、セグメンテーションサンプルに対してハミング距離に基づくグリーディクラスタリングアルゴリズムを適用する。この手法により、複数のサンプルにわたって一貫したクラスラベルを持つピクセルをグループ化する。
• 後処理として、不審な小さな領域を除去するエロージョンベースのフィルタリングと、境界付近のアーチファクトを解消するメジャリティラベル置換を実施する。
• エンドツーエンドで訓練されるモデルは、再構成損失およびKLダイバージェンス損失を用い、階層的潜在変数が空間的およびスケール特有の変動を分離してモデリングすることを可能にする。

実験結果

リサーチクエスチョン

• RQ1 階層的潜在空間は、曇りや曇りのある医療画像における、生成および再構成されたセグメンテーションマップの再現性を向上させることができるか？
• RQ2 本モデルは、構造的出力予測タスクにおいて、複数の空間スケールにまたがる独立した変動を学習し、分離できるか？
• RQ3 階層的cVAEアーキテクチャは、隠蔽やあいまいさが生じる複雑なシーンにおいて、グローバル潜在変数モデルに比べてより優れたインスタンスセグメンテーションを実現できるか？
• RQ4 本モデルは、特に部分的に遮蔽された領域において、一貫性があり多様なセグメンテーション仮説を生成できるか？
• RQ5 スケール間での変動要因の分離能力が、下流の診断タスクにおける性能向上にどの程度寄与するか？

主な発見

• HPU-Netは、標準的な確率的U-Netに比べて、特に細部の構造的特徴を保持する点で顕著に高い再構成精度およびサンプリング再現性を達成している。
• 神経生物学的画像（SNEMI3D）および自然画像（Cityscapes）の両データセットにおいて、本モデルは一貫したインスタンスセグメンテーションを生成し、1枚の画像あたり5つの異なるインスタンスIDを学習している。
• Cityscapesの車両インスタンスセグメンテーションタスクにおいて、本モデルはサンプル間で曇っている領域（例：道路/歩道、トラック/バス）を適切に反転させることができ、学習された不確実性を反映している。
• 最もグローバルなスケールからのみサンプリングを行うと粗い低解像度マスクが得られ、逆に最も局所的なスケールからのみサンプリングを行うと高解像度だがノイズの多い予測が得られる。これは、階層的潜在変数が補完的な役割を果たしていることを示している。
• ブラックアウトされた領域においても、不確実性を伴った一貫したセグメンテーションを予測できており、時系列医療画像における4D疾患進行モデリングへの応用可能性が示唆されている。
• グリーディハミング距離クラスタリングアルゴリズムは、ピクセルを一貫したインスタンスに効果的にグループ化できており、後処理によりセグメンテーション境界におけるアーチファクトが除去されている。