[Paper Review] Unimodal-uniform Constrained Wasserstein Training for Medical Diagnosis

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

Medical Image Segmentation Techniques4 citations

TL;DR

This paper proposes a unimodal-uniform constrained Wasserstein training loss for medical diagnosis with ordinal labels, leveraging a ground metric based on increasing functions of Euclidean distance to encode class correlations. It achieves state-of-the-art performance on Diabetic Retinopathy and Ultrasound BIRADS datasets by combining smoothed target labels with exact closed-form solutions, outperforming CE loss, regression, and prior ordinal methods.

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.

Motivation & Objective

• To address the limitation of standard cross-entropy loss in ordinal medical diagnosis, where inter-class similarity and sequential label structure are ignored.
• To model label noise in clinical data by incorporating inlier and outlier uncertainty via a unimodal-uniform mixture distribution.
• To develop a differentiable, efficient loss function that explicitly encodes ordinal relationships using Wasserstein distance with a structured ground metric.
• To provide closed-form solutions for both one-hot and smoothed target labels, enabling fast and stable training.
• To demonstrate superior performance on real-world medical datasets with discrete, successive risk levels.

Proposed method

• Uses the 1st Wasserstein distance as a loss function to measure the optimal transport cost between predicted and target label distributions.
• Defines a ground metric as an increasing function of Euclidean distance along a line, encoding ordinal similarity between classes.
• Introduces a unimodal-uniform mixture distribution to model label noise, where true labels are perturbed by nearby or distant classes.
• Derives exact closed-form solutions for the Wasserstein loss under both one-hot and smoothed target label settings, enabling efficient optimization.
• Applies convex and concave increasing functions of the Euclidean distance as ground metrics to explore optimization behavior.
• Employs a smoothed target label distribution to improve robustness and generalization without additional training cost.

Experimental results

Research questions

• RQ1 Can a Wasserstein-based loss with a structured ground metric improve performance in ordinal medical diagnosis compared to standard cross-entropy and regression?
• RQ2 How does modeling label noise via a unimodal-uniform mixture distribution affect model robustness and accuracy?
• RQ3 What is the impact of using exact closed-form solutions versus approximate optimization in Wasserstein training for ordinal classification?
• RQ4 How do different forms of the ground metric (e.g., linear, convex, concave) affect model performance and convergence?
• RQ5 Does the proposed method generalize beyond Diabetic Retinopathy to other ordinal medical datasets like US-BIRADS?

Key findings

• The proposed Wasserstein loss with a convex increasing function of Euclidean distance (D²) achieved the highest performance on the Diabetic Retinopathy dataset, with 80.6% accuracy and 0.716 QWK.
• The unimodal-uniform smoothed target label improved performance across all baselines, especially in noisy settings like US-BIRADS, where it contributed significantly to state-of-the-art results.
• The exact closed-form solution consistently outperformed its approximate counterpart, demonstrating faster convergence and better generalization.
• The method achieved 80.6% accuracy and 0.716 QWK on the Diabetic Retinopathy dataset, surpassing previous methods including multi-task, stick-breaking, and Poisson-based approaches.
• On the US-BIRADS dataset, the method achieved 80.6% accuracy and 0.716 QWK, with the smoothed label providing a more substantial improvement due to higher label noise.
• Hyperparameter sensitivity analysis showed that QWK remained stable under variations in τ and ξ, indicating robustness to parameter tuning.