[Paper Review] Wasserstein-Fisher-Rao Document Distance

Zihao Wang, Datong Zhou+3 more|arXiv (Cornell University)|Apr 23, 2019

Topic Modeling31 references4 citations

TL;DR

This paper proposes the Wasserstein-Fisher-Rao (WFR) document distance as a robust, interpretable alternative to Word Mover's Distance (WMD) for measuring semantic similarity between documents of varying lengths. By leveraging unbalanced optimal transport theory, WFR automatically re-weights transport plans based on semantic similarity, reducing WMD's overestimation of dissimilarity in length-varying texts. The method achieves significant improvements in KNN classification accuracy across eight datasets, especially on those with high word count variance.

ABSTRACT

As a fundamental problem of natural language processing, it is important to measure the distance between different documents. Among the existing methods, the Word Mover's Distance (WMD) has shown remarkable success in document semantic matching for its clear physical insight as a parameter-free model. However, WMD is essentially based on the classical Wasserstein metric, thus it often fails to robustly represent the semantic similarity between texts of different lengths. In this paper, we apply the newly developed Wasserstein-Fisher-Rao (WFR) metric from unbalanced optimal transport theory to measure the distance between different documents. The proposed WFR document distance maintains the great interpretability and simplicity as WMD. We demonstrate that the WFR document distance has significant advantages when comparing the texts of different lengths. In addition, an accelerated Sinkhorn based algorithm with GPU implementation has been developed for the fast computation of WFR distances. The KNN classification results on eight datasets have shown its clear improvement over WMD.

Motivation & Objective

• To address the overestimation of semantic dissimilarity by Word Mover's Distance (WMD) when comparing documents of different lengths.
• To develop a document distance metric that maintains WMD's interpretability while improving robustness to length variation.
• To apply the newly developed Wasserstein-Fisher-Rao (WFR) metric from unbalanced optimal transport theory to document similarity.
• To evaluate the WFR document distance in both standalone and integrated frameworks, such as Word Mover's Embedding (WME).
• To demonstrate that WFR enables more effective semantic matching in tasks like KNN classification and document retrieval.

Proposed method

• The WFR document distance is defined using the Wasserstein-Fisher-Rao metric, which extends classical Wasserstein distance by allowing unbalanced transport with a penalty for mass imbalance.
• The method models documents as normalized bag-of-words (nBOW) distributions and computes transport costs between word vectors in embedding space.
• Unlike WMD, which enforces mass conservation, WFR introduces a penalty term for mass imbalance, enabling adaptive re-weighting of transport plans based on semantic proximity.
• The WFR transport plan is computed efficiently using WFR Sinkhorn iterations with GPU acceleration, enabling scalable inference.
• The framework is integrated into the Word Mover's Embedding (WME) pipeline, replacing WMD with WFR to improve embedding estimation.
• Hyperparameters such as neighborhood size (k) and regularization (η) are optimized via 5-fold cross-validation on each dataset.

Experimental results

Research questions

• RQ1 Does the WFR document distance outperform WMD in classifying documents of varying lengths?
• RQ2 Can WFR mitigate the overestimation of dissimilarity caused by WMD in length-varying text comparisons?
• RQ3 How does WFR perform relative to supervised baselines like S-WMD, especially when the latter relies on global re-weighting?
• RQ4 Can WFR improve performance in existing frameworks such as Word Mover's Embedding (WME)?
• RQ5 Is the WFR metric efficient and scalable enough for practical use in NLP applications?

Key findings

• WFR document distance achieves lower KNN classification error rates than WMD across all eight datasets, with a significant margin on datasets with high variance in document length (e.g., BBCSPORTS, AMAZON, 20NEWS).
• On BBCSPORTS, WFR reduces the error rate to 2.7±1.0 compared to WMD’s 3.5±0.7, demonstrating robustness to length variation.
• WFR outperforms S-WMD on six out of eight datasets, even though S-WMD uses supervised re-weighting, indicating WFR’s superior, text-specific re-weighting mechanism.
• In the WME framework, WFR with 512 Monte Carlo samples achieves performance comparable to WMD with 4096 samples, significantly reducing computational cost.
• The WFR-based WME framework achieves a test error of 29.8 on the 20NEWS dataset, outperforming WME+WMD (30.7) under the same sampling conditions.
• Numerical experiments confirm that WFR is both effective and efficient, with GPU-accelerated Sinkhorn iterations enabling scalable deployment.