[Paper Review] Optimal Transport for Stationary Markov Chains via Policy Iteration

Kevin O’Connor, Kevin McGoff+1 more|arXiv (Cornell University)|Jun 14, 2020

Markov Chains and Monte Carlo Methods65 references4 citations

TL;DR

This paper proposes a novel optimal transport framework for stationary finite-state Markov chains, termed Optimal Transition Coupling (OTC), by leveraging Markov decision processes and policy iteration to compute couplings that minimize long-term average cost while preserving Markovian and stationary structure. The key contribution is a provably convergent exact algorithm (ExactOTC) and a fast, entropy-regularized variant (EntropicOTC) with nearly-linear per-iteration complexity, enabling efficient computation on large state spaces with statistical consistency guarantees.

ABSTRACT

We study the optimal transport problem for pairs of stationary finite-state Markov chains, with an emphasis on the computation of optimal transition couplings. Transition couplings are a constrained family of transport plans that capture the dynamics of Markov chains. Solutions of the optimal transition coupling (OTC) problem correspond to alignments of the two chains that minimize long-term average cost. We establish a connection between the OTC problem and Markov decision processes, and show that solutions of the OTC problem can be obtained via an adaptation of policy iteration. For settings with large state spaces, we develop a fast approximate algorithm based on an entropy-regularized version of the OTC problem, and provide bounds on its per-iteration complexity. We establish a stability result for both the regularized and unregularized algorithms, from which a statistical consistency result follows as a corollary. We validate our theoretical results empirically through a simulation study, demonstrating that the approximate algorithm exhibits faster overall runtime with low error. Finally, we extend the setting and application of our methods to hidden Markov models, and illustrate the potential use of the proposed algorithms in practice with an application to computer-generated music.

Motivation & Objective

• To develop a constrained optimal transport framework that respects the dynamics of stationary Markov chains, enabling alignment and comparison of stochastic processes with temporal dependence.
• To address the computational challenge of solving optimal transport problems under Markovian and stationarity constraints, especially for large state spaces.
• To extend the framework to hidden Markov models (HMMs) by coupling hidden chains using emission-based costs.
• To establish theoretical convergence and stability of the proposed algorithms, ensuring statistical consistency from empirical data.
• To design a scalable, efficient algorithm with nearly-linear per-iteration complexity suitable for practical applications.

Proposed method

• Formulates the Optimal Transition Coupling (OTC) problem as a constrained optimal transport problem where transport plans must preserve the transition dynamics and stationarity of two Markov chains.
• Establishes a connection between the OTC problem and Markov decision processes (MDPs), enabling the use of policy iteration to solve it.
• Proposes ExactOTC (Algorithm 1), a policy iteration-based algorithm that converges to the optimal solution in finite time for irreducible Markov chains.
• Introduces EntropicOTC (Algorithm 6), a regularized variant that improves computational efficiency by adding entropy to the objective, enabling faster iterations.
• Derives nearly-linear per-iteration complexity bounds for EntropicOTC, matching state-of-the-art computational OT methods.
• Uses stability analysis of the OTC problem to derive statistical consistency of plug-in estimators from finite samples.

Experimental results

Research questions

• RQ1 Can optimal transport be effectively adapted to stationary Markov chains by enforcing transition dynamics and stationarity constraints?
• RQ2 How can the OTC problem be solved efficiently, especially for large state spaces?
• RQ3 What is the computational complexity of the proposed regularized algorithm, and does it scale favorably with state space size?
• RQ4 Is the solution robust to estimation errors in empirical transition matrices, and does it converge consistently from data?
• RQ5 Can the framework be extended to hidden Markov models by coupling hidden states and incorporating emission costs?

Key findings

• ExactOTC converges to the optimal solution of the OTC problem in a finite number of iterations for irreducible Markov chains, as proven in Theorem 7.
• EntropicOTC achieves nearly-linear per-iteration complexity in the dimension of the coupling matrix, with bounds established in Theorems 9 and 12.
• The OTC problem is stable under perturbations of the input transition matrices, as formalized in Theorem 13.
• Statistical consistency of the plug-in estimator for the optimal transition coupling and its expected cost follows as a corollary to the stability result (Corollary 14).
• Empirical validation shows that EntropicOTC achieves significantly faster runtime with low error compared to exact methods.
• The framework successfully extends to HMMs, with a practical application demonstrated in computer-generated music generation.