【论文笔记】HF-UNet: Learning Hierarchically Inter-Task Relevance in Multi-Task U-Net

2020

Title：HF-UNet:Learning Hierarchically Inter-Task Relevance in Multi-Task U-Net for Accurate Prostate Segmentation in CT images

---

摘要

       本文中，作者通过两阶段网络来解决CT图像中前列腺分割的挑战性任务，其中第一阶段是快速定位，第二是第二阶段是为了准确地分割前列腺。为了在第二阶段精确分割前列腺，将前列腺分割制定为一个多任务学习框架，该框架包括分割前列腺的主要任务和勾画前列腺边界的辅助任务。在这里，第二项任务是为CT图像中不清楚的前列腺边界提供额外的指导。多任务采用分层融合的U-Net结构（HF-UNet），具有用于两个任务的两个互补分支，其中新颖的基于注意力的任务一致性学习模块可在两个解码分支之间的每个级别进行通信。作者在大型计划CT图像数据集上进行了广泛的评估，实验结果表明HF-UNet优于传统的多任务网络体系结构和最新方法。

---

贡献

（1）提出了一种分层融合的多任务架构，它可以更好地分层学习这两个任务的互补功能。 具体来说，提出的TCL块旨在实现层次特征交互以解决两个任务。
（2）通过将轮廓感知纳入网络来解决前列腺分割任务，并通过多任务学习来解决它。通过回归有关前列腺边界及其背景的特殊设计轮廓敏感标签来引导网络。
（3）在339位患者的大型CT图像数据集上进行的评估表明，与最新方法相比，本文方法在前列腺分割上的性能得到了改善。

---

一、解决的问题

       将常规FCN直接用于前列腺分割是有挑战性的，因为由于组织对比度低，无法在CT图像中轻易区分前列腺边界，以及CT图像中的器官形状和外观通常显示出较大不同个体之间的差异性。而且由于常规FCN中使用的降采样操作，通常会破坏位于前列腺边界的体素之间的空间相关性，从而导致最终的分割无法分割完整的前列腺轮廓。现有的普遍解决方案是在多任务学习框架中利用情境感知策略自动定义器官边界。但常规的多任务深度网络通常在所有任务之间共享大多数参数（即特征表示），这可能会限制其数据拟合能力，因为不可避免地会忽略不同任务的特殊性。

二、实验方法

       作者提出了一个两阶段网络，一个阶段用于定位前列腺结构的位置所在，以去缩小所占面积过大的嘈杂背景区域；另一个阶段是在前者的基础上进行前列腺结构的精准分割。

Detection
       采用U-Net架构来自动定位前列腺区域，裁剪完全包括整个前列腺区域的子图像。然后，基于粗略预测的分割图来确定前列腺的中心。再基于获得的器官中心裁剪大小为128×128×128的区域，以完全覆盖前列腺器官。最后将裁剪完成的补丁馈入到后续网络（HF-UNet）。

HF-UNet
       该网络中包含了两项任务——前列腺分割（Branch a）和前列腺边界描绘（Branch b）。后者作为辅助任务，为网络提供关键的形态学指导，对抗降采样对前列腺边界的体素之间的空间相关性的破坏。

（1） 轮廓边界划分。常用的方法是将二分类，将器官的轮廓边界体素作为正样本，而其他体素作为负样本。因为前列腺的边界难以区分，因而这种方法并不适用，于是作者将此问题看作回归问题，并考虑了轮廓附近的像素作为子地面真相。将每个轮廓点及其周围的像素公式化为高斯分布，其核为$\sigma$（即$\sigma$= 5）。 形式上，将loss层的输入表示为（ξX，y），其中$\xi$X是网络的输出特征，而y表示前列腺轮廓的地面标签，并使用均方误差来计算两个任务之间的差异。 然后是轮廓感知的目标输入（X，y）和权重$\theta$可以写成，

$\varphi (X,\theta )$代表轮廓感知分支（Branch b）的输出，$S(y_{ij})$是高斯滤波器的聚集，其中$y$代表位置$(i,j)$，可以公式化为，

其中$(K,L)$表示距离小于$\sigma$的$(i,j)$的邻集。


（2）任务一致性学习模块（TCL）。从全局角度对两个分支学习到的知识进行建模。作者构建了两种TCL块，一种是加权残差TCL块来评估共享信息和任务特定信息的有效性，另一种是基于注意力TCL块自适应学习上述两种信息的比率。TCL模块的结构如图所示，

这个模块的总损失函数是各级的TCL模块的损失之和（使用TCL块与两个分支（输出）之间的特征对的均方误差聚合）。整个TCL的流程就是将两个分支的信息通过加和的操作进行融合，再其后设置瓶颈层对抗增加的噪声信息，然后再经过数个卷积层后，将TCL生成的特征反馈回每个分支。在反馈方法中，作者定义了两个连接，private连接（子网块之间的连接）和public连接（子网块和TCL块之间的连接）。下一级的各分支的输入为上一级中private连接中的特征图与public连接中的特征图以某个权重值进行加和的融合图。然后作者引入了特征图的通道信息和位置信息具现化权重值，最后可得，

多任务过程
       总损失函数如下，由分割损失、回归损失、tcl损失和正则化的参数值组成。其中，分割损失采用的是交叉熵损失。另外，由于回归损失的值普遍较大，因而在冷启动阶段训练分割损失，并且将回归分支和tcl模块中的参数值固定。此外，还引用了权重值约束。

三、实验结果

1.实验数据集

由339位患者组成的大型计划CT图像数据集 ，包含339张图像，图像大小为512×512×(61~508)。

2.实验评估指标

Dice相似系数（DSC），平均表面距离（ASD），灵敏度（SEN），阳性预测值（PPV）。

3.结果分析

       作者先验证了提出的TCL模块的性能。作者将HF-UNet-n(n代表TCL块的数量)与传统的U-Net，后期分支网络（LB），早期分支网络（EB）进行了比较。其中，FB代表融合块的数量（共享参数块+TCL块）。实验结果表明，HF-UNet-6在HF-UNet-1~6中达到了最佳的性能，即通过增加TCL块的数量可以提高分割性能，分层共享策略可以使特征融合收益。

       作者还对private连接上的信息和public连接上的信息权重值进行了研究。实验结果表明，就DSC而言，相对于$\alpha$= 0.8的HF-UNet在DSC中达到0.880±0.028的最佳性能。相对于$\alpha$= 0.4的HF-UNet在ASD下达到1.339±0.500的最佳性能。但是，当同时考虑DSC和ASD时，相对于$\alpha$= 0.2的HF-UNet可获得最佳性能。对于前列腺分割任务而言，private信息比率更低时能获得更好的分割效果。

       最后，作者将自己的方法与几种最新的分割方法进行了比较，实验结果表明就DSC和ASD而言，HF-UNet6-dAtt始终都能达到最佳性能，且ASD的改进幅度较大，这表明所提出方法的有效性。

原文献链接: link.
https://arxiv.org/abs/2005.10439