Learning to solve inverse problems using Wasserstein loss

1 Introduction

In inverse problems the goal is to determine model parameters from indirect noisy observations. Example of such problems arise in many different fields in science and engineering, e.g., in X-ray CT natterer2001mathematical , electron tomography oktem2015mathematics , and magnetic resonance imaging brown2014magnetic

. Machine learning has recently also been applied in this area, especially in imaging applications. Using supervised machine-learning to solve inverse problems in imaging requires training data where ground truth images are paired with corresponding noisy indirect observations. The learning provides a mapping that associates observations to corresponding images. However, in several applications there are difficulties in obtaining the ground truth, e.g., in many cases it may have undergone a distortion. For example, a recent study showed that MRI images may be distorted by up to 4 mm due to, e.g., inhomogeneities in the main magnetic field

Walker2014 . If these images are used for training, the learned MRI reconstruction will suffer in quality. Similar geometric inaccuracies arise in several other imaging modalities, such as Cone Beam CT and full waveform inversion in seismic imaging.

This work seeks to provide a scheme for learning a reconstruction scheme for an ill-posed inverse problem with a Wasserstein loss by leveraging upon recent advances in efficient solutions of optimal transport cuturi2013sinkhorn ; karlsson2016generalized and learned iterative schemes for inverse problems adler2017learned . The proposed method is demonstrated on a computed tomography example, where we show a significant improvement compared to training the same network using mean squared error loss. In particular, using the Wasserstein loss instead of standard mean squared error gives a result that is more robust against potential miss-alignment in training data.

2 Background

2.1 Inverse problems

In inverse problems the goal is to reconstruct an estimate of the signal

f_{true} \in X

from noisy indirect measurements (data)

g \in Y

assuming

g = T (f_{true}) + δ g .

(1)

In the above $X$ and $Y$ are referred to as the reconstruction and data space, respectively. Both are typically Hilbert or Banach spaces. Moreover $T : X \to Y$ denotes the forward operator, which models how a given signal gives rise to data in absence of noise. Finally, $δ g \in Y$ is the noise component of data. Many inverse problems of interest are ill-posed, meaning that there is no uniques solution to (1) and hence there is no inverse to $T$ . Typically reconstructions of $f_{true}$ are sensitive to the data and small errors gets amplified. One way to mitigate these effects is to use regularization engl2000regularization .

Variational regularization

In variational regularization one formulates the reconstruction problem as an optimization problem. To this end, one introduces a data discrepancy functional $f \mapsto L (T (f), g)$ , where $L : Y \times Y \to R$ , that quantifies the miss-fit in data space, and a regularization functional $S : X \to R$ that encodes a priori information about $f_{true}$ by penalizing undesirable solutions. For a given $g \in Y$ , this gives an optimization problem of the form

min f \in X L (T (f), g) + λ S (f) .

(2)

Here, $λ$ acts as a trade-off parameter between the data discrepancy and regularization functional. In many cases $L$ is taken to be the negative data log-likelihood, e.g., $L (T (f), g) = ∥ T (f) - g ∥_{2}^{2}$ in the case of additive white Gaussian noise. Moreover, a typical choice for regularization functional is total variation (TV) regularization, $S (f) = ∥ \nabla f ∥_{1}$ rudin1992nonlinear . These regularizers typically give rise to large scale and non-differentiable optimization problems, which requires advanced optimization algorithms.

Learning for inverse problems

In many applications, and so also in inverse problems, data driven approaches have shown dramatic improvements over the state-of-the-art LeBeHi15

. Using supervised learning to solve an inverse problem amounts to finding a parametrized operator

T_{Θ}^{†} : Y \to X

where the parameters

Θ

are selected so that

g = T (f_{true}) + δ g ⟹ T_{Θ}^{†} (g) \approx f_{% t r u e} .

For inverse problems in image processing, such as denoising and deblurring, we have $Y = X$

and it is possible to apply a wide range of widely studied machine learning techniques, such as fully convolutional deep neural networks with various architectures, including fully convolutional networks

NIPS2008_3506 and denoising auto-encoders NIPS2012_4686 .

However, in more complicated inverse problems as in tomography, the data and reconstruction spaces are very different, e.g., their dimension after discretization may differ. For this reason, learning a mapping from $Y$ to $X$ becomes nontrivial, and classical architectures that map, e.g., images to images using convolutional networks cannot be applied as-is. One solution is to use fully-connected layers as in PaGiKaLoMa04 for very small scale tomographic reconstruction problems. A major disadvantage with such a fully learned approach is that the parameters space has to be very high dimensional in order to be able to learn both the prior and the data model, which often renders it infeasible due to training time and lack of training data.

A more successful approach is to first apply some crude reconstruction operator $T^{†} : Y \to X$ and then use machine learning to post process the result. This separates the learning from the complications of mapping between spaces since the operator $T^{†}$ can be applied off-line, prior to training. Such an approach has been demonstrated for tomographic reconstruction in PeBa13 ; WuGhChMa16 . Its drawback for ill posed inverse problems is that information is typically lost by using $T^{†}$ , and this information cannot be recovered by post processing.

Finally, another approach is to incorporate the forward operator $T$ and its adjoint $T^{*}$ into the neural network. In these learned iterative schemes, classical neural networks are interlaced with applications of the forward and backward operator, thus allowing for the learned reconstruction operator to work directly from data without having to learn the data model. For example, in YaSuLiXu16 an ADMM-like scheme for Fourier inversion is learned and PuWe17 consider solving inverse problems typically arising in image restoration by a learned gradient-descent scheme. In adler2017solving this later approach is shown to be applicable to large scale tomographic inversion. Finally, in adler2017learned they apply learning in both spaces $X$ and $Y$ , yielding a Learned Primal-Dual scheme, and show that it outperforms learned post-processing for reconstruction of medical CT images.

Loss functions for learning

Once the $Θ$ parametrization of $T_{Θ}^{†}$ is set, the parameters are typically chosen by minimization of some loss functional $L$ . Without doubt, the most common loss function is the mean squared error, also called $L_{2}$ loss, given by

L (Θ) = E_{f, g} [∥ T_{Θ}^{†} (g) - f ∥_{2}^{2}] .

(3)

It has however been noted that it is sub-optimal for imaging, and a range of other loss functions have been investigated. These include the classical $ℓ_{p}$ norms and the structural similarity index (SSIM) Zhao2017 , as well as more complex losses such as perceptual losses Johnson2016 and adversarial networks DBLP:journals/corr/MardaniGCVZATHD17 .

Recently, optimal mass transport has also been considered as loss function for classification frogner2015learning and generative models arjovsky2017wasserstein . In this work we consider using optimal transport for training a reconstruction scheme for ill-posed inverse problems.

2.2 Optimal mass transport and Sinkhorn iterations

In optimal mass transport the aim is to transform one distribution into another by moving the mass in a way that minimizes the cost of the movement. For an introduction and overview of the topic, see, e.g., the monograph villani2008optimal . Lately, the area has attracted a lot of research cuturi2013sinkhorn ; cuturi2015asmoothed ; chizat2015unbalanced with applications to, e.g., signal processing haker2004optimal ; georgiou2009metrics ; jiang2012geometric ; engquist2014application and inverse problems benamou2015iterative ; karlsson2016generalized .

The optimal mass transport problem can be formulated as follows: let $Ω \subset R^{d}$ be a compact set, and let $μ_{0}$ and $μ_{1}$ be two measures, defined on $Ω$ , with the same total mass. Given a cost $c : Ω \times Ω \to R_{+}$ that describes the cost for transporting a unit mass from one point to another, find a (mass preserving) transference plan $M$ that is as cheap as possible. Here, the transference plan characterizes how to move the mass of $μ_{0}$ in order to deform it into $μ_{1}$ . Letting the transference plan be a nonnegative measure $d M$ on the space $Ω \times Ω$

yields a linear programming problem in the space of measures:

$T (μ_{0}, μ_{1}) = min d M \geq 0$	$\int_{(x_{0}, x_{1}) \in Ω \times Ω} c (x_{0}, x_{1}) d M (x_{0}, x_{1})$	(4)
subject to	$μ_{0} (x_{0}) d x_{0} = \int_{x_{1} \in Ω} d M (x_{0}, x_{1}),$
	$μ_{1} (x_{1}) d x_{1} = \int_{x_{0} \in Ω} d M (x_{0}, x_{1}) .$

Although this formulation is only defined for measures $μ_{0}$ and $μ_{1}$ with the same total mass, it can also be extended to handle measures with unbalanced masses georgiou2009metrics ; chizat2015unbalanced . Moreover, under suitable conditions one can define the Wasserstein metrics $W_{p}$ using $T$ , by taking $c (x_{0}, x_{1}) = d (x_{0}, x_{1})^{p}$ for $p \geq 1$ and where $d$ is a metric on $Ω$ , and $W_{p} (μ_{0}, μ_{1}) := T (μ_{0}, μ_{1})^{1 / p}$ (villani2008optimal, , Definition 6.1). As the name indicates, $W_{p}$ is a metric on the set of nonnegative measures on $Ω$ with fixed mass (villani2008optimal, , Theorem 6.9), and $T$ is weak $^{*}$ continuous on this set. One important property is that $T$ (and thus also $W_{p}$ ) does not only compare objects point by point, as standard $L^{p}$ metrics, but instead quantifies how the mass is moved. This makes optimal transport natural for quantifying uncertainty and modelling deformations jiang2012geometric ; karlsson2013uncertainty .

One way to solve the optimal transport problem in applications is to discretize $Ω$

and solve the corresponding finite-dimensional linear programming problem. In this setting the two measures are represented by point masses on the discretization grid, i.e., by two vectors

μ_{0}, μ_{1} \in R_{+}^{n}

where the element

[μ_{k}]_{i}

corresponds to the mass in the point

x_{(i)} \in Ω

for

i = 1, \dots, n

and

k = 0, 1

. Moreover, a transference plan is represented by a matrix

M \in R_{+}^{n \times n}

where the value

m_{i j} := [M]_{i j}

denotes the amount of mass transported from point

x_{(i)}

x_{(j)}

. The associated cost of a transference plan is

\sum_{i, j = 1}^{n} c_{i j} m_{i j} = t r a c e (C^{T} M)

, where

[C]_{i j} = c_{i j} = c (x_{(i)}, x_{(j)})

is the transportation cost from

x_{(i)}

x_{(j)}

, and by discretizing the constraints we get that

M

is a feasible transference plan from

μ_{0}

μ_{1}

if the row sums of

M

μ_{0}

and the column sums of

M

μ_{1}

. The discrete version of (4) thus takes the form

	$T (μ_{0}, μ_{1}) = min M \geq 0$	$t r a c e (C^{T} M)$		(5)
	subject to	$μ_{0} = M 1_{n} μ_{1} = M^{T} 1_{n},$

where $M \geq 0$ denotes element-wise non-negativity of the matrix. However, even though (5) is a linear programming problem it is in many cases computationally infeasible due to the vast number of variables. Since $M \in R_{+}^{n \times n}$ the number of variables is $n^{2}$ , and thus if one seek to solve the optimal transport problem between two $512 \times 512$ images this results in more than $6 \cdot 10^{10}$ variables.

One approach for addressing this problem was proposed by Cuturi cuturi2013sinkhorn that introduces an entropic regularizing term $D (M) = \sum_{i, j = 1}^{n} (m_{i j} log (m_{i j}) - m_{i j} + 1)$ for approximating the transference plan, so the resulting perturbed optimal transport problem reads as

	$min M \geq 0$	$t r a c e (C^{T} M) + ε D (M)$		(6)
	subject to	$μ_{0} = M 1_{n} μ_{1} = M^{T} 1_{n} .$

One can show that an optimal solution to (6) is of the form

M = d i a g (u) K d i a g (v),

(7)

where $K = exp (- C / ε)$ (point-wise exponential) is known, and $u, v \in R_{+}^{n}$ are unknown. This shows that the solution is parameterized by only $2 n$ variables. Moreover, the two vectors can be computed iteratively by so called Sinkhorn iterations, i.e., alternatingly compute $u$ and $v$ that matches $μ_{0}$ and $μ_{1}$ respectively. This is summarizied in Algorithm 1 where $⊙$ denotes elementwise multiplication and $. /$ elementwise division. The procedure has been shown to have a linear convergence rate, see cuturi2013sinkhorn and references therein.

Moreover, when the underlying cost $c (x_{0}, x_{1})$ is translation invariant the discretized cost matrix $C$ , and thus also the transformation $K$ , gets a Toeplitz-block-Toeplitz structure. This structure can be used in order to compute $K v$ and $K^{T} u$

efficiently using the fast Fourier transform in

O (n log n)

, instead of naive matrix-vector multiplication in

O (n^{2})

karlsson2016generalized . This is crucial for applications in imaging since for images of size

512 \times 512

pixels one would have to explicitly store and multiply with matrices of size

262144 \times 262144

1:Input

C, ε, μ_{0}, μ_{1}

2:initialize

v_{0} > 0

and

K = exp (- C / ε)

3:for

i = 1, \dots, N

u_{i} \leftarrow μ_{0} . / (K v_{i - 1})

v_{i} \leftarrow μ_{1} . / (K^{T} u_{i})

6:return

u_{N}^{T} (K ⊙ C) v_{N}

Algorithm 1 Sinkhorn iterations for computing entropy-regularized optimal transport cuturi2013sinkhorn

3 Learning a reconstruction operator using Wasserstein loss

In this work we propose to use entropy regularized optimal transport (6) to train a reconstruction operator, i.e., to select the parameters as

Θ^{*} \in arg min Θ E_{f, g} [T (T_{Θ}^{†} (g), f)] .

(8)

This should give better results when data $g$ is not aligned with the ground truth $f$ . To see this, consider the case when $f$ is a point mass. In that case training the network with the $L_{2}$ loss (3) will (in the ideal case) result in a perfect reconstruction composed with a convolution that “smears” the reconstruction over the area of possible miss-alignment. On the other hand since optimal mass transport does not only compare objects point-wise, the network will (in the ideal case) learn a perfect reconstruction combined with a movement of the object to the corresponding barycenter (centroid) of the miss-alignment. These statements are made more precise in the following propositions. Formal definitions and proofs are deferred to the appendix.

Proposition 1.

Let $g \in L_{2} (R^{n})$ , let $τ$ be a $R^{n}$

-valued random variable with probability measure

d P (t)

, and let

g_{τ} (x) := g (x - τ)

. Then there exists a function

f \in L_{2} (R^{n})

that minimizes

E_{τ} [∥ f - g_{τ} ∥_{2}^{2}]

, and this

f

has the form

f (x) = (d P * g) (x) := \int_{R^{n}} g (x - t) d P (t) .

Proposition 2.

Let $δ (x)$ be the Dirac delta function on $R^{n}$ , let $τ$ be a $R^{n}$ -valued random variable with probability measure $d P (t)$ , and let $δ_{τ} (x) := δ (x - τ)$ . Then

E_{τ} [T (δ_{τ}, μ)] = \int_{R^{n}} (\int_{R^{n}} c (t, x) d P (t)      := F (x)) d μ (x) for μ \geq 0 and \int_{R^{n}} d μ (x) = 1

and $E_{τ} [T (δ_{τ}, μ)] = \infty$ otherwise. Furthermore, finding a $μ$ that minimizes $E_{τ} [T (δ_{τ}, μ)]$ is equivalent to finding the global minimizers to $F (x)$ . In particular, if (i) the probability measure $d P$ is symmetric around its mean, (ii) the underlying cost $c$ is of the form $c (t, x) = d (x - t)$ , where $d$ is convex and symmetric, and (iii) $d P$ and $d$ are such that

\int_{R^{n}} d (x - t) d P (t) and \int_{R^{n}} \partial d (x - t) d P (t)

both exist and are finite for all $x \in R^{n}$ , then $μ (x) = δ (x - E [τ])$ is an optimal solution. Furthermore, if $d$ is also strictly convex, then this is the unique minimizer.

To illustrate Propositions 1 and 2 we consider the following example.

Example 3.

Let $τ$

be uniformly distributed on

[- 1, 1]

, and let

c (x_{0}, x_{1}) = (x_{0} - x_{1})^{2}

. This gives

F (x) = \frac{1}{2} \int_{- 1}^{1} (x - t)^{2} d t = \frac{1}{3} + x^{2},

which has minimum $x = 0$ , and hence the (unique) minimizer to $E_{τ} [T (δ_{τ}, μ)]$ is $μ (x) = δ (x)$ . For the $L_{2}$ case with the uniform distribution, the minimizer of $E_{τ} [∥ f - g_{τ} ∥_{2}^{2}]$ is the smoothed function $g * \frac{1}{2} χ_{[- 1, 1]}$ .

The most common choice of distance $c$ is to use the squared norm $c (x_{0}, x_{1}) = ∥ x_{0} - x_{1} ∥^{2}$ , as in the previous example. In this case the result of Proposition 2 can be strengthened, as shown in the following example.

Example 4.

Let $τ$ be a $R^{n}$ -valued random variable with probability measure $d P (t)$

with finite first and second moments, and let

c (x_{0}, x_{1}) = ∥ x_{0} - x_{1} ∥^{2}

. This gives

F (x) = \int_{R^{n}} (x - t)^{2} d P (t) = x^{2} - 2 x E [τ] + E [τ^{2}],

which has a unique global minimum in $x = E [τ]$ and hence $μ (x) = δ (x - E [τ])$ .

4 Implementation and evaluation

We use the recently proposed learned primal-dual structure in adler2017learned for learning a reconstruction operator $T_{Θ}^{†}$ for solving the inverse problem in (1). In this algorithm, a sequence of small blocks work alternatingly in the data (dual) space $Y$ and the reconstruction (primal) space $X$ and are connected using the forward operator $T$ and its adjoint $T^{*}$ . The algorithm works with any differentiable operator $T$ , but we state the version for linear operators for simplicity in algorithm 2.

1:Initialize

f_{0} \in X^{N_{primal}}, h_{0} \in U^{N_{dual}}

2:for

i = 1, \dots, I

h_{i} \leftarrow Γ_{Θ_{i}^{d}} (h_{i - 1}, T (f_{i - 1}^{(2)}), g)

f_{i} \leftarrow Λ_{Θ_{i}^{p}} (f_{i - 1}, T^{*} (h_{i}^{(1)}))

T_{Θ}^{†} (g) : = f_{I}^{(1)}

Algorithm 2 Learned Primal-Dual reconstruction algorithm

Figure 1: Network architecture used to solve the inverse problem. Dual and primal iterates are in blue and red boxes, respectively. Several arrows pointing to the same box indicates concatenation. The initial values

f_{0}, h_{0}

enter from the left, while the data

g

is supplied to the dual iterates.

The method was implemented using ODL adler2017ODL , ASTRA vanaerle2016fast

, and TensorFlow

abadi2016tensorflow . We used the reference implementation¹¹1https://github.com/adler-j/learned_primal_dual with default parameters, i.e., the number of blocks in the primal and dual space was

I = 10

, and the number of primal and dual variables was set to

N_{primal} = N_{dual} = 5

. Moreover, the blocks used a residual structure and had three layers of

3 \times 3

convolutions with

32

filters. PReLU nonlinearities were used. Thus, this corresponds to a residual CNN with convolutional depth of

10 \cdot 2 \cdot 3 = 60

, as shown in graphical format in fig. 1. We used zero initial values,

f_{0} = h_{0} = 0

We compare a learned reconstruction operator of this form when trained using $L_{2}$ loss (3) and using optimal transport loss (8). Moreover, the evaluation is done on a problem similar to the evaluation problem in adler2017learned ; adler2017solving , i.e., on a problem in computed tomography. More specifically, training is done on data that consists of randomly generated circles on a domain of 512 ×512, and the forward operator $T$ is the ray transform natterer2001mathematical . What makes this an ill-posed problem is that the data acquisition is done from only 30 views with 727 parallel lines. Moreover, the source of noise is two-fold in this set-up: (i) the pairs $(g_{i}, f_{i})$ of data sets and phantoms are not aligned, meaning that the data is computed from a phantom with a random change in position. This random change is independent for the different circles, and for each circle it is a shift which is uniformly distributed over $[- 40, 40]$ pixels, both in up-down and left-right direction. (ii) on the data computed from the shifted phantom, 5% additive Gaussian noise was added. For an example, see fig. 2.

The optimal mass transport distance computed with Sinkhorn iterations was used as loss function, where we used the transport cost

c (x_{1}, x_{2}) = (1 - e^{- ∥ x_{1} - x_{2} ∥^{4} / 80^{4}}) .

This was chosen since it heavily penalizes large movements, while not diverging to infinity which causes numerical instabilities. Moreover, $c (x_{1}, x_{2})^{1 / 4}$ is in fact a metric on $R^{2}$ (see lemma 6 in the appendix) and thus $W_{4} (μ_{0}, μ_{1}) := T (μ_{0}, μ_{1})^{1 / 4}$ gives rise to a Wasserstein metric on the space of images, where $T (μ_{0}, μ_{1})$ is the optimal mass transport distance with the transport cost $c (x_{1}, x_{2})$ . Since this cost is translation invariant, the matrix-vector multiplications $K u$ and $K^{T} v$ can be done with fast Fourier transform, as mentioned in section 2.2, and this was implemented in Tensorflow. We used 10 Sinkhorn iterations with entropy regularization $ε = 10^{- 3}$ , to approximate the optimal mass transport. Automatic differentiation was used to back-propagate the result during training.

Since the optimal mass transport function (6) is only finite for marginals $μ_{0}$ and $μ_{1}$ with the same total mass, in the training we normalize the output of the reconstruction $T_{Θ}^{†} (g)$ with $m a s s (f) / m a s s (T_{Θ}^{†} (g))$ . This makes $T_{Θ}^{†} (g)$ invariant with respect to the total mass, which is undesirable. To compensate for this, a small penalization on the error in total mass was added to the loss function.

The training also followed adler2017learned closely. In particular, we used $2 \cdot 10^{4}$ batches of size $1$ , using the ADAM optimizer kingma2014adam with default values except for $β_{2} = 0.99$ . The learning rate (step length) used was cosine annealing loshchilov2016sgdr with initial step length $10^{- 3}$ . Moreover, in order to improve training stability we performed gradient norm clipping pascanu2012understanding with norms limited to 1. The convolution parameters were initialized using Xavier initialization glorot2010understanding , and all biases were initialized to zero. The training took approximately 3 hours using a single Titan X GPU. The source code used to replicate these experiments are available online ²²2https://github.com/adler-j/wasserstein_inverse_problems

Results are presented in fig. 3. As can be seen, the reconstruction using $L_{2}$ loss “smears” the reconstruction to an extent where the shape is impossible to recover. On the other hand, the reconstruction using the Wasserstein loss retains the over-all global shape of the object, although relative and exact positions of the circles are not recovered.

5 Conclusions and future work

In this work we have considered using Wasserstein loss to train a neural network for solving ill-posed inverse problems in imaging where data is not aligned with the ground truth. We give a theoretical motivation for why this should give better results compared to standard mean squared error loss, and demonstrate it on a problem in computed tomography. In the future, we hope that this method can be applied to other inverse problems and to other problems in imaging such as segmentation.

Appendix: Deferred definition and proofs

Proof of Proposition 1.

To show that $f (x) = (d P * g) (x) \in L_{2} (R^{n})$ minimizes $E_{τ} [∥ f - g_{τ} ∥_{2}^{2}]$ we expand the expression and use Fubini’s theorem to get

E_{τ} [∥ f - g_{τ} ∥_{2}^{2}] = \int_{R^{n}} (\int_{R^{n}} | f (x) - g_{t} (x) |^{2} d x) d P (t) = \int_{R^{n}} (\int_{R^{n}} (f (x) - g_{t} (x))^{2} d P (t)) d x .

Rearranging terms and using that $\int_{R^{n}} d P (t) d t = 1$ , this can be written as

E_{τ} [∥ f - g_{τ} ∥_{2}^{2}] = \int_{R^{n}} {(f (x) - \int_{R^{n}} g_{t} (x) d P (t))}_{R^{n}}^{2} d x + c,

where $c$ is a constant. Using this it follows that the minimizing $f$ is of the form

f (x) = \int_{R^{n}} g_{t} (x) d P (t) .

To see that $f \in L_{2} (R^{n})$ we note that, by using Fubini’s theorem, we have

	$∥ f ∥_{2}^{2}$	$= \int_{R^{n}} {(\int_{R^{n}} g_{t} (x) d P (t))}_{R^{n}}^{2} d x = \int_{R^{n}} \int_{R^{n}} (\int_{R^{n}} g_{s} (x) g_{t} (x) d x) d P (s) d P (t)$
		$\leq \int_{R^{n}} \int_{R^{n}} (\frac{1}{2} \int_{R^{n}} g_{s} (x)^{2} + g_{t} (x)^{2} d x) d P (s) d P (t) = ∥ g ∥_{2}^{2} < \infty$

where the first inequality is the arithmetic-geometric mean inequality. This completes the proof. ∎

Definition 5.

Let $h : R^{n} \to R$ . A subgradient to $h$ in a point $y$ is a vector $v$ so that

h (x) \geq h (y) + ⟨ v, x - y ⟩, \forall x \in R^{n} .

The set of all subgradients in a point $y$ is called the subdifferential of $h$ at $y$ , and is denoted by $\partial h (y)$ . This is a set-valued operator, and for any measure $d ν$ on $R^{n}$ we define $(d ν * \partial h) (y) := \int_{R^{n}} \partial h (t) d ν (y - t)$ to be the set-valued operator

y \mapsto {\int_{R^{n}} v (t) d ν (y - t) \in R^{n} ∣ v (t) \in \partial h (t)} .

Proof of Proposition 2.

We consider finding the marginal $μ$ that minimize $E_{τ} [T (δ_{τ}, μ)]$ . Without loss of generality we assume that $τ$ is zero-mean, since otherwise we simply consider $τ - E [τ]$ which is a zero-mean random variable. First we note that $T (δ_{t}, μ)$ is only finite for nonegative measures $μ$ with total mass $1$ , and hence $E_{τ} [T (δ_{τ}, μ)]$ is only finite for such measures. Second, for such a $μ$ we have

T (δ_{t}, μ) = \int_{R^{n}} c (t, x) d μ (x),

since one needs to transport all mass in $μ$ into the point $t$ where $δ_{t}$ has its mass. Using this and expanding the expression for the expectation gives that

E_{τ} [T (δ_{τ}, μ)] = \int_{R^{n}} T (δ_{t}, μ) d P (t) = \int_{R^{n}} (\int_{R^{n}} c (t, x) d P (t)) d μ (x),

where we have used Fubini’s theorem in the last step. This completes the first half of the statement.

To prove the second half of the statement, note that the optimal $μ$ have support only in the global minimas of the function

F (x) := \int_{R^{n}} d (x - t) d P (t) = \int_{R^{n}} d (t) d P (x - t),

which by assumption exists and is finite. Now, since $d$ is convex we have that

d (x) \geq d (y) + ⟨ \partial d (y), x - y ⟩, for all x, y \in R^{n},

(9)

and convolving this inequality with $d P$ gives the inequality

F (x) \geq F (y) + ⟨ \int_{R^{n}} \partial d (y - t) d P (t), x - y ⟩, for all x, y \in R^{n},

(10)

where all terms exist and are bounded by assumption. This shows that

\int_{R^{n}} \partial d (y - t) d P (t) \subset \partial F (y) .

Now, since $d$ is symmetric we have that $\partial d$ is anti-symmetric, i.e., that $\partial d (x) = - \partial d (- x)$ , since

d (x) = d (- x) \geq d (- y) + ⟨ \partial d (- y), - x + y ⟩ = d (y) + ⟨ - \partial d (- y), x - y ⟩ .

Therefore

\partial F (0) \supset \int_{R^{n}} \partial d (- t) d P (t) ∋ 0,

where the last inclusion follows since $d P$ is symmetric and $\partial d$ is anti-symmetric. Now, since $0 \in \partial F (0)$ we have that $x = 0$ is a global minimizer to $F (x)$ (bauschke2011convex, , Theorem 16.2), and thus one optimal solution to the problem is $μ (x) = δ (x)$ . Now, if $d$ is strictly convex, the inequality (9) is strict for $x \neq y$ , and thus (10) is also strict, which shows that the optimal solution is unique. ∎

Lemma 6.

Let $∥ \cdot ∥$ be a norm on $R^{m}$ . Then

d (x_{1}, x_{2}) = (1 - e^{- ∥ x_{1} - x_{2} ∥^{n}})^{\frac{1}{n}} .

is a metric on $R^{m}$ for $n \geq 1$ .

Proof.

It is easily seen that $d (x_{1}, x_{2})$ is symmetric, nonnegative, and equal to zero if only if $x_{1} = x_{2}$ . Thus we only need to verify that the triangle inequality holds. To this end we note that if

(1 - e^{- (a + b)^{n}})^{\frac{1}{n}} \leq (1 - e^{- a^{n}})^{\frac{1}{n}} + (1 - e^{- b^{n}})^{\frac{1}{n}}, for all a, b \geq 0,

(11)

for all $n \geq 1$ , then by taking $a = ∥ x_{1} - x_{2} ∥$ , $b = ∥ x_{2} - x_{3} ∥$ , and using the triangle inequality for the norm $∥ \cdot ∥$ we have that

	$d (x_{1}, x_{3})$	$= (1 - e^{- ∥ x_{1} - x_{3} ∥^{n}})^{\frac{1}{n}} \leq (1 - e^{- (∥ x_{1} - x_{2} ∥ + ∥ x_{2} - x_{3} ∥)^{n}})^{\frac{1}{n}}$
		$\leq (1 - e^{- ∥ x_{1} - x_{2} ∥^{n}})^{\frac{1}{n}} + (1 - e^{- ∥ x_{2} - x_{3} ∥^{n}})^{\frac{1}{n}} = d (x_{1}, x_{2}) + d (x_{2}, x_{3}) .$

Therefore we will show that (11) holds for all $n \geq 1$ , and to do so we will

show that if a function $g : R_{+} \to R_{+}$ fulfills $g (0) = 0$ , $g (x)^{'} \geq 0$ , $g^{''} (x) \leq 0$ for all $x \in R_{+}$ , then $g (x_{1} + x_{2}) \leq g (x_{1}) + g (x_{2})$ ,
show that for $x \geq 0$ the map $x \mapsto (1 - e^{- x^{n}})^{\frac{1}{n}}$ fulfills the assumptions in (i) for any $n \geq 1$ .

To show (i) we note that

	$g (x_{1} + x_{2})$	$= \int_{0}^{x_{1} + x_{2}} g^{'} (t) d t = \int_{0}^{x_{1}} g^{'} (t) d t + \int_{x_{1}}^{x_{1} + x_{2}} g^{'} (t) d t$
		$\leq \int_{0}^{x_{1}} g^{'} (t) d t + \int_{0}^{x_{2}} g^{'} (t) d t = g (x_{1}) + g (x_{2}),$

where the inequality uses that $g^{'} (t) \geq g^{'} (x + t)$ for any $x, t \geq 0$ since $g^{''} (x) \leq 0$ for all $x \geq 0$ .

To show (ii), let $g (x) := (1 - e^{- x^{n}})^{\frac{1}{n}}$ and observe that $g (0) = 0$ . Differentiating $g$ twice gives

	$g^{'} (x) = \frac{e^{- x^{n}} (1 - e^{- x^{n}})^{\frac{1}{n}} x^{n - 1}}{1 - e^{- x^{n}}}$
	$g^{''} (x) = - \frac{(1 - e^{- x^{n}})^{\frac{1}{n}} x^{n - 2} (=: h_{n} (x^{n})      n e^{x^{n}} x^{n} - x^{n} + e^{x^{n}} - n e^{x^{n}} + n - 1)}{(e^{x^{n}} - 1)^{2}} .$

For $x \geq 0$ we see that $g^{'} (x) \geq 0$ for all $n \geq 1$ . Moreover, for $x \geq 0$ we see that $g^{''} (x) \leq 0$ for all $x \geq 0$ and for all $n \geq 1$ if and only if $h_{n} (x^{n}) \geq 0$ . With the change of variable $x^{n} = y$ , we thus want to show that $h_{n} (y) \geq 0$ for all $y \geq 0$ and all $n \geq 1$ . To see this we note that $h_{n} (0) = 0$ and that

h_{n}^{'} (y) = n e^{y} y + e^{y} - 1 \geq 0 for all y \geq 0 and n \geq 1.

This shows (ii), and hence completes the proof. ∎

References

[1] M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. Corrado, A. Davis, J. Dean, M. Devin, et al. Tensorflow: Large-scale machine learning on heterogeneous distributed systems. arXiv preprint arXiv:1603.04467, 2016.
[2] J. Adler, H. Kohr, and O. Öktem. ODL - a Python framework for rapid prototyping in inverse problems. 2017. In preparation, KTH Royal Institute of Technology. Code and documentation available online: https://github.com/odlgroup/odl.
[3] J. Adler and O. Öktem. Learned primal-dual reconstruction. arXiv preprint arXiv:1707.06474, 2017.
[4] J. Adler and O. Öktem. Solving ill-posed inverse problems using iterative deep neural networks. Inverse Problems, 2017.
[5] M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein GAN. arXiv preprint arXiv:1701.07875, 2017.
[6] H.H. Bauschke and P.L. Combettes. Convex analysis and monotone operator theory in Hilbert spaces. Springer, New York, 2011.
[7] J.-D. Benamou, G. Carlier, M. Cuturi, L. Nenna, and G. Peyré. Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111–A1138, 2015.
[8] R.W. Brown, Y.-C. N. Cheng, E.M. Haacke, M.R. Thompson, and R. Venkatesan. Magnetic Resonance Imaging: Physical Principles and Sequence Design. John Wiley & Sons Ltd, 2014.
[9] L. Chizat, G. Peyré, B. Schmitzer, and F.-X. Vialard. Unbalanced optimal transport: geometry and Kantorovich formulation. arXiv preprint arXiv:1508.05216, 2015.
[10] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, pages 2292–2300, 2013.
[11] M. Cuturi and G. Peyré. A smoothed dual approach for variational Wasserstein problems. SIAM Journal on Imaging Sciences, pages 320–343, 2016.
[12] H.W. Engl, M. Hanke, and A. Neubauer. Regularization of inverse problems. Kluwer Academic Publisher, 2000.
[13] B. Engquist and B.D. Froese. Application of the Wasserstein metric to seismic signals. Communications in Mathematical Sciences, 12(5), 2014.
[14] C. Frogner, C. Zhang, H. Mobahi, M. Araya, and T.A. Poggio. Learning with a Wasserstein loss. In Advances in Neural Information Processing Systems, pages 2053–2061, 2015.
[15] T.T. Georgiou, J. Karlsson, and M.S. Takyar. Metrics for power spectra: an axiomatic approach. IEEE Transactions on Signal Processing, 57(3):859–867, 2009.
[16] X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. In
Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics
, pages 249–256, 2010.
[17] S. Haker, L. Zhu, A. Tannenbaum, and S. Angenent. Optimal mass transport for registration and warping.
International Journal of computer vision
, 60(3):225–240, 2004.
[18] V. Jain and S. Seung. Natural image denoising with convolutional networks. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems 21, pages 769–776. Curran Associates, Inc., 2009.
[19] X. Jiang, Z.-Q. Luo, and T.T. Georgiou. Geometric methods for spectral analysis. IEEE Transactions on Signal Processing, 60(3):1064–1074, 2012.
[20] J. Johnson, A. Alahi, and L. Fei-Fei.
Perceptual Losses for Real-Time Style Transfer and Super-Resolution
, pages 694–711. Springer International Publishing, Cham, 2016.
[21] J. Karlsson and T.T. Georgiou. Uncertainty bounds for spectral estimation. IEEE Transactions on Automatic Control, 58(7):1659–1673, 2013.
[22] J. Karlsson and A. Ringh. Generalized Sinkhorn iterations for regularizing inverse problems using optimal mass transport. arXiv preprint arXiv:1612.02273, 2016.
[23] D. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[24] Y. LeCun, Y. Bengio, and G. Hinton. Deep learning. Nature, 521:436—444, 2015.
[25] I. Loshchilov and F. Hutter. SGDR: stochastic gradient descent with restarts. arXiv preprint arXiv:1608.03983, 2016.
[26] M. Mardani, E. Gong, J.Y. Cheng, S. Vasanawala, G. Zaharchuk, M.T. Alley, N. Thakur, S. Han, W.J. Dally, J.M. Pauly, and L. Xing. Deep generative adversarial networks for compressed sensing automates MRI. CoRR, abs/1706.00051, 2017.
[27] F. Natterer and F. Wübbeling. Mathematical methods in image reconstruction. SIAM, 2001.
[28] O. Öktem. Mathematics of electron tomography. In O. Scherzer, editor, Handbook of Mathematical Methods in Imaging, pages 937–1031. Springer, New York, NY, 2015.
[29] R. Pascanu, T. Mikolov, and Y. Bengio. Understanding the exploding gradient problem. CoRR, abs/1211.5063, 2012.
[30] P. Paschalis, N. D. Giokaris, A. Karabarbounis, G. K. Loudos, D. Maintas, C. N. Papanicolas, V. Spanoudaki, Ch. Tsoumpas, and E. Stiliaris. Tomographic image reconstruction using artificial neural networks. Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment, 527(1–2):211–215, 2004.
[31] D. M. Pelt and K. J. Batenburg. Fast tomographic reconstruction from limited data using artificial neural networks. IEEE Transactions on Image Processing, 22(12):5238–5251, 2013.
[32] P. Putzky and M. Welling. Recurrent inference machines for solving inverse problems. Submitted to ICLR 2017, Toulon, France, April 24–26, 2017. Report available from https://openreview.net/pdf?id=HkSOlP9lg, 2017.
[33] L.I. Rudin, S. Osher, and E. Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: Nonlinear Phenomena, 60(1):259–268, 1992.
[34] W. van Aarle, W.J. Palenstijn, J. Cant, E. Janssens, F. Bleichrodt, A. Dabravolski, J. De Beenhouwer, K.J. Batenburg, and J. Sijbers. Fast and flexible X-ray tomography using the ASTRA toolbox. Optics express, 24(22):25129–25147, 2016.
[35] C. Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
[36] A. Walker, G. Liney, P. Metcalfe, and L. Holloway. MRI distortion: considerations for MRI based radiotherapy treatment planning. Australasian Physical & Engineering Sciences in Medicine, 37(1):103–113, Mar 2014.
[37] T. Würfl, F. C. Ghesu, V. Christlein, and A. Maier. Deep learning computed tomography. In S. Ourselin, L. Joskowicz, M. Sabuncu, G. Unal, and W. Wells, editors, MICCAI 2016: Medical Image Computing and Computer-Assisted Intervention – MICCAI 2016, volume 9902 of Lecture Notes in Computer Science, pages 432–440. Springer-Verlag, 2016.
[38] J. Xie, L. Xu, and E. Chen. Image denoising and inpainting with deep neural networks. In F. Pereira, C. J. C. Burges, L. Bottou, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 25, pages 341–349. Curran Associates, Inc., 2012.
[39] Y. Yang, J. Sun, H. Li, and Z. Xu. Deep ADMM-Net for compressive sensing MRI. In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, pages 10–18. Curran Associates, 2016. Report available from http://papers.nips.cc/paper/6406-deep-admm-net-for-compressive-sensing-mri.pdf.
[40] H. Zhao, O. Gallo, I. Frosio, and J. Kautz. Loss functions for image restoration with neural networks. IEEE Transactions on Computational Imaging, 3(1):47–57, March 2017.

Learning to solve inverse problems using Wasserstein loss

Authors

How many samples are needed to reliably approximate the best linear estimator for a linear inverse problem?

Learning the optimal regularizer for inverse problems

Spatial Frequency Loss for Learning Convolutional Autoencoders

Regularized Variational Data Assimilation for Bias Treatment using the Wasserstein Metric

Minimax Estimation of Conditional Moment Models

Residual Entropy

Error Analysis of Douglas-Rachford Algorithm for Linear Inverse Problems: Asymptotics of Proximity Operator for Squared Loss

1 Introduction

2 Background

2.1 Inverse problems

Variational regularization

Learning for inverse problems

Loss functions for learning

2.2 Optimal mass transport and Sinkhorn iterations

3 Learning a reconstruction operator using Wasserstein loss

Proposition 1.

Proposition 2.

Example 3.

Example 4.

4 Implementation and evaluation

5 Conclusions and future work

Appendix: Deferred definition and proofs

Proof of Proposition 1.

Definition 5.

Proof of Proposition 2.

Lemma 6.

Proof.

References

Learning to solve inverse problems using Wasserstein loss

Authors

How many samples are needed to reliably approximate the best linear estimator for a linear inverse problem?

Learning the optimal regularizer for inverse problems

Spatial Frequency Loss for Learning Convolutional Autoencoders

Regularized Variational Data Assimilation for Bias Treatment using the Wasserstein Metric

Minimax Estimation of Conditional Moment Models

Residual Entropy

Error Analysis of Douglas-Rachford Algorithm for Linear Inverse Problems: Asymptotics of Proximity Operator for Squared Loss

This week in AI

1 Introduction

2 Background

2.1 Inverse problems

Variational regularization

Learning for inverse problems

Loss functions for learning

2.2 Optimal mass transport and Sinkhorn iterations

3 Learning a reconstruction operator using Wasserstein loss

Proposition 1.

Proposition 2.

Example 3.

Example 4.

4 Implementation and evaluation

5 Conclusions and future work

Appendix: Deferred definition and proofs

Proof of Proposition 1.

Definition 5.

Proof of Proposition 2.

Lemma 6.

Proof.

References