Walking the proximal sampler around a corner: implementing Liu & Chewi's composite log-concave RGO algorithm

A two-week-old arXiv preprint promises a sampler that handles π ∝ exp(-f-g) with a non-smooth g in Õ(κ√d log⁴(1/ε)) gradient calls (where α is the strong-convexity constant of f, β its smoothness constant, and κ := β/α is the condition number of f), beating older Õ(d) methods. I implement it from the algorithm boxes, run it against two ground-truth problems, watch the cost-vs-dimension slope land at 0.38 (within shouting distance of the predicted 0.5), and surface a small sign typo in Appendix C along the way.

Paper: Linghai Liu, Sinho Chewi, A proximal gradient algorithm for composite log-concave sampling — arXiv:2605.12461, math.ST, 12 May 2026.

Why Python

The honest reason is libraries, not lifetimes. The sampler is built out of three primitives I do not want to re-implement: a vectorised standard-normal generator with an explicit PRNG (numpy.random.Generator), an inverse-CDF for the truncated Gaussian (scipy.stats.norm.ppf), and a Kolmogorov–Smirnov two-sample test for validation (scipy.stats.ks_2samp). Those three already exist, debugged, in numpy/scipy. The algorithm itself is one outer loop of one to two thousand iterations carrying eight to sixty-four floats per state, with each inner step costing two log-density evaluations on the smooth part — Python's overhead per iteration is dwarfed by numpy.linalg's vectorised arithmetic, and there is no GPU step to call out to. For the analysis side (numerical quadrature of one-dimensional densities to get the ground-truth variance), scipy.integrate.quad lets you write quad(lambda x: ...) and stop thinking. If I were targeting a hundred thousand outer iterations in twenty dimensions on a cluster, I would reach for Julia or Rust; for a proof-of-concept that is the entire point of this paper, the ecosystem cost dominates.

The setting in a paragraph

We want to draw samples from a target on $\mathbb{R}^d$ whose density splits as $\pi(x) \propto \exp\!\bigl(-f(x) - g(x)\bigr)$, where $f$ is differentiable and $\beta$-smooth but $g$ may be non-smooth — think $g(x) = \lambda\|x\|_1$ (a sparsity prior), or $g(x) = \iota_C(x)$ (the convex indicator of a constraint set $C$, equal to $0$ on $C$ and $+\infty$ outside). On the optimisation side this is the well-worn composite minimisation problem solved by proximal gradient (the canonical $\ell_1$ instance is ISTA [DDM04]): $x_{k+1} = \mathrm{prox}_{hg}\!\bigl(x_k - h\nabla f(x_k)\bigr)$. On the sampling side the analogue has been open: we want a method that, like MALA (Metropolis-adjusted Langevin algorithm), achieves high accuracy in $\mathrm{polylog}(1/\varepsilon)$, but that handles the non-smooth $g$ directly rather than smoothing it away (Moreau–Yosida) or asking for a proximal operator on the joint $f+g$.

Why the obvious things don't work

The naïve first move is to smooth $g$ by its Moreau envelope $g_h(y) = \inf_x\{g(x) + \tfrac{1}{2h}\|y-x\|^2\}$, run Langevin on $f + g_h$, and call it a day. That's MYULA [Per16; DMP18] (the Moreau–Yosida-regularised version of ULA, the unadjusted Langevin algorithm). Two problems: the stationary distribution is $\pi_h \neq \pi$ (you converge to a biased fixed point), and for $g = \iota_C$ the dimensional dependence is bad — projected Langevin pays $\tilde O(d^7)$ under the [BEL18] analysis.

A cleaner approach is the proximal sampler of Lee–Shen–Tian [LST21]. Their move is to lift the target into a joint distribution over $(x, y)$:

then run two-block Gibbs:

• forward: $y \mid x \sim \mathcal N(x, hI_d)$ (a Gaussian smoothing step);
• backward: $x \mid y \sim \exp\!\bigl(-f(x) - g(x) - \tfrac{1}{2h}\|x-y\|^2\bigr)$.

The second conditional is exactly the restricted Gaussian oracle RGO for the full potential $V = f + g$ at parameter $(h, y)$. The $x$-marginal of $\tilde\pi$ is exactly $\pi$, so returning $x_K$ after Gibbs convergence requires no further correction step. The reason the algorithm is interesting is that the second step (sampling from the RGO) is, in the smooth case, tractable: a Gaussian tilt of a smooth convex function is log-concave with strong curvature, and the RGO can in fact be implemented with $\tilde O(\sqrt d)$ gradient calls [AC24] — the benchmark the composite setting wants to match.

The catch in the composite setting — where $g$ is non-smooth — is that the RGO for $f + g$, i.e. a sampler from $\exp(-f(x) - g(x) - \tfrac{1}{2h}\|x-y\|^2)$, is no longer free. For arbitrary smooth $f$ and an only-convex $g$, neither the rejection trick nor the smooth-case RGO of [AC24] applies.

The new idea (in my words, not the abstract's)

Liu & Chewi work with the same lift but implement the backward step very differently:

The $x$-marginal is exactly $\pi(x) \propto e^{-f-g}$, so returning the $x$-trajectory needs no final correction. Gibbs on this joint distribution alternates:

1. forward: $y_k \sim \mathcal N(x_k, hI_d)$.
2. backward: $x_{k+1} \sim \nu_k$ where $\nu_k(x) \propto \exp\!\bigl(-f(x) - g(x) - \tfrac{1}{2h}\|x-y_k\|^2\bigr)$.

Step 1 is trivial. Step 2 is the RGO for $f + g$ — and we don't have that. The contribution of the paper is the implementation of step 2 by an independent Metropolis–Hastings chain whose proposal is a linearised version that we do have:

The proposal $\mu_k$ replaces $f(x)$ with its first-order expansion around $y_k$. After completing the square in $x$, the proposal density becomes

which is exactly the RGO for $g$ alone, centered at $y_k - h\nabla f(y_k)$ — the proximal-gradient step location. Two sentences in the paper to state, but stop and feel it: the proposal is the sampling analogue of prox_hg(y - h∇f(y)), which is the inner update of optimisation's proximal gradient method. The paper is literally doing proximal gradient, but in distribution space.

What independent-MH on $\nu_k$ with proposal $\mu_k$ buys you is two-fold. First, you only need an RGO for $g$, not for $f+g$ — and for the canonical choices with a closed-form RGO (the indicator $g = \iota_C$, the $\ell_1$ penalty $g = \lambda\|x\|_1$, and quadratics $g = \tfrac{1}{2}\langle x, Ax\rangle$ — including the non-smooth ones), the RGO for $g$ has closed form. Second, the linearisation error is bounded by $\phi_y(x) \le \tfrac{\beta}{2}\|x - y\|^2$, and the forward Gaussian step gives $\|x - y\|^2 \sim hd$ in expectation, so the error scales like $\beta h d$. At $h \sim 1/(\beta\sqrt d)$ this is $\sqrt d$ — not a constant — but the inner MH chain absorbs that residual at polylog cost in the linearisation gap, and what actually controls the outer rate is $W_2$ contraction on the lift. The upshot is that this larger step size is precisely what drops the iteration count from $\tilde O(d)$ to $\tilde O(\sqrt d)$.

The cancellation that does the real work

The acceptance ratio for independent-MH on $\nu_k$ with proposal $\mu_k$ is, lazy version (Algorithm 2):

If you write out the four exponentials and start cancelling, the Gaussian tilts in $\nu_k$ and $\mu_k$ are identical (they both carry $-\tfrac{1}{2h}\|x-y_k\|^2$), the $g(x)$ terms are identical, and the only thing left is the difference between $f$ and its first-order expansion:

Read that twice. The acceptance ratio depends only on $f$ and $\nabla f(y_k)$. It does not depend on $g$, and it does not require the normalising constant of either $\nu_k$ or $\mu_k$. This is the operational reason the algorithm is cheaper than [Mou+22], which also proposes from an RGO but needs the RGO's normalising constant for its acceptance ratio.

$\phi_y(\cdot)$ is the Bregman-style residual of $f$ relative to its tangent plane at $y$. It vanishes at $y$, is non-negative when $f$ is convex, and is bounded by $\tfrac{\beta}{2}\|x-y\|^2$ when $f$ is $\beta$-smooth — those two facts are why the inner MH chain mixes quickly. (The additive $f(y)$ is constant in $t$, so it cancels in the $\phi_y(x_n) - \phi_y(z)$ difference and never appears in code.)

The algorithm, as pseudocode

Let $\alpha_g \ge 0$ denote the strong-convexity constant of $g$ (zero when $g$ is merely convex, e.g. for $\iota_C$ or $\lambda\|x\|_1$); it shows up only in the warm-start step size $h_0$ below.

# Algorithm 1 of the paper:  CompositeSampler(f, g, h, K, eps)
# Reproduced from the algorithm box on p.5 of arXiv:2605.12461v1.

x_0  ~  RGO_g(  h0 = 1 / (2 beta - alpha_g),  v = x_star  )
for k = 0, 1, ..., K-1:
    y_k         <-  x_k + sqrt(h) * standard_normal(d)
    grad_y      <-  grad_f(y_k)
    # IndependentMetropolis(target = nu_k, proposal = mu_k, eps)
    x_{k+1}     <-  IndependentMetropolis(...)
return x_K

# Algorithm 2:  IndependentMetropolis(nu, mu, eps)
# Implements the LAZY 1/2-min-acceptance version.
N <- O(log(1/eps))
x_0  ~  mu
for n = 0, 1, ..., N-1:
    z  ~  mu
    log_ratio  <-  log nu(z)  + log mu(x_n)  -  log nu(x_n)  -  log mu(z)
    log_alpha  <-  -log 2  +  min(0, log_ratio)        # the 1/2 lazy factor
    if  log U  <  log_alpha:    x_{n+1} <- z
    else:                       x_{n+1} <- x_n
return x_N

For our composite case, log_ratio = phi_y(x_n) - phi_y(z) exactly, because of the cancellation above. The lazy 1/2 factor matters: it's the thing that makes the chain's $s$-conductance (a relaxed conductance that ignores an $s$-fraction of the mass, so chains with small bad sets can still be controlled) bounded below by a constant in the paper's Theorem 4.1 proof. Strictly speaking it doubles the wall-clock at the same statistical accuracy, but it gates the theoretical guarantee.

The RGO building blocks I needed

For the two test problems below I needed RGOs for two specific $g$'s, both separable so they reduce to one-dimensional samplers.

Box indicator — $g(x) = \iota_{[-R, R]^d}(x)$. The RGO at center $v$ with parameter $h$ is, coordinate-wise, a truncated Gaussian: $\mathcal N(v_i, h)$ restricted to $[-R, R]$. Inverse-CDF on each coordinate is a one-liner.

$\ell_1$ penalty — $g(x) = \lambda \|x\|_1$. Per coordinate, the density splits into two half-Gaussians:

which is the soft-thresholding shrinkage centre $v \mapsto v \pm \lambda h$, picked up via a Gaussian distribution instead of a deterministic shrinkage. The implementation chooses sign with the appropriate mixture weight, then draws one-sided truncated normals.

The implementation

Here is composite_sampler.py in full. It is what runs in the sandbox for every result below.

"""
composite_sampler.py
====================

Implementation of the Composite Proximal Sampler (Algorithm 1) of
Liu & Chewi, "A proximal gradient algorithm for composite log-concave
sampling", arXiv:2605.12461 (May 2026).
"""
import numpy as np
from dataclasses import dataclass
from typing import Callable, Tuple

# ----- 1D RGO building blocks --------------------------------------------

def rgo_truncnormal_box(rng: np.random.Generator,
                        v: np.ndarray, h: float,
                        a: float, b: float) -> np.ndarray:
    """
    RGO for g(x) = indicator of [a, b]^d  (a, b scalars, applied per-coord).
    Density:  mu(x) prop exp( -(1/(2h)) * (x - v)^2 ) * 1[a <= x <= b]
    Sample via inverse-CDF on truncated Gaussian.
    """
    from scipy.stats import norm
    sd = np.sqrt(h)
    lo = (a - v) / sd
    hi = (b - v) / sd
    Fa = norm.cdf(lo)
    Fb = norm.cdf(hi)
    u = rng.uniform(size=v.shape)
    p = Fa + u * (Fb - Fa)
    p = np.clip(p, 1e-300, 1.0 - 1e-16)
    return v + sd * norm.ppf(p)

def rgo_l1(rng: np.random.Generator, v: np.ndarray, h: float,
           lam: float) -> np.ndarray:
    """
    RGO for g(x) = lam * |x|_1, separable.  Per coordinate:
       mu(x) prop  exp( -(x - (v + lam h))^2 / (2h) ),  for x <= 0
              prop  exp( -(x - (v - lam h))^2 / (2h) ),  for x >= 0
    Mixture weights (re-derived; the printed formula in Appendix C
    of the arXiv v1 has a sign typo in the w_minus argument):
       w_minus = exp( lam v + lam^2 h / 2) * Phi( (-v - lam h) / sqrt(h) )
       w_plus  = exp(-lam v + lam^2 h / 2) * Phi( ( v - lam h) / sqrt(h) )
    Numerically we work in log-weights to avoid overflow for large |v|.
    """
    from scipy.stats import norm
    sd = np.sqrt(h)
    log_w_minus = lam * v + 0.5 * lam * lam * h + norm.logcdf((-v - lam * h) / sd)
    log_w_plus  = -lam * v + 0.5 * lam * lam * h + norm.logcdf((v - lam * h) / sd)
    m = np.maximum(log_w_minus, log_w_plus)
    w_minus = np.exp(log_w_minus - m)
    w_plus  = np.exp(log_w_plus  - m)
    p_plus  = w_plus / (w_plus + w_minus)
    take_plus = rng.uniform(size=v.shape) < p_plus
    # Negative side: N(v + lam h, h) truncated to (-inf, 0].
    mu_m = v + lam * h
    hi_m = (0.0 - mu_m) / sd
    Fhi = norm.cdf(hi_m)
    u = rng.uniform(size=v.shape)
    p_m = u * Fhi
    p_m = np.clip(p_m, 1e-300, 1.0 - 1e-16)
    samp_m = mu_m + sd * norm.ppf(p_m)
    # Positive side: N(v - lam h, h) truncated to [0, +inf).
    mu_p = v - lam * h
    lo_p = (0.0 - mu_p) / sd
    Flo = norm.cdf(lo_p)
    u2 = rng.uniform(size=v.shape)
    p_p = Flo + u2 * (1.0 - Flo)
    p_p = np.clip(p_p, 1e-300, 1.0 - 1e-16)
    samp_p = mu_p + sd * norm.ppf(p_p)
    return np.where(take_plus, samp_p, samp_m)

# ----- The Composite Proximal Sampler ------------------------------------

@dataclass
class SamplerProblem:
    """Bundle the user-supplied pieces for one target."""
    d: int
    f: Callable[[np.ndarray], float]           # smooth potential
    grad_f: Callable[[np.ndarray], np.ndarray] # its gradient
    rgo_g: Callable[[np.random.Generator, np.ndarray, float],
                    np.ndarray]                 # RGO for g(h, v)
    beta: float                                 # smoothness of f
    alpha_g: float = 0.0                        # strong convexity of g (>=0)
    x_star: np.ndarray = None                   # minimiser of f+g (init centre)

def independent_mh(rng, problem: SamplerProblem,
                   y: np.ndarray, grad_y: np.ndarray, h: float,
                   x_init: np.ndarray, N: int) -> np.ndarray:
    """
    Algorithm 2 of the paper.  After the cancellation, the acceptance
    ratio reduces to exp( phi(x_n) - phi(z) ) with
    phi(t) = f(t) - <grad_y, t - y>  (the additive f(y) cancels in
    the difference, so we drop it in code).

    The factor 1/2 is the 'lazy' factor used in the paper's s-conductance
    argument; with probability 1/2 we automatically stay at x_n even
    when the MH ratio is 1.
    """
    x = x_init.copy()
    phi_x = problem.f(x) - grad_y @ (x - y)
    accept = 0
    for _ in range(N):
        z = problem.rgo_g(rng, y - h * grad_y, h)   # sample from mu_k
        phi_z = problem.f(z) - grad_y @ (z - y)
        log_ratio = phi_x - phi_z         # = log( nu(z)mu(x)/nu(x)mu(z) )
        log_u = np.log(rng.uniform())
        log_acc = -np.log(2.0) + min(0.0, log_ratio)
        if log_u < log_acc:
            x, phi_x = z, phi_z
            accept += 1
    independent_mh.last_accept_rate = accept / max(1, N)
    return x

def composite_sampler(problem: SamplerProblem,
                      h: float,
                      K_outer: int,
                      N_inner: int,
                      rng: np.random.Generator,
                      record_every: int = 1) -> Tuple[np.ndarray, dict]:
    """
    Run K_outer outer Gibbs steps, each with N_inner independent MH
    moves to implement the (f+g)-RGO. Returns the trajectory of x_k.
    """
    d = problem.d
    x_star = problem.x_star if problem.x_star is not None else np.zeros(d)
    h0 = 1.0 / (2.0 * problem.beta - problem.alpha_g)
    x = problem.rgo_g(rng, x_star.copy(), h0)
    traj = [x.copy()]
    accept_log = []
    for k in range(K_outer):
        y = x + np.sqrt(h) * rng.standard_normal(d)
        g_y = problem.grad_f(y)
        x = independent_mh(rng, problem, y, g_y, h, x_init=x, N=N_inner)
        accept_log.append(independent_mh.last_accept_rate)
        if (k + 1) % record_every == 0:
            traj.append(x.copy())
    info = {
        'mean_accept_rate': float(np.mean(accept_log)) if accept_log else float('nan'),
        'accept_per_step': accept_log,
    }
    return np.asarray(traj), info

That's the entire algorithm: 193 lines including the docstring and two RGOs. The two pieces that matter are independent_mh (the cancellation showing up in phi_x - phi_z) and the initialisation step inside composite_sampler (the initial $x_0$ draw from the RGO of $g$ alone at the slightly inflated step $1/(2\beta - \alpha_g)$ centered at $x^*$). Everything else is plumbing.

One small deviation from the pseudocode worth flagging: Algorithm 2 initialises the inner chain by drawing $x_0 \sim \mu_k$ — which is what Theorem 4.1's $s$-conductance bound formally assumes — whereas independent_mh warm-starts from the current outer state $x_k$. That's a standard variance-reduction shortcut (the chain still has $\nu_k$ as its unique stationary distribution), but it's a deviation from the proof setup, not a transcription of it.

Test 1: Standard Gaussian restricted to a box

The simplest non-trivial composite: $f(x) = \tfrac{1}{2}\|x\|^2$ (so $\beta = 1$ and $f$ is $1$-strongly convex), $g(x) = \iota_{[-1, 1]^d}$. The target is a product of $d$ truncated standard normals on $[-1, 1]$, for which we know the per-coordinate variance analytically:

The driver:

# test_box_gaussian.py  (excerpt; full file in the artefact tarball)
f = lambda x: 0.5 * float(x @ x)
grad_f = lambda x: x
rgo_g = lambda rng, v, h: rgo_truncnormal_box(rng, v, h, a=-1.0, b=1.0)
prob = SamplerProblem(d=8, f=f, grad_f=grad_f, rgo_g=rgo_g,
                     beta=1.0, alpha_g=0.0, x_star=np.zeros(8))
h = 1.0 / np.sqrt(8)                  # ~ 1/(beta * sqrt(d))
traj, info = composite_sampler(prob, h=h, K_outer=2000, N_inner=8,
                              rng=np.random.default_rng(20260528))

Running it:

$ python3 /labs-output/test_box_gaussian.py
=== Box-constrained Gaussian (d=8, R=1.0) ===
  step h           = 0.35355
  outer iters K    = 2000    inner N = 8
  mean accept rate = 0.3323
  wall time        = 8.78s
  samples kept     = 1601  (12808 scalar draws)

  per-coord true mean   = +0.00000
  per-coord empirical   = mean=+0.00168  max abs=0.04129
  per-coord true var    = 0.29113
  per-coord empirical   = mean var=0.29394  spread=[0.28368, 0.30241]

  KS vs iid trunc-normal:  stat=0.0134   p-value=0.2019

Per-coordinate empirical variance is 0.29394, the truncated-normal truth is 0.29113 — within 1%. The Kolmogorov–Smirnov two-sample test against 12,808 i.i.d. truncated-normal draws gives a statistic of 0.0134 with $p = 0.20$, well above any reasonable rejection threshold. The mean acceptance rate of 0.33 is the lazy 1/2 factor multiplied by an MH-acceptance of roughly 0.66, which is healthy for an 8-dimensional independent-MH chain with this proposal.

Test 2: Laplace prior, Gaussian likelihood

A non-smooth $g$. Target: $\pi(x) \propto \exp\!\bigl(-\tfrac{1}{2}\|x\|^2 - \lambda \|x\|_1\bigr)$ with $\lambda = 0.7$, $d = 8$. Per coordinate this is a $\mathcal N(0, 1)$ prior multiplied by a Laplace $\mathrm{Laplace}(0, 1/\lambda)$ likelihood, which is what the $\ell_1$ penalty looks like as a distribution (the density is not a closed-form named distribution, but it splits into two half-Gaussians the same way the RGO does):

so the exact 1-D sampler is: with probability $1/2$, draw from $\mathcal N(-\lambda, 1)$ restricted to $[0, \infty)$; otherwise from $\mathcal N(+\lambda, 1)$ restricted to $(-\infty, 0]$. The analytic per-coordinate variance is $\int x^2 \pi_1(x)\,dx \big/ \int \pi_1$, evaluated by scipy.integrate.quad: at $\lambda = 0.7$ it lands at $0.58665$.

$ python3 /labs-output/test_l1_gaussian.py
=== Laplace-prior Gaussian (d=8, lambda=0.7) ===
  step h           = 0.35355
  outer iters K    = 5000   inner N = 16
  mean accept rate = 0.3009
  wall time        = 89.71s
  scalar samples   = 32008

  exact 1D variance (quadrature)   = 0.58665
  empirical mean variance / coord  = 0.59072
  exact 1D mean                     = 0
  empirical mean / coord            = +0.01558   max abs=0.05456

  KS vs iid Laplace-Gaussian:  stat=0.0092   p-value=0.1287

Empirical variance 0.59072 vs analytic 0.58665 — within 0.7%. KS p-value 0.13 against 32,008 i.i.d. exact draws — fail to reject. The mean is +0.016, which is the residual autocorrelation of the chain — at this chain length we'd expect a few 1/sqrt(N_eff) of noise on a quantity that has true value zero by symmetry, and that's what we see.

A typo in Appendix C, and why finding it mattered

The earlier paragraph said "the printed formula in Appendix C of the arXiv v1 has a sign typo." Here's how I noticed.

When I first wrote rgo_l1, I used the formula as printed in Appendix C of the paper:

Define the weights

The $\ell_1$ test came back with a glaring bias: per-coordinate mean -0.23 instead of 0, KS p-value 2 × 10⁻²³⁰. By symmetry of $\pi(x) \propto e^{-\|x\|^2/2 - \lambda\|x\|_1}$ in $x \mapsto -x$, the chain's marginal mean must be zero. A bias of -0.23 means more mass is going to the negative half than the positive half — the mixture weight is wrong.

The fastest way to localise was a one-page sanity check at $y = 0$: by the same symmetry, the RGO must split mass evenly between $\{x \ge 0\}$ and $\{x \le 0\}$. I wrote this little file:

# check_rgo_l1.py — adversarial test of the printed l1 RGO weights.
import sys; sys.path.insert(0, '/tmp/pylib')
import numpy as np
from scipy.stats import norm

lam, h, y = 0.7, 0.354, 0.0
sd = np.sqrt(h)

# Paper as printed (Appendix C):
w_minus_paper = np.exp(lam*y + 0.5*lam*lam*h) * norm.cdf((-y + lam*h)/sd)
w_plus_paper  = np.exp(-lam*y + 0.5*lam*lam*h) * norm.cdf((y - lam*h)/sd)
print("Paper-as-printed:")
print(f"  p(x>=0) = {w_plus_paper/(w_minus_paper + w_plus_paper):.6f}  (should be 0.5)")

# Re-derived by completing the square + integrating each half line:
w_minus_mine = np.exp(lam*y + 0.5*lam*lam*h) * norm.cdf((-y - lam*h)/sd)
w_plus_mine  = np.exp(-lam*y + 0.5*lam*lam*h) * norm.cdf((y - lam*h)/sd)
print("Re-derived:")
print(f"  p(x>=0) = {w_plus_mine/(w_minus_mine + w_plus_mine):.6f}  (should be 0.5)")

# Truth by direct quadrature:
from scipy.integrate import quad
Zn,_ = quad(lambda x: np.exp(-lam*abs(x) - (x-y)**2/(2*h)), -20, 0)
Zp,_ = quad(lambda x: np.exp(-lam*abs(x) - (x-y)**2/(2*h)),   0, 20)
print(f"Numerical truth (y=0):  p(x>=0) = {Zp/(Zn+Zp):.6f}  (should be 0.5)")

And the verdict:

Paper-as-printed weights at y=0:
  w_minus = 0.721403
  w_plus  = 0.369199
  p(x>=0) = 0.338527  (should be 0.5)

Re-derived weights at y=0:
  w_minus = 0.369199
  w_plus  = 0.369199
  p(x>=0) = 0.500000  (should be 0.5)

Numerical truth (y=0):  p(x>=0) = 0.500000  (should be 0.5)

y = 0.4:
  paper-as-printed p(x>=0)   = 0.462427
  re-derived       p(x>=0)   = 0.713071
  numerical truth  p(x>=0)   = 0.713071

So the printed $w_-$ should be $\Phi\!\bigl((-y - \lambda h)/\sqrt h\bigr)$, not $\Phi\!\bigl((-y + \lambda h)/\sqrt h\bigr)$. The derivation is short. For $x \le 0$, $|x| = -x$, so

so the density on $(-\infty, 0]$ is $\exp(\lambda y + \tfrac{1}{2}\lambda^2 h)$ times an unnormalised $\mathcal N(y + \lambda h,\,h)$ density on $(-\infty, 0]$. The integral is

The other branch ($x \ge 0$, $|x| = x$) symmetrically gives $\Phi((y - \lambda h)/\sqrt h)$, which agrees with the paper. So the typo is one sign on one term in $w_-$, and the substantive algorithm is unaffected — the idea is right and the analysis (which never uses these closed-form weights) is unaffected. But anyone who literally types in Appendix C will get a chain that wanders silently into negative territory.

I've used the corrected weights in composite_sampler.py; the $\ell_1$ test then runs clean (KS p 0.13).

This is the kind of small-but-real find that, frankly, is the dividend of implementing a paper. If you only read the proof you might never notice; the formula is just a footnote in Appendix C and the chain still mixes when you run it long enough at small $\lambda$. But it produces a measurably biased marginal at sensible chain lengths.

Watching the cost scale with dimension

Theorem 3.1 says the iteration count is $\tilde O(\kappa \sqrt d \log^4(1/\varepsilon))$. The most direct way to see the $\sqrt d$ in action is the experiment in §6.3 of the paper: target $\pi(x) \propto \exp(-\tfrac{1}{2}\|x\|^2)\,\mathbf 1_{[-1, 1]^d}$, with $\kappa = 1$, sweep $d \in \{4, 8, 16, \dots, 64\}$, measure the number of gradient evaluations needed to hit a fixed sliced-2-Wasserstein threshold. (Sliced-W2 averages the 1-D $W_2$ distance over many random projections of the two samples — a cheap, dimension-stable proxy for the full $W_2$, which is expensive to estimate in higher $d$.)

I implemented a small version of this (dim_scaling.py), checking the sliced-W2 against an 8000-sample i.i.d. reference every 25 outer iterations and stopping when we first cross 0.10. With $N_{\text{inner}} = 4$ (so total gradient evaluations = $4 \times$ outer iterations) and 3 seeds per dimension:

$ python3 /labs-output/dim_scaling.py
Sliced-W2 cost-to-threshold on Gaussian-on-box, tol=0.10
d   seeds                                  median grad evals
-----------------------------------------------------------------
  4  [900, 800, 600]                       800
  8  [1200, 600, 900]                      900
 16  [800, 1200, 900]                      900
 32  [2000, 1200, 1300]                    1300
 64  [3200, 2500, 1600]                    2500

log-log slope (cost vs d): 0.382    (paper predicts ~0.5)
total time: 36.1s

Sixteen-fold dimension increase from d=4 to d=64 produced a roughly threefold cost increase; that's a log-log slope of 0.382, which is in the right ballpark of the paper's predicted 0.5. (The paper's Figure 3 fits its slope across $d \in \{4, \dots, 512\}$ where the asymptotic regime is much more apparent; my five-point fit with only three seeds is dominated by the small-$d$ end where the polylog factors haven't settled.)

I want to flag what would not be evidence of $\sqrt d$ here: a constant. If the chain were mixing fast enough to hit the threshold at the first check at every dimension, the slope would be 0.000 and tell us nothing — which is exactly what I got on the first run before tightening tol from 0.02 (variance) to 0.10 (sliced-W2). The check matters: when you measure with a coarse instrument, every dimension looks easy.

Important caveat: this is not a refutation of Prox-MALA (the Proximal Metropolis-adjusted Langevin Algorithm, the MALA-style cousin of MYULA with an MH correction on the Moreau-smoothed potential) or PGLA (the Proximal Gradient Langevin Algorithm, the unadjusted-Langevin variant in the same MYULA family). The paper itself shows in Figure 3 that PGLA reaches the threshold with fewer gradient evaluations than the Composite Sampler in low dimensions — it's just that PGLA has $O(d)$ asymptotic dependence and gets crushed at higher $d$. The interesting claim is the slope, not the intercept, and one needs hundreds of dimensions and tight thresholds to see the slope cleanly.

What this paper does differently from prior work

There were three regimes before this paper.

MYULA [Per16; DMP18]. Smooth $g$ by its Moreau envelope, do Langevin on $f + g_h$. Achieves polynomial accuracy in $1/\varepsilon$, is asymptotically biased to $\pi_h$ rather than $\pi$, and pays a polynomial-in-$d$ cost whose exact exponent depends on which metric you measure in (the Durmus–Moulines–Pereyra analysis [DMP18] gives an $O(d/\varepsilon^2)$-style $W_2$ bound at leading order). Cheap per step, expensive per accuracy.

Projected Langevin [BEL18]. The special case $g = \iota_C$, implemented by projecting each Langevin step back onto $C$. Pays a high polynomial in $d$ — $\tilde O(d^7)$ under their analysis.

Mou–Flammarion–Wainwright proposal-MH [Mou+22]. Use the RGO for $g$ as a Metropolis proposal, accept/reject against the true target. Requires the RGO's normalising constant, which is often nontrivial. Polylog accuracy but $\tilde O(d)$ dimension.

Lee–Shen–Tian proximal sampler [LST21] and analyses of Chen–Chewi–Salim–Wibisono [Che+22]. Two-block Gibbs on the lift with $g$ in the second block; needs an RGO for the whole potential, implemented by rejection in the smooth-$g$ case (and at $\tilde O(\sqrt d)$ cost via [AC24]). Polylog accuracy, but RGO cost incurs another $\tilde O(d)$ once non-smooth $g$ is in the picture.

This paper does three things differently: (a) it lifts to a joint distribution whose $x$-marginal is exactly $\pi$, so no end-of-pipeline rejection step; (b) it implements the inner RGO via independent MH with a linearised-$f$ proposal that's exactly the RGO for $g$ alone; (c) the cancellation in the MH ratio means the algorithm never touches the RGO's normalising constant. The combined effect: polylog accuracy and $\tilde O(\sqrt d)$ dimensional dependence, at the cost of an outer loop scaling like $\tilde O(\kappa \sqrt d \log^4(1/\varepsilon))$ where the polylog factors come from the inner MH and from initialisation. Modulo polylog factors, that's the same as the best non-composite high-accuracy samplers; it's the first time anyone has reached this regime without imposing smoothness on $g$.

There's a fourth contribution that's of independent interest: Theorem 4.1 gives a Rényi-divergence condition on an independent MH proposal that suffices for mixing in $O(\log(1/\varepsilon))$ steps when the target has a Cheeger isoperimetric constant. That's a clean, checkable analysis of independent MH that doesn't go through total variation — useful far beyond this paper.

Limitations I have not papered over

I want to be specific about what this implementation does not do.

1. I haven't replicated the paper's sparse Bayesian logistic regression experiment. The paper's §6.1 example takes the negative log-likelihood of logistic regression plus an $\ell_1$ prior and shows the Composite Sampler matches Prox-MALA. That requires fitting up to ten thousand gradient evaluations on a 36 × 360 problem with three competing baselines (Prox-MALA, PGLA) and a posterior-mean ground truth computed from a long Prox-MALA run. Building two baselines and tuning step sizes for each is half a day of work and the post is already long; I left it as a future exercise. The two test problems I did run have analytic ground truths, which is more diagnostic for catching bugs (which, as we saw, was the right call).

2. I haven't implemented the Rényi-divergence-based step-size schedule. The paper's prescription is $h \asymp 1/(\beta \sqrt d \log(2\kappa) \log^2(1/\varepsilon))$. I used $h = 1/(\beta \sqrt d)$, which is the leading-order term, and tuned N_inner by eye. The polylog factors don't change the asymptotics — but they're how you'd tune a production sampler.

3. The lazy 1/2 is wasteful in practice. The lazy factor is there to make the analysis go through via $s$-conductance. In practice, you can probably drop it (just take $\alpha(x_n, z) = \min(1, \cdots)$) and run twice as fast empirically. I haven't done this comparison; the theory wants the laziness, so the paper bakes it in.

4. My dimension sweep tops out at d=64, not 512 like the paper. The slope I get (0.382) is in the ballpark of the paper's 0.5 but visibly low — five points and three seeds is not enough for a confident fit, and the polylog factors that the paper's theorem includes have not washed out at d=64.

5. My RGO derivation differs from Appendix C in one sign. Documented above. The substantive algorithm is unchanged.

What I'd reuse, what I'd rewrite

The 193 lines of composite_sampler.py are the parts I'd keep. The SamplerProblem dataclass is a thin enough wrapper that, if you want to drop in a different smooth $f$ and a different RGO, you write two lambdas and a closure and you're done. That's not a lot of code for an algorithm whose theoretical claims are this specific.

What I'd rewrite if I were going further: the independent_mh inner loop is Python-with-numpy and spends most of its time in the Python interpreter, not in numerical work — about 90 µs per inner step for d=8. For a real production sampler I'd put the whole inner loop in a compiled language (the obvious move: numba.njit over independent_mh, which keeps the numpy calling convention; or jax.jit if I wanted to vmap across chains for the dimension sweep). At d=512 and tight thresholds the Python overhead starts to matter. For this writeup it's irrelevant.

Reproducing the runs

Inside the sandbox these three commands reproduce every table I quoted:

export PYTHONPATH=/tmp/pylib:/labs-output
python3 /labs-output/test_box_gaussian.py
python3 /labs-output/test_l1_gaussian.py
python3 /labs-output/dim_scaling.py

The seeds are fixed (20260528 for the validation tests, 1, 2, 3 for the dimension sweep), so the numbers in this post are byte-identical with what your machine prints. The whole battery runs in about three minutes on a single CPU core.

Artefacts

The companion tarball ships the four Python files referenced above: composite_sampler.py, test_box_gaussian.py, test_l1_gaussian.py, dim_scaling.py, plus the small diagnostic check_rgo_l1.py that surfaced the Appendix C typo.

References

@misc{LiuChewi2026,
  title  = {A proximal gradient algorithm for composite log-concave sampling},
  author = {Liu, Linghai and Chewi, Sinho},
  year   = {2026},
  month  = {May},
  eprint = {2605.12461},
  archivePrefix = {arXiv},
  primaryClass  = {math.ST},
  url    = {https://arxiv.org/abs/2605.12461},
}

Related work cited inline:

• [BEL18] Bubeck, S.; Eldan, R.; Lehec, J. Sampling from a log-concave distribution with projected Langevin Monte Carlo. Discrete & Computational Geometry, 2018. — the projected-Langevin dimension bound referenced when ruling out the naïve MYULA approach.
• [LST21] Lee, Y. T.; Shen, R.; Tian, K. Structured logconcave sampling with a restricted Gaussian oracle. COLT 2021. — the proximal sampler.
• [Mou+22] Mou, W.; Flammarion, N.; Wainwright, M. J.; Bartlett, P. L. An efficient sampling algorithm for non-smooth composite potentials. JMLR 2022. — Metropolis with an RGO proposal; uses the normalising constant.
• [Che+22] Chen, Y.; Chewi, S.; Salim, A.; Wibisono, A. Improved analysis for a proximal algorithm for sampling. COLT 2022. — the analysis the paper leans on.
• [AC24] Altschuler, J.; Chewi, S. Faster high-accuracy log-concave sampling via algorithmic warm starts. 2024. — the $\sqrt d$ RGO implementation in the smooth case.
• [DDM04] Daubechies, I.; Defrise, M.; De Mol, C. An iterative thresholding algorithm for linear inverse problems with a sparsity constraint. Comm. Pure Appl. Math. 2004. — ISTA, the canonical $\ell_1$ instance of the proximal-gradient iteration that this whole paper rhymes with.
• [Per16] Pereyra, M. Proximal Markov chain Monte Carlo algorithms. Statistics and Computing, 2016. — the original MYULA paper.
• [DMP18] Durmus, A.; Moulines, É.; Pereyra, M. Efficient Bayesian computation by proximal Markov chain Monte Carlo: when Langevin meets Moreau. SIAM J. Imaging Sciences, 2018. — the $W_2$ analysis of MYULA referenced for the $d/\varepsilon^2$ leading-order bound.