remarkRemark hypothesisHypothesis claimClaim THE COLLINEAR TRIANGULAR MATRIX
AND FS-2 ALGORITHMQ. LIAO, Z. WANG, J. CHEN, B. BAI, S. JIN AND H. WU
[][nocite]ex_supplement

# Fast Sinkhorn II: Collinear Triangular Matrix and Linear Time Accurate
Computation of Optimal Transport††thanks: Submitted to the editors DATE. This
work was supported by National Natural Science Foundation of People’s Republic
of China Grant Nos.11871297 and 12031013, Tsinghua University Initiative
Scientific Research Program, and Shanghai Municipal Science and Technology
Major Project 2021SHZDZX0102.

Qichen Liao Department of Mathematical Sciences, Tsinghua University, Beijing
100084, People’s Republic of China; Theory Lab, Central Research Institute,
2012 Labs, Huawei Technologies Co. Ltd., Hong Kong SAR, People’s Republic of
China<EMAIL_ADDRESS>Zihao Wang Department of Computer
Science and Engineering, Hong Kong University of Science and Technology, Clear
Water Bay, Hong Kong SAR, People’s Republic of China<EMAIL_ADDRESS>Jing Chen School of Physical and Mathematical Sciences, Nanyang Technological
University, Singapore 639798<EMAIL_ADDRESS>Bo Bai Theory Lab,
Central Research Institute, 2012 Labs, Huawei Technologies Co. Ltd., Hong Kong
SAR, People’s Republic of China<EMAIL_ADDRESS>Shi Jin School of
Mathematical Sciences, Institute of Natural Sciences, and MOE-LSC Shanghai
Jiao Tong University, Shanghai 200240, People’s Republic of China ().
<EMAIL_ADDRESS>Hao Wu Corresponding author. Department of Mathematical
Sciences, Tsinghua University, Beijing 100084, People’s Republic of China ()
<EMAIL_ADDRESS>
###### Abstract

In our previous work [arXiv:2202.10042], the complexity of Sinkhorn iteration
is reduced from $O(N^{2})$ to the optimal $O(N)$ by leveraging the special
structure of the kernel matrix. In this paper, we explore the special
structure of kernel matrices by defining and utilizing the properties of the
Lower-ColLinear Triangular Matrix (L-CoLT matrix) and Upper-ColLinear
Triangular Matrix (U-CoLT matrix). We prove that (1) L/U-CoLT matrix-vector
multiplications can be carried out in $O(N)$ operations; (2) both families of
matrices are closed under the Hadamard product and matrix scaling. These
properties help to alleviate two key difficulties for reducing the complexity
of the Inexact Proximal point method (IPOT), and allow us to significantly
reduce the number of iterations to $O(N)$. This yields the Fast Sinkhorn II
(FS-2) algorithm for accurate computation of optimal transport with low
algorithm complexity and fast convergence. Numerical experiments are presented
to show the effectiveness and efficiency of our approach.

###### keywords:

Optimal Transport, Wasserstein-1 metric, Sinkhorn algorithm, IPOT method, FS-2
algorithm

49M25; 65K10

##  1 Introduction

The Wasserstein metric, broadly used in optimal transport theory with
applications in many fields including machine learning, quantifies the
dissimilarity between two probabilistic distributions. Many methods have been
proposed to compute the Wasserstein metrics directly, such as the linear
programming methods [30, 22, 36], combinatorial methods [33], solving the
Monge-Ampère equations [15, 14, 3], via Benamou-Brenier formulation [2, 21]
and the proximal splitting methods [8, 28]. In recent years, several
approximation techniques in optimal transport for high-dimensional
distributions have also been proposed [26, 27].

The Sinkhorn algorithm [10, 34] is a popular $O(N^{2})$ algorithm to
approximate the Wasserstein metric [31] by minimizing the entropy regularized
optimal transport (OT) problem. In [24], the FS-1 algorithm is proposed to
solve entropy regularized OT in $O(N)$ time by leveraging the special
structure of the Sinkhorn kernel matrix of the Wasserstein-1 metric. The
solution of entropy regularized OT approximates the accurate OT solution only
if the regularization parameter is sufficiently small. However, small
regularization parameters lead to numerical instability and excessive
iterations [13]. This causes the slow convergence of the Sinkhorn algorithm.

The Inexact Proximal point method [35] for the Optimal Transport problem
(IPOT) has been proposed to address this challenge. It regularizes the
original OT by introducing the proximal point term and solves a series of
successive subproblems. Only fairly mild regularization parameters are
required to ensure the method’s fast convergence to the accurate OT solution
in an $O(N^{2})$ algorithm. The goal of this paper is to construct a new
method to accurately compute OT solutions with good convergence behavior and
$O(N)$ algorithm complexity by combing the IPOT method and the FS-1 algorithm.
Note the two key steps in the IPOT method make it hard to reduce the
complexity to $O(N)$: the matrix Hadamard product (Algorithm 1, line 4) and
the matrix scaling (Algorithm 1, line 8). For general matrices, the complexity
of the above operations is both $O(N^{2})$. Moreover, these operations may
destroy the special structure of the kernel matrix [24], making it impossible
for us to implement matrix-vector multiplication with $O(N)$ cost.

We will explore the special structure of kernel matrices by defining and
exploiting the properties of the Lower-ColLinear Triangular Matrix (L-CoLT
matrix) and Upper-ColLinear Triangular Matrix (U-CoLT matrix). For these
matrices, we can realize the matrix-vector multiplication with $O(N)$ cost by
using the idea of dynamic programming similar to [24]. Next, we show that each
L/U-CoLT matrix can be represented by two vectors of dimension $N$.
Furthermore, we prove the closure of families of L/U-CoLT matrices to matrix
Hadamard product and matrix scaling. This means that the special structure of
the kernel matrix is preserved by matrix Hadamard product and matrix scaling,
so we can still implement matrix-vector multiplication (Algorithm 1, lines
6-7) with $O(N)$ cost. On the other hand, by updating two representation
vectors of the L/U-CoLT matrix, we can also implement matrix Hadamard product
(Algorithm 1, line 4) and matrix scaling (Algorithm 1, line 8) with $O(N)$
cost. Consequently, the Fast Sinkhorn II (FS-2) algorithm is developed, which
integrates the advantages of both IPOT and FS-1. Moreover, we also find that
the FS-2 algorithm has the advantage in reducing the space complexity since
all the matrices can be represented by vectors. Due to these benefits, one can
expect that our FS-2 could be applied in various fields, e.g., machine
learning [26, 27, 16, 25], image processing [32, 29], inverse problems [6, 12,
37, 18], density function theory[19, 5, 9].

The rest of the paper is organized as follows. In Section 2 , the basics of
the Wasserstein-1 metric and the IPOT method are briefly reviewed. After
presenting the definition, properties, and fast matrix-vector multiplications
of the L/U-CoLT matrix in Section 3, we apply them to accelerate the IPOT
method, thus developing the FS-2 algorithm in Section 4. In Section 5, the
FS-2 algorithm is extended to high dimensions. The numerical experiments are
performed to verify our conclusions in Section 6. We conclude the paper in
Section 7.

##  2 The Wasserstein-1 metric and the IPOT method

Given two unit discrete distributions $\boldsymbol{u}$ and $\boldsymbol{v}$,

$\boldsymbol{u}=(u_{1},u_{2},\cdots
u_{N})^{\top}\in\mathbb{R}^{N},\quad\boldsymbol{v}=(v_{1},v_{2},\cdots,v_{N})^{\top}\in\mathbb{R}^{N},$

where $u_{i}\geq 0$, $v_{j}\geq 0$, and $\sum_{i}u_{i}=\sum_{j}v_{j}=1$. The
Wasserstein-1 distance between them is defined as [31]

(1)
$W_{1}\left(\boldsymbol{u},\boldsymbol{v}\right)=\min_{\Gamma\mathbf{1}=\boldsymbol{u},\Gamma^{T}\mathbf{1}=\boldsymbol{v}}\left\langle
C,\;\Gamma\right\rangle,$

where $C=[c_{ij}]\in\mathbb{R}^{N\times N}$ is the cost matrix. The element
$c_{ij}=\|\boldsymbol{x}_{i}-\boldsymbol{y}_{j}\|_{1}$ represents the cost of
transporting the unit mass from position $\boldsymbol{x}_{i}$ to position
$\boldsymbol{y}_{j}$ and the variable
$\Gamma=[\gamma_{ij}]\in\mathbb{R}^{N\times N}$ to be optimized is the
transport plan. Here, the Frobenius inner product $\left\langle
A,B\right\rangle=\sum_{i,j}a_{ij}b_{ij}$, where $A=[a_{ij}],B=[b_{ij}]$ are
real-valued matrices.

The Sinkhorn algorithm [10, 34] solves an entropy regularized OT problem to
obtain an approximate result of (1). However, the small regular parameter
required by good approximation leads to a slow convergence rate and numerical
instability. To avoid this problem, the proximal point iteration (2) is
developed to solve (1) accurately [35]. It begins with a transport map
$\Gamma^{(0)}$ and iteratively solves the following minimization problem

(2)
$\Gamma^{\left(t+1\right)}=\mathop{\arg\min}\limits_{\Gamma\mathbf{1}=\boldsymbol{u},\Gamma^{T}\mathbf{1}=\boldsymbol{v}}\left\langle
C,\;\Gamma\right\rangle+\delta^{\left(t\right)}D_{h}\left(\Gamma,\;\Gamma^{\left(t\right)}\right),$

where $D_{h}$ is Bregman divergence, taken in the form of the KL divergence in
[35],

$D_{h}\left(A,\;B\right)=\sum_{i,j}\left(a_{ij}\ln\frac{a_{ij}}{b_{ij}}-a_{ij}+b_{ij}\right)$

and $\delta^{\left(t\right)}$ is the regular parameter. The Lagrangian of the
above equation writes

$L(\Gamma,\boldsymbol{\alpha},\boldsymbol{\beta})=\left\langle
C,\;\Gamma\right\rangle+\delta^{\left(t\right)}D_{h}\left(\Gamma,\;\Gamma^{\left(t\right)}\right)+\boldsymbol{\alpha}^{T}\left(\Gamma\mathbf{1}-\boldsymbol{u}\right)+\boldsymbol{\beta}^{T}\left(\Gamma^{T}\mathbf{1}-\boldsymbol{v}\right).$

Taking derivative of the Lagrangian with respect of $\gamma_{ij}$ directly
leads to

$\gamma_{ij}=e^{-\alpha_{i}/\delta^{\left(t\right)}}Q_{ij}^{(t)}e^{-\beta_{j}/\delta^{\left(t\right)}},\;\;\textrm{where}\;\;Q_{ij}^{(t)}=\gamma^{\left(t\right)}_{ij}e^{-c_{ij}/\delta^{\left(t\right)}}>0.$

Denoting $\odot$ as the Hadamard product, $Q^{(t)}=K\odot\Gamma^{(t)}$ and
$K=[e^{-c_{ij}/\delta^{(t)}}]\in\mathbb{R}^{N\times N}$ is the kernel matrix.
Letting $\phi_{i}=e^{-\alpha_{i}/\delta^{\left(t\right)}}$,
$\psi_{j}=e^{-\beta_{j}/\delta^{(t)}}$, and vectors
$\boldsymbol{\phi}=(\phi_{i})$ and $\boldsymbol{\psi}=(\psi_{j})$, one obtains

(3) ${\rm diag}(\boldsymbol{\phi})Q^{(t)}{\rm
diag}(\boldsymbol{\psi})\mathbf{1}=\boldsymbol{u},\quad{\rm
diag}(\boldsymbol{\psi})Q^{(t)\top}{\rm
diag}(\boldsymbol{\phi})\mathbf{1}=\boldsymbol{v}.$

By iteratively updating vectors $\boldsymbol{\phi}$ and $\boldsymbol{\psi}$

(4)
$\boldsymbol{\psi}^{(t,\ell+1)}=\boldsymbol{v}\oslash(Q^{(t)\top}\boldsymbol{\phi}^{(t,\ell)}),\quad\boldsymbol{\phi}^{(t,\ell+1)}=\boldsymbol{u}\oslash(Q^{(t)}\boldsymbol{\psi}^{(t,\ell+1)}),$

one can obtain an accurate solution for the original OT problem (1). Here
$\oslash$ represents pointwise division, $t$ is the proximal iteration step
(outer iteration) and $\ell$ it the Sinkhorn-type iteration step (inner
iteration). The pseudo-code of IPOT is shown in Algorithm 1.

Algorithm 1 IPOT

Input: $\boldsymbol{u},\;\boldsymbol{v}\in\mathbb{R}^{N}$;
$K=e^{-C/\delta}\in\mathbb{R}^{N\times N}$;
$L,\mathrm{itr\\_max}\in\mathbb{N}^{+}$
Output: $W_{1}(\boldsymbol{u},\boldsymbol{v})$

1:$\boldsymbol{\phi},\;\boldsymbol{\psi}\leftarrow\frac{1}{N}\mathbf{1}_{N}$

2:$\Gamma=\mathbf{1}_{N}\mathbf{1}^{T}_{N}$

3:for $t=1\;:\;\mathrm{itr\\_max}$ do

4: $Q\leftarrow K\odot\Gamma$

5: for $\ell=1\;:\;L$ do

6:
$\boldsymbol{\psi}\leftarrow\boldsymbol{v}\oslash\left(Q^{T}\boldsymbol{\phi}\right)$

7:
$\boldsymbol{\phi}\leftarrow\boldsymbol{u}\oslash\left(Q\boldsymbol{\psi}\right)$

8:
$\Gamma\leftarrow\text{diag}\left(\boldsymbol{\phi}\right)Q\text{diag}\left(\boldsymbol{\psi}\right)$
return $W_{1}(\boldsymbol{u},\boldsymbol{v})$

##  3 The Collinear Triangular Matrix

###  3.1 Definition and Fast Matrix-Vector Multiplication

######  Definition 3.1 (Lower/Upper-Collinear Triangular Matrix).

A lower triangular matrix is called a Lower-Collinear Triangular Matrix(L-CoLT
matrix) if its corresponding entries on any two rows (columns) have the
same–column (row) independent– ratio except those dividing by $0$.
Specifically, the N-dimensional L-CoLT matrix set is defined as follows:

(5) $\mathcal{C}_{L}^{N}\\!=\\!\left\\{M\in\mathbb{R}^{N\times
N}\;|\;m_{\\!i+1,j\\!}/m_{\\!i,j\\!}=r_{i},\;j\leq
i;\;m_{\\!i,j\\!}=0,\;i<j,\;\boldsymbol{r}\in\left(\mathbb{R}\backslash\\{0\\}\right)^{N-1}\right\\}.$

Similarly, we define Upper-Collinear Triangular Matrix(U-CoLT matrix), which
is a strictly upper triangular matrix:

(6) $\mathcal{C}_{U}^{N}\\!=\\!\left\\{M\in\mathbb{R}^{N\times
N}\;|\;m_{\\!i-1,j\\!}/m_{\\!i,j\\!}=r^{\prime}_{i-1},\;i<j;\;m_{\\!i,j\\!}=0,\;i\geq
j,\;\boldsymbol{r}^{\prime}\in\left(\mathbb{R}\backslash\\{0\\}\right)^{N-2}\right\\}.$

We call the vectors $r$ and $r^{\prime}$ in (5)-(6) the ratio vectors of the
collinear triangular matrix.

The matrices introduced in Definition 3.1 are termed as collinear triangular
matrices (CoLT), due to the following collinearity between columns:

$m_{i,j}/m_{i,j+1}=m_{k,j}/m_{k,j+1}\iff m_{i,j}/m_{k,j}=m_{i,j+1}/m_{k,j+1}.$

######  Theorem 3.2 (Vector Representation of Collinear Triangular Matrix).

Any L-CoLT matrix $M_{L}$ can be represented by its diagonal elements
$\boldsymbol{\gamma}$ and the ratio vector $\boldsymbol{r}$ in Equation (5).
Any U-CoLT matrix $M_{U}$ can be represented by its superdiagonal elements
$\boldsymbol{\gamma}^{\prime}$ and the ratio vector $\boldsymbol{r}^{\prime}$
in (6).

######  Proof 3.3.

For any L-CoLT matrix $M_{L}\in\mathcal{C}_{L}^{N}$, if its corresponding
$\boldsymbol{\gamma}$ and $\boldsymbol{r}$ are given, then
$m_{i,j}=\gamma_{j}\prod\limits_{k=i}^{N-1}r_{k}$. The proof of U-CoLT is
similar.

In the following, we use
$\textsc{L-CoLT}(\boldsymbol{\gamma},\;\boldsymbol{r})$,
$\boldsymbol{\gamma}\in\mathbb{R}^{N},\;\boldsymbol{r}\in\mathbb{R}^{N-1}$ and
$\textsc{U-CoLT}(\boldsymbol{\gamma}^{\prime},\;\boldsymbol{r}^{\prime})$,
$\boldsymbol{\gamma}^{\prime}\in\mathbb{R}^{N-1},\;\boldsymbol{r}^{\prime}\in\mathbb{R}^{N-2}$
to represent a L-CoLT matrix and a U-CoLT matrix, respectively. A specific
correspondence of L-CoLT and U-CoLT is shown as follow:

For the L-CoLT matrix $M_{L}$

$M_{L}=\textsc{L-CoLT}(\boldsymbol{\gamma},\;\boldsymbol{r})=\begin{pmatrix}\gamma_{1}&&&&\\\
\gamma_{1}r_{1}&\gamma_{2}&&&\\\
\gamma_{1}r_{1}r_{2}&\gamma_{2}r_{2}&\gamma_{3}&&\\\
\vdots&\vdots&\vdots&\ddots&\\\
\gamma_{1}\\!\prod\limits_{i=1}^{\\!N\\!-\\!1\\!}\\!r_{i}&\gamma_{2}\\!\prod\limits_{i=2}^{\\!N\\!-\\!1\\!}\\!r_{i}&\gamma_{3}\\!\prod\limits_{i=3}^{\\!N\\!-\\!1\\!}\\!r_{i}&\cdots&\;\;\gamma_{N}\end{pmatrix}$

Similarly, for the U-CoLT matrix $M_{U}$

$M_{U}=\textsc{U-CoLT}(\boldsymbol{\gamma}^{\prime},\;\boldsymbol{r}^{\prime})=\begin{pmatrix}0\quad&\;\gamma^{\prime}_{1}&\;\gamma^{\prime}_{2}r^{\prime}_{1}&\;\cdots&\;\gamma^{\prime}_{\\!N\\!-\\!1\\!}\\!\prod\limits_{i=1}^{\\!N\\!-\\!2\\!}\\!r^{\prime}_{i}\\\
\;&\;0&\;\gamma^{\prime}_{2}&\;\cdots&\;\gamma^{\prime}_{\\!N\\!-\\!1\\!}\\!\prod\limits_{i=2}^{\\!N\\!-\\!2\\!}\\!r^{\prime}_{i}\\\
\;&\;&\quad\ddots&\quad\ddots&\;\vdots\\\ \;&\;&\;&\quad
0&\;\gamma^{\prime}_{\\!N\\!-\\!1\\!}\\\ \;&\;&\;&\;&0\;\\\ \end{pmatrix}$

The special nature of the L-CoLT and U-CoLT matrices allows us to compute
matrix-vector multiplications in $O(N)$ operations.

For any $M_{L}=\textsc{L-CoLT}(\boldsymbol{\gamma},\;\boldsymbol{r})$ and
vector $\boldsymbol{y}\in\mathbb{R}^{N}$, the matrix-vector multiplication
$M_{L}\boldsymbol{y}$ is written as

(7) $M_{L}\boldsymbol{y}=\begin{pmatrix}\gamma_{1}y_{1}&+&0&+&0&\cdots&+&0\\\
\gamma_{1}r_{1}y_{1}&+&\gamma_{2}y_{2}&+&0&\cdots&+&0\\\
\gamma_{1}r_{1}r_{2}y_{1}&+&\gamma_{2}r_{2}y_{2}&+&\gamma_{3}y_{3}&\cdots&+&0\\\
\vdots&\vdots&\vdots&\vdots&\vdots&\ddots&\vdots&\vdots\\\
\gamma_{1}\prod\limits_{i=1}^{N-1}r_{i}y_{1}&+&\gamma_{2}\prod\limits_{i=2}^{N-1}r_{i}y_{2}&+&\gamma_{3}\prod\limits_{i=3}^{N-1}r_{i}y_{3}&\cdots&+&\gamma_{N}y_{N}\end{pmatrix}.$

Denote $\boldsymbol{p}_{k}$ as the summation of the $k$-th row in (7), then
one has

$\displaystyle p_{1}=\gamma_{1}y_{1},\quad
p_{k}=r_{k-1}p_{k-1}+\gamma_{k}y_{k},\quad k=2,\cdots,N.$

Based on this recursion formula, a fast implementation is proposed in
Algorithm 2.

Algorithm 2 Fast L-CoLT Matrix-Vector multiplication

Input: input vector $\boldsymbol{y}$ of size $N$, input matrix
$M_{L}=\textsc{L-CoLT}(\boldsymbol{\gamma},\;\boldsymbol{r})$
Output: $\boldsymbol{p}=M_{L}\boldsymbol{y}$

1:procedure LCMV($\boldsymbol{y},\;\boldsymbol{\gamma},\;\boldsymbol{r}$)

2: $p_{1}=\gamma_{1}y_{1}$

3: for $i=1\;:\;N-1$ do

4: $p_{i+1}=r_{i}p_{i}+\gamma_{i+1}y_{i+1}$ return $\boldsymbol{p}$

Similarly, the fast matrix-vector multiplication for U-CoLT matrices is shown
in Algorithm 3.

Algorithm 3 Fast U-CoLT Matrix-Vector multiplication

Input: input vector $\boldsymbol{y}$ of size $N$, input matrix
$M=\textsc{U-CoLT}(\boldsymbol{\gamma}^{\prime},\;\boldsymbol{r}^{\prime})$
Output: $\boldsymbol{q}=M_{U}\boldsymbol{y}$

1:procedure
UCMV($\boldsymbol{y},\;\boldsymbol{\gamma}^{\prime},\;\boldsymbol{r}^{\prime}$)

2: $q_{N}=0$, $q_{N-1}=\gamma^{\prime}_{N-1}y_{N}$

3: for $i=2\;:\;N-1$ do

4: $q_{N-i}=r^{\prime}_{N-i}q_{N-i+1}+\gamma^{\prime}_{N-i}y_{N-i+1}$ return
$\boldsymbol{q}$

Next, we denote the set $\mathcal{C}^{N}$ as the direct sum of
$\mathcal{C}_{L}^{N}$ and $\mathcal{C}_{U}^{N}$, defined by

######  Definition 3.4.

(8)
$\mathcal{C}^{N}=\mathcal{C}_{L}^{N}+\mathcal{C}_{U}^{N}=\left\\{A+B\;|\;A\in\mathcal{C}_{L}^{N},\;B\in\mathcal{C}_{U}^{N}\right\\}.$

Due to the linearity of matrix-vector multiplication, we can further develop
the fast matrix-vector multiplication algorithm for matrices in
$\mathcal{C}^{N}$, which is given in Algorithm 4.

Algorithm 4 CoLT Matrix-Vector multiplication

Input: input vector $\boldsymbol{x}$ of size $N$, diagonal elements
$\boldsymbol{\gamma},\;\boldsymbol{\gamma}^{\prime}$ and the ratio vector
$\boldsymbol{r},\;\boldsymbol{r}^{\prime}$
Output: $\boldsymbol{p}+\boldsymbol{q}=M\boldsymbol{y}$

1:procedure
CMV($\boldsymbol{y},\;\mathbf{c}^{L},\;\mathbf{c}^{U},\;\boldsymbol{\gamma},\;\boldsymbol{\gamma}^{\prime}$)

2:
$\boldsymbol{p}=\textsc{LCMV}(\boldsymbol{y},\;\boldsymbol{\gamma},\;\boldsymbol{r})$

3:
$\boldsymbol{q}=\textsc{UCMV}(\boldsymbol{y},\;\boldsymbol{\gamma}^{\prime},\;\boldsymbol{r}^{\prime})$
return $\boldsymbol{p}+\boldsymbol{q}$

The space and time complexities of these algorithms are $O(N)$, which is much
better than the original matrix-vector multiplication.

###  3.2 Some Basic Properties

In this subsection, we justify some basic properties of those matrices
involved, which will be used in our algorithm.

######  Theorem 3.5.

$\left(\mathcal{C}_{L}^{N},\;\odot\right)$ and
$\left(\mathcal{C}_{U}^{N},\;\odot\right)$ are Abelian groups, where $\odot$
is the Hadamard product.

######  Proof 3.6.

We only prove the theorem for $\left(\mathcal{C}_{L}^{N},\;\odot\right)$. It
suffices to show that $\left(\mathcal{C}_{L}^{N},\;\odot\right)$ has the
following properties:

Closure: For any two matrices
$A=\textsc{L-CoLT}(\hat{\boldsymbol{\gamma}},\;\hat{\boldsymbol{r}})$ and
$B=\textsc{L-CoLT}(\tilde{\boldsymbol{\gamma}},\;\tilde{\boldsymbol{r}})$, we
set $D=A\odot B$. Since $d_{i,j}=a_{i,j}b_{i,j}$, one has

(9) $\displaystyle
d_{i,j}/d_{i+1,j}=\left(a_{i,j}b_{i,j}\right)/\left(a_{i+1,j}b_{i+1,j}\right)=\hat{r}_{i}\tilde{r}_{i},$
$\displaystyle\text{ }j=1,2,\cdots,i,$

and the strictly upper triangle part of $D$ is obviously 0, which means
$D=\textsc{L-CoLT}(\hat{\boldsymbol{\gamma}}\odot\tilde{\boldsymbol{\gamma}},\;\hat{r}\odot\tilde{r})\in\mathcal{C}_{L}^{N}$.

Identity and Inverses: Let

(10) $E=\begin{pmatrix}1&&&&\\\ 1&\quad 1&&&\\\ 1&\quad 1&\quad 1&&\\\
\vdots&\quad\vdots&\quad\vdots&\quad\ddots&\\\ 1&\quad 1&\quad
1&\quad\cdots&\quad 1\end{pmatrix}\in\mathcal{C}_{L}^{N},$

then for any $A=\textsc{L-CoLT}(\boldsymbol{\gamma},\;\boldsymbol{r})$,
$A\odot E=E\odot A=A$, which means $E$ is the identity element. Let

$B=\begin{pmatrix}1/a_{11}&&&&\\\ 1/a_{21}&\quad 1/a_{22}&&&\\\ 1/a_{31}&\quad
1/a_{32}&\quad 1/a_{33}&&\\\ \vdots&\quad\vdots&\quad\vdots&\quad\ddots&\\\
1/a_{n1}&\quad 1/a_{n2}&\quad 1/a_{n3}&\quad\cdots&\quad
1/a_{nn}\end{pmatrix}.$

Since

$b_{i,j}/b_{i+1,j}=a_{i+1,j}/a_{i,j}=1/r_{i},\;j=1,2,\cdots,i,$

then $B\in\mathcal{C}_{L}^{N}$. And obviously, $A\odot B=B\odot A=E$, which
means $B$ is the inverse of $A$.

Commutativity and Associativity: The commutativity and associativity can be
derived from the commutative and associative law of real number
multiplication.

Based on the above theorem, we can deduce directly

######  Corollary 3.7.

$\left(\mathcal{C}^{N},\;\odot\right)$ is an abelian group with identity
element $\mathbf{1}_{N}\mathbf{1}_{N}^{T}$.

######  Theorem 3.8.

For any vector
$\boldsymbol{x}\in\left(\mathbb{R}\backslash\\{0\\}\right)^{N}$,
$f_{\boldsymbol{x}}\left(M\right)=\left(\text{diag}\left(\boldsymbol{x}\right)\right)M$
and
$g_{\boldsymbol{x}}\left(M\right)=M\left(\text{diag}\left(\boldsymbol{x}\right)\right)$
are permutations in $\mathcal{C}_{L}^{N}$ and $\mathcal{C}_{U}^{N}$.

######  Proof 3.9.

We only prove the theorem for $\mathcal{C}_{L}^{N}$.

Closure: For any vector
$\boldsymbol{x}\in\left(\mathbb{R}\backslash\\{0\\}\right)^{N}$, and
$M_{L}=\textsc{L-CoLT}(\boldsymbol{\gamma},\;\boldsymbol{r})$, let $E$ be the
one defined in Equation (10). Since

$E_{1}=\left(\text{diag}\left(\boldsymbol{x}\right)\right)E=\begin{pmatrix}x_{1}&&&&\\\
x_{2}&\;x_{2}&&&\\\ x_{3}&\;x_{3}&\;x_{3}&&\\\
\vdots&\;\vdots&\;\vdots&\;\ddots&\\\
x_{n}&\;x_{n}&\;x_{n}&\;\cdots&\;x_{n}\end{pmatrix}\in\mathcal{C}_{L}^{N},$
$E_{2}=E\left(\text{diag}\left(\boldsymbol{x}\right)\right)=\begin{pmatrix}x_{1}&&&&\\\
x_{1}&\;x_{2}&&&\\\ x_{1}&\;x_{2}&\;x_{3}&&\\\
\vdots&\;\vdots&\;\vdots&\;\ddots&\\\
x_{1}&\;x_{2}&\;x_{3}&\;\cdots&\;x_{n}\end{pmatrix}\in\mathcal{C}_{L}^{N},$

we have

(11) $\displaystyle
f_{\boldsymbol{x}}\left(M\right)=\left(\text{diag}\left(\boldsymbol{x}\right)\right)M=\left(\text{diag}\left(\boldsymbol{x}\right)\right)E\odot
M=E_{1}\odot M\in\mathcal{C}_{L}^{N};$ $\displaystyle
g_{\boldsymbol{x}}\left(M\right)=M\left(\text{diag}\left(\boldsymbol{x}\right)\right)=M\odot
E\left(\text{diag}\left(\boldsymbol{x}\right)\right)=M\odot
E_{2}\in\mathcal{C}_{L}^{N}.$

The last set membership can be derived by the closure of
$\left(\mathcal{C}_{L}^{N},\;\odot\right)$ proved in Theorem 3.5. Hence,
$f_{\boldsymbol{x}}$ and $g_{\boldsymbol{x}}$ are maps from
$\mathcal{C}_{L}^{N}$ to itself.

Injectiveness: For any two matrices
$A=\textsc{L-CoLT}(\hat{\boldsymbol{\gamma}},\;\hat{r})$ and
$B=\textsc{L-CoLT}(\tilde{\boldsymbol{\gamma}},\;\tilde{\boldsymbol{r}})$, let
$D_{1}$ be the inverse of $E_{1}$ and $D_{2}$ be the inverse of $E_{2}$. If
$f_{\boldsymbol{x}}\left(A\right)=f_{\boldsymbol{x}}\left(B\right)$, then

$A=D_{1}\odot E_{1}\odot A=D_{1}\odot
f_{\boldsymbol{x}}\left(A\right)=D_{1}\odot
f_{\boldsymbol{x}}\left(B\right)=D_{1}\odot E_{1}\odot B=B.$

If $g_{\boldsymbol{x}}\left(A\right)=g_{\boldsymbol{x}}\left(B\right)$, then

$A=A\odot E_{2}\odot D_{2}=g_{\boldsymbol{x}}\left(A\right)\odot
D_{2}=g_{\boldsymbol{x}}\left(B\right)\odot D_{2}=B\odot E_{2}\odot D_{2}=B,$

which means $f_{\boldsymbol{x}}\left(\cdot\right)$ and
$g_{\boldsymbol{x}}\left(\cdot\right)$ are injective functions.

Surjectiveness: For any $M\in\mathcal{C}_{L}^{N}$, let $Q_{1}=D_{1}\odot M$
and $Q_{2}=M\odot D_{2}$, then

$\displaystyle f_{\boldsymbol{x}}\left(Q_{1}\right)=E_{1}\odot D_{1}\odot
M=E\odot M=M;$ $\displaystyle g_{\boldsymbol{x}}\left(Q_{2}\right)=M\odot
D_{2}\odot E_{2}=M\odot E=M,$

which means $f_{\boldsymbol{x}}\left(\cdot\right)$ and
$g_{\boldsymbol{x}}\left(\cdot\right)$ are surjective functions.

######  Corollary 3.10.

$f_{\boldsymbol{x}}\left(\cdot\right)$ and
$g_{\boldsymbol{x}}\left(\cdot\right)$ are permutations in $\mathcal{C}^{N}$.

##  4 The Fast Sinkhorn II

In this section, we will discuss the implementation details to accelerate
IPOT. In Algorithm 1, three parts lead to $O(N^{2})$ algorithm complexity,
i.e., the matrix Hadamard product (line 4), the matrix-vector multiplication
(lines 6-7), and the matrix scaling (line 8). They all rely on the
representation and manipulation of L/U CoLT matrices.

For two discrete distributions on a 1D uniform mesh grid with a grid spacing
of $h$, by introducing the notation $\lambda=e^{-h/\delta}$, the kernel matrix
$K$ is written as

(12) $K=\begin{pmatrix}1&\lambda&\lambda^{2}&\cdots&\lambda^{N-1}\\\
\lambda&1&\lambda&\cdots&\lambda^{N-2}\\\
\lambda^{2}&\lambda&1&\cdots&\lambda^{N-3}\\\
\vdots&\vdots&\vdots&\ddots&\vdots\\\
\lambda^{N-1}&\lambda^{N-2}&\lambda^{N-3}&\cdots&1\end{pmatrix}\in\mathcal{C}^{N}.$

Below we discuss step by step of IPOT (Algorithm 1) to reduce the complexity:

  * •

line 2: $\Gamma=\mathbf{1}_{N}\mathbf{1}^{T}_{N}\in\mathcal{C}^{N}$, we only
need four vectors
$(\boldsymbol{\gamma},\boldsymbol{\gamma}^{\prime},\boldsymbol{r},\boldsymbol{r}^{\prime})$
to represent $\Gamma$ according to Theorem 3.2.

  * •

line 4: the matrix Hadamard product $Q=K\odot\Gamma\in\mathcal{C}^{N}$
according to Theorem 3.5 and Corollary 3.7. By updating the four
representation vectors
$(\boldsymbol{\gamma},\boldsymbol{\gamma}^{\prime},\boldsymbol{r},\boldsymbol{r}^{\prime})$,
we can obtain $Q$ with $O(N)$ cost.

  * •

lines 6-7: the matrix-vector multiplication $Q^{T}\boldsymbol{\phi}$ and
$Q\boldsymbol{\psi}$ can be implemented with $O(N)$ cost according to
Algorithm 4.

  * •

line 8: the matrix scaling
$\Gamma=\text{diag}(\boldsymbol{\phi})Q\text{diag}(\boldsymbol{\psi})\in\mathcal{C}^{N}$
according to Theorem 3.8 and Corollary 3.10. By updating the four
representation vectors
$(\boldsymbol{\gamma},\boldsymbol{\gamma}^{\prime},\boldsymbol{r},\boldsymbol{r}^{\prime})$,
we can obtain $\Gamma$ with $O(N)$ cost.

Based on the above discussions, we proposed the FS-2 algorithm with $O(N)$
complexity. The pseudo-code is presented in Algorithm 5.

Algorithm 5 1D FS-2 Algorithm

Input: $\boldsymbol{u},\;\boldsymbol{v}\in\mathbb{R}^{N}$;
$L,\mathrm{itr\\_max}\in\mathbb{N}^{+}$; $h,\delta\in\mathbb{R}$
Output: $W_{1}(\boldsymbol{u},\boldsymbol{v})$

1:$\lambda\leftarrow e^{-h/\delta}$;
$\boldsymbol{\phi},\boldsymbol{\psi}\leftarrow\frac{1}{N}\mathbf{1}_{N}$;
$\boldsymbol{r},\boldsymbol{s}\leftarrow\mathbf{0}_{N}$

2:$\boldsymbol{\alpha}^{L},\;\boldsymbol{\beta}^{L},\;\boldsymbol{\alpha}^{U},\;\boldsymbol{\beta}^{U}\leftarrow\lambda\mathbf{1}_{N-1}$;
$\boldsymbol{\gamma}\leftarrow\mathbf{1}_{N}$;
$\boldsymbol{\gamma}^{\prime}\leftarrow\lambda\mathbf{1}_{N-1}$

3:for $t=1\;:\;\mathrm{itr\\_max}$ do

4: for $\ell=1\;:\;L$ do

5:
$\boldsymbol{r}\leftarrow\textsc{CMV}(\boldsymbol{\phi},\boldsymbol{\beta}^{L},\boldsymbol{\beta}^{U},\boldsymbol{\gamma},\boldsymbol{\gamma}^{\prime})$

6: $\boldsymbol{\psi}\leftarrow\boldsymbol{v}\oslash\boldsymbol{r}$

7:
$\boldsymbol{s}\leftarrow\textsc{CMV}(\boldsymbol{\psi},\boldsymbol{\alpha}^{L},\boldsymbol{\alpha}^{U},\boldsymbol{\gamma},\boldsymbol{\gamma}^{\prime})$

8: $\boldsymbol{\phi}\leftarrow\boldsymbol{u}\oslash\boldsymbol{s}$

9: for $i=1\;:\;N-1$ do

10:
$\alpha^{L}_{i}\leftarrow\lambda\alpha^{L}_{i}\left(\phi_{i+1}/\phi_{i}\right)$,
$\beta^{L}_{i}\leftarrow\lambda\beta^{L}_{i}\left(\psi_{i+1}/\psi_{i}\right)$

11:
$\gamma^{\prime}_{i}\leftarrow\lambda\gamma^{\prime}_{i}\phi_{i}\psi_{i+1}$

12: for $i=1\;:\;N-2$ do

13:
$\alpha^{U}_{i}\leftarrow\lambda\alpha^{U}_{i}\left(\phi_{i}/\phi_{i+1}\right)$,
$\beta^{U}_{i}\leftarrow\lambda\beta^{U}_{i}\left(\psi_{i}/\psi_{i+1}\right)$

14:
$\boldsymbol{\gamma}\leftarrow\boldsymbol{\phi}\odot\boldsymbol{\psi}\odot\boldsymbol{\gamma}$
return $W_{1}(\boldsymbol{u},\boldsymbol{v})$

There is a minor flaw in the above algorithm. The computational cost of
$W_{1}(\boldsymbol{u},\boldsymbol{v})$ is still $O(N^{2})$ in the last step.
This was also ignored in our previous paper [24]. Now, we would like to
discuss this issue. The computation of
$W_{1}\left(\boldsymbol{u},\boldsymbol{v}\right)=\left\langle
C,\;\Gamma\right\rangle$ can be regarded as the summation of all elements of
the following matrix.

(13) $\small
C\odot\Gamma=\begin{pmatrix}0&h\gamma^{\prime}_{1}&2h\gamma^{\prime}_{2}r^{\prime}_{1}&\cdots&(\\!N\\!-\\!1\\!)h\gamma^{\prime}_{\\!N\\!-\\!1\\!}\\!\prod\limits_{i=1}^{\\!N\\!-\\!2\\!}\\!r^{\prime}_{i}\\\
h\gamma_{1}r_{1}&0&h\gamma^{\prime}_{2}&\cdots&(\\!N\\!-\\!2\\!)h\gamma^{\prime}_{\\!N\\!-\\!1\\!}\\!\prod\limits_{i=2}^{\\!N\\!-\\!2\\!}\\!r^{\prime}_{i}\\\
2h\gamma_{1}r_{1}r_{2}&h\gamma_{2}r_{2}&0&\cdots&(\\!N\\!-\\!3\\!)h\gamma^{\prime}_{\\!N\\!-\\!1\\!}\\!\prod\limits_{i=3}^{\\!N\\!-\\!2\\!}\\!r^{\prime}_{i}\\\
\vdots&\vdots&\vdots&\ddots&\vdots\\\
\left(\\!N\\!-\\!1\\!\right)h\gamma_{1}\\!\prod\limits_{i=1}^{\\!N\\!-\\!1\\!}\\!r_{i}&\;\;\left(\\!N\\!-\\!2\\!\right)h\gamma_{2}\\!\prod\limits_{i=2}^{\\!N\\!-\\!1\\!}\\!r_{i}&\;\;\left(\\!N\\!-\\!3\\!\right)h\gamma_{3}\\!\prod\limits_{i=3}^{N\\!-\\!1}\\!r_{i}&\cdots&0\end{pmatrix}.$

We separate the summation of the matrix to the lower and strictly upper
triangular parts. Thus, the $k$-th line summation of two parts can be written
as

$p_{k}=\sum\limits_{i=1}^{k}\omega_{ki},\quad
q_{k}=\sum\limits_{i=k+1}^{N}\omega_{ki}.$

We can consider the following recursive computation

(14) $\displaystyle p_{1}=0,\quad p_{2}=h\gamma_{1}r_{1},\quad
p^{\prime}_{2}=h\gamma_{1}r_{1}+h\gamma_{2},$ $\displaystyle
p_{i}=r_{i-1}\left(p_{i-1}+p^{\prime}_{i-1}\right),\quad
p^{\prime}_{i}=r_{i-1}p^{\prime}_{i-1}+h\gamma_{i},\quad i=3,\;4,\cdots,\;N.$
$\displaystyle q_{N}=0,\quad q_{N-1}=h\gamma^{\prime}_{N-1},\quad
q^{\prime}_{N-1}=h\gamma^{\prime}_{N-1}+h,$ $\displaystyle
q_{j}=r^{\prime}_{j}\left(q_{j+1}+q^{\prime}_{j+1}\right),\quad
q^{\prime}_{j}=r^{\prime}_{j}q^{\prime}_{j+1}+h,\quad j=1,\;2,\cdots,\;N-2.$

Thus, the Wasserstein-1 metric can be finally obtained with $O(N)$ cost

$W_{1}(\boldsymbol{u},\boldsymbol{v})=\left\langle
C,\;\Gamma\right\rangle=\sum_{i=1}^{N}\left(p_{i}+q_{i}\right).$

##  5 Extension to high dimension

In this section, we illustrate how the FS-2 algorithm generalizes to higher
dimensions using the two-dimensional case as an example.

###  5.1 Block Collinear Triangular Matrix

Hereinafter, for $A\in\mathbb{R}^{MN\times MN}$, we break it into $M^{2}$
uniform blocks with size $N\times N$:

$A=\left(\begin{array}[]{c|c|c|c|c}A_{1,1}&A_{1,2}&A_{1,3}&\cdots&A_{1,M}\\\
\hline\cr A_{2,1}&A_{2,2}&A_{2,3}&\cdots&A_{2,M}\\\
\hline\cr\vdots&\vdots&\vdots&\ddots&\vdots\\\ \hline\cr
A_{M,1}&A_{M,2}&A_{M,3}&\cdots&A_{M,M}\end{array}\right).$

And for vectors $\boldsymbol{x}\in\mathbb{R}^{kN}$, we break it into $k$
uniform blocks as
$\left(\boldsymbol{x}_{1},\boldsymbol{x}_{2},\cdots,\boldsymbol{x}_{k}\right)^{T}$,
in which

$\boldsymbol{x}_{i}=\left(x_{1+(i-1)N},\;x_{2+(i-1)N},\;\cdots,\;x_{iN}\right)^{T},\quad
i=1,2,\cdots,k.$

To carry out the fast implementation of the matrix-vector multiplication of
the block matrix above, we generalize the definition of $\mathcal{C}^{N}$ in
(8) to the block case as

######  Definition 5.1.

$\begin{aligned} \mathcal{C}^{N,M}=\\{A\in&\mathbb{R}^{MN\times
MN}|A_{k,k}\in\mathcal{C}^{N};\;\boldsymbol{r}^{L},\boldsymbol{r}^{U}\in\left(\mathbb{R}\backslash\\{0\\}\right)^{\left(M-1\right)N};\\\
&A_{i+1,j}=\left(\text{diag}\left(\boldsymbol{r}^{L}_{i}\right)\right)A_{i,j},\;j\leq
i;\;A_{i-1,j}=\left(\text{diag}\left(\boldsymbol{r}^{U}_{i-1}\right)\right)A_{i,j},\;i\leq
j\\}\end{aligned}.$

Since $\mathcal{C}^{N,M}$ is a generalization of $\mathcal{C}^{N}$, we can
also use the strategy of Algorithm 4 in blocks to reduce the computational
cost of matrix-vector multiplications. For a vector
$\boldsymbol{x}\in\mathbb{R}^{NM}$, the matrix-vector multiplication
$A\boldsymbol{x}$ is written as

(15)
$A\boldsymbol{x}=\begin{pmatrix}A_{1,1}\boldsymbol{x}_{1}&+&A_{1,2}\boldsymbol{x}_{2}&+&A_{1,3}\boldsymbol{x}_{3}&\cdots&+&A_{1,M}\boldsymbol{x}_{M}\\\
A_{2,1}\boldsymbol{x}_{1}&+&A_{2,2}\boldsymbol{x}_{2}&+&A_{2,3}\boldsymbol{x}_{3}&\cdots&+&A_{2,M}\boldsymbol{x}_{M}\\\
A_{3,1}\boldsymbol{x}_{1}&+&A_{3,2}\boldsymbol{x}_{2}&+&A_{3,3}\boldsymbol{x}_{3}&\cdots&+&A_{3,M}\boldsymbol{x}_{M}\\\
\vdots&\vdots&\vdots&\vdots&\vdots&\ddots&\vdots&\vdots\\\
A_{M,1}\boldsymbol{x}_{1}&+&A_{M,2}\boldsymbol{x}_{2}&+&A_{M,3}\boldsymbol{x}_{3}&\cdots&+&A_{M,M}\boldsymbol{x}_{M}\end{pmatrix}.$

We separate the summation of row $k$ to the lower triangular part
$\boldsymbol{p}_{k}$ and the strictly upper triangular part
$\boldsymbol{q}_{k}$. Then computing $A\boldsymbol{x}$ is formulated as

$A\boldsymbol{x}=\boldsymbol{p}+\boldsymbol{q},\quad\boldsymbol{p}_{k}=\sum_{i=1}^{k}A_{k,i}\boldsymbol{x}_{i},\quad\boldsymbol{q}_{k}=\sum_{i=k+1}^{M}A_{k,i}\boldsymbol{x}_{i},\quad
k=1,\cdots,M.$

If $A$ is in $\mathcal{C}^{N,M}$ with $\boldsymbol{r}^{L}$ and
$\boldsymbol{r}^{U}$, instead of directly calculating $\boldsymbol{p}_{k}$ and
$\boldsymbol{q}_{k}$, a successive computation is used

(16)
$\displaystyle\boldsymbol{p}_{1}=A_{1,1}\boldsymbol{x}_{1},\quad\boldsymbol{p}_{k}=\boldsymbol{r}^{L}_{k-1}\odot\boldsymbol{p}_{k-1}+A_{k,k}\boldsymbol{x}_{k},\quad
k=2,\cdots,M,$
$\displaystyle\boldsymbol{q}_{M}=\mathbf{0}_{N},\quad\boldsymbol{q}_{k-1}=\boldsymbol{r}^{U}_{k-1}\odot(\boldsymbol{q}_{k}+A_{k,k}\boldsymbol{x}_{k}),\quad
k=M,M-1,\cdots,2.$

Since the computation of $A_{k,k}\boldsymbol{x}_{k}$ can be carried out with
$O(N)$ complexity by using Algorithm 4, the whole computation is of $O(NM)$
complexity.

Similar to Theorem 3.5 and Theorem 3.8, $\mathcal{C}^{N,M}$ is closed under
the Hadamard product and matrix scaling.

######  Theorem 5.2.

$\left(\mathcal{C}^{N,M},\odot\right)$ is an Abelian group; Matrix scaling
operations are permutations in $\mathcal{C}^{N,M}$.

######  Proof 5.3.

For any $A\in\mathcal{C}^{N,M}$, since the diagonal blocks of $A$ are in
$\mathcal{C}^{N}$, by Corollary 3.10, all blocks in $A$ are in
$\mathcal{C}^{N}$. Then the two properties can be proved in a similar way as
in Section 3.2.

###  5.2 The 2D FS-2 Algorithm

Consider two discretized probabilistic distributions

$\displaystyle\boldsymbol{u}$
$\displaystyle=\left(u_{11},u_{21},\cdots,u_{N1},u_{12},\cdots,u_{i_{1}j_{1}},\cdots,u_{NM}\right),$
$\displaystyle\boldsymbol{v}$
$\displaystyle=\left(v_{11},v_{21},\cdots,v_{N1},v_{12},\cdots,v_{i_{2}j_{2}},\cdots,v_{NM}\right),$

on a uniform 2D mesh of size $N\times M$ with a vertical spacing of $h_{1}$
and a horizontal spacing of $h_{2}$. The corresponding kernel matrix is
written as

$K=\left(\begin{array}[]{c|c|c|c|c}K_{0}&\lambda_{2}K_{0}&\lambda_{2}^{2}K_{0}&\cdots&\lambda_{2}^{M-1}K_{0}\\\
\hline\cr\lambda_{2}K_{0}&K_{0}&\lambda_{2}K_{0}&\cdots&\lambda_{2}^{M-2}K_{0}\\\
\hline\cr\vdots&\vdots&\vdots&\ddots&\vdots\\\
\hline\cr\lambda_{2}^{M-1}K_{0}&\lambda_{2}^{M-2}K_{0}&\lambda_{2}^{M-3}K_{0}&\cdots&K_{0}\end{array}\right),$

where the sub-matrix

$K_{0}=\begin{pmatrix}1&\lambda_{1}&\cdots&\lambda_{1}^{N-1}\\\
\lambda_{1}&1&\cdots&\lambda_{1}^{N-2}\\\ \vdots&\vdots&\ddots&\vdots\\\
\lambda_{1}^{N-1}&\lambda_{1}^{N-2}&\cdots&1\end{pmatrix},$

and

$\lambda_{1}=e^{-h_{1}/\delta},\quad\lambda_{2}=e^{-h_{2}/\delta}.$

Obviously, the original 2D kernel $K$ contains blocks which are multiples of
the 1D kernel, hence belongs to $\mathcal{C}^{N,M}$. By an analysis similar to
that in Section 4 and using Theorem 5.2, the matrices $Q$ and $\Gamma$ are in
$\mathcal{C}^{N,M}$ throughout the course of the iteration, which means that
all the matrix-vector multiplications can be carried out by using recursion
(16). Thus, the total cost of matrix-vector multiplication of our FS-2
algorithm for 2D Wasserstein-1 metric is reduced to $O(NM)$.

In the 2D FS-2 algorithm, we use ‘ $\widehat{}$ ’ to distinguish whether it is
a coefficient of the block or the inner sub-matrix, and update them
simultaneously after an inner loop. The pseudo-code is presented in Algorithm
6. Updating the coefficients of inner sub-matrices and computation of
$W_{1}\left(\boldsymbol{u},\boldsymbol{v}\right)$ are omitted in the pseudo-
code since they are similar to the 1D case, which we have described in detail.

Algorithm 6 2D FS-2 Algorithm

Input: $\boldsymbol{u},\;\boldsymbol{v}\in\mathbb{R}^{NM}$;
$L,\mathrm{itr\\_max}\in\mathbb{N}^{+}$; $h_{1},h_{2},\delta\in\mathbb{R}$
Output: $W_{1}(\boldsymbol{u},\boldsymbol{v})$

1:$\lambda_{1}\leftarrow e^{-h_{1}/\delta}$; $\lambda_{2}\leftarrow
e^{-h_{2}/\delta}$;
$\boldsymbol{\phi}^{(0)},\boldsymbol{\psi}^{(0)}\leftarrow\frac{1}{NM}\mathbf{1}_{NM}$;
$\boldsymbol{p}^{(0)},\boldsymbol{r}^{(0)},\boldsymbol{q}^{(0)},\boldsymbol{s}^{(0)}\leftarrow\mathbf{0}_{NM}$

2:$\boldsymbol{\gamma}\leftarrow\mathbf{1}_{NM},\;\boldsymbol{\gamma}^{\prime},\;\boldsymbol{\alpha}^{L},\;\boldsymbol{\beta}^{L},\lambda_{1}\mathbf{1}_{(N-1)M},\;\boldsymbol{\alpha}^{U},\;\boldsymbol{\beta}^{U}\leftarrow\lambda_{1}\mathbf{1}_{(N-2)M}$

3:$\widehat{\boldsymbol{\alpha}}^{L},\;\widehat{\boldsymbol{\beta}}^{L},\;\widehat{\boldsymbol{\alpha}}^{U},\;\widehat{\boldsymbol{\beta}}^{U}\leftarrow\lambda\mathbf{1}_{N(M-1)}$

4:while $t=1\;:\;\mathrm{itr\\_max}$ do

5: for $\ell=1\;:\;L$ do

6:
$\boldsymbol{r}_{1}\leftarrow\textsc{CMV}(\boldsymbol{\phi}_{1},\;\boldsymbol{\beta}_{1})$,
$s_{M}\leftarrow\mathbf{0}_{N}$

7: for $i=1\;:\;M-1$ do

8:
$\boldsymbol{r}_{i+1}\leftarrow\widehat{\boldsymbol{\beta}}^{L}_{i}\odot\boldsymbol{r}_{i}+\textsc{CMV}(\boldsymbol{\phi}_{i+1},\boldsymbol{\beta}_{i+1}^{L},\boldsymbol{\beta}_{i+1}^{U},\boldsymbol{\gamma}_{i+1},\boldsymbol{\gamma}^{\prime}_{i+1})$

9:
$\boldsymbol{s}_{\\!N-i\\!}\\!\leftarrow\\!\widehat{\boldsymbol{\beta}}^{U}_{\\!N-i\\!}\\!\odot\\!(\boldsymbol{s}_{\\!N-i+1\\!}+\textsc{\\!CMV\\!}(\boldsymbol{\phi}_{\\!N-i+1\\!},\boldsymbol{\beta}_{\\!N-i+1\\!}^{L},\boldsymbol{\beta}_{\\!N-i+1\\!}^{U},\boldsymbol{\gamma}_{\\!N-i+1\\!},\boldsymbol{\gamma}^{\prime}_{\\!N-i+1\\!}))$

10:
$\boldsymbol{\psi}\leftarrow\boldsymbol{v}\oslash\left(\boldsymbol{r}+\boldsymbol{s}\right)$

11:
$\boldsymbol{p}_{1}\leftarrow\textsc{CMV}(\boldsymbol{\psi}_{1},\;\boldsymbol{\alpha}_{1})$,
$\boldsymbol{q}_{M}\leftarrow\mathbf{0}_{N}$

12: for $i=1\;:\;M-1$ do

13:
$\boldsymbol{p}_{i+1}\leftarrow\widehat{\boldsymbol{\alpha}}^{L}_{i}\odot\boldsymbol{p}_{i}+\textsc{CMV}(\boldsymbol{\psi}_{i+1},\boldsymbol{\alpha}_{i+1}^{L},\boldsymbol{\alpha}_{i+1}^{U},\boldsymbol{\gamma}_{i+1},\boldsymbol{\gamma}^{\prime}_{i+1})$

14:
$\boldsymbol{q}_{\\!N-i\\!}\\!\leftarrow\\!\widehat{\boldsymbol{\alpha}}^{U}_{\\!N-i\\!}\\!\odot\\!(\boldsymbol{q}_{\\!N-i+1\\!}+\textsc{\\!CMV\\!}(\boldsymbol{\psi}_{\\!N\\!-\\!i\\!+\\!1\\!},\boldsymbol{\alpha}^{L}_{\\!N-i+1\\!},\boldsymbol{\alpha}^{U}_{\\!N-i+1\\!},\boldsymbol{\gamma}_{\\!N-i+1\\!},\boldsymbol{\gamma}^{\prime}_{\\!N-i+1\\!}))$

15:
$\boldsymbol{\phi}\leftarrow\boldsymbol{u}\oslash\left(\boldsymbol{p}+\boldsymbol{q}\right)$

16: for $i=1\;:\;M-1$ do

17:
$\widehat{\boldsymbol{\alpha}}^{L}_{i}\leftarrow\lambda_{2}\left(\boldsymbol{\phi}_{i+1}/\boldsymbol{\phi}_{i}\right)\odot\widehat{\boldsymbol{\alpha}}^{L}_{i}$,
$\widehat{\boldsymbol{\beta}}^{L}_{i}\leftarrow\lambda_{2}\left(\boldsymbol{\psi}_{i+1}/\boldsymbol{\psi}_{i}\right)\odot\widehat{\boldsymbol{\beta}}^{L}_{i}$

18:
$\widehat{\boldsymbol{\alpha}}^{U}_{i}\leftarrow\lambda_{2}\left(\boldsymbol{\phi}_{i}/\boldsymbol{\phi}_{i+1}\right)\odot\widehat{\boldsymbol{\alpha}}^{U}_{i}$,
$\widehat{\boldsymbol{\beta}}^{U}_{i}\leftarrow\lambda_{2}\left(\boldsymbol{\psi}_{i}/\boldsymbol{\psi}_{i\\!+\\!1}\right)\odot\widehat{\boldsymbol{\beta}}^{U}_{i}$

19: Update
$\boldsymbol{\gamma},\;\boldsymbol{\gamma}^{\prime},\;\boldsymbol{\alpha}^{L},\;\boldsymbol{\beta}^{L},\;\boldsymbol{\alpha}^{U},\;\boldsymbol{\beta}^{U}$
return $W_{1}(\boldsymbol{u},\boldsymbol{v})$

##  6 Numerical Experiments

In this section, we carry out three numerical experiments to evaluate the FS-2
algorithm, including one 1D example and two 2D examples. The true Wasserstein
metric $W_{LP}$ is obtained by solving the original OT (1) using interior-
point methods [11, 20]. In our experiments, for both IPOT and FS-2, the number
of inner loops is set as $L=20$ and the regularization parameter
$\delta^{(t)}$ is set to $1$. The number of iterations here is the total
number of loops: $\\#\mathrm{iteration}=\mathrm{itr\\_max}\times L$. In order
to deal with the difficulties caused by zeros, we utilize the rescaling method
in [23]:

(17)
$D(f,g)=W_{1}\left(\frac{\frac{f}{\left\|f\right\|}+\eta}{1+N\eta},\frac{\frac{g}{\left\|g\right\|}+\eta}{1+N\eta}\right),$

In the following, we refer to formula (17) for numerical stability with
$\eta=10^{-5}$. All the experiments are conducted on a platform with 128G RAM,
and one Intel(R) Xeon(R) Gold 5117 CPU @2.00GHz with 14 cores.

###  6.1 1D Gaussian distributions

We consider the Wasserstein-1 metric between two mixtures of 1D Gaussian
distributions:
$0.4\mathcal{N}\left(60,64\right)+0.6\mathcal{N}\left(40,36\right)$ and
$0.5\mathcal{N}\left(35,81\right)+0.5\mathcal{N}\left(70,81\right)$, which is
the experiment setting in [35]. Input vectors $\boldsymbol{u}$ and
$\boldsymbol{v}$ are generated by integration on the uniform discretization of
interval $[0,100]$ with node size $N$.

We first compare the convergence of FS-1 and FS-2 for $N=1000$. We tested
$100$ experiments, and each experiment was performed for $10,000$ iterations.
In Fig. 1, the differences of the Wasserstein-1 metric between the true
solution $W_{LP}$ and the numerical solutions generated by FS-1 and FS-2 are
depicted. As expected, as $\varepsilon$ decreases, the error of FS-1 decreases
gradually after the iterations converge. We can observe this for
$\varepsilon=1/20$ and $\varepsilon=1/80$. However, for $\varepsilon=1/320$,
we can not observe convergence. In fact, the error does not drop over $10,000$
iterations. This is because $\varepsilon$ is too small, making updates
extremely slow. In fact, after $20,000$ iterations, the result of
$\varepsilon=1/320$ will continue to drop, and the final error is smaller than
that of $\varepsilon=1/20$ and $\varepsilon=1/80$. However, in any case, the
results of FS-1 are far inferior to those of FS-2, both in terms of accuracy
and convergence rate.

Figure 1: The 1D Gaussian distribution problem. The errors between the
numerical results generated by FS-1 or FS-2 and the true Wasserstein-1 metric
w.r.t. number of iterations.

The averaged computational time of the IPOT method and the FS-2 algorithm is
given in Table 1 and Fig. 2 (left). Apparently, the FS-2 algorithm has
achieved an overwhelming advantage over the IPOT method in terms of
computational speed, and ensures that the transport plans of the two are
almost the same. This replicates the advantages of the FS-1 algorithm over the
Sinkhorn algorithm. According to the data fitting results, the empirical
complexity of the FS-2 algorithm is $O(N^{1.07})$, which is much smaller than
the $O(N^{2.40})$ complexity of the IPOT method. At last, we show the
computational time required to reach the absolute error of the Wasserstein-1
metric for $N=1,000$ in Fig. 2 (right). Clearly, the FS-2 algorithm has an
advantage of two orders of magnitude in computational time compared to the
IPOT method.

Figure 2: The 1D Gaussian distribution problem. Left: The comparison of computational time between the FS-2 algorithm and the IPOT method with different numbers of grid points $N$. Right: The computational time required to reach the absolute error of the Wasserstein-1 metric. Table 1: The 1D Gaussian distribution problem. The comparison between the IPOT method and the FS-2 algorithm with the different number of grid points $N$. Columns 2-4 are the averaged computational time of the two algorithms and the speed-up ratio of the FS-2 algorithm. Column 5 is the Frobenius norm of the difference between the transport plan computed by the two algorithms. N | Computational time (s) | Speed-up ratio | $\left\|P_{FS}-P\right\|_{F}$
---|---|---|---
FS-2 | IPOT
500 | $1.25\times 10^{-2}$ | $1.22\times 10^{0}$ | $9.76\times 10^{1}$ | $2.09\times 10^{-15}$
2000 | $4.95\times 10^{-2}$ | $3.73\times 10^{1}$ | $7.52\times 10^{2}$ | $6.65\times 10^{-16}$
8000 | $2.32\times 10^{-1}$ | $8.41\times 10^{2}$ | $3.63\times 10^{3}$ | $8.57\times 10^{-16}$

###  6.2 2D Random distributions

Next, we compute the Wasserstein-1 metric between two $N\times N$ dimensional
random vectors whose elements obey the uniform distribution on $(0,1)$.
Without loss of generality, we set $h_{1}=h_{2}=0.1$. We also tested $100$
experiments, and each experiment was performed for $10,000$ iterations. We
hope to test the performance of the FS-2 algorithm in 2D through this example.
The differences in the Wasserstein-1 metric between the true solution $W_{LP}$
and the numerical solutions generated by FS-1 and FS-2 are shown in Fig. 3.
From this, we can observe that FS-1 converges quickly for $\varepsilon=1/20$,
but the error is large. When $\varepsilon=1/80$, the iteration converges at
about $5,000$ steps. The error keep decreasing even after $10,000$ steps for
$\varepsilon=1/320$. However, their errors and convergence speed are not as
good as FS-2.

Figure 3: The 2D random distribution problem. The errors between the numerical
results generated by FS-1 or FS-2 and the true Wasserstein-1 metric w.r.t.
number of iterations.

The averaged computational time of the IPOT method and the FS-2 algorithm is
given in Table 2 and Fig. 4 (left). According to the data fitting results, the
empirical complexity of the FS-2 algorithm is $O(N^{2.05})$, which is much
smaller than the $O(N^{4.98})$ complexity of the IPOT method. The
computational time required to reach the absolute error of the Wasserstein-1
metric for $N\times N=32\times 32$ is also presented in Fig. 4 (right).
Similar to the previous subsection, we can also observe the huge computational
efficiency of the FS-2 algorithm over the IPOT method.

Figure 4: The 2D random distribution problem. Left: The comparison of computational time between the FS-2 algorithm and the IPOT method with different numbers of grid points $N$. Right: The computational time required to reach the absolute error of the Wasserstein-1 metric. Table 2: The 2D random distribution problem. The comparison between the IPOT method and the FS-2 algorithm with different total number of grid nodes $N\times N$. Columns 2-4 are the averaged computational time of the two algorithms and the speed-up ratio of the FS-2 algorithm. Column 5 is the Frobenius norm of the difference between the transport plan computed by the two algorithms. N$\times$N | Computational time (s) | Speed-up ratio | $\left\|P_{FS}-P\right\|_{F}$
---|---|---|---
FS-2 | IPOT
20$\times$20 | $2.24\times 10^{-2}$ | $4.88\times 10^{-1}$ | $2.18\times 10^{1}$ | $2.40\times 10^{-16}$
40$\times$40 | $7.74\times 10^{-2}$ | $1.34\times 10^{1}$ | $1.73\times 10^{2}$ | $1.52\times 10^{-16}$
80$\times$80 | $3.38\times 10^{-1}$ | $4.79\times 10^{2}$ | $1.42\times 10^{3}$ | $1.40\times 10^{-16}$
160$\times$160 | $1.42\times 10^{0}$ | $1.46\times 10^{4}$ | $1.03\times 10^{4}$ | $8.51\times 10^{-17}$
320$\times$320 | $6.31\times 10^{0}$ | $-$ | $-$ | $-$

###  6.3 Image matching problem

The final experiment tests the performance of our FS-2 algorithm for high-
resolution image matching. This is a successful application of the Optimal
Transport [4, 17, 7]. We select two images from the DIV2K dataset [1]. Through
a process similar to Subsection 5.4 in the manuscript [24], we compute the
Wasserstein-1 metric between the two images. The differences in the
Wasserstein-1 metric between the true solution $W_{LP}$ and the numerical
solutions generated by FS-1 and FS-2 are depicted in Fig. 6. We also present
the averaged computational time of the IPOT method and the FS-2 algorithm in
Table 3. Moreover, the computational time required to reach the absolute error
of the Wasserstein-1 metric for $N\times N=32\times 32$ is shown in Fig. 7.
From these results, we can get the same conclusion as before, that is, the
FS-2 algorithm seems to be the numerical algorithm with the fastest
convergence and the lowest complexity for computing the Wasserstein-1 metric.

Figure 5: The image matching problem. Illustration of images. Figure 6: The image matching problem. The errors between the numerical results generated by FS-1 or FS-2 and the true Wasserstein-1 metric w.r.t. number of iterations. Table 3: The image matching problem. The comparison between the IPOT method and the FS-2 algorithm with the different total number of grid nodes $N\times N$. Columns 2-4 are the averaged computational time of the two algorithms and the speed-up ratio of the FS-2 algorithm. Column 5 is the Frobenius norm of the difference between the transport plan computed by the two algorithms. $N\times N$ | Computational time (s) | Speed-up ratio | $\left\|P_{FS}-P\right\|_{F}$
---|---|---|---
FS-2 | IPOT
100$\times$100 | $5.16\times 10^{-1}$ | $1.44\times 10^{3}$ | $2.79\times 10^{3}$ | $6.44\times 10^{-17}$
200$\times$200 | $2.31\times 10^{0}$ | $3.69\times 10^{4}$ | $1.60\times 10^{4}$ | $4.65\times 10^{-17}$
400$\times$400 | $9.69\times 10^{0}$ | $-$ | $-$ | $-$
800$\times$800 | $4.18\times 10^{1}$ | $-$ | $-$ | $-$
Figure 7: The image matching problem. The computational time required to reach
the absolute error of the Wasserstein-1 metric.

##  7 Conclusion

As the follow-up of the FS-1 paper, we generalize the result of matrix-vector
multiplication at $O(N)$ costs for the special matrix to the more general
L/U-CoLT matrix. We illustrate that only two vectors are required to represent
any L/U-CoLT matrix. Moreover, we also prove the closure of families of
L/U-CoLT matrices to matrix Hadamard product and matrix scaling. Therefore,
the above matrix operations are essentially updating the representation
vectors, which reduce both time and space complexity to $O(N)$. These results
can be directly applied to the Inexact Proximal point method for Optimal
Transport problem and reduce the overall computational complexity to O(N).
From this, we develop the Fast Sinkhorn II algorithm. It does not seem to be
an overstatement that for the computation of the Wasserstein-1 metric, we have
probably obtained the most competitive method, both in terms of convergence
speed and computational complexity.

## References

  * [1] E. Agustsson and R. Timofte, Ntire 2017 challenge on single image super-resolution: Dataset and study, in The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, July 2017.
  * [2] J.-D. Benamou and Y. Brenier, A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem, Numer. Math., 84 (2000), pp. 375–393.
  * [3] J.-D. Benamou, B. D. Froese, and A. M. Oberman, Numerical solution of the optimal transportation problem using the Monge–Ampère equation, J. Comput. Phys., 260 (2014), pp. 107–126.
  * [4] M. Burger, M. Franek, and C.-B. Schönlieb, Regularized regression and density estimation based on optimal transport, Appl. Math. Res. Express, 2012 (2012), pp. 209–253.
  * [5] G. Buttazzo, L. De Pascale, and P. Gori-Giorgi, Optimal-transport formulation of electronic density-functional theory, Phys. Rev. A, 85 (2012), p. 062502.
  * [6] J. Chen, Y. Chen, H. Wu, and D. Yang, The quadratic Wasserstein metric for earthquake location, J. Comput. Phys., 373 (2018), pp. 188–209.
  * [7] P. Clarysse, B. Delhay, M. Picq, and J. Pousin, Optimal extended optical flow subject to a statistical constraint, J. Comput. Appl. Math., 234 (2010), pp. 1291–1302.
  * [8] P. L. Combettes and J.-C. Pesquet, Proximal splitting methods in signal processing, in Fixed-Point Algorithms for Inverse Problems in Science and Engineering, Springer, 2011, pp. 185–212.
  * [9] C. Cotar, G. Friesecke, and C. Klüppelberg, Density functional theory and optimal transportation with Coulomb cost, Comm. Pure Appl. Math., 66 (2013), pp. 548–599.
  * [10] M. Cuturi, Sinkhorn distances: Lightspeed computation of optimal transport, in Advances in Neural Information Processing Systems, vol. 26, 2013, pp. 2292–2300.
  * [11] I. Dikin, Iterative solution of problems of linear and quadratic programming, Dokl. Akad. Nauk, 174 (1967), pp. 747–748.
  * [12] B. Engquist, K. Ren, and Y. Yang, The quadratic Wasserstein metric for inverse data matching, Inverse Problems, 36 (2020), p. 055001.
  * [13] J. Franklin and J. Lorenz, On the scaling of multidimensional matrices, Linear Algebra Appl., 114 (1989), pp. 717–735.
  * [14] B. D. Froese, Numerical methods for the elliptic Monge-Ampère equation and optimal transport, PhD thesis, Simon Fraser University, Burnaby, BC, Canada, 2012.
  * [15] B. D. Froese and A. M. Oberman, Convergent finite difference solvers for viscosity solutions of the elliptic Monge–Ampère equation in dimensions two and higher, SIAM J. Numer. Anal., 49 (2011), pp. 1692–1714.
  * [16] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio, Generative adversarial nets, in Advances in Neural Information Processing Systems, vol. 27, 2014.
  * [17] S. Haker, L. Zhu, A. Tannenbaum, and S. Angenent, Optimal mass transport for registration and warping, Int. J. Comput. Vis., 60 (2004), pp. 225–240.
  * [18] H. Heaton, S. W. Fung, A. T. Lin, S. Osher, and W. Yin, Wasserstein-based projections with applications to inverse problems, arXiv preprint arXiv:2008.02200, (2020).
  * [19] Y. Hu, H. Chen, and X. Liu, A global optimization approach for multi-marginal optimal transport problems with Coulomb cost, arXiv preprint arXiv:2110.07352, (2021).
  * [20] N. Karmarkar, A new polynomial-time algorithm for linear programming, in Proceedings of the sixteenth annual ACM symposium on Theory of computing, 1984, pp. 302–311.
  * [21] W. Li, E. K. Ryu, S. Osher, W. Yin, and W. Gangbo, A parallel method for earth mover’s distance, J. Sci. Comput., 75 (2018), pp. 182–197.
  * [22] X. Li, D. Sun, and K.-C. Toh, An asymptotically superlinearly convergent semismooth Newton augmented Lagrangian method for linear programming, SIAM J. Optim., 30 (2020), pp. 2410–2440.
  * [23] Z. Li, Y. Tang, J. Chen, and H. Wu, The quadratic Wasserstein metric with squaring scaling for seismic velocity inversion, arXiv preprint arXiv:2201.11305, (2022).
  * [24] Q. Liao, J. Chen, Z. Wang, B. Bai, S. Jin, and H. Wu, Fast Sinkhorn I: An $O(N)$ algorithm for the Wasserstein-1 metric, Commun. Math. Sci., accepted, (2022).
  * [25] A. T. Lin, W. Li, S. Osher, and G. Montúfar, Wasserstein proximal of GANs, in International Conference on Geometric Science of Information, Springer, 2021, pp. 524–533.
  * [26] C. Meng, Y. Ke, J. Zhang, M. Zhang, W. Zhong, and P. Ma, Large-scale optimal transport map estimation using projection pursuit, in Advances in Neural Information Processing Systems, vol. 32, 2019, pp. 8118–8129.
  * [27] C. Meng, J. Yu, J. Zhang, P. Ma, and W. Zhong, Sufficient dimension reduction for classification using principal optimal transport direction, in Advances in Neural Information Processing Systems, vol. 33, 2020.
  * [28] L. Métivier, R. Brossier, Q. Merigot, É. Oudet, and J. Virieux, An optimal transport approach for seismic tomography: Application to 3D full waveform inversion, Inverse Problems, 32 (2016), p. 115008.
  * [29] O. Museyko, M. Stiglmayr, K. Klamroth, and G. Leugering, On the application of the Monge–Kantorovich problem to image registration, SIAM J. Imaging Sci., 2 (2009), pp. 1068–1097.
  * [30] O. Pele and M. Werman, Fast and robust earth mover’s distances, in 2009 IEEE 12th International Conference on Computer Vision, IEEE, 2009, pp. 460–467.
  * [31] G. Peyré, M. Cuturi, et al., Computational optimal transport: With applications to data science, Found. Trends Mach. Learn., 11 (2019), pp. 355–607.
  * [32] Y. Rubner, C. Tomasi, and L. J. Guibas, The earth mover’s distance as a metric for image retrieval, Int. J. Comput. Vis., 40 (2000), pp. 99–121.
  * [33] F. Santambrogio, Optimal transport for applied mathematicians: Calculus of variations, pdes, and modeling, Progr. Nonlinear Differential Equations Appl., Birkäuser, Basel, 2015.
  * [34] R. Sinkhorn, Diagonal equivalence to matrices with prescribed row and column sums, Amer. Math. Monthly, 74 (1967), pp. 402–405.
  * [35] Y. Xie, X. Wang, R. Wang, and H. Zha, A fast proximal point method for computing exact Wasserstein distance, in Uncertainty in artificial intelligence, PMLR, 2020, pp. 433–453.
  * [36] L. Yang, J. Li, D. Sun, and K.-C. Toh, A fast globally linearly convergent algorithm for the computation of Wasserstein barycenters., J. Mach. Learn. Res., 22 (2021), pp. 1–37.
  * [37] Y. Yang, B. Engquist, J. Sun, and B. F. Hamfeldt, Application of optimal transport and the quadratic Wasserstein metric to full-waveform inversion, Geophysics, 83 (2018), pp. R43–R62.


* Frankle and Carbin (2019) Frankle, J. and Carbin, M. (2019). The lottery ticket hypothesis: Finding sparse, trainable neural networks.
  * Gal and Ghahramani (2016) Gal, Y. and Ghahramani, Z. (2016). Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of the 33rd International Conference on Machine Learning (ICML), p. 1050–1059.
  * Gelman et al. (1995) Gelman, A., Carlin, J. B., Stern, H. S. and Rubin, D. B. (1995). Bayesian Data Analysis. Chapman & Hall/CRC, 1st edn.
  * Havasi et al. (2020) Havasi, M., Jenatton, R., Fort, S., Liu, J. Z., Snoek, J., Lakshminarayanan, B., Dai, A. M. and Tran, D. (2020). Training independent subnetworks for robust prediction.
  * Hinton and Van Camp (1993) Hinton, G. E. and Van Camp, D. (1993). Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the sixth annual conference on Computational learning theory, pp. 5–13.
  * Kingma and Ba (2014) Kingma, D. P. and Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv e-prints, arXiv:1412.6980.
  * Lázaro-Gredilla and Figueiras-Vidal (2010) Lázaro-Gredilla, M. and Figueiras-Vidal, A. R. (2010). Marginalized neural network mixtures for large-scale regression. IEEE transactions on neural networks, 21 (8), 1345–1351.
  * Liu et al. (2020) Liu, J. Z., Lin, Z., Padhy, S., Tran, D., Bedrax-Weiss, T. and Lakshminarayanan, B. (2020). Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. arXiv preprint arXiv:2006.10108.
  * Lizotte (2008) Lizotte, D. J. (2008). Practical Bayesian Optimization. Ph.D. thesis, CAN, aAINR46365.
  * MacKay (1992) MacKay, D. J. (1992). A practical bayesian framework for backpropagation networks. Neural computation, 4 (3), 448–472.
  * Metropolis et al. (1953) Metropolis, N., Rosenbluth, A. W., Rosenbluth, M. N., Teller, A. H. and Teller, E. (1953). Equation of state calculations by fast computing machines. The Journal of Chemical Physics, 21 (6), 1087–1092.
  * Močkus (1975) Močkus, J. (1975). On bayesian methods for seeking the extremum. In G. I. Marchuk (ed.), Optimization Techniques IFIP Technical Conference Novosibirsk, July 1–7, 1974, Berlin, Heidelberg: Springer Berlin Heidelberg, pp. 400–404.
  * Neal (2012) Neal, R. M. (2012). Bayesian learning for neural networks, vol. 118. Springer Science & Business Media.
  * Ober and Rasmussen (2019) Ober, S. W. and Rasmussen, C. E. (2019). Benchmarking the neural linear model for regression. In Advances in Approximate Bayesian Inference (AABI).
  * Orloff and Bloom (2014) Orloff, J. and Bloom, J. (2014). The frequentist school of statistics, class 17.
  * Pedregosa et al. (2011) Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapeau, D., Brucher, M., Perrot, M. and Duchesnay, E. (2011). Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12, 2825–2830.
  * Pinsler et al. (2019) Pinsler, R., Gordon, J., Nalisnick, E. and Hernández-Lobato, J. M. (2019). Bayesian batch active learning as sparse subset approximation. In Advances in Neural Information Processing Systems 32, pp. 6359–6370.
  * Ranganath et al. (2013) Ranganath, R., Gerrish, S. and Blei, D. M. (2013). Black box variational inference.
  * Rasmussen and Williams (2006) Rasmussen, C. and Williams, C. (2006). Gaussian Processes for Machine Learning. Adaptive Computation and Machine Learning, Cambridge, MA, USA: MIT Press.
  * Riquelme et al. (2018) Riquelme, C., Tucker, G. and Snoek, J. (2018). Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. In International Conference on Learning Representations (ICLR).
  * Ritter et al. (2018) Ritter, H., Botev, A. and Barber, D. (2018). A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018-Conference Track Proceedings, International Conference on Representation Learning, vol. 6.
  * Ross et al. (2018) Ross, A. S., Pan, W. and Doshi-Velez, F. (2018). Learning qualitatively diverse and interpretable rules for classification. In International Conference on Machine Learning (ICML): Workshop on Human Interpretability in Machine Learning.
  * Shi et al. (2019) Shi, J., Khan, M. E. and Zhu, J. (2019). Scalable training of inference networks for gaussian-process models. In International Conference on Machine Learning, PMLR, pp. 5758–5768.
  * Snelson and Ghahramani (2006) Snelson, E. and Ghahramani, Z. (2006). Sparse gaussian processes using pseudo-inputs. In Advances in neural information processing systems, pp. 1257–1264.
  * Snoek et al. (2015) Snoek, J., Rippel, O., Swersky, K., Kiros, R., Satish, N., Sundaram, N., Patwary, M. M. A., Prabhat and Adams, R. P. (2015). Scalable bayesian optimization using deep neural networks. In Proceedings of the 32nd International Conference on Machine Learning (ICML), pp. 2171–2180.
  * Springenberg et al. (2016) Springenberg, J. T., Klein, A., Falkner, S. and Hutter, F. (2016). Bayesian optimization with robust bayesian neural networks. In Advances in Neural Information Processing Systems 29, pp. 4134–4142.
  * Sun et al. (2019) Sun, S., Zhang, G., Shi, J. and Grosse, R. (2019). Functional variational bayesian neural networks. In International Conference on Learning Representations.
  * Wang et al. (2019) Wang, K., Pleiss, G., Gardner, J., Tyree, S., Weinberger, K. Q. and Wilson, A. G. (2019). Exact gaussian processes on a million data points. In Advances in Neural Information Processing Systems, pp. 14648–14659.
  * Wilson and Nickisch (2012) Wilson, A. and Nickisch, H. (2012). Kernel interpolation for scalable structured gaussian processes (kiss-gp). In 32nd International Conference on Machine Learning (ICML), PCMLR, vol. 32.
  * Wilson et al. (2015) Wilson, A. G., Dann, C. and Nickisch, H. (2015). Thoughts on massively scalable gaussian processes.
  * Wilson et al. (2016) —, Hu, Z., Salakhutdinov, R. and Xing, E. P. (2016). Deep kernel learning. In Artificial intelligence and statistics, pp. 370–378.
  * Yao et al. (2019) Yao, J., Pan, W., Ghosh, S. and Doshi-Velez, F. (2019). Quality of uncertainty quantification for bayesian neural network inference. In International Conference on Machine Learning (ICML): Workshop on Uncertainty & Robustness in Deep Learning.
  * Zhou and Precioso (2019) Zhou, W. and Precioso, F. (2019). Adaptive Bayesian linear regression for automated machine learning. Tech. Rep. arXiv:1904.00577 [cs.LG], ArXiV.

##  Appendix A Appendix to Chapter 3

In proposing a novel framework, it is important to understand the effect of
hyperparameters. Here, the effect of various hyperparamters, and architectures
are explored, as well as the effect of random restarts. It is important to
note this is not an exhaustive exploration.

###  1 Neural Linear Model

This experiment was run with a 2-layer ReLU network with 50 and 20 neurons in
the first and second layers respectively (20 features). We used MAP training.
With no regularization and very noisy priors, the NLM is able to model in-
between uncertainty, albeit inconsistently, seen in Figure 1. With
regularization, we see the priors are not expressive enough and the NLM fails
to ever capture in-between uncertainty.

Figure 1: NLM is inconcistent in capturing the in-between uncertainty.

This experiment was run with a 2-layer ReLU network with 50 neurons in the
first layer and without regularization. We used MAP training. The number of
neurons in the second layer correspond to the number of features. We see
clearly as model capacity increases NLM better fits the data, seen in Figure
2. However, this increased capacity still fails to consistently model in-
between uncertainty.

Figure 2: NLM is inconsistent in capturing hte in-between uncertainty even
with increase capacity.

This experiment was run with a 2-layer ReLU network with 50 neurons in the
first hidden layers and 20 neurons in the last hidden layer (20 features). We
used MAP training. We see that NLM is able to capture more complex
relationships as capacity increases, but this increased capacity does not lead
to consistent in-between uncertainty in Figure 3. Additionally, these added
layers make NLM susceptible to overfitting.

Figure 3: Deeper NLMs are still inconsistent in capturing in-between
uncertainty.

This experiment was run with a 2-layer ReLU network with 50 and 20 neurons in
the first and second layers respectively (20 features). We used marginal
likelihood training and a smaller $\gamma=0.1$. We see that this NLM is able
to capture higher in-between uncertainty when $\alpha$ is high enough in
Figure 5, but is inconsistent in doing so.

Figure 4: NLM is inconsistent in capturing the in-between uncertainty, even
with higher $\alpha$.

This experiment was run with a 2-layer ReLU network with 50 and 20 neurons in
the first and second layers respectively (20 features). We used marginal
likelihood training and a larger $\gamma=1.0$. We see that this NLM is unable
to capture higher in-between uncertainty even when $\alpha$ is high in Figure
5.

Figure 5: NLM is inconsistent in capturing the in-btween uncertainty, even
with higher $\alpha$.

###  2 LUNA

This experiment was run with a 2-layer ReLU network with 50 and 20 neurons in
the first and second layers respectively (20 features). We used MAP training.
We see that LUNA is able to capture uncertainty with diverse regressors. The
regressors extrapolate neatly away from the data when there are 2 in Figure 6.
The regressors form a very diverse basis, with some fitting the data and
extrapolating away from each other, and others approximating the means on each
side of the gap.

Figure 6: LUNA’s auxiliary regressors consistently fit the data well and
extrapolate differently away from the data.

###  3 TUNA

This experiment was run with a 2-layer ReLU network with 50 and 20 neurons in
the first and second layers respectively (20 features). We used MAP training
and an RBF kernel with length scale $l$. We see that TUNA is able to
consistently capture uncertainty in the data scarce region when the auxiliary
regressors are diverse. When the auxiliary regressors are not diverse, we see
that TUNA fails to capture in-between uncertainty in Figure 7. 7.

Figure 7: TUNA is able to consistently capture in-between uncertainty with
appropriate kernel length scale.

##  Appendix B Appendix to Chapter 4

###  1 Experimental Results

####  1 Toy Experiments

Additional toy data experiments were run to showcase TUNA’s predictive
performance when compared to GPs.

####  2 UCI Gap

The raw values for epistemic uncertainty are given in Table 1. We see that
LUNA tends to have larger values of epistemic uncertainty in both the gap and
not gap regions. This is expected because we are explicitly penalizing the
auxiliary regressors to be diverse. All models have epistemic uncertainty
within the same order of magnitude, with BBVI tending to have the largest.

Avg. Epistemic Uncertainty Yacht - FROUDE Concrete - CEMENT Concrete - SUPER
Boston - RM Boston - LSTAT Boston - PTRATIO Not Gap Gap Not Gap Gap Not Gap
Gap Not Gap Gap Not Gap Gap Not Gap Gap Ens. Boot. 0.73 $\pm$ 0.11 0.58 $\pm$
0.03 3.63 $\pm$ 0.24 5.10 $\pm$ 0.15 4.46 $\pm$ 0.39 8.32 $\pm$ 0.23 1.54
$\pm$ 0.18 1.38 $\pm$ 0.04 1.59 $\pm$ 0.20 1.67 $\pm$ 0.08 1.61 $\pm$ 0.17
1.58 $\pm$ 0.08 Ensemble 0.39 $\pm$ 0.03 0.49 $\pm$ 0.02 2.10 $\pm$ 0.23 4.58
$\pm$ 0.21 1.92 $\pm$ 0.20 5.01 $\pm$ 0.15 0.92 $\pm$ 0.12 0.85 $\pm$ 0.02
0.92 $\pm$ 0.10 1.09 $\pm$ 0.03 0.90 $\pm$ 0.14 1.23 $\pm$ 0.04 Anc. Ens. 0.60
$\pm$ 0.23 0.63 $\pm$ 0.19 2.76 $\pm$ 0.38 5.25 $\pm$ 0.29 2.53 $\pm$ 0.36
6.23 $\pm$ 0.49 1.18 $\pm$ 0.16 1.06 $\pm$ 0.05 1.10 $\pm$ 0.10 1.20 $\pm$
0.08 1.13 $\pm$ 0.19 1.44 $\pm$ 0.07 NLM 0.12 $\pm$ 0.01 0.15 $\pm$ 0.02 0.77
$\pm$ 0.05 0.73 $\pm$ 0.04 0.74 $\pm$ 0.05 0.78 $\pm$ 0.04 0.37 $\pm$ 0.03
0.34 $\pm$ 0.03 0.90 $\pm$ 0.07 0.85 $\pm$ 0.06 0.76 $\pm$ 0.07 0.73 $\pm$
0.04 GP 0.94 $\pm$ 0.13 1.59 $\pm$ 0.10 3.04 $\pm$ 0.25 5.31 $\pm$ 0.07 2.76
$\pm$ 0.27 6.01 $\pm$ 0.13 1.88 $\pm$ 0.14 1.57 $\pm$ 0.02 1.76 $\pm$ 0.15
1.71 $\pm$ 0.06 1.93 $\pm$ 0.25 2.23 $\pm$ 0.08 MCD 1.61 $\pm$ 0.23 0.75 $\pm$
0.08 1.36 $\pm$ 0.07 1.45 $\pm$ 0.03 1.29 $\pm$ 0.08 1.44 $\pm$ 0.03 0.79
$\pm$ 0.08 0.65 $\pm$ 0.02 0.80 $\pm$ 0.07 0.70 $\pm$ 0.04 0.81 $\pm$ 0.07
0.77 $\pm$ 0.03 SNGP 0.07 $\pm$ 0.01 0.10 $\pm$ 0.05 0.33 $\pm$ 0.03 0.31
$\pm$ 0.03 0.28 $\pm$ 0.04 0.27 $\pm$ 0.04 0.29 $\pm$ 0.04 0.26 $\pm$ 0.04
0.56 $\pm$ 0.05 0.52 $\pm$ 0.06 0.29 $\pm$ 0.03 0.27 $\pm$ 0.02 BBVI 1.60
$\pm$ 0.11 1.43 $\pm$ 0.02 2.79 $\pm$ 0.16 2.28 $\pm$ 0.09 6.34 $\pm$ 0.35
7.00 $\pm$ 0.70 2.48 $\pm$ 0.20 2.13 $\pm$ 0.05 2.50 $\pm$ 0.18 1.70 $\pm$
0.08 5.70 $\pm$ 0.38 6.84 $\pm$ 0.45 LUNA 0.44 $\pm$ 0.10 0.68 $\pm$ 0.14 1.32
$\pm$ 0.15 2.05 $\pm$ 0.37 1.45 $\pm$ 0.31 7.29 $\pm$ 2.61 1.12 $\pm$ 0.09
1.00 $\pm$ 0.07 1.81 $\pm$ 0.25 2.30 $\pm$ 0.22 1.14 $\pm$ 0.09 1.83 $\pm$
0.52

Table 1: Average Epistemic Uncertainty in the Not Gap and Gap regions.

We do see that none of the models are able to detect the gap region for the
Boston - RM data set. We see in Figure 1 that the gap region is quite thin
relative to the rest of the data. The width of the gap region is approximately
0.5, where as the data has a range of 6. Additionally, the gap occurs in a
region of fairly steep signal. The combination of these two factors makes
detecting the gap region much more difficult than in other data sets, even
from the same parent data set, such as Boston - LSTAT, where the gap is
substantially wider, or Concrete - CEMENT, where the gap occurs in a
relatively flat region.

Figure 1: The Boston - RM UCI Gap Data Set is uniquely challenging for
detecting the gap region.

###  2 Experimental Setup

####  1 Synthetic Data

###### Cubic Gap Example

Following the set-up in (Ober and Rasmussen, 2019), we construct a synthetic
1-D dataset comprising 100 train and 100 test pairs $(x,y)$, where $x$ is
sampled uniformly in the range $[-4,-2]\cup[2,4]$ and $y$ is generated as
$y=x^{3}+\epsilon,\epsilon\sim\mathcal{N}(0,3^{2})$. All models were hand
tuned using an architecture of two hidden layers with 50 nodes and ReLU
activation.

###### Transfer Learning

The transfer learning experiment was done on the squiggle gap function seen in
Section 2. 100 data points were sampled from the not gap region for train,
test, and validation, split using an 80-10-10 split. 10 random samples were
taken for both the gap and not gap sections, and models were run on each
sample with three random restarts. Data was normalized for training and
prediction, reported results were then unnormalized. One hidden layer of 50
nodes, and one hidden layer of $N$ nodes, were used for both LUNA and NLM. One
layer of 50 nodes and $N$ Random Fourier Features were used for SNGP. $N$ is
the number of features, plotted along the x-axis in Section 2. The Adam
optimizer was used with a learning rate of $10^{-3}$ and 20000 epochs.
Hyperparameters were chosen over a grid with maximum log-likelihood on
validation data used as the model selection criterion for both random restarts
and on validation data. No minibatching was used.

For LUNA, prior variance was selected over $\log\alpha\in[-1,0,0]$,
regularization was selected over $\log\gamma\\{-6,-5,-4\\}$, diversity penalty
was selected over $\log\lambda\in\\{-4,-3\\}$, or the sigmoid annealing
schedule, perturbation standard deviation was selected over
$\log\sigma\in\\{-2,-1\\}$, and $N/2$ auxiliary regressors were used. For NLM,
prior variance was selected over $\log\alpha\in\\{-1,0,1\\}$, and
regularization was selected over $\log\gamma\in\\{-6,-5,-4\\}$. For SNGP,
prior variance was selected over $\log\alpha\in\\{-1,0,1\\}$, regularization
was selected over $\log\gamma\in\\{-6,-5,-4\\}$, and the normalizing factor
was selected over $c\in\\{1,5\\}$.

###### Radial Uncertainty Benchmark

Hyperparameters were selected for all models by first using a BayesOpt
procedure, then hand tuning each parameter, for each model, in each iteration
of the experiment. In one dimension, 10 random initializations with 20
iterations of BayesOpt were used. In two and three dimensions, five random
initializations and 10 iterations of BayesOpt were used. All models were
trained with 20,000 epochs using the Adam optimizer Kingma and Ba (2014) with
a learning rate of $10^{-}3$. All NN-based models had an architecture of two
layers of 50 hidden nodes with ReLU activation. The gold-standard GP used the
sum of a Matern Kernel starting at length scale 1 and a White Kernel with the
noise level set to the sampled data noise in one and two dimensions. In three
dimensions, the Matern kernel started with a length scale of 0.1. 10 random
restarts were used with the optimizer, as part of Scikit-Learn’s built-in
functionality Pedregosa et al. (2011). No minibatching was used.

For LUNA, a sigmoid diversity annealing schedule was used. The scaling factor
of the annealing, that is, $C$ in the equation
$f_{sigmoid}=C/(1+exp(-6x/N+3))$, was searched for in log space in the range
$\log C\in[-5,1]$. Regularization was searched for in the range
$\log\gamma\in[-10,2]$. Prior variance was searched for in the range
$\alpha\in[10^{-7},10^{2}]$. Perturbation standard deviation was searched for
in the range $\log\sigma\in[-2,0]$, and 25 auxiliary regressors were used. For
NLM, the prior variance was searched for in the range
$\alpha\in[10^{-8},10^{2}]$ and the regularization parameter was searched for
in the range $\log\gamma\in[-10,2]$. For TUNA, the prior variance was searched
for in the range $\alpha\in[10^{-7},10^{2}]$, the regularization was searched
for in the range $\log\gamma\in[-10,2]$, the kernel scaling factor was
searched for in the range $\log c\in[-7,2]$, the perturbation standard
deviation was searched for in the range $\log\sigma\in[-3,0]$, and the length
scale was searched for in the range $\log l\in[-3,1]$. For SNGP, the prior
variance was searched for in the range $[10^{-8},10^{2}]$, the regularization
was searched for in the range $\log\gamma\in[-10,2]$, and the normalizing
factor was searched for in the range $c\in[0,2]$. For Anchored Ensembles, 5
networks were used, with the initial variance, the data variance, and anchor
variance all searched for in log space over the range $[-9,1]$. For
Bootstrapped Ensembles, 20 networks were used and the regularizaion was
searched for in the range $\log\gamma\in[-9,1]$. For BBVI, the weight variance
was searched for in the range $\log\sigma^{2}\in[-9,1]$, and 1000 posterior
samples were used for prediction. For MCDropout, the regularization was
searched for in the range $\log\gamma\in[-9,1]$, the dropout was searched for
in the range $[0,1]$, and 100 output samples were used for prediction.

###### BayesOpt Example

For all models a three layer, 50 hidden node architecture with Tanh activation
was used unless noted otherwise. For the Branin and Hartmann6 functions, we
simply were searching for the global minimum. For the logistic regression and
SVM experiments, we first searched for hyperparameters on the validation set,
then reported accuracy on the test set. An 80-10-10 train-test-validation
split was used. Each experiment was run with 10 random restarts of three
different random initial guesses. A grid search over hyperparameters were
chosen once at the beginning of each experiment. The Adam optimizer with a
learning rate of $10^{-3}$ and 1000 epochs was used in each experiment. No
minibatching was used.

For LUNA, the prior variance was searched for in log space over
$\log\alpha\in[-1,0,1]$, the regularization was searched for in log space over
$\log\gamma\in[-5,-4,-3]$, the diversity penalty was searched over
$\lambda=0.1$, tanh, sigmoid, square root annealing schedules specified in
section 2, the perturbation standard deviation was searched for in
$\sigma\in[0.1,0.01]$, and 50 auxiliary regressors were used. Additional hand-
tuning was done for LUNA on the Hartmann6 experiment. For NLM, the prior
variance was searched for in log space over $\log\alpha\in[-1,0,1]$, the
regularization was searched for in $\log\lambda\in[-5,-4,-3]$. The GP length
scale was searched for in $l\in[0.001,0.01,0.1,1,5,10]$ with an RBF kernel.
SNGP prior variance was searched for over the range $\log\alpha\in[-1,0,1]$,
the normalizing factor was searched for in the range $c\in[1,5]$, no dropout,
and 200 Random Fourier Features were used. For Anchored Ensembles, 5 models
were used, and the anchor variance, data variance, and initial variance were
all searched for over the log space $[-1,0,1]$. For both bootstrap and regular
ensembles, 25 networks were used and regularization was searched for over the
range $\log\gamma\in[-5,-4,-3,-2,-1]$. For MCDropout, 5 forward passes were
used due to the multiple evaluations needed during acquisition function
optimization, dropout was searched for over the range $[0.01,0.05,0.1]$, and
regularization was searched for over the range $\log\gamma\in[-4,-3,-2]$. For
BBVI, a reduced architecture of one layer of 50 nodes was used due to
computational complexity and memory overhead, a learning rate of $10^{-2}$ was
also used,, dropout rate was selected in the range $\\{0.01,0.05,0.1,0.2\\}$
and regularization was automatically set to $1$ 1000 posterior samples were
used, weight variance was searched for over the range
$\sigma^{2}_{w}\in[1,10,50]$, and prior variance was searched for in the range
$\log\alpha\in[-1,0,1]$.

####  2 Real Data

###### UCI Regression

We used six UCI regression data sets to benchmark our models. Those are the
Boston Housing, Concrete Compressive Strength, Yacht Hydrodynamics, Energy
Efficiency, Abalone (Kin8nm), and Red Wine Quality datasets. For each
experiment, the data was split into 90% train data and 10% test data. The
train data was then split using the first 80% train, last 20% validation sets.
The Adam optimizer was used with 20,000 epochs and a learning rate of
$10^{-3}$. A two hidden layer, 50 node architecture with ReLU activation was
used for each NN-based model. Data noise was hand selected based on visual
examination of the data. Hyperparameters were selected using grid search, with
minimum RMSE being used to select the best random restart and model for
Anchored Ensembles, Bootstrap Ensembles, and Ensembles. Maximum log-liklihood
for Bayesian models, and minimum RMSE for non-Bayesian models, was used for
random restart and model selection for each of the other models. A minibatch
size of 32 was used unless noted otherwise.

For LUNA, prior variance was selected over $\log\alpha\in\\{-1,0,1\\}$,
regularization was selected over $\log\gamma\in\\{-5,-4,-3,-2,-1\\}$,
diversity penalty was selected over $\log\lambda\in\\{-8,-7,-6\\}$ and the
sigmoid annealing schedule, perturbation standard deviation was selected over
$\log\sigma\in\\{-2,-1\\}$, minibatch size of 128 was used. Additional hand
tuning was done for the Concrete and Boston data sets. For NLM, prior variance
was selected over $\log\alpha\in\\{-2,-1,0,1,2\\}$, regularization was
selected over $\log\gamma\in\\{-5,-4,-3,-2,-1\\}$. For GP, and RBF Kernel was
used and length scale was searched for in $\log l\in\\{-2,-1,0,1\\}$. For
SNGP, prior variance was selected over $\log\alpha\in\\{-2,-1,0,1,2\\}$,
regularization was selected over $\log\gamma\in\\{-5,-4,-3\\}$, normalizing
factor was searched for over $c\in\\{1,5,10\\}$, 200 Random Fourier Features
were used. For Anchored Ensembles, 5 models were used, and the anchor
variance, data variance, and initial variance were all searched for over the
log space $\\{-2,-1,0\\}$. For both Bootstrap Ensembles and Ensembles, 20
networks were used and regularization was searched for over the range
$\log\gamma\in\\{-5,-4,-3,-2,-1\\}$. For BBVI, 1000 posterior samples were
used, weight variance was searched for over the range
$\sigma^{2}_{w}\in[1,10,50]$, and prior variance was searched for in the range
$\log\alpha\in[-1,0,1]$. For MCDropout, dropout rate was selected in the range
$p\in\\{0.01,0.05,0.1,0.2\\}$ and regularization was automatically set to
$1/\sigma^{2}_{data}$.

###### UCI Gap

We used 3 standard UCI (Dua and Graff, 2017) regression data sets and modify
them to create 6 “gap data-sets”, wherein we purposefully created a gap in the
data where we can test our model’s in-between uncertainty (i.e. we train our
model on the non-gap data and test the model’s epistemic uncertainty on the
gap data). We adapt the procedure from Foong et al. (2019b) to convert these
UCI data sets into UCI gap data sets. For a selected input dimension, we (1)
sort the data in increasing order in that dimension, and (2) remove middle
$1/3$ to create a gap. We specifically selected input dimensions that have
high correlation with the output in order to ensure that the learned model
should have epistemic uncertainty in the gap; that is, if we select a
dimension that is not useful for prediction, any model need not have increased
uncertainty in the gap. The features we selected are:

  * •

Boston Housing: “Rooms per Dwelling” (RM), “Percentage Lower Status of the
Population” (LSTAT), and “Parent Teacher Ratio” (PTRATIO)

  * •

Concrete Compressive Strength: “Cement” and “Superplasticizer”

  * •

Yacht Hydrodynamics: “Froude Number”

The not gap region of the data was then split into 10 different 80-10-10
train,test, validation splits. Final results are computed as the mean and
standard deviation of the predicions over all of the splits. The same
architecture was used here as in the UCI regression experiments above. Data
noise was hand tuned by maximizing validation log-likelihood of an NLM with
prior variance $0.1$ and regularization $10^{-5}$. Maximum log-liklihood for
Bayesian models, and minimum RMSE for non-Bayesian models, was used for random
restart and model selection for each of the other models. All models used a
batch size of 32 unless noted otherwise.

For LUNA, prior variance was selected over $\log\alpha\in\\{-1,0,1\\}$,
regularization was selected over $\log\gamma\in\\{-5,-4,-3,-2,-1\\}$,
diversity penalty was selected over $\log\lambda\in\\{-2,-1\\}$, the sigmoid,
and tanh annealing schedules, perturbation standard deviation was set to
$\log\sigma=-2$, minibatch size of 128 was used. Model selection was done by
first selecting the top 90th percentile runs based on validation log-
likelihood, then selecting the lowest validation diversity penalty from these
runs. For NLM, prior variance was selected over $\log\alpha\in\\{-1,0,1\\}$,
regularization was selected over $\gamma\in\\{-5,-4,-3,-2,-1\\}$. For GP, and
RBF Kernel was used and length scale was searched for in
$l\in\\{10^{-2},5\cdot 10^{-2},10^{-1},5\cdot^{-1},0,1,5,10,50,500\\}$. For
SNGP, prior variance was selected over $\log\alpha\in\\{-1,0,1\\}$,
regularization was selected over $\log\gamma\in\\{-5,-4,-3,-2,-1\\}$,
normalizing factor was searched for over $c\in\\{1,5\\}$, 200 Random Fourier
Features were used. For Anchored Ensembles, 5 models were used, and the anchor
variance, data variance, and initial variance were all searched for over the
log space $\\{-2,-1,0\\}$. For both Bootstrap Ensembles and Ensembles, 20
networks were used and regularization was searched for over the range
$\log\gamma\in\\{-5,-4,-3,-2,-1\\}$. For BBVI, 1000 posterior samples were
used, weight variance was searched for over the range
$\sigma^{2}_{w}\in\\{1,10,50\\}$, and prior variance was searched for in the
range $\log\alpha\in[-1,0,1]$. 50,000 training epochs and no random restarts
were used due to increased computational time. For MCDropout, dropout rate was
selected in the range $p\in\\{0.005,0.01,0.05,0.1\\}$ and regularization was
automatically set to $1/\sigma^{2}_{data}$.