[Paper Review] Sampling with Riemannian Hamiltonian Monte Carlo in a Constrained Space

Yunbum Kook, Yin Tat Lee+2 more|arXiv (Cornell University)|Feb 3, 2022

Bioinformatics and Genomic Networks20 citations

TL;DR

CRHMC integrates constraints directly into Riemannian HMC to sample ill-conditioned, high-dimensional constrained distributions efficiently, preserving sparsity and achieving mixing independent of condition numbers; outperforms CHRR/CDHR by orders of magnitude on real datasets.

ABSTRACT

We demonstrate for the first time that ill-conditioned, non-smooth, constrained distributions in very high dimension, upwards of 100,000, can be sampled efficiently $ extit{in practice}$. Our algorithm incorporates constraints into the Riemannian version of Hamiltonian Monte Carlo and maintains sparsity. This allows us to achieve a mixing rate independent of smoothness and condition numbers. On benchmark data sets in systems biology and linear programming, our algorithm outperforms existing packages by orders of magnitude. In particular, we achieve a 1,000-fold speed-up for sampling from the largest published human metabolic network (RECON3D). Our package has been incorporated into the COBRA toolbox.

Motivation & Objective

• Motivate constrained high-dimensional sampling in ill-conditioned settings common in systems biology and linear programming.
• Develop a constrained Riemannian HMC (CRHMC) that preserves sparsity and respects linear constraints.
• Provide a practical, scalable implementation with theoretical guarantees.
• Demonstrate empirical speedups over existing methods on benchmark metabolic and LP datasets.

Proposed method

• Extend RHMC to enforce linear constraints via a constrained Hamiltonian with a carefully chosen M(x) whose range matches the constraint null space.
• Use a self-concordant barrier-based local metric g(x) to define M(x)=Q(x)ᵀ g(x) Q(x) that preserves sparsity and constraint feasibility.
• Employ an implicit midpoint integrator for discretization to maintain symplectic, reversible dynamics and enable a Metropolis correction.
• Derive efficient formulas to avoid pseudo-inverse/pseudo-determinant by expressing M(x)† and log pdet M(x) in terms of g(x) and Dc(x); leverage sparse linear solvers.
• Simplify subspace constraints to reduce computation, and compute leverage scores efficiently via sparse Cholesky decompositions to avoid dense updates.
• Provide specialized updates for the case c(x)=Ax−b to further simplify dynamics and avoid dense matrix operations.

Experimental results

Research questions

• RQ1 Can we sample from e^−f(x) subject to Ax=b and x ∈ K efficiently in very high dimensions?
• RQ2 Does CRHMC achieve mixing times independent of the condition number in ill-conditioned, constrained, non-smooth settings?
• RQ3 How can one maintain sparsity and feasibility throughout the constrained Hamiltonian dynamics without incurring prohibitive computational costs?
• RQ4 What practical speedups and scalability can CRHMC achieve on real-world datasets (systems biology, LPs) compared to existing constrained samplers?

Key findings

• CRHMC achieves sub-quadratic mixing and sampling time and can scale to very large models (up to ~100k variables) where CHRR/CDHR struggle.
• On multiple real datasets (metabolic networks and NETLIB LPs), CRHMC outperforms CHRR and CDHR by orders of magnitude in both mixing rate and time per effective sample.
• CRHMC enables sampling from large constrained polytopes (e.g., Recon3D) with substantial speedups (up to 1,000× faster) over the largest published models.
• Structured experiments show CRHMC scales to half-million dimensions for hypercubes and simplices, and up to 10^5 for Birkhoff polytopes, with reasonable effective sample sizes.
• The uniformity tests indicate that CRHMC samples closely approximate the uniform distribution over the polytopes examined.