---
abstract: 'We study an extension of optimal transport techniques to stationary Markov chains from a computational perspective. In this context, naively applying optimal transport to the stationary distributions of the Markov chains of interest would not capture the Markovian dynamics. Instead, we study a new problem, called the optimal transition coupling problem, in which the optimal transport problem is constrained to the set of stationary Markovian couplings satisfying a certain transition matrix condition. After drawing a connection between this problem and Markov decision processes, we prove that solutions can be obtained via the policy iteration algorithm. For settings with large state spaces, we also define a regularized problem, propose a faster, approximate algorithm, and provide bounds on the computational complexity of each iteration. Finally, we validate our theoretical results empirically, demonstrating that the approximate algorithm exhibits faster overall runtime with low error in a simulation study.'
author:
- |
    Kevin O’Connor\
    UNC-Chapel Hill\
    `koconn@live.unc.edu`\
    Kevin McGoff\
    UNC-Charlotte\
    `kmcgoff1@uncc.edu`\
    Andrew B. Nobel\
    UNC-Chapel Hill\
    `nobel@email.unc.edu`
bibliography:
- 'references.bib'
title: |
    Optimal Transport for Stationary Markov Chains\
    via Policy Iteration
---

Introduction {#sec:intro}
============

The application and computation of optimal transport plans has recently received a great deal of attention within the machine learning community. Recent applications of optimal transport include a variety of problems in generative modeling [@arjovsky2017wasserstein; @genevay2017learning; @salimans2018improving; @kolouri2018sliced; @deshpande2018generative] and supervised learning [@frogner2015learning; @janati2019wasserstein; @Luise2018DifferentialPO]. In this paper, we study the optimal transport (OT) problem in the case where the objects of interest are stationary Markov chains. In particular, we provide algorithms for computing solutions to a constrained form of the OT problem by combining existing ideas from computational OT with techniques from Markov decision processes.

[The principled extension of computational OT techniques to distributions capturing stochastic structure, for example processes with serial dependence, is an important problem in computational OT.]{} Indeed, several recent applications of OT, including modeling the growth of cell populations over time [@schiebinger2019optimal] and embedding natural language [@xu2018distilled], fit naturally into the framework of dependent processes. Furthermore, such extensions of OT open the door to applications in climate science, finance, epidemiology and other fields where observations possess temporal structure. The case of stationary Markov chains considered in this paper constitutes a step towards rigorous extensions of computational OT to dependent [processes]{}.

#### Main contributions

\(1) We propose the optimal transport problem for stationary Markov chains (OTC) in terms of transition couplings. (2) [We recast the OTC problem for Markov chains as a multichain, average-cost Markov decision process.]{} (3) We prove that the standard policy iteration algorithm converges to a global solution of this problem. (4) We leverage regularization techniques for computational optimal transport to derive a faster, approximate algorithm. (5) We provide initial validation for our theoretical results through a simulation study.

#### Notation

Let ${\mathbb{R}}_+$ be the non-negative reals and $\Delta_n = \{u \in \mathbb{R}^n_+ | \sum_{i=1}^n u_i = 1\}$ denote the probability simplex in ${\mathbb{R}}^n$. Given a metric space ${\mathcal{U}}$, let ${\mathcal{M}}({\mathcal{U}})$ denote the set of Borel probability measures on ${\mathcal{U}}$. For a vector $u \in {\mathbb{R}}^n$, let $\|u\|_\infty = \max_i |u_i|$ and $\|u\|_1 = \sum_i |u_i|$. [Occasionally we will treat matrices in ${\mathbb{R}}^{n\times n}$ as vectors in ${\mathbb{R}}^{n^2}$.]{}

Constraining the optimal transport problem {#sec:background_on_otc}
==========================================

The optimal transport problem is defined in terms of couplings. Given probability measures $\mu \in {\mathcal{M}}({\mathcal{U}})$ and $\nu \in {\mathcal{M}}({\mathcal{V}})$ defined on [metric spaces]{} ${\mathcal{U}}$ and ${\mathcal{V}}$, a *coupling* of $\mu$ and $\nu$ is defined as a probability measure $\pi \in {\mathcal{M}}({\mathcal{U}}\times{\mathcal{V}})$ such that $\pi(A\times {\mathcal{V}}) = \mu(A)$ and $\pi({\mathcal{U}}\times B) = \nu(B)$ for every measurable $A \subset {\mathcal{U}}$ and $B \subset {\mathcal{V}}$. Letting $\Pi(\mu, \nu)$ denote the set of couplings of $\mu$ and $\nu$, the optimal transport problem with respect to a cost function $\tilde{c}: {\mathcal{U}}\times{\mathcal{V}}\rightarrow {\mathbb{R}}$ is defined as $$\label{eq:ot_problem}
\mbox{minimize} \,\, \int \tilde{c}\, d\pi \,\,  \mbox{subject to} \,\, \pi \in \Pi(\mu, \nu).$$ [As stated, the problem makes no particular assumptions about what the measures $\mu$ and $\nu$ describe. In most existing applications, $\mu$ and $\nu$ represent static quantities such as images or measurements of gene expression. However, in other application areas, the measures $\mu$ and $\nu$ may represent dependent processes. For example, $\mu$ and $\nu$ may correspond to sequences of words, the symptoms of a patient over time, or daily high temperatures. In these cases, additional care is needed in order to study .]{}

As a step toward computational OT for general dependent processes, we consider the case when $\mu$ and $\nu$ represent stationary Markov chains $X = (X_0, X_1, ...)$ and $Y = (Y_0, Y_1, ...)$ taking values in finite spaces ${\mathcal{X}}$ and ${\mathcal{Y}}$, respectively. Unlike general processes, which may exhibit infinitely long-range dependence, stationary Markov chains have a relatively simple dependence structure and are thus especially conducive to computation. Without loss of generality, we assume that ${\mathcal{X}}$ and ${\mathcal{Y}}$ both contain $d$ points. [Let $P, Q \in {\mathbb{R}}_+^{d \times d}$ be the respective transition matrices, and let $p, q \in \Delta_d$ be the respective stationary distributions, of the chains $X$ and $Y$.]{} For a review of basic Markov chain theory, we refer the reader to [@levin2017markov]. We also suppose that a non-negative cost function $c: {\mathcal{X}}\times {\mathcal{Y}}\rightarrow {\mathbb{R}}_+$ has been specified. We remark that this setting mirrors that of standard OT, where a cost or metric is defined *a priori* on the sample space.

[One may naively apply the standard formulation of the optimal transport problem in this setting, by taking ${\mathcal{U}}= {\mathcal{X}}$, ${\mathcal{V}}= {\mathcal{Y}}$, $\tilde{c}(x, y) = c(x,y)$, and considering the optimal coupling of the stationary distributions $p$ and $q$. Note that this marginal approach does not capture the dependence structure of $X$ and $Y$, and can lead to misleading conclusions. Consider for example the case when ${\mathcal{X}}= {\mathcal{Y}}= \{0, 1\}$, $c(x,y) = \delta(x \neq y)$, $$P= \kbordermatrix{&0&1 \\ 0 & \nicefrac{1}{2} & \nicefrac{1}{2} \\ 1 & \nicefrac{1}{2} & \nicefrac{1}{2}}, \quad\mbox{and}\quad Q = \kbordermatrix{&0&1 \\ 0 & 0& 1 \\ 1 & 1 & 0}.$$ Even though $X$ is iid and $Y$ is deterministic, their optimal transport distance is zero since $p = q$. Furthermore, the optimal coupling only specifies a distribution on ${\mathcal{X}}\times {\mathcal{Y}}$: it does not provide a means of generating a joint process having $X$ and $Y$ as marginals. We seek a variation of that captures and preserves the stochastic structure (stationarity, Markovity) of $X$ and $Y$.]{}

[As an alternative to the marginal approach, one may consider instead the full measures ${\mathbb{P}}\in {\mathcal{M}}({\mathcal{X}}^{\mathbb{N}})$ and ${\mathbb{Q}}\in {\mathcal{M}}({\mathcal{Y}}^{\mathbb{N}})$ of the processes $X$ and $Y$.]{} In particular, ${\mathbb{P}}$ is the unique probability measure on ${\mathcal{X}}^{\mathbb{N}}$ such that for any cylinder set $[a_i^j] := \{{\mathbf{x}}\in {\mathcal{X}}^{\mathbb{N}}: x_k = a_k, i \leq k \leq j\}$, $${\mathbb{P}}([a_i^j]) := p(a_i) \prod\limits_{k=i+1}^j P(a_{k-1}, a_k),$$ and ${\mathbb{Q}}$ is defined similarly in terms of $q$ and $Q$. Then one may let ${\mathcal{U}}= {\mathcal{X}}^{\mathbb{N}}$, ${\mathcal{V}}= {\mathcal{Y}}^{\mathbb{N}}$, ${\mathbf{x}}= (x_0, x_1, ...)$, ${\mathbf{y}}= (y_0, y_1, ...)$, $\tilde{c}({\mathbf{x}}, {\mathbf{y}}) = c(x_0, y_0)$ and couple ${\mathbb{P}}$ and ${\mathbb{Q}}$, obtaining a probability measure on the joint sequence space $({\mathcal{X}}\times{\mathcal{Y}})^{\mathbb{N}}$. However, such a coupling need not be stationary or Markovian. [In order to capture the dynamics of $X$ and $Y$, we might restrict the feasible set to couplings of ${\mathbb{P}}$ and ${\mathbb{Q}}$ that have the same dependence structure as $X$ and $Y$, namely the family $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$ of stationary Markovian couplings. While this is a natural choice, the minimum expected cost over this set may violate the triangle inequality even when the cost is a metric [@ellis1976thedj; @ellis1978distances]. Moreover, the family $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$ is not characterized by a simple set of constraints [@boyle2009hidden].]{} For the sake of computational tractability, we require a subset of $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$ which admits a sufficiently simple, computationally tractable representation. To alleviate these issues, we further constrain the set of couplings to the subset of $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$ whose transition distributions are couplings of those of $X$ and $Y$. Note that, to reduce notation when considering vectors and matrices indexed by elements of ${\mathcal{X}}\times{\mathcal{Y}}$, we will indicate only the cardinality of the index set and adopt an indexing convention whereby a vector $u \in \mathbb{R}^{d^2}$ is indexed as $u(x, y)$ and a matrix $R \in {\mathbb{R}}^{d^2\times d^2}$ is indexed as $R((x,y),(x',y'))$ for $(x,y),(x',y') \in {\mathcal{X}}\times{\mathcal{Y}}$. Note further that we regard vectors of the form $R((x,y),\cdot)$ as row vectors.

\[def:transitioncouplings\] A paired chain in $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$ with transition matrix $R$ is called a **transition coupling** of $X$ and $Y$ if, for every $(x,y) \in {\mathcal{X}}\times {\mathcal{Y}}$, $R((x,y),\cdot) \in \Pi(P(x,\cdot), Q(y,\cdot))$. We denote the set of transition couplings of $X$ and $Y$ by ${\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}$ and, abusing notation slightly, the set of transition matrices satisfying the condition above by ${\Pi(P, Q)}$.

The couplings defined in Definition \[def:transitioncouplings\] are referred to as “Markovian couplings” in the literature [@levin2017markov] and have been used, for example, to study diffusions [@banerjee2018coupling; @banerjee2016coupling; @banerjee2017rigidity]. We refer to such couplings as “transition couplings” in order to distinguish them from elements of $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$. Note that ${\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}\neq \emptyset$ since it contains the independent coupling, that is, the distribution obtained by coupling $X$ and $Y$ independently.

A key advantage of considering ${\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}$ over $\Pi_M({\mathbb{P}}, {\mathbb{Q}})$ is that the constraints defining this set are linear and thus computationally tractable. In the case that $X$ and $Y$ are irreducible, this set of transition matrices actually characterizes the set of transition couplings.

[prop]{}[transmatchar]{} \[prop:transmat\_char\] Let $X$ and $Y$ be irreducible stationary Markov chains with transition matrices $P$ and $Q$, respectively. Then any stationary Markov chain with a transition matrix contained in ${\Pi(P, Q)}$ is a transition coupling of $X$ and $Y$.

For brevity, we will also use “transition couplings” to refer to elements of ${\Pi(P, Q)}$. Defining $\tilde{c}: ({\mathcal{X}}\times {\mathcal{Y}})^{\mathbb{N}}\rightarrow {\mathbb{R}}_+$ such that $\tilde{c}({\mathbf{x}}, {\mathbf{y}}) = c(x_0, y_0)$, we define the *optimal transition coupling (OTC) problem* by $$\label{eq:otc_problem_rv}
\mbox{minimize}\,\, \int \tilde{c} \,  d\pi \,\, \mbox{subject to} \,\, \pi \in {\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}.$$ As shown in Appendix \[app:existence\], the minimum in is achieved by an element of ${\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}$ under our assumptions. Problem involves the minimization of a linear objective over the non-convex set ${\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}$ and thus poses a significant computational challenge. However, Proposition \[prop:transmat\_char\] allows us to optimize instead over the convex polyhedron ${\Pi(P, Q)}$. Informally, can be restated as the minimization of ${\mathbb{E}}c(X_0', Y'_0)$ over $R \in {\Pi(P, Q)}$, where $(X', Y')$ is a stationary Markov chain generated by $R$. However, this reformulation of has a non-convex objective, so some care is needed in order to obtain global solutions.

Related Work {#sec:related_work}
------------

Stationary couplings of stationary processes, known as a *joinings*, were first studied in [@furstenberg1967disjointness]. Distances between processes based on joinings have been proposed in the ergodic theory literature [@ornstein1973application; @gray1975generalization], but have been explored primarily as a theoretical tool: no tractable algorithms have been proposed for computing such distances exactly. In the context of Markov chains, coupling methods have been widely used as a tool to establish rates of convergence [@griffeath1976coupling; @lindvall2002lectures]. Optimal Markovian couplings of Markov processes are studied in [@ellis1976thedj; @ellis1978distances; @ellis1980conditions; @ellis1980kamae]. Despite the theoretical progress, little work has been done to develop tractable algorithms for computing optimal couplings of Markov chains. In [@mufa1994optimal; @zhang2000existence; @song2016measuring] the authors consider a different, computationally simpler, form of the optimal transition coupling problem studied here, in which one minimizes the expected cost of the next step. We also remark that the optimal transition coupling problem appears in an unpublished manuscript of Aldous and Diaconis.

Other work has considered modifications of the Wasserstein distance for time series. The work [@cazelles2019wasserstein] studies the Wasserstein-Fourier distance, which is the Wasserstein distance between normalized power spectral densities, while [@muskulus2011wasserstein] suggest using the optimal transport cost between the $k$-block empirical measures constructed from observed samples. For general observed sequences, [@su2018order] consider only couplings that do not disturb the ordering of the two sequences too much, as quantified by the inverse difference moment. In contrast to these approaches, we seek a more direct modification of the optimal transport problem itself that best captures the Markovian dynamics.

Connection to Markov decision processes {#sec:connection_to_mdp}
=======================================

In the remainder of the paper, we focus on developing tractable algorithms for solving the OTC problem . We begin by making a connection between the OTC problem and Markov decision processes (MDP), which allows us to build upon existing techniques and algorithms in the MDP literature.

A Markov decision process is characterized by a 4-tuple $({\mathcal{S}}, {\mathcal{A}}, {\mathcal{P}}, c')$ consisting of a state space $\mathcal{S}$, an action space ${\mathcal{A}}= \bigcup_s {\mathcal{A}}_s$, a set of transition distributions ${\mathcal{P}}= \{p(\cdot | s, a): s \in {\mathcal{S}}, a \in {\mathcal{A}}\}$, and a cost function $c': \mathcal{S}\times {\mathcal{A}}\rightarrow\mathbb{R}$. At each time step the process occupies a state $s \in \mathcal{S}$ and an agent chooses an action $a \in {\mathcal{A}}_s$; the process then moves to a new state according to the distribution $p(\cdot | s, a)$, incurring a cost $c'(s, a)$. Informally, the goal of the agent is to choose actions in order to incur minimum average cost. The behavior of an agent is described by a family $\gamma = \{ \gamma_s(\cdot) : s \in  \mathcal{S} \}$ of distributions $\gamma_s(\cdot) \in {\mathcal{M}}(\mathcal{A}_s)$ on the set of admissible actions, which is known as a [*policy*]{}. An agent following policy $\gamma$ chooses her next action according to $\gamma_s(\cdot)$ whenever the system is in state $s$, independently of her previous actions.

It is easy to see that, in conjunction with the transition distributions ${\mathcal{P}}$, every policy $\gamma$ induces a stationary Markov chain on the state space $\mathcal{S}$ of the MDP. In the average-cost MDP problem the goal is to identify a policy for which the induced Markov chain minimizes the limiting average cost, namely a policy $\gamma$ minimizing $$\label{eq:mdp_average_cost}
\overline{c}_\gamma(s) := \lim\limits_{T \rightarrow\infty} \frac{1}{T}\sum\limits_{t=1}^T \mathbb{E}_\gamma \left[c'(s_t, a_t) \bigg| s_0 = s\right],$$ for each $s \in {\mathcal{S}}$. Note that the expectation in is taken with respect to the Markov chain induced by $\gamma$, and that the limit exists by the ergodic theorem. In general, the limiting average cost $\overline{c}_\gamma(s)$ will depend on the initial state $s$, but if $\gamma$ induces an ergodic chain then the average cost will be constant. If all policies induce ergodic Markov chains, the MDP is referred to as “unichain”; otherwise the MDP is classified as “multichain”. We refer the reader to [@puterman2005markov] for more details on MDP’s.

The OTC problem may readily be recast as an MDP. In detail, let the state space $\mathcal{S} = {\mathcal{X}}\times {\mathcal{Y}}$. Furthermore, letting $s = (x,y)$ denote an element of $\mathcal{S}$, define the set of admissible actions in $s$ as ${\mathcal{A}}_s = {\Pi(P(x,\cdot), Q(y,\cdot))}$, the transition distributions $p(\cdot | s, r_s) := r_s(\cdot)$ for $r_s \in {\mathcal{A}}_s$, and the cost function $c'(s, r_s) = c(x,y)$. Note that the cost function $c'$ is independent of the action $r_s$. We refer to this MDP as TC-MDP. Any policy $\gamma$ for TC-MDP specifies distributions over ${\Pi(P(x,\cdot), Q(y,\cdot))}$ for each $(x,y) \in {\mathcal{X}}\times{\mathcal{Y}}$ and thus corresponds to a single distribution over ${\Pi(P, Q)}$ that governs the random actions of the agent. In TC-MDP it suffices to consider only deterministic policies $\gamma$, namely those such that for each state $s = (x,y)$ the distribution $\gamma_s(\cdot)$ is a point mass at unique element of ${\mathcal{A}}_{s} = {\Pi(P(x,\cdot), Q(y,\cdot))}$.

[prop]{}[deterministicpolicy]{} \[prop:deterministic\_policy\] Let $\gamma$ be a policy for TC-MDP. Then there exists a deterministic policy $\tilde{\gamma}$ such that $\overline{c}_\gamma(s) = \overline{c}_{\tilde{\gamma}}(s)$ for every $s \in {\mathcal{S}}$.

As such, optimization over deterministic policies and over ${\Pi(P, Q)}$ are equivalent. Going forward, we refer to $R \in {\Pi(P, Q)}$ directly instead of the equivalent deterministic policy $\tilde{\gamma}$ in our notation. We briefly note that, even when $X$ and $Y$ are ergodic, the same may not be true of the stationary Markov chain induced by a transition coupling matrix $R \in {\Pi(P, Q)}$ (see Appendix \[sec:redicible\_tc\]). Specifically, a single element of ${\Pi(P, Q)}$ may have multiple stationary distributions and thus give rise to multiple stationary Markov chains depending on the initial state $s \in {\mathcal{S}}$. Thus TC-MDP is classified as multichain.

Finally, supposing that $X$ and $Y$ are irreducible, note the equivalence of the objective functions in TC-MDP and : For every $R \in {\Pi(P, Q)}$ and $s \in {\mathcal{S}}$, let $\pi_R^s \in {\mathcal{M}}(({\mathcal{X}}\times{\mathcal{Y}})^{\mathbb{N}})$ denote the measure of the stationary Markov chain generated by $R$ with stationary distribution arising from the initial state $s$. Note that by Proposition \[prop:transmat\_char\], $\pi_R^s \in {\Pi_{\tiny TC}({\mathbb{P}}, {\mathbb{Q}})}$. Then for every $s \in {\mathcal{S}}$, $$\overline{c}_R(s) = \lim\limits_{T\rightarrow\infty} \frac{1}{T} \sum\limits_{t=1}^T \sum\limits_{s_t} R^t(s,s_t) c(s_t) = \int \tilde{c} \, d\pi_R^s.$$ If $R^* \in {\Pi(P, Q)}$ is optimal in TC-MDP, then $\overline{c}_{R^*}(s) = \min_{R \in {\Pi(P, Q)}} \overline{c}_{R}(s)$ for every $s \in {\mathcal{S}}$. Letting $s^* \in \operatorname*{argmin}_{s \in {\mathcal{S}}} \overline{c}_{R^*}(s)$, $\pi_{R^*}^{s^*}$ is thus an optimal transition coupling of $X$ and $Y$. So any solution to TC-MDP necessarily yields a solution to .

Policy iteration for optimal transition couplings {#sec:policy_iteration}
=================================================

Now that we have shown that the OTC problem can be viewed as an MDP, we can leverage existing algorithms for MDP’s to obtain solutions. To this end, we propose to use the policy iteration algorithm [@howard1960dynamic] because of its favorable convergence properties and ease of implementation. To facilitate our discussion, in what follows, we regard the cost function $c$ and limiting average cost $\overline{c}_R$ as vectors in ${\mathbb{R}}_+^{d^2}$. For each $R \in {\Pi(P, Q)}$, standard results [@puterman2005markov] guarantee that the limit $\overline{R} := \lim_{T\rightarrow\infty} \nicefrac{1}{T} \sum_{t=0}^{T-1} R^t$ exists. When $R$ is aperiodic and irreducible, the Perron-Frobenius theorem implies that $\overline{R} = \lim_{T\rightarrow\infty} R^T$ and the rows of $\overline{R}$ are equal to the stationary distributions of $R$.

[r]{}[0.50]{}

\[alg:pia\] $R_0 \leftarrow P \otimes Q$, $n \leftarrow 0$

The policy iteration algorithm repeatedly [evaluates]{} and [improves]{} policies. For a given policy $R \in {\Pi(P, Q)}$ the evaluation step computes the average cost (*gain*) vector $g = \overline{R} \, c$ and the total extra cost (*bias*) vector $h = \sum_{t=0}^\infty R^t(c - g)$. In practice, $g$ and $h$ may be obtained by solving a linear system of equations rather than evaluating infinite sums (see Algorithm \[alg:exact\_pe\]). The policy improvement step selects a new transition coupling matrix $R'$ that minimizes $R' \, g$ or $R' \, h$ element-wise (see Algorithm \[alg:exact\_pi\]). In particular, one may select $R'$ such that for each $(x,y)$ the row $r = R'((x,y),\cdot)$ minimizes $r g$ or $r h$ over $r \in {\Pi(P(x,\cdot), Q(y,\cdot))}$ . Once a fixed point in the evaluation and improvement process is reached, the procedure terminates. The resulting policy iteration algorithm will be referred to as `ExactOTC` (see Algorithm \[alg:pia\]). We initialize Algorithm \[alg:pia\] to the independent transition coupling $P \otimes Q$, which satisfies $P\otimes Q((x,y),(x',y')) = P(x,x')Q(y,y')$.

\[alg:exact\_pe\] Solve for $(g, h, w)$ such that $$\left[\begin{array}{ccc} I - R & 0 &0 \\ I & I - R & 0 \\ 0 & I & I - R \end{array}\right] \left[\begin{array}{c} g \\ h \\ w \end{array}\right] = \left[\begin{array}{c} 0 \\ c \\ 0 \end{array}\right]$$

\[alg:exact\_pi\] $R' \leftarrow \operatorname*{argmin}_{R \in \Pi} R g$

For finite state and action spaces, policy iteration is known to yield an optimal policy for the average-cost MDP in a finite number of steps [@puterman2005markov]. While policy iteration may fail to converge for general compact action spaces [@dekker1987counter; @schweitzer1985undiscounted; @puterman2005markov], as is the case for TC-MDP, we may exploit the polyhedral structure of ${\Pi(P, Q)}$ to establish the following convergence result.

[thm]{}[convergenceofpia]{} \[thm:convergence\_of\_pia\] Algorithm \[alg:pia\] converges to a solution $(g^*, h^*, R^*)$ of TC-MDP in a finite number of iterations. Moreover, if $X$ and $Y$ are irreducible, $R^*$ is an optimal transition coupling of $X$ and $Y$.

Recall from the discussion in Section \[sec:connection\_to\_mdp\] that a solution to TC-MDP necessarily yields a solution to . Thus Theorem \[thm:convergence\_of\_pia\] ensures that a solution to the OTC problem can be obtained from Algorithm \[alg:pia\] in a finite number of iteration. A proof of this result can be found in Appendix \[app:pia\_convergence\].

Fast approximate policy iteration {#sec:fast_approx_pia}
=================================

The simplicity of Algorithm \[alg:pia\] along with Theorem \[thm:convergence\_of\_pia\] make it an appealing method for solving the OTC problem when $d$ is small. However, each call to Algorithm \[alg:exact\_pe\] involves solving a system of $3d^2$ linear equations, requiring a total of ${\mathcal{O}}(d^6)$ operations. Furthermore, each call to Algorithm \[alg:exact\_pi\] entails solving $d^2$ linear programs each with ${\mathcal{O}}(d)$ constraints, which can be accomplished in a total time of ${\mathcal{O}}(d^5 \log d)$. For even moderate $d$, this may be too slow for practical use. A similar dependence on the dimension of each coupling is observed in exact OT algorithms, such as the Network Simplex Algorithm [@peyre2019computational]. To alleviate the poor scaling with $d$, one may use entropic regularization, whereby a negative entropy term is added to the OT objective. Cuturi [@cuturi2013sinkhorn] showed that solutions to the regularized OT problem may be obtained efficiently via Sinkhorn’s algorithm [@sinkhorn1967diagonal]. More recently, [@altschuler2017near] proved that Sinkhorn’s algorithm yields an approximation of the OT cost with error bounded by $\varepsilon$ in near-linear time with respect to the dimension of the couplings under consideration. This represents the state of the art in terms of dependence on dimension for arbitrary discrete measures. Analogously, one may hope that a similar dependence on the size of elements of ${\Pi(P, Q)}$ may be achievable for each policy iteration when solving the OTC problem.

In this section, we explore the extension of entropic regularization techniques to the OTC problem and provide an approximate algorithm that runs in $\tilde{{\mathcal{O}}}(d^4)$ time per iteration, where $\tilde{{\mathcal{O}}}(\cdot)$ hides poly-logarithmic factors. This complexity is nearly-linear in the dimension $d^4$ of the transition couplings. Mirroring the derivation of entropic OT, we first introduce a constrained OTC problem in which we consider transition couplings that are close to the independent transition coupling. We show that this type of constraint induces beneficial regularity properties among the transition couplings in the constrained set. We also propose a truncation-based approximation of the `ExactPE` algorithm, which we call `ApproxPE`. Using the regularity of the constrained set, we show that one can obtain approximations of the gain and bias from `ApproxPE` with error bounded by $\varepsilon$ in $\tilde{{\mathcal{O}}}(d^4 \log \varepsilon^{-1})$ time. Furthermore, we propose an entropy-regularized approximation to the `ExactPI` algorithm, called `ApproxPI`. We perform a new analysis of the Sinkhorn algorithm to show that `ApproxPI` yields an improved transition coupling with error bounded by $\varepsilon$ in $\tilde{{\mathcal{O}}}(d^4 \varepsilon^{-4})$ time. Combining these two algorithms, we obtain the `FastEntropicOTC` algorithm, which runs in $\tilde{{\mathcal{O}}}(d^4 \varepsilon^{-4})$ time per iteration. Finally, we provide empirical support for our theoretical results through a simulation study. We find that the improved efficiency at each iteration of `FastEntropicOTC` leads to a much faster runtime in practice as compared to `ExactOTC`. Our experiments also show that `FastEntropicOTC` yields an expected cost that closely approximates the unregularized OTC cost.

Constrained Optimal Transition Coupling Problem
-----------------------------------------------

We begin by defining a constrained set of transition couplings. Let ${\mathcal{K}}(\cdot \| \cdot)$ be the Kullback-Leibler (KL) divergence, where for any $u, v \in \Delta_{d^2}$ with $u \ll v$, ${\mathcal{K}}(u \| v) = \sum_{s} u(s) \log(u(s) / v(s))$ (letting $0 \log(0/0) = 0$). Then for every $\eta > 0$ and $(x,y) \in {\mathcal{X}}\times {\mathcal{Y}}$, define the set $${\Pi_\eta(P(x,\cdot), Q(y,\cdot))}= \left\{ r \in {\Pi(P(x,\cdot), Q(y,\cdot))}: \, {\mathcal{K}}\big(r \| P\otimes Q((x,y),\cdot)\big) \leq \eta\right\},$$ and the subset of transition coupling matrices ${\Pi_\eta(P,Q)}= \{ R \in {\Pi(P, Q)}: R((x,y),\cdot) \in {\Pi_\eta(P(x,\cdot), Q(y,\cdot))}, \, \forall (x,y) \in {\mathcal{X}}\times {\mathcal{Y}}\}$. Thus, elements of ${\Pi_\eta(P,Q)}$ have rows that are close in KL-divergence to the rows of the independent transition coupling $P \otimes Q$. When $P$ and $Q$ are aperiodic and irreducible, we find that $P\otimes Q$ is as well. We prove that, for appropriate $\eta$, the proximity of each element in ${\Pi_\eta(P,Q)}$ to $P\otimes Q$ in KL-divergence is enough to establish two beneficial regularity properties for the entire set.

[prop]{}[propertiesofentropictc]{} \[prop:properties\_of\_entropic\_tc\] Let $P$ and $Q$ be aperiodic and irreducible. Then for $\eta$ small enough, every $R \in {\Pi_\eta(P,Q)}$ is aperiodic and irreducible and thus has a unique stationary distribution $\lambda_R \in \Delta_{d^2}$. Moreover, there exist constants $M < \infty$ and $\alpha \in (0, 1)$ such that for any $t \geq 1$, $$\max\limits_{R\in {\Pi_\eta(P,Q)}} \max\limits_{s \in {\mathcal{X}}\times{\mathcal{Y}}} \|R^t(s, \cdot) - \lambda_{R}\|_1 \leq M \alpha^t.$$

We give an explicit choice of $\eta$ in the proof of Proposition \[prop:properties\_of\_entropic\_tc\]. Now, let ${\Pi_{\tiny TC}^\eta({\mathbb{P}}, {\mathbb{Q}})}$ be the set of transition couplings with transition matrices in ${\Pi_\eta(P,Q)}$ and define the constrained OTC problem, $$\label{eq:entropic_otc}
\mbox{minimize} \,\, \int \tilde{c} \, d\pi \,\, \mbox{subject to} \,\, \pi \in {\Pi_{\tiny TC}^\eta({\mathbb{P}}, {\mathbb{Q}})}.$$ For completeness, we prove that a solution to exists in Appendix \[app:existence\]. Since the function $R \mapsto {\mathcal{K}}(R(s,\cdot) \| P\otimes Q(s,\cdot))$ is uniformly bounded over $R \in {\Pi(P, Q)}$ for each $s \in {\mathcal{X}}\times {\mathcal{Y}}$, coincides with the unconstrained OTC problem for large enough $\eta$. As such, we expect that, for large $\eta$, a solution to will be close in some sense to a solution of . While we do not provide a proof of this fact, we give some empirical evidence that this is the case in Section \[sec:experiments\]. Finally, note that corresponds to an MDP in the same way that does but with a constrained set of policies. In the rest of the section, we seek fast approximations of Algorithms \[alg:exact\_pe\] and \[alg:exact\_pi\] for this constrained MDP. From now on, we assume that $P$ and $Q$ are aperiodic and irreducible and fix $\eta > 0$ such that Proposition \[prop:properties\_of\_entropic\_tc\] holds.

Fast Approximate Policy Evaluation
----------------------------------

[r]{}[0.28]{}

\[alg:fast\_pe\] $\tilde{g} \leftarrow (\nicefrac{1}{d^2} (R^L c)^\top \mathbbm{1})\mathbbm{1}$ $\tilde{h} \leftarrow \sum_{t=0}^T R^t(c - \tilde{g})$

Next, we propose a fast approximation of Algorithm \[alg:exact\_pe\]. By our choice of $\eta$ and Proposition \[prop:properties\_of\_entropic\_tc\], all elements of ${\Pi_\eta(P,Q)}$ are aperiodic and irreducible. Thus, the gain vector is constant and may be written as $g = g_0 \mathbbm{1}$ for a scalar $g_0$. Fixing $R \in {\Pi_\eta(P,Q)}$ and $L, T \geq 1$, we approximate the gain $g$ by $\tilde{g} := (\nicefrac{1}{d^2} (R^L c)^\top \mathbbm{1})\mathbbm{1}$ and the bias vector $h$ as $\tilde{h} := \sum_{t=0}^T R^t(c - \tilde{g})$. The resulting algorithm, which we refer to as `ApproxPE`, is detailed in Algorithm \[alg:fast\_pe\]. Note that $\tilde{g}$ and $\tilde{h}$ can be computed in $\mathcal{O}(L d^4)$ and ${\mathcal{O}}(T d^4)$ time, respectively. Since $g$ and $h$ are equal to the limits of $\tilde{g}$ and $\tilde{h}$ as $L, T \rightarrow \infty$, we expect that larger $L$ and $T$ will yield better approximations. One must ensure that the $L$ and $T$ that are required for a good approximation do not grow too quickly with $d$. Using Proposition \[prop:properties\_of\_entropic\_tc\], we show that this is the case.

[prop]{}[policyevaluationcomplexity]{} \[prop:policy\_evaluation\_complexity\] Let $P$ and $Q$ be aperiodic and irreducible transition matrices, $R \in {\Pi_\eta(P,Q)}$ and $\varepsilon > 0$. Furthermore, let $g \in {\mathbb{R}}^{d^2}$ and $h \in {\mathbb{R}}^{d^2}$ be the gain and bias of $R$, respectively. Then for appropriate choice of $L$ and $T$, $\emph{\texttt{ApproxPE}}(R, L, T)$ yields $(\tilde{g}, \tilde{h})$ such that $\|\tilde{g} - g\|_\infty \leq \varepsilon$ and $\|\tilde{h} - h\|_1 \leq \varepsilon$ in $\tilde{\mathcal{O}}(d^4 \log \varepsilon^{-1})$ time.

In particular, `ApproxPE` does approximate `ExactPE` in time scaling like $\tilde{{\mathcal{O}}}(d^4)$. Explicit choices of $L$ and $T$ are given in the proof of Proposition \[prop:policy\_evaluation\_complexity\], which may be found in Appendix \[app:complexity\].

Entropic Policy Improvement {#sec:entropic_regularization}
---------------------------

Next we describe a means of approximating Algorithm \[alg:exact\_pi\]. Note that since the gain vector for any element of ${\Pi_\eta(P,Q)}$ is constant, we need only improve policies with respect to the bias vector. For the constrained MDP, exact policy improvement can be performed by calling `ExactPI` with $\Pi = {\Pi_\eta(P,Q)}$. However, no computation time is saved by doing this. Instead, we settle for an algorithm that yields approximately improved transition couplings with better efficiency. To find such an approximation, we reconsider the linear optimization problems that comprise the policy improvement step. Namely, for each $s = (x,y) \in {\mathcal{X}}\times {\mathcal{Y}}$, $$\label{eq:ent_policy_improvement}
\mbox{minimize} \,\, \sum\limits_{s'} r(s') h(s') \,\, \mbox{subject to} \,\, r \in {\Pi_\eta(P(x,\cdot), Q(y,\cdot))}.$$ Recognizing that is in fact a constrained OT problem, it is equivalent to $$\label{eq:penalized_ot}
\mbox{minimize} \,\, \sum\limits_{s'} r(s') h(s') + \frac{1}{\xi_s} \sum_{s'} r(s') \log r(s') \,\, \mbox{subject to} \,\, r\in {\Pi(P(x,\cdot), Q(y,\cdot))},$$

[r]{}[0.52]{}

\[alg:approx\_entropic\_pi\]

for some $\xi_s < \infty$ depending on $\eta$ and $h$ [@dessein2018regularized]. Problem is an instance of an entropy-regularized OT problem. In order to solve , we use the `ApproxOT` algorithm of [@altschuler2017near], detailed in Appendix \[app:complexity\]. Using `ApproxOT` instead of solving exactly, we obtain the `ApproxPI` algorithm detailed in Algorithm \[alg:approx\_entropic\_pi\]. It was shown in [@altschuler2017near] that `ApproxOT` yields an approximation of the OT cost in near-linear time with respect to the size of the couplings of interest. However, in order to control the the approximation error of `ApproxPI`, we rely on a different analysis showing that one can obtain an approximation of the regularized optimal coupling in near-linear time. To the best of our knowledge, this result does not exist in the literature, so we provide a proof in Appendix \[app:complexity\]. Using this result, we show the following complexity bound.

[prop]{}[policyimprovementcomplexity]{} \[prop:policy\_improvement\_complexity\] Let $P$ and $Q$ be aperiodic and irreducible, $h \in {\mathbb{R}}^{d^2}$ and $\varepsilon > 0$. Then there exist finite constants $(\xi_s)$ such that $\emph{\texttt{ApproxPI}}(h, (\xi_{s}), \varepsilon)$ returns $\hat{R} \in {\Pi(P, Q)}$ with $\max_s \|\hat{R}(s,\cdot) - R^*(s,\cdot)\|_1 \leq \varepsilon$ for some $R^* \in \operatorname*{argmin}_{R' \in {\Pi_\eta(P,Q)}} R'h$ in $\tilde{{\mathcal{O}}}(d^4 \varepsilon^{-4})$ time.

To summarize, this result states that `ApproxPI` approximates `ExactPI` in $\tilde{{\mathcal{O}}}(d^4)$ time rather than $\tilde{{\mathcal{O}}}(d^5)$ as previously discussed. In practice, further speedups are possible by utilizing the fact that the $d^2$ entropic OT problems to be solved are decoupled and thus may be computed in parallel.

`FastEntropicPIA`
-----------------

[r]{}[0.48]{}

\[alg:fastentropic\_pia\] $n \leftarrow 0$

Finally, using Algorithms \[alg:fast\_pe\] and \[alg:approx\_entropic\_pi\], we define the `FastEntropicOTC` algorithm, detailed in Algorithm \[alg:fastentropic\_pia\]. Essentially, `FastEntropicOTC` is defined by taking `ExactOTC` for ${\Pi_\eta(P,Q)}$ and replacing `ExactPE` and `ExactPI` by their approximations, `ApproxPE` and `ApproxPI`. In practice, `ApproxPI` returns transition couplings in the relative interior of ${\Pi(P, Q)}$, so the iterates of `FastEntropicOTC` are not restricted to the finite set of extreme points of ${\Pi(P, Q)}$. Thus, convergence for Algorithm \[alg:fastentropic\_pia\] must be assessed differently than in Algorithm \[alg:pia\]. In our simulations we found that the element-wise inequality $\tilde{g}_{n+1} \geq \tilde{g}_n$ works well as an indicator of convergence.

[0.48]{} ![A comparison of total runtimes between `ExactOTC` and `FastEntropicOTC` and approximation errors for a range of $d$ and $\xi$. Error bars show the maximum and minimum values over five simulations. []{data-label="fig:exp_results"}](runtime_diff_plot_05-29-20_09-42-41.png "fig:")

[0.48]{} ![A comparison of total runtimes between `ExactOTC` and `FastEntropicOTC` and approximation errors for a range of $d$ and $\xi$. Error bars show the maximum and minimum values over five simulations. []{data-label="fig:exp_results"}](error_plot_05-29-20_09-42-41.png "fig:")

Simulation study {#sec:experiments}
================

In order to validate the use of Algorithm \[alg:fastentropic\_pia\] as a fast approximation of Algorithm \[alg:pia\], we performed a simulation study to compare the two. In each simulation, we generated random $P$, $Q$ and $c$ and ran both `ExactOTC` and `FastEntropicOTC` until convergence. We used a range of parameters, letting $d \in \{10, 20, ... 100\}$, $L = 100$, $T = 1000$, and $\xi \in \{75, 100, 200\}$, where $\xi_s = \xi$ for all $s \in {\mathcal{X}}\times {\mathcal{Y}}$. Complete implementation details may be found in Appendix \[app:experimental\_details\]. The resulting runtimes and errors are reported in Figure \[fig:exp\_results\]. In our simulations, we found that the time savings at each iteration from `FastEntropicOTC` resulted in substantial time savings over the entire runtime of the algorithm without substantial loss of accuracy. Moreover, weakening the regularization by increasing $\xi$ reduces the error of `FastEntropicOTC` with little additional runtime. This suggests that `FastEntropicOTC` may be a more efficient alternative to `ExactOTC` when $d$ is large.

Broader impact {#broader-impact .unnumbered}
==============

Optimal transport has been used in several applications including image analysis and generation, domain adaptation, modeling cell development, and embedding natural language. Our work extends the optimal transport problem to a setting in which the objects of interest are stationary Markov chains Markov chains can be used for modeling a variety of phenomena including population dynamics, gene sequences, text, and music. We anticipate that our work can be applied to problems arising in these areas.

The authors would like to thank Quoc Tran-Dinh for helpful discussions and Jason Altschuler for contributions to the proof of Lemma \[lemma:bound\_on\_sinkhorn\_error\]. K.O. and A.N were supported in part by NIH Grant R01 HG009125-01. K.M. gratefully acknowledges the support of NSF Grant DMS-1847144. K.M. and A.N. were supported in part by NSF Grant DMS-1613261. A.N. was supported in part by NSF Grant DMS-1613072.
