\documentclass[10pt]{article}
\renewcommand{\baselinestretch}{1.5}
\usepackage{epsfig}
\usepackage{subfigure}
\usepackage{latexsym}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{amsthm}
\usepackage{epsfig}
\usepackage{bm}

\newcommand\s{{\bf s}}
\newcommand\x{{\bf x}} 
\newcommand\y{{\bf y}} 
\newcommand\z{{\bf z}} 
\newcommand\p{{\bf p}} 
\newcommand\e{{\bf e}} 
\renewcommand\d{{\bf d}} 
\renewcommand\a{{\bf a}} 
\renewcommand\b{{\bf b}} 
\renewcommand\c{{\bf c}} 
\newcommand\R{\mathbb{R}} 


\setlength{\textheight}{22cm} \setlength{\textwidth}{6.5in}
\setlength{\topmargin}{0pt} \setlength{\evensidemargin}{1pt}
\setlength{\oddsidemargin}{1pt} \setlength{\headsep}{10pt}
\parskip=2mm
\parindent=8mm

\newtheorem{theorem}{Theorem}
\newtheorem{corollary}[theorem]{Corollary}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{example}[theorem]{Example}
\newtheorem{observation}[theorem]{Observation}
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{claim}[theorem]{Claim}
\newtheorem{fact}[theorem]{Fact}
\newtheorem{assumption}[theorem]{Assumption}
\newtheorem{remark}[theorem]{Remark}
\begin{document}
\title{CME307\&E311 Course Project IV: Computing Wasserstein Barycenter via Linear Programming}
\author{Yinyu Ye}
\maketitle

In this project, we study the computation the Wasserstein barycenter of a set of discrete probability measures. Given support points of probability measures in a metric space and
a transportation cost function (e.g. the Euclidean distance), Wasserstein distance defines a distance between two measures as the minimal transportation cost between them. Given a set of measures in the same space, the $p$-Wasserstein barycenter is defined as the measure minimizing the sum of $p$-Wasserstein distances to all measures in the set. Note that computing the barycenter of a set of discrete measures can be formulated by linear programming. 

In this project, we focus on the case of $p=2$ and compare the performance of different interior-point methods in solving the barycenter problem. We refer the notations and model setup to \cite{DHZY}. Instead of running experiments on MNIST dataset, we first restrict our attention to the algorithmic side and consider the following way in specifying the distributions $\mathcal{P}^{(t)}$.
\begin{itemize}
\item Generate $m_t$ samples from normal distribution $\mathcal{N}(\mu_t, \sigma_t^2)$ and construct $\mathcal{P}^{(t)}$ as the empirical distribution on the $m_t$ samples.
\item In the following experiments, you should vary the choice of the number of samples $m_t$, the number of distributions $N,$ and the parameters $(\mu_t, \sigma_t^2).$
\end{itemize}
In this way, the objective becomes finding the Wasserstein barycenter of $N$ normal distributions. Additionally, we first focus on the case of \textit{Pre-specified Support Problem} (See \cite{DHZY}) and choose the support of the barycenter distribution $\mathcal{P}$ be the union of the supports of $\mathcal{P}^{(t)}$'s.

\textbf{Question 1}: Implement central-path method for this problem. Alternatively, you may also implement steepest descent method instead. Clearly state the linear program, include pseudo-code for your implementation, and report the parameters you use in the optimization algorithm such as step size. 

\textbf{Question 2:} Implement the predictor-corrector interior point method \cite{STY} for this problem with the single low-rank regularization method (SLRM) and double low-rank regularization method
(DLRM) in \cite{DHZY}. These two methods aim to reduce the cost of solving the Newton equations in interior point method. How does this algorithm compare to the vanilla implementation of interior point method in Question 1? Plot the barycenter distribution $\mathcal{P}$ for two to three problem instances (specifications of $N$, $m_t$, etc.). 

Now we consider the general \textit{Free Support Problem} where the support of distribution $\mathcal{P}$ is also a decision variable. In \cite{MC} and \cite{MA}, the authors proposed an entropy-smoothed version of Wasserstein distance for both regularization and computation purpose. The new distance replaces the summand $\langle D^{(t)}, \Pi^{(t)} \rangle$ in (3) and (4) in \cite{DHZY} with 
$$\langle D^{(t)}, \Pi^{(t)} \rangle -\frac{1}{\lambda} h(\Pi^{(t)})$$
where $h(\cdot)$ is the entropy function. Intuitively, the entropy function will encourage the dispersion of the distribution $\Pi^{(t)}$ and avoid concentrations on a few points. Computationally, this new formulation enables a cheap computation of the gradients with respect to both the probability distribution parameters $(a_1,...,a_m)$ and the support $\bm{X} = (\bm{q}_1,...,\bm{q}_m)$ (in the language of \cite{DHZY}). 

\textbf{Question 3}: Implement Algorithm 3 in \cite{MA} with different choice of the regularization parameter $\gamma$. How does the resultant barycenter distribution compare with the ones obtained from interior-point methods? Note that Algorithm 3 in \cite{MA} considers a free-support setting while the two proceeding questions consider a pre-specified support. For a fair comparison, you may implement the free-support version of the interior-point method in \cite{DHZY}. Specifically, it will alternate between solving a linear program for the distribution parameter $(a_1,...,a_m)$ and solving a quadratic program for the support $\bm{X} = (\bm{q}_1,...,\bm{q}_m)$. The quadratic program features for an analytical solution as (7) in \cite{DHZY}.

As noted in these papers, the free support problem is then a non-convex problem. Now we are interested in how the gradient-based algorithm (Algorithm 3 in \cite{MA}) compares to the interior point method in respect with escaping saddle points and local minima. 

\textbf{Question 4}: Implement MAAIPM algorithm in \cite{DHZY} and compare it against Algorithm 3 in \cite{MA} in respect with the original objective function value
$$\sum_{t=1}^N \langle D^{(t)}, \Pi^{(t)} \rangle.$$
While implementing Algorithm 3 in \cite{MA}, you may want to periodically increase the regularization parameter $\gamma$ to mitigate the effect of the additional penalty term. Please report the runtime, the number of iterations, and the objective value under both algorithms.

\textbf{Question 5}: Now you may migrate the experiments to the MNIST dataset\footnote{http://yann.lecun.com/exdb/mnist/} and Fashion MNIST dataset\footnote{https://www.kaggle.com/zalando-research/fashionmnist}. The advantage of using these datasets is that it can provide a more meaningful visualization of the barycenter distribution. 


\begin{thebibliography}{1}
\bibitem{DHZY} Dongdong Ge, Haoyue Wang, Zikai Xiong, Yinyu Ye.
\newblock Interior-Point Methods Strike Back: Solving the Wasserstein Barycenter Problem.
\newblock https://arxiv.org/abs/1905.12895.


\bibitem{STY}  Mizuno, Shinji, Michael J. Todd, and Yinyu Ye. 
\newblock On adaptive-step primal-dual interior-point algorithms for linear programming.
\newblock Mathematics of Operations research 18.4 (1993): 964-981.


\bibitem{MC} Marco Cuturi.
\newblock Sinkhorn distances: Lightspeed computation of optimal transport.
\newblock Advances in neural information processing systems, 2013.

\bibitem{MA} Marco Cuturi, Arnaud Doucet. 
\newblock Fast computation of Wasserstein barycenters.
\newblock Proceedings of the International Conference on Machine Learning 2014.

\end{thebibliography}

\end{document}

