Learning to solve inverse problems using Wasserstein loss

TL;DR

This paper introduces the use of Wasserstein loss in inverse problem training to improve reconstruction quality, especially under data miss-alignments, demonstrated through CT imaging experiments.

Contribution

It presents a novel application of Wasserstein loss in learned primal-dual schemes for inverse problems, addressing data misalignments that degrade standard MSE-based training.

Findings

Wasserstein loss corrects for miss-alignments in training data.

Training with Wasserstein loss yields sharper reconstructions than MSE.

Empirical validation on CT data shows improved results with Wasserstein loss.

Abstract

We propose using the Wasserstein loss for training in inverse problems. In particular, we consider a learned primal-dual reconstruction scheme for ill-posed inverse problems using the Wasserstein distance as loss function in the learning. This is motivated by miss-alignments in training data, which when using standard mean squared error loss could severely degrade reconstruction quality. We prove that training with the Wasserstein loss gives a reconstruction operator that correctly compensates for miss-alignments in certain cases, whereas training with the mean squared error gives a smeared reconstruction. Moreover, we demonstrate these effects by training a reconstruction algorithm using both mean squared error and optimal transport loss for a problem in computerized tomography.