[论文解读] 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（HPU-Net），一种将U-Net架构与条件变分自编码器（cVAE）相结合的模型，通过使用空间结构化的分层潜在变量图来建模医学图像与自然图像分割中的多尺度模糊性。该方法在学习跨尺度的复杂结构化分布的同时，实现了高保真度的重建与采样，显著优于先前模型，在实例分割与不确定性感知预测任务中表现更优。

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 该模型在不同尺度上分离变化因子的能力，在多大程度上提升了其在下游诊断任务中的性能？

主要发现

• 与标准概率U-Net相比，HPU-Net在保留细粒度结构细节方面，显著提升了重建与采样保真度。
• 在神经生物学图像（SNEMI3D）与自然图像（Cityscapes）数据集中，该模型均能生成一致的实例分割，每幅图像学习到5个独立的实例ID。
• 在Cityscapes的汽车实例分割任务中，模型成功在样本间翻转模糊区域（如街道/人行道、卡车/公共汽车），反映出学习到的不确定性。
• 仅从最全局尺度采样时，模型生成粗略、低分辨率的掩码；仅从最局部尺度采样时，生成高分辨率但噪声较大的预测，表明分层潜在变量在功能上具有互补性。
• 在被遮挡区域，模型能预测一致的分割结果并体现不确定性，表明其在时序医学影像中4D疾病进展建模中的适用性。
• 基于汉明距离的贪婪聚类算法能有效将像素聚类为连贯的实例，后处理可有效去除分割边界处的伪影。