[论文解读] Unimodal-uniform Constrained Wasserstein Training for Medical Diagnosis

Xiaofeng Liu, Xu Han+3 more|arXiv (Cornell University)|Nov 3, 2019

Medical Image Segmentation Techniques被引用 4

一句话总结

本文提出了一种用于具有有序标签的医学诊断的单模态-均匀约束Wasserstein训练损失，利用基于欧几里得距离递增函数的基度量来编码类别相关性。通过结合平滑目标标签与精确的闭式解，该方法在糖尿病视网膜病变（Diabetic Retinopathy）和超声BIRADS（Ultrasound BIRADS）数据集上实现了最先进性能，优于交叉熵损失、回归方法以及先前的有序方法。

ABSTRACT

The labels in medical diagnosis task are usually discrete and successively distributed. For example, the Diabetic Retinopathy Diagnosis (DR) involves five health risk levels: no DR (0), mild DR (1), moderate DR (2), severe DR (3) and proliferative DR (4). This labeling system is common for medical disease. Previous methods usually construct a multi-binary-classification task or propose some re-parameter schemes in the output unit. In this paper, we target on this task from the perspective of loss function. More specifically, the Wasserstein distance is utilized as an alternative, explicitly incorporating the inter-class correlations by pre-defining its ground metric. Then, the ground metric which serves as a linear, convex or concave increasing function w.r.t. the Euclidean distance in a line is explored from an optimization perspective. Meanwhile, this paper also proposes of constructing the smoothed target labels that model the inlier and outlier noises by using a unimodal-uniform mixture distribution. Different from the one-hot setting, the smoothed label endues the computation of Wasserstein distance with more challenging features. With either one-hot or smoothed target label, this paper systematically concludes the practical closed-form solution. We evaluate our method on several medical diagnosis tasks (e.g., Diabetic Retinopathy and Ultrasound Breast dataset) and achieve state-of-the-art performance.

研究动机与目标

• 解决标准交叉熵损失在有序医学诊断中的局限性，即忽略类别间相似性与顺序标签结构。
• 通过单模态-均匀混合分布建模临床数据中的标签噪声，引入内点与外点不确定性。
• 开发一种可微且高效的损失函数，利用带有结构化基度量的Wasserstein距离显式编码有序关系。
• 为单热编码与平滑目标标签提供闭式解，实现快速且稳定的训练。
• 在具有离散、连续风险等级的真实世界医学数据集上展示优越性能。

提出的方法

• 使用一阶Wasserstein距离作为损失函数，衡量预测标签分布与目标标签分布之间的最优传输成本。
• 将基度量定义为沿直线上欧几里得距离的递增函数，以编码类别间的有序相似性。
• 引入单模态-均匀混合分布以建模标签噪声，其中真实标签被附近或遥远的类别所扰动。
• 推导出在单热编码与平滑目标标签设置下Wasserstein损失的精确闭式解，实现高效优化。
• 采用欧几里得距离的凸函数与凹函数作为基度量，以探索优化行为。
• 采用平滑目标标签分布以在不增加额外训练成本的情况下提升鲁棒性与泛化能力。

实验结果

研究问题

• RQ1 与标准交叉熵损失和回归方法相比，基于Wasserstein的损失结合结构化基度量是否能提升有序医学诊断的性能？
• RQ2 通过单模态-均匀混合分布建模标签噪声，对模型鲁棒性与准确率有何影响？
• RQ3 在有序分类的Wasserstein训练中，使用精确闭式解与近似优化相比有何影响？
• RQ4 不同形式的基度量（如线性、凸函数、凹函数）对模型性能与收敛性有何影响？
• RQ5 所提出方法是否能超越糖尿病视网膜病变数据集，在其他有序医学数据集（如US-BIRADS）上也实现泛化？

主要发现

• 在糖尿病视网膜病变数据集上，采用欧几里得距离的凸递增函数（D²）作为基度量的Wasserstein损失取得了最高性能，准确率达到80.6%，QWK为0.716。
• 单模态-均匀平滑目标标签在所有基线方法中均提升了性能，尤其在US-BIRADS等噪声较高的设置中贡献显著，达到最先进结果。
• 精确闭式解始终优于其近似版本，表现出更快的收敛速度与更好的泛化能力。
• 在糖尿病视网膜病变数据集上，该方法实现了80.6%的准确率与0.716的QWK，优于先前的多任务、分段构建与泊松基方法。
• 在US-BIRADS数据集上，该方法实现了80.6%的准确率与0.716的QWK，由于标签噪声更高，平滑标签带来了更显著的性能提升。
• 超参数敏感性分析表明，QWK在τ与ξ变化下保持稳定，表明对参数调优具有鲁棒性。