[論文レビュー] 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 訓練損失を提案する。クラスの相関関係を表現するために、ユークリッド距離の増加関数に基づく基底度量を活用する。滑らかなターゲットラベルと正確な閉形式解を組み合わせることで、CE損失、回帰、および先行研究の順序付き手法を上回る、最先端の性能を達成している。Diabetic Retinopathy および 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 距離と構造化された基底度量を用いて、順序関係を明示的に符号化する微分可能で効率的な損失関数を開発すること。
• ワンホットおよびスムージング済みターゲットラベルの両方に対して、閉形式解を提供することで、高速かつ安定した訓練を可能にすること。
• 離散的で連続するリスクレベルを持つ実世界の医療データセットにおいて、優れた性能を示すこと。

提案手法

• 予測ラベル分布とターゲットラベル分布の間の最適輸送コストを測るため、1番目の Wasserstein 距離を損失関数として用いる。
• クラス間の順序的類似性を符号化するため、直線上のユークリッド距離の増加関数として基底度量を定義する。
• 真のラベルが近隣または遠くのクラスに摂動されるように、一様なモード混合分布を導入してラベルノイズをモデル化する。
• ワンホットおよびスムージング済みターゲットラベル設定の両方において、Wasserstein 損失の正確な閉形式解を導出する。これにより、効率的な最適化が可能になる。
• 最適化挙動を調査するため、ユークリッド距離の凸関数および凹関数を基底度量として用いる。
• 追加の訓練コストをかけずに、一般化性能とロバストネスを向上させるために、スムージング済みターゲットラベル分布を採用する。

実験結果

リサーチクエスチョン

• RQ1 構造化された基底度量を備えた Wasserstein 基盤損失は、標準的な交差エントロピー損失や回帰手法と比較して、順序付き医療診断の性能向上に寄与するか？
• RQ2 一様なモード混合分布によるラベルノイズのモデル化は、モデルのロバストネスと精度にどのような影響を与えるか？
• RQ3 Wasserstein 訓練における正確な閉形式解と近似最適化の比較において、性能にどのような差が生じるか？
• RQ4 基底度量の異なる形（例：線形、凸、凹）は、モデルの性能と収束性にどのような影響を与えるか？
• RQ5 提案手法は Diabetic Retinopathy を超えて、US-BIRADS などの他の順序付き医療データセットにも一般化可能か？

主な発見

• ユークリッド距離の凸関数（D²）を基底度量として用いた本手法は、Diabetic Retinopathy データセットで最高の性能を達成し、80.6% の正解率と 0.716 の QWKを記録した。
• 一様なモードスムージング済みターゲットラベルは、すべてのベースラインで性能を向上させ、特に US-BIRADS のようなノイズの多い環境では顕著な向上を示した。
• 正確な閉形式解は、近似解を常に上回り、収束が速く、一般化性能に優れていることが示された。
• Diabetic Retinopathy データセットでは、80.6% の正解率と 0.716 の QWKを達成し、マルチタスク、スティックブレイキング、Poissonベースの手法をすべて上回った。
• US-BIRADS データセットでも、80.6% の正解率と 0.716 の QWKを達成し、スムージング済みラベルが高ノイズ環境下でより顕著な向上をもたらした。
• ハイパーパramータ感度分析の結果、τ および ξ の変動に対しても QWK が安定しており、パラメータチューニングに対してロバストであることが示された。