Unfolding with a Wasserstein Loss

TL;DR

This paper introduces a Wasserstein loss-based approach for data unfolding, providing a robust alternative to classical methods like Richardson-Lucy, with proven convergence and improved performance in physics applications.

Contribution

It establishes conditions for optimizer uniqueness, develops a convergent Sinkhorn algorithm, and demonstrates improved robustness over traditional deconvolution methods.

Findings

Wasserstein loss yields more robust unfolding than KL-based methods.

The Sinkhorn algorithm scales efficiently with data size.

Numerical experiments show superior performance in physics-inspired problems.

Abstract

Data unfolding -- the removal of noise or artifacts from measurements -- is a fundamental task across the experimental sciences. Of particular interest are applications in physics, where the dominant approach is Richardson-Lucy (RL) deconvolution. The classical RL approach aims to find denoised data that, once passed through the noise model, is as close as possible to the measured data in terms of Kullback-Leibler (KL) divergence. This requires that the support of the measured data overlaps with the output of the noise model, a hypothesis typically enforced by binning, which introduces numerical error. As a counterpoint, the present work studies an alternative formulation using a Wasserstein loss. We establish sharp conditions for existence and uniqueness of optimizers, answering open questions of Li, et al., regarding necessary conditions for uniqueness in the case of transport map noise models. We then develop a provably convergent generalized Sinkhorn algorithm to compute approximate optimizers. Our algorithm requires only empirical observations of the noise model and measured data and scales with the size of the data, rather than the ambient dimension. Numerical experiments on one- and two-dimensional problems inspired by jet mass unfolding in particle physics demonstrate that the optimal transport approach offers robust, accurate performance compared to classical RL deconvolution, particularly when binning artifacts are significant.