Optimal Preconditioning and Fisher Adaptive Langevin Sampling
Authors: Michalis Titsias
Venue: NeurIPS 2023 poster (2023)
Keywords: MCMC
Langevin diffusion
preconditioning
Fisher information
adaptive MCMC
score function

Abstract: We define  an optimal preconditioning for the Langevin diffusion by analytically optimizing the expected squared jumped distance. This yields as the optimal preconditioning an inverse Fisher information covariance matrix, where the covariance matrix is computed as the outer product of log target gradients averaged under the target.  We apply this result to the Metropolis adjusted Langevin algorithm (MALA)  and derive a computationally efficient adaptive MCMC scheme that learns the preconditioning from the history of gradients produced as the algorithm runs. We show in several experiments that the proposed algorithm is very robust in high dimensions and significantly outperforms other methods, including a closely related adaptive MALA scheme that learns the preconditioning with standard adaptive MCMC as well as the position-dependent  Riemannian manifold MALA sampler.

Optimal Preconditioning and Fisher Adaptive Langevin Sampling
Michalis K. Titsias Google DeepMind mtitsias@@oogle..om com
Abstract
We define an optimal preconditioning for the Langevin diffusion by analytically optimizing the expected squared jumped distance. This yields the optimal pre- conditioning an inverse Fisher information covariance matrix, where the covariance matrix is computed as the outer product of log target gradients averaged under the target. We apply this result to the Metropolis adjusted Langevin algorithm (MALA) and derive computationally efficient adaptive MCMC scheme that learns the preconditioning from the history of gradients produced as the algorithm runs. We show in several experiments that the proposed algorithm is very robust in high dimensions and significantly outperforms other methods, including a closely related adaptive MALA scheme that learns the preconditioning with standard adaptive MCMC as well as the position-dependent Riemannian manifold MALA sampler.
Introduction
Markov chain Monte Carlo (MCMC) general framework for simulating from arbitrarily complex distributions, and it has shown to be useful for statistical inference in a wide range of problems [18, 11j. The main idea an MCMC algorithm quite simple. Given a complex target π (z)) Markov chain is constructed using x-invariant transition kernel that allows to simulate dependent realizations æ ...... that eventually converge to samples from æ. These samples can used for Monte Carlo integration by forming ergodic averages. A general way to define r-invariant transition accept-reject mechanism in which the chain moves from state on
kernels the Catropolis-Hastingg to the next state xn+1 by first generating candidate state yn from proposal distribution ( (xn) and then it sets In+1 = Yo with probability a aIIn yn))
(yn) ((In()n) afEn, yn) = min(1, an T(xn) q(vn (x,)''
or otherwise rejects Yn and sets In+1 In. The choice of the proposal distribution g (x,,) is crucial because determines the mixing of the chain, i.e. the dependence of samples across time. For example, "slowly mixing" chain even after convergence may not be useful for Monte Carlo integration since it will output a highly dependent set of samples producing ergodic estimates of very high variance. Different ways of defining ((.)) lead common algorithms such as random walk Metropolis (RWM), Metropolis-adjuste
Langevin algorithm (MALA) [40, 35] and Hamiltonian Monte Carlo (HMC) [15, 28]. Within each class of these algorithms adaptation of parameters of the proposal distribution, such as a step size, is also important and this has been widely studied in the literature by producing optimal scaling results [34, 35, 36, 22, 6, 8, 38, 7, 39, 9]. and also by developing adaptive MCMC algorithms [21, 5, 37, 19, 2, 1, 4, 24]. The standard adaptive MCMC procedure in [21] uses the history of the chain to recursively compute empirical covariance of the target π and build a multivariate Gaussian proposal distribution. However, this type of covariance adaptation can be too slow and not robust in high dimensional settings [38, 3].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this paper, we derive fast and very robust adaptive MCMC technique in high dimensions that learns preconditioning matrix for the MALA method, which is the standard gradient- craddeet-based MCMC algorithm obtained by first-order discretization of the continuous-time Langevin diffusion. Our first contribution to define an optimal preconditioning by analytically optimizing criterion on the Langevin diffusion. The criterion is the well-known expected squared jumped distance [30] which at optimum yields as a preconditioner the inverse matrix I of the following Fisher information covariance matrix I = Ex(x)) V log π (xx) V log [. This contradicts the common belief in adaptive MCMC that the covariance of π is the best preconditioner. While this is a surprising result show that I connects with certain quantity appearing in optimal scaling of RWM [34, 36).
Having recognized I=! as the optimal preconditioning we derive an easy to implement and com- putationally efficient adaptive MCMC algorithm that learns from the history of gradients produced as MALA runs. This method sequentially updates empirical inverse Fisher estimate I,-² using a recursion having quadratic cost ((f) (d is the dimension of x) per iteration. In practice, since for sampling we need a square root matrix of 2,-1 we implement the recursions over a square root matrix by adopting classical results from Kalman filtering [31, 10J. We compare our method against MALA that learns the preconditioning with standard adaptive MCMC [21]. a
Riemannian manifold MALA [20] well as simple MALA (without preconditioning) and HMC. In several experiments we show that the proposed algorithm significantly outperforms all other methods.
2 Background
We consider an intractable target distribution π (z)) with x E R. known up some normalizing constant, and we assume that V log (<x) := v, log (2) is well defined. A continuous time process with stationary distribution ı is the overdamped Langevin diffusion
dæ
AV AV + VAdB.,
(2)
2
where B, denotes d-dimensional Brownian motion. This is stochastic differential equation (SDE) that generates sample paths such that for large 1, ~ т. We also incorporate a preconditioning matrix
A, which is a symmetric positive definite covariance matrix, while v is such that VAVA A.
Simulating from the SDE in (2) is intractable and the standard approach is to use a first-order Euler- Maruyama discretization combined with a Metropolis-Hastings adjustment. This leads to the so Langevin algorithm (MALA) where at each iteration given
called preconditioned Metropolis the current state Tn (where n = .... we sample yn from the proposal distribution
c(ym|rr.) Nyml + 0²
AV log (In), o² A):
where the step size o² O appears due time discretization. We accept yo with probability (IIn Bni) min (1, ax) where an follows the form in (1). The obvious way to compute a, is
Tyn) q( Bn) π (Bn) exp( -/ n yn 5 AV log T (3,, 113 4 }} an 7 (xx)) ((yn x (xx)) exp(
2 Un Zn 2/2 AV log π (Ix, .3'
where ||=||-- A However, in some cases that involve high dimensional targets, this can be costly since in the ratio of proposal densities both the preconditioning matrix A and its inverse A appear. In turns out that can avoid A-1 and simplify the computation as stated below.
Proposition 1. For preconditioned MALA with proposal density given by (3) the ratio of proposals in the M-H acceptance probubility can be written as
((Znlyn) exp(tilan, Un) - hly,, {n)), 122, 2
o² -- AV log (())))))))))))))))))
V log ((c).
(yyn
This expression does not depend on the inverse A-1, and this leads to computational gains and simplified implementation that we exploit in the adaptive MCMC algorithm presented in Section 4.
The motivation behind the use of preconditioned MALA is that with suitable preconditioner the mixing of the chain can be drastically improved, especially for very anisotropic target distributions.

---

A very general way to specify A is by applying an adaptive MCMC algorithm, which learns A online. To design such an algorithm is useful to first specify a notion of optimality. common argument in the literature, that is used for both RWM and MALA, is that a suitable A is the unknown covariance matrix Σ [21, 38, 24] of the target я. This means that we should learn A so that to approximate E. However, this argument is rather heuristic since it is not based on an optimality criterion. One of our contributions is to specify an optimal A based on an optimization procedure, that we describe in Section 3. This A will turn out to be not the covariance matrix of the target but an inverse Fisher information matrix.
3 Optimal preconditioning using expected squared jumped distance
Preconditioning aims to improve sampling when different directions (or individual variables x:;) in the state space can have different scalings under the target æ. Here, we develop method for selecting the preconditioning through the optimization of an objective function. This method uses the observation that an effective preconditioning correlates with large values of the global step size o² in MALA, i.e. o² is allowed to increase when preconditioning becomes effective as shown in the sampling efficiency scores in Table and the corresponding estimated step sizes reported in Appendix E.
In our analysis we consider the rejection-free or unadjusted Langevin sampler where we discretize the time continuous Langevin diffusion in (2) with a small finite := 0² > 0 so that
1 2₁ AV log T (2;) + VA(B++S Be), where Bi+5 B, ~ N(
(4)
We will use the expected squared jumped distance (ô, A) = 10[ - x=ll²1 computed follows. Proposition 2. If ae ~ (2x)) the vector J+45 - Ie defined by (4) has zero mean and covariance
8²
B[(a++a - (x248 ])
AB AB [ log (lgg((a; log (xx) A 8A.
(5)
Further, tr (E((31445 - (2) - x,)
E|(r ((I2+5 - 2x - x4) T)) = E1| x2.1%]
which shows that J(ō, A) is the trace of the covariance matrix in (5).
To control discretization error we impose upper bound constraint (6, A) S for small 0. A preconditioning that "symmetrizes" the target can be obtained by maximizing the discretization step size 5 subject to J(ō. A) c. Since J (6, A) monotonically increases with 8, the maximum 5* satisfies minA [ō"", A) = e. This means that the optimal preconditioning A is obtained by minimizing [6, A) under some global scale constraint A, stated next.
Proposition 3. Suppose A is a symmetric positive definite matrix satisfying tr( = c, with > 0 constant. Then the objective J for any 5 > 0, is miminized for A given by
A"= kI
k
==Br(w) [V log क (<) log π (x)
(6)
Lim1 pc
where μ&s are the eigenvalues of 'I assumed to satisfy 0 < 14 < co.
The positive multiplicative scalar & in (6) is not important since the specific value 0 is arbitrary, e.g. if we choose e = Σi= s4; then
1 and A' I-1 In other words, what matters is that the optimal A is proportional to the inverse matrix I-1 so it follows the curvature of I-1, For multivariate Gaussian π (z) (zrlu,,) 2)) it holds I 2,, so the optimal preconditioner coincides with the covariance matrix of x. More generally though, for non-Gaussian targets this will not hold.
Connection with classical Fisher information matrix. The matrix I is positive definite since it is the covariance of the gradient V log (ιz) := log (2x) where Ex VV log [(x) = 0. Also, I is similar to the classical Fisher information matrix. To illustrate some differences suppose that the target T (I) is a Bayesian posterior (0) x p(Y|0)p ()) = p(Y, 0) where Y are the observations are the random parameters. The classical Fisher information is a frequentist quantity where we fix some parameters and compute G()) = E ( Vo log p (0)) V log p\ p(V by averaging over data. In contrast, I = Ep(ry) Vo log PY 0) pYY, ( 8) [ is more like Bayesian quantity where we fix the data Y and average over the parameters 0. Importantly, I is not a function of 0 while G(O) is. Similarly to the classical Fisher information, I also satisfies the following standard property: Given that log (ar) twice differentiable and V2 log (xx) (rr) the Hessian matrix, I from (6) is also written as I (E) (z) 2 log (x)). Next, refer to I as the Fisher matrix.
and 0
3

---

Connection with optimal scaling. The Fisher matrix I connects also with the optimal scaling result for the RWM algorithm [34, 36]. Specifically, for targets of the form x(x) f(.. the RWM proposal [Vvn (())) = (yn, In, (o²/d)La) and as d - oo the optimal parameter o² is o² 2.38 where J f [( log [()
di )²] is the (univariate) Fisher information for the univariate density fz), and the preconditioning involves as in our case the inverse Fisher I = Id. This result has been generalized also for heterogeneous targets in [36] where again the inverse Fisher information matrix (having now a more general diagonal form) appears as the optimal preconditioner.
Fisher information adaptive MALA
Armed with the previous optimality result, we wish develop adaptive MCMC algorithm to optimize the proposal in (3) by learning online the global variance o² and the preconditioner A. For a² we follow the standard practice to tune this parameter in order to reach average acceptance rate around 0.574 as suggested by optimal scaling results [35, 36]. For the matrix A we want to adapt it that approximately it becomes proportional to the inverse Fisher Z-1 from (6). We also incorporate parametrization that helps the adaptation of o² to be more independent from the one of A. Specifically, we remove the global scale from by defining the overall proposal as
qyn (an) =N -N(1n (Vonle
o²
AV log (In), 2 ttr(A) A)
3 (r)(A)
where o² is normalized by ,,,((A), i.e. the average eigenvalue of A. Another way to view this is that the effective preconditioner is A/ A/()r((A) which has average eigenvalue equal to The proposal (7) is invariant to any scaling of A, i.c. if A is replaced by kA (with h > 0) the proposal remains the same. Also, note that when A is the identity matrix Ia (or multiple of identity) then (tr(Ta) = 1 and the above proposal reduces to standard MALA with isotropic step size o².
is straightforward to adapt 6² towards an average acceptance rate 0.574; see pseudocode in A Algorithm 1. Thus our main focus next is to describe the learning update for A, in fact eventually not for A itself but for square root matrix A which is what need to sample from the proposal in (7).
To start with, let us simplify notation by writing the score function the n-th MCMC iteration as 8n = Vzn om log क (. We introduce the n-sample empirical Fisher estimate
1,
S
la. n
(8)
n2
where À > is fixed damping parameter. Given that certain conditions apply [21, that the chain converges and ergodic averages converge exact expected values, I,, is consistent estimator satisfying lim,,+00 In I since as n + so the damping part À Ia vanishes. Including the damping is very important since it offers Tikhonov-like regularization, similar to ridge regression, and it ensures that for any finite the eigenvalues of I, are strictly positive. An estimate then for the preconditioner An can be set to be proportional to the inverse of the empirical Fisher 1,, i.e.
1
T2 Σ
An
8
Ia n
n
s;S
Ald
n1
Since any positive multiplicative sealar in front of A, plays no role, ignore the scalar and define A, = (Σi= 868, + -1a) Then, as MCMC iterates we can adapt A, in (d2)) cost per iteration based the recursion
Initialization: A = (sis) + Aaa) Aaa) = 1 la
8188
(10)
A A+
An-18 A
Iteration: A, = (A + 8,8 = A 1 S A
(11)
where we applied Woodbury matrix identity. This estimation in the limit can give the optimal preconditioning in the sense that under the ergodicity assumption, lim An z-1
In
a(A... mz

---

practice we do not need to compute directly the matrix An but square root matrix R₂ := VAn. such that R, An. since we need square root matrix draw samples from the proposal in (7). To express the corresponding recursion for R. we will rely on technique that dates back to the early days of Kalman filtering [31, 10], which applied to our case gives the following result.
An can be computed recursively in Proposition square root matrix Ryg such that R.R
(dd) time per iteration as follows:
Initialization: R1
( aT- + VA
-
1
5181
11
(12)
(Rx (Rn-1dn)0n
1
Iteration: Rr Rn-1 In
On Rm- Sn, On
(13)
+ 2x On
1+
1+0,, Do
A way to generalize the above recursive estimation of square root for the inverse Fisher matrix is to consider the stochastic approximation framework [33]. This requires to write an online learning update for the empirical Fisher of the form
in in-1 (3 - 1i, initialized I₁ = ssss + Ald:
(14)
where the leamning rates "tn satisfy the standard conditions En= 7n - 0o, En=1 7/ oc. Then, is straightforward to generalize the recursion for the square root matrix in Proposition 4 to account for this more general case; see Appendix B. The recursion in Proposition 4 is special case when The = L. In our simulations we did not observe significant improvement by using more general learning rate sequences, and therefore in all our experiments in Section 5 we use the standard learning rate 7n- = !.. Note that this learning rate is also used by other adaptive MCMC methods [21].
An adaptive algorithm that learns online from the score function vectors Sn can work well in some cases, but still it can be unstable in general. One reason is that 8n, = V log (xx,) will not have zero expectation when the chain transient and states T. are not yet draws from the stationary distribution т. To analyze this, note that the learning signal 8,, enters in the empirical Fisher estimator Z, through the outer product 8, as shown by Eqs. (8) and (14). However, in the transient phase 8,,8% will be biased since the expectation E[s,,,, T [s,,,, [[]ss,)ss,, E(s,,))] + [[s,]]][s,,]]]]]]]]]]][]]]] * I. where the expectations are taken under the marginal distribution of the chain at time n. In practice the mean vector [[(s,,, can take large absolute values, which can introduce significant bias through the additive term [ss,]][[,,]] T Thus, to reduce some bias we could track the empirical mean 8n, E 8; and center the signal 8n 8, so that the Fisher matrix is estimated by the empirical covariance [==(ss (s; sn.) The recursive estimation becomes similar to standard adaptive MCMC [21] where we recursively propagate an online empirical estimate for the mean of 8, and incorporate into the online empirical estimate of the covariance matrix (in our case the inverse Fisher matrix); see Eq. (17) in Section 5 for the standard adaptive MCMC recursion [21] and Appendix C for our Fisher method. While this can make learning quite stable we experimentally discovered that there is another scheme, presented next in Section 4.1, that is significantly better and stable especially for very anisotropic high dimensional targets; see detailed results in Appendix E.3.
4.1 Adapting to score function increments
An MCMC algorithm updates at each iteration its state according Tn+1 In + I(u, < (I yn (1n, 2n)) where (I,,,, ya) is the M-H probability, un ~ UO, is an uniform ran- dom number and I(-) is the indicator function. This sets In+1 either the proposal yn, or the previous state n, based on the binary value I(u, < a(In, Bn)). Similarly, we can consider the update of the score function s(x) V log (α)) and conveniently re-arrange it as increment,
5 s(xn+1) - s(x,)) = I(un < a(zn. yn (s(y.) s(xx))
(15)
While both 8n and S have zero expectation when In is from stationarity, i.e. In ~ ग,, the increment sf (unlike s,)) tends in practice to be more centered and close to zero even when the chain is transient, e.g. note that s is zero when yn is rejected. Further, since the difference s = a(xx+1) - s(x..) conveys information about the covariance of the score function
use it in the recursion of
5

---

Algorithm 1 Fisher adaptive MALA (blue lines are ommitted when not adapting (R,
Input:
Log target log π (2r); gradient V Vlog π (4); ;
(default X = 10); a, -0.574
Initialize 1 and o² by running simple MALA (i.e. with (vvl + (a² /2) log (xz), (2) T)) for no
(default 500) iterations where o² is adapted towards acceptance rate ax
Initialize square root matrix R = La and compute (log π (Ix),, V log (z1))
Initialize R = x²
# placeholder for the normalized step size o² 3 (
for For n = 1, 2,,,, 3
do
Propose Un Vn + ( ( V log x (In)) Rn NO(0 Ta)
Compute (log x (ym), V log T (yn))
..)) by using Proposition 1 Compute adaptation signal 5 Valan, Un (V) log π. (Un) - V log 7 (xx))
Compute (z vn) = min (1, elog ((yn
Use sy to adapt R based on Proposition 4 (if n = 1 use (12) and if n > 1 use (13))
Adapt step size o² + o² [1 Pn (a(zn; Uo) - o.)] # default const learning rate p.. 0.015
Normalize step size o = o² / (1((RR
(r(RRR) = sum(R. R R) which (dP) Accept/reject yn with probability a(En, Ba) to obtain (In+1, log (2n+1), V log (Ix+1))
end for
Proposition 4 to learn the preconditioner A, where we simply replace Sn by sh As shown in the experiments this leads to remarkably fast and effective adaptation of the inverse Fisher matrix I-1 without observable bias, or at least no observable for Gaussian targets where the true I-1 known. We can further apply Rao-Blackwe noo-Blackwlllization to reduce some variance of o Since & enters into the estimation of the empirical Fisher, see Eq. (8) or (14), through the outer product (s))) = I(u, < a( yn (s(m)) s(xn)) - s(xx)) we can marginalize out the n.v. a(xn, Un (s(x)))))
&
Un which yields E,,, [s (s)) Blackwellization an alternative vector to use for adaptation is
((xx))
s(x,)) T After this Rao-
$ = Va(2zn Un) - s(zn)),
(16)
which depends on the square root v aa Fn) of the M-H probability. As long as alI, Mn) > the learning signal in (16) depends on the proposed sample Un even when it is rejected.
express the full algorithm for Fisher information adaptive MALA as outlined by Algorithm 1, which adapts by using the tho-Blackwellized score function increments from Eq. (16). Note that, while Algorithm uses Sñ from Eq. (16), the initial signal from Eq. (15) works equally well; see Appendix E.3. Also, the algorithm includes an initialization phase where simple MALA runs for few iterations to move away from the initial state, discussed further in Section 5.
Finally,
5 Experiments
5.1 Methods and experimental setup
We apply the Fisher information adaptive MALA algorithm (FisherMALA) to high dimensional problems and we compare it with the following other samplers. (i) The simple MALA sampler with proposal N (1 Ix,, + (o² / log 7 (Ix,), o²I) which adapts only step size x² without having a preconditioner. (ii) A preconditioned adaptive MALA (AdaMALA) where the proposal follows exactly the from in (7) but where the preconditioning matrix is learned using standard adaptive MCMC based on the well-known recursion from [21]:
12
n 2
2 n -n
1 + - (Ix, pn-1)(In - Tn-1)
pn
Σ. En.
(17)
7
71
7
where the recursion is initialized at μ1 I1 and 2n [ u1 (1) - TT)T + AL, and À >> is the damping parameter that plays the same role as in FisherMALA. (iii) The Riemannian manifold MALA (mMALA) [20] which uses high dimensions runs slower than other schemes since the computation of A(z) may involve second derivatives and requires matrix decomposition that costs (dd)) per iteration. (iv) Finally, we include in the comparison the Hamiltonian Monte Carlo (HMC) sampler with fixed number of 10 leap frog steps and identity mass matrix. We leave the possibility to learn with our method a preconditioner in HMC for future work since this is more involved; see discussion at Section 7.
preconditioning matrix A(z). mMALA in

---

For all experiments and samplers we consider 2 x 10* burn-in iterations and 2 x 10 iterations for collecting samples. We set À 10 in FisherMALA and AdaMALA. Adaptation of the proposal distributions, i.e. the parameter o², the preconditioning or the step size of HMC, occurs only during burn-in and at collection of samples stage the proposal parameters are kept fixed. For all three MALA schemes the global step size o² is adapted to achieve an acceptance rate around 0.574 (see Algorithm 1) while the corresponding parameter for HMC is adapted towards 0.651 rate [9]. In MALA from the 2 x 10a burn-in iterations the first 500 iterations are used as the initialization phase in Algorithm I where samples are generated by just MALA with adaptable o². Thus, only the last 1.95 x 10t burn-in iterations used to adapt the preconditioner. For AdaMALA this initialization scheme proved to be unstable and we used a more elaborate scheme, as described in Appendix D.
We compute effective sample size (ESS) scores for each method by using the 2 x 10 samples from the collection phase. We estimate ESS across each dimension of the state vector xr and we report maximum, median and minimum values, by using the built-in method in TensorFlow Probability Python package. Also, we show visualizations that indicate sampling efficiency or effectiveness in estimating the preconditioner (when the ground truth preconditioner is known).
5.2 Gaussian targets
We consider three examples of multivariate Gaussian targets of the form T(x) (N(xa ju, 2), where the optimal preconditioner (up to any positive scaling) is the covariance matrix > since the inverse Fisher is --- Σ. For such case the Riemannian manifold sampler mMALA [20] is the optimal MALA sampler since it uses precisely Σ as the preconditioning. In contrast to mMALA which somehow has access to the ground-truth oracle, both FisherMALA and AdaMALA use adaptive recursive estimates of the preconditioner that should converge to the optimal Σ, and thus the question is which of them learns faster. To quantify this we compute the Frobenius norm || - S|لF across adaptation iterations n, where B denotes the matrix normalized by the average trace, i.c. ("LB)) ), for either An given by FisherMALA or A, := Σn, given by AdaMALA and where E is the optimal normalized preconditioner. The faster the Frobenius norm goes to zero the more effective is the corresponding adaptive scheme. For all three Gaussian targets the mean vector μ was taken to be the vector of ones and samplers were initialized by drawing from standard normal. The first example is a two-dimensional Gaussian target with covariance matrix Σ = 0.995; 0.995 1]. Both FisherMALA and AdaMALA perform almost the same (FisherMALA has faster convergence) in this low dimensional example as shown by Frobenius norm Figure la; see also Figure the Appendix for visualizations of the adapted preconditioners. The following two examples involve 100-dimensional targets.
B B/(( d
Gaussian process correlated target. We consider a Gaussian process to construct a 100- dimensional Gaussian where the covariance matrix is obtained by a non-stationary covariance function comprising the product of linear and squared exponential kernels plus small white noise,
i.e. [E]a_ 8:8j expl (as (da-)))² + 0.0018.; where the scalar inputs 8; form regular grid in 0.09
the range [1, 2]. Figure 1b shows the evolution of the Frobenius norms and panels d,c depict as 100 x 100 images the true covariance matrix and the preconditioner estimated by FisherMALA, For AdaMALA see Figure 6 in the Appendix. Clearly, FisherMAAA learns much faster and achieves more accurate estimates of the optimal preconditioner. Further Table I shows that FASA achieves significantly better ESS than AdaMALA and reaches the same performance with mMALA.
Inhomogeneous Gaussian target. In the last example we follow [28, 39] and we consider a Gaus- sian target with diagonal covariance matrix Σ (iag..r......................... 2oo) where the standard deviation values o, take values in the grid {0.01, o.. 1). This target is challenging because the different scaling across dimensions means that samplers with single step size, i.c. without preconditioning, will adapt to the smallest dimension of the state while the chain at the higher dimensions, such as T100, will be moving slowly exhibiting high autocorrelation. Note that FisherMALA and AdaMALA run without knowing that the optimal preconditioner is diagonal matrix, i.e. they learn a full covariance matrix. Figure 2a shows the ESS scores for all 100 dimensions of x for four samplers (except mMALA which has the same performance with FisherMALA), where we can observe that only FiL is able to achieve high ESS uniformly well across all dimensions. contrast, MALA and HMC that use single step size cannot achieve high sampling efficiency and their
7

---

AGAMALA
Rroprius Norm
50
AdaMALA RiseeMAAA
True
FisherMALA
1.0
0.0
10000
iteraioons
(a)
(b)
(c)
(d)
Figure 1: Panel (a) shows the Frobenius norm across burn-in iterations for the 2-D Gaussian and (b) for the GP target. The exact GP covariance matrix is shown in (c) and the estimated one by FisherMALA in (d).
1.0
0.8
FisherMALA
Frebentus Norm
0
100
AdaMALA
Std values
0.6
40
AdaMALA
FisherMALA MALA
20
0.0
20
40
80
100
20
40
80
100
5000
15000 Iterations
10000
20000
Dimension
Dimension
(a)
(b)
(c)
Figure 2: Results in the inhomogeneous Gaussian target.
ESS for dimensions close to T100 drops significantly. The same holds for AdaMALA due to its inability to learn fast the preconditioner, as shown by the Frobenius norm values in Figure 2c and the estimated standard deviations in Figure 2b. AdaMALA can eventually get very close to the optimal precondtioner but requires hundred thousands of adaptive steps, while FisherMALA learns with only few thousand steps.
5.3 Bayesian logistic regression
We consider Bayesian logistic regression distributions of the form π (0]Y, Z) x p(Y|0, 2)p(0) with data (Y, Z) {y), 24 I== where Rd is the imput vector and y/ the binary label. The likelihood is P( |, Z) = I1 (0 2, 2.) ( 2,)) 1-", where 0 E IRd are the random parameters assigned the prior p(0) = -N La). We consider six binary classification datasets (Australian Credit, Heart, Pima Indian, Ripley, German Credit and Caravan) with number of data ranging from 250 to n 5822 and dimensionality of the 0 ranging from 3 to 87. We also consider a much higher 785-dimensional example MNIST for classifying the digits 5 and 6, that has 11339 training examples. To make the inference problems more challenging, in the first six examples we do not standardize the inputs 2, which creates very anisotropic posteriors over 9. For the MNIST data, which initially are grey-scale images in [0, 255], we simply divide by the maximum pixel value, i.e. 255, to bring the images in [0, 1]. In Table we report the ESS for the low 7-dimensional Pima Indians dataset, the medium -7-iimension Caravan dataset and the higher 785-dimensional MNIST dataset, while the results for the remaining datasets shown in Appendix E. Further, Figure 3 shows the evolution of the unnormalized log target density log(p(Y Z (0)) for the best four samplers in
AdaMALA
HMC
MMALA
RisherAAL
-1240
Log target
1260
1280
-1300
20000
20000
Iterations
Figure 3: dataset.
The evolution of the log-target across iterations for the best four algorithms Caravan

---

Table 1: ESS scores are averages after repeating the simulations 10 times under different random initializations.
Max ESS
Median ESS
Min ESS
GP target (d 100)
MALA
15.235 = 4.246
6.326 = 1.934
3.619 = 0.956
AdaMALA
845.100 = 80.472
662.978 = 81.127
552.377 = 74.441
HMC
17.625 = 6.706
6.680 + 2.205
4.315 + 0.889
mMALA
2109.441 + 101.553
2007.640- + 104.867
1841.978 + 114.266
FisherMALA
2096.259 94.751
1923.753 + 95.820
1784.962 + 104.440
Pima Indian (d=7)
MALA
106.668 + 29.601
1 3.821
4.061 + 1.587 1.587
AdaMALA
211.948 + 133.363
52.277 + 26.566
6.401 + 3.344
HMC
1624.773 = 544.777
337.100 = 212.158
6.052 2.062
mMALA
6086.948 1: 117.241
5690.967 118.401
5297.835 1 160.084
FisherMALA
6437.4119 207.548
5981.960 + 156.072
5628.541 168.425
Caravan (d 87)
MALA
27.247 = 7.554
5.890 = 0.398
2.906 = 0.150
AdaMALA
41.522 = 9.343
7.144 0.663
3.144 = 0.135
HMC
787.901 + 173.863
37.303 + 6.808
4.238 + 0.532
mMALA
179.899 + 61.502
121.867 41.801
51.114 + 24.995
FisherMALA
2257.737 45.289
1920.903 55.821
498.016 + 96.692
MNIST (d 785)
MALA
34.074 4.977
7.589 0.149
2.944 + 0.066
AdaMALA
62.301 = 9.203
8.188 + 0.399
2.985 + 0.089
HMC
889.386 + 118.050
303.345 = 10.976
114.439 = 20.965
mMALA
51.589 = 3.447
20.222 - 1.259
5,240 0.492
FisherMALA
1053.455 35.680
811.522 19.165
439.580 1. 52.800
Caravan dataset which visualizes chain autocorrelation. From all these results we can conclude that FisherMALA is better than all other samplers, and remarkably it outperforms significantly the position-dependent mMALA, especially in the high dimensional Caravan and MNIST datasets.
6
Related Work
There exist works that use some form of global preconditioning in gradient-based samplers for specialized targets such as latent Gaussian models [13, 44], which make use the tractable Gaussian prior. Our method differs, since it is more agnostic to the target and learns a preconditioning from the history of gradients, analogously how traditional adaptive MCMC learns from states [21, 38].
Several research works use dosition-dependens preconditioning A(x) within gradient-basse samplers, such as MALA. This is for example the idea behind Riemannian manifold MALA [20] and extensions [45]. Similar to Riemannian manifold methods there are approaches inspired by second order optimization that use the Hessian matrix, or some estimate of the Hessian, for sampling in a MALA- style manner [32, 17, 27J. Recently, such samplers and their time-continuous diffusion limits have been theoretically analyzed by obtaining convergence guarantees [12, 26). All such methods form position-dependent preconditioning and not the preconditioning we use in this paper, e.g. note that Z-1 we consider here requires an expectation under the target and thus it is always a global preconditioner rather than a position-dependent one. Another difference is that our method has quadratic cost, while position-dependent preconditioning methods have cubic cost and they require computationally demanding quantities like the Hessian matrix. Therefore, in order for these methods to run faster some approximation may be needed, e.g. low rank [27] or quasi-Newton type [46, 24]. Furthermore, the Bayesian logistic results in Table (see also Figure 3) show that the proposed FisherMALA method significantly outperforms manifold MALA [20] in Caravan and MNIST examples, despite the fact that manifold MALA preconditions with the exact negative inverse Hessian matrix of the log target. This could suggest that deeenenndeependent preconditioning may be less effective in certain type of high--dimensional and log-concave problems.

---

Finally, there is recent work for learning flexible MCMC proposals by using neural networks [42, 25, 23, 41] and by adapting parameters using differentiable objectives [25, 29, 43, 14]. Our method differs, since it does not use objective functions (which have extra cost because they require an optimization to run in parallel with the MCMC chain), but instead adapts similarly to traditionally MCMC methods by accumulating information from the observed history of the chain.
7 Conclusion
We derived an optimal preconditioning for the Langevin diffusion by optimizing the expected squared jumped distance, and subsequently developed an adaptive MCMC algorithm that approximates the optimal preconditioning by applying an efficient quadratic cost recursion. Some possible topics for future research are: Firstly, would be useful to investigate whether the score function differences that we use as the adaptation signal introduce any bias in the estimation of the inverse Fisher matrix. Secondly, would be interesting to extend our method learn the preconditioning for other gradient- based samplers such as Hamiltonian Monte Carlo (HMC), where such a matrix there is referred to as the inverse mass matrix. For HMC this is more complex since both the mass matrix and its inverse are needed in the iteration. Finally. it could be interesting to investigate adaptive schemes of the inverse Fisher matrix by using multiple parallel and interacting chains, similarly to ensemble covariance matrix estimation for Langevin diffusions [16].
Acknowledgments and Disclosure of Funding
We are grateful to the reviewers for their comments. Also, we wish to thank Arnaud Doucet, Sam Power, Francisco Ruiz, Jiaxin Shi, Yee Whye Teh, Siran Liu, Kazuki Osawa and James Martens for useful discussions.
References
[1] Christophe Andrieu and Yves Atchade. On the efficiency of adaptive meme algorithms. Elec- tronic Communications in Probability, 12:336-349, 2007.
[2] Christophe Andrieu and fric Moulines. On the ergodicity properties of some adaptive meme algorithms. The Annals of Applied Probability, 16(3):1462-1555 1505, 2006.
[3] Christophe Andrieu and Johannes Thoms. A tutorial on adaptive memc. Statistics and computing, 18(4):343 101133333333,-22 2008
[4] Yves Atchade, Gersende Fort, Eric Moulines, and Pierre Priouret. Adaptive markov chain monte carlo: theory and methods. Preprint, 2009.
[5] Yves F Atchadé and Jeffrey Rosenthal. On adaptive markov chain monte carlo algorithms. Bernoulli, 1155):815-828, 828, 2005.
[6] Myléne Bédard. Weak convergence of metropolis algorithms for non-iid target distributions. The Annals of Applied Probability, pages 1222-1244 1244. 2007.
[7] Mylène Bédard. Efficient sampling using metropolis algorithms: Applications of optimal scaling results. Journal of Computational and Graphical Statistics, 17(2):312 2008.
[8] Mylene Bedard. Optimal acceptance rates for metropolis algorithms: Moving beyond 0.234. Stochastic Processes and their Applications, 118(12)2198-2222, 2008
[9] Alexandros Beskos, Natesh Pillai, Gareth Roberts, Jesus-Maria Sanz-Serna, and Andrew Stuart. Optimal tuning of the hybrid monte carlo algorithm. Bernoulli, 1955A)1501-1534, 2013.
[10] Gerald J. Bierman. Factorization Methods for Discrete Sequential Estimation, volume 128. 1977.
[11] Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of Markov Chain Monte Carlo. CRC press, 2011.
10

---

[12] Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Exponential ergodicity of lanror-langevin diffusions. In Proceedings of the 34th International Conference on Neural Information Processing Systems, 2020.
[13] S. L. Cotter, G. O. Roberts, A. M. Stuart, and D. White. MCMC methods for functions: modifying old algorithms to make them faster. Statistical Science, 28(3):224446 2013.
[14] Ameer Dharamshi, Vivian Ngo, and Jeffrey S. Rosenthal. Sampling by divergence minimization. Communications in Statistics -Simulation and Computation, 2023.
[15] Simon Duane, A. D. Kennedy, Brian J. Pendleton, and Duncan Roweth. Hybrid monte carlo. Physics Letters B, 195(2):166- - 222, 1987.
[16] Alfredo Garbuno-Inigo, Franca Hoffmann, Wuchen Li, and Andrew M. Stuart. Interacting langevin diffusions: Gradient structure and ensemble kalman sampler. SIAM Journal on Applied Dynamical Systems, 19(1): 41112444,4 2020.
[17] John Geweke and Hisashi Tanizaki. Note on the sampling distribution for the metropolis- hastings algorithm. Communications in Statistics Theory and Methods, 32(4):775-789, 2003.
[18) W.R. Gilks, S. Richardson, and D. Spiegelhalter. Markov Chain Monte Carlo in Practice. Chapman & Hall/CRC Interdisciplinary Statistics. Taylor & Francis, 1995.
[19] Paolo Giordani and Robert Kohn. Adaptive independent metropolis-hastingg by fast estimation of mixtures of normals. Journal of Computational and Graphical Statistics, 19(2):243-259, 2010.
[20] Mark Girolami and Ben Calderhead. Riemann manifold Langevin and Hamiltonian Monte Carlo methods, Journal of the Royal Statistical Society: Series (Statistical Methodology). 73(2):123-2144 2011.
[21] Heikki Haario, Eero Saksman, and Johanna Tamminen. An adaptive metropolis algorithm. Bernoulli, 7(2):223-242, 2001.
[22] Heikki Haario, Eero Saksman, and Johanna Tamminen. Componentwise adaptation for high dimensional meme. Computational Statistics, 20(2):265-273, 2005.
[23] Raza Habib and David Barber. Auxiliary variational meme. To appear ICLR 2019, 2019.
[24] Benedict Leimkuhler, Charles Matthews, and Jonathan Weare. Ensemble preconditioning for markov chain monte carlo simulation. Statistics and Computing, 28, 03 2018
[25]
Daniel Levy, Matt D. Hoffman, and Jascha Sohl-Dickstein. Generalizing hamiltonian monte carlo with neural networks. In International Conference on Learning Representationss, 2018.
[26] Ruilin Li, Molei Tao, Santosh S. Vempala, and Andre Wibisono. The mirror langevin algorithm converges with vanishing bias. In Sanjoy Dasgupta and Nika Haghtalab, editors, Proceedings of The 33rd International Conference on Algorithmic Learning Theory, volume 167 of Proceedings of Machine Learning Research, pages 718-742. PMLR, 29 Mar-01 Apr 2022.
[27] James Martin, Lucas C. Wilcox, Carsten Burstedde, and Omar Ghattas. A stochastic newton meme method for large-scale statistical inverse problems with application to seismic inversion. SIAM Journal on Scientific Computing, 1460--4746-- 2012.
[28] Radford M. Neal. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 54:113-163 2010.
[29] Kirill Neklyudov, Pavel Shvechikov, and Dmitry Vetrov. Metropolis-hastings view on variational inference and adversarial training. arXiv preprint arXiv:/810.07151, 2018.
[30] Cristian Pasarica and Andrew Gelman. Adaptively scaling the metropolis algorithm using expected squared jumped distance. Statistica Sinica, pages 343-364, 2010.
[31] James E. Potter and Robert G. Stern. Statistical filtering of space navigation measurements. 1963.
11

---

[32] Yuan Qi and Thomas P. Minka. Hessian-based Markov Chain Monte-cane carto Algorithms. In First Cape Cod Workshop on Monte Carlo Methods, Cape Cod, Mass, 2002.
[33] H. Robbins and S. Monro, stochastic approximation method. Annals of Mathematical Statistics, 22:400-407, 1951.
[34] Gareth O Roberts, Andrew Gelman, and Walter R Gilks. Weak convergence and optimal scaling of random walk metropolis algorithms. The annals of applied probability, 7(1):110-122) 1997.
[35] Gareth O Roberts and Jeffrey S Rosenthal. Optimal scaling of discrete approximations to langevin diffusions. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 60(1):255-268, 1998.
[36] Gareth O Roberts and Jeffrey S Rosenthal. Optimal scaling for various cettrpolis-hastinng algorithms. Statistical Science, pages 351-367, 2001.
[37] Gareth O Roberts and Jeffrey S Rosenthal. Coupling and ergodicity of adaptive markov chain monte carlo algorithms. Journal of applied probability, 44(2):458 475, 2007.
[38] Gareth O Roberts and Jeffrey S Rosenthal. Examples of adaptive meme. Journal of Computa- tional and Graphical Staristics, 18(2):349-366 2009.
[39] Jeffrey S Rosenthal. Optimal proposal distributions and adaptive memc. In Handbook of Markov Chain Monte Carlo, pages 114-132. Chapman and Hall/CRC, 2011.
[40] J. Rossky, J. Doll, and H. L. Friedman. Brownian dynamics as smart Monte Carlo simulation. The Journal of Chemical Physics, (1110):668-46333 November 1978.
[41] T. Salimans, D. P. Kingma, and M. Welling. Markov chain Monte Carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, 2015.
[42] Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-nice-me: Adversarial training for meme. In Advances Neural Information Processing Systems, pages 5140-5150, 2017.
[43] Michalis Titsias and Petros Dellaportas. Gradient-based adaptive markov chain monte carlo. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché-Bue, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[44] Michalis Titsias and Omiros Papaspiliopoulos. Auxiliary gradient- based sampling algorithms. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 80, 10 2016.
[45] Tatiana Xifara, Chris Sherlock, Samuel Livingstone, Simon Byrne, and Mark Girolami. Langevin diffusions and the metropolis addjupoli--ajjusee langevin algorithm. Statistics & Probability Letters, 04 2014.
[46] Yichuan Zhang and Charles A. Sutton. Quasi-newton methods for markov chain monte carlo. In Advances in Neural Information Processing Systems 24: 25th Annual Conference on Neural Information Processing Systems 201 Proceedings of meeting held 12- 14 December 201 Granada, Spain., pages 2393-2401, 2011.
12


Reviews:
[{"review_id": "zBFf5eYCFL", "reviewer": "Reviewer_bdsU", "summary": "This is a nice paper that develops a preconditioner for improved sampling of complex target distributions using gradient history.  The algorithm is well presented and is placed in context of other recent proposals from the literature.   There are several numerical experiments presented on small to medium sized examples that support the idea.  \n\nA nice feature is that the paper is very readable which will enhance the uptake of the idea.\n\nThe numerical experimentation is good but not exhaustive, nonetheless convincing to this reader.", "questions": "What were wall clock times for all the different computations?  In particular they should be reported for the larger models.\n\nYou give many results in the table for Bayesian logistic regression.  How did you parameterize the various methods?\n\nHow does the performance of the method presented in the paper scale up in a parallel computing setting?\n", "limitations": "The authors miss the chance to introduce the work by indicating the importance of sampling in large scale data science in providing improved robustness and the potential for uncertainty quantification.  I would suggest that the introduction be modified to develop this theme.   By its very nature, this work has potential to strengthen the foundations of AI and this should at least be mentioned.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Reasonable overview of literature.   Careful presentation.  Readable style.  ", "weaknesses": "Limited experimentation in high dimensional settings.     Limited ablation study.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": " \"What were wall clock times for all the different computations? In particular they should be reported for the larger models.\" \n\nIn the Table below we report the overall running time for all samplers in the largest MNIST Bayesian logistic regression problem. These running times were obtained by running all algorithms using a v100 GPU. We can observe the following. All MALA variants have similar running times (besides mMALA) and this is because the gradient computation of the log target dominates the running time and the quadratic cost needed for FisherMALA is somehow insignificant compared to the gradient computation. mMALA is slower because it requires computation of the Hessian plus the Hessian matrix decomposition which adds an extra cost, while the slowest is HMC because it requites multiple gradient evaluations per iteration. \n\n**Running times (in seconds) for MNIST**\n\n**MALA**          $1900.086 \\pm 60.186$ \n\n**AdaMALA**    $1927.065 \\pm 190.311$ \n\n**HMC**            $8046.461 \\pm 272.085$ \n\n**mMALA**        $4442.111 \\pm 207.497$ \n\n**FisherMALA**  $1910.230 \\pm 112.168$ \n\n\"You give many results in the table for Bayesian logistic regression. How did you parameterize the various methods?\"\n\nSection 5.1 provides details about how adaptation is done for all the different samplers and for all targets (including all Bayesian logistic regression examples in the Appendix, i.e. all runs follow Section 5.1), such as how the global $\\sigma^2$ is adapted, number of burn-in iterations etc. Please let us know if this information is not sufficient.   \n\n\"How does the performance of the method presented in the paper scale up in a parallel computing setting?\" \n\nIf we use multiple parallel MCMC chains then we would like perhaps these chains to communicate information and adapt a preconditioner that uses gradients from all chains. How to do this efficiently is an interesting research topic, but we believe that this is beyond the scope of the current work and it is suitable for future research.     \n\n\"The authors miss the chance to introduce the work by indicating the importance of sampling in large scale... \"\n\nThank you for this comment. We will try to follow reviewer's suggestion and add a related sentence in the first paragraph of the introduction.  \n "}, {"review_id": "Ht44EoY10Z", "reviewer": "Reviewer_NzaV", "summary": "Preconditioned Metropolis-adjusted Langevin Algorithm is considered. Optimality of the preconditioning matrix, with respect to expected squared jump distance, is found to be proportional to an object related to the inverse Fisher information. In contrast to the Fisher, the expectation is taken with respect to model parameters, giving an object that is not position dependent (in contrast to related Riemannian manifold MALA, and RMHMC). An online algorithm, based on classical tools from Kalman filtering, is given for efficiently computing the inverse of a Monte Carlo approximation to the Fisher Information, leading to a novel adaptive preconditioned MALA method. Debiasing in the transient phase is discussed, and a Rao-Blackwellization scheme is given to reduce variance of the estimator. Experiments are conducted on a range of standard testbed tasks.", "questions": "Numerical performance of the non-centred version of FisherMALA are not reported, yet there is a burn-in period before the Fisher is computed that is usually used to mitigate the influence of the transient phase. How much effect does centering the scores have on performance, and how do the empirical means change throughout sampling?\n\nThe RB vs no-RB versions of your method appear to have performed very similarly, however with an acceptance rate tuned to <0.6, it seems natural to ask - how do they compare in terms of ESS/s? If the quadratic step is only being computed less than 60% of the time, does this negate the perceived/marginal benefits of using RB? This may also help demonstrate your performance advantage over mMALA.", "limitations": "Limitations are addressed in weaknesses/questions. There is little view for negative societal impact in this type of work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper provides a well principled approach to preconditioning MALA that shows improvements over both adaptive and Riemannian MCMC methods. In particular, it provides a quadratic approach to a Fisher-informed preconditioner that does not require cubic inversion/decomposition at every step, instead only needing quadratic updates. This approach outperforms the more expensive mMALA on nontrivial tasks.\n\nThe paper is communicated in an accessible way, and the experimental section covers a reasonable range of problems for samplers of this type.", "weaknesses": "The main body of the text claims that the Rao-Blackwellized version is significantly better, but this is not reflected in the numerics in Appendix E. This is not a weakness of the method itself, but appears misleading. When appendix E is mentioned, maybe \"detailed comparison between the proposed methods\" is clearer than \"detailed results\".\n\nThere are numerous grammatical errors and awkward phrasings, and the document would benefit from another read through with this in mind.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Regarding the Rao-Blackwellized version the reviewer is correct. Our numerical results in the Appendix E indeed show that the Rao-Blackwellized version is not significantly different than the non Rao-Blackwellized version, and both methods work well. Thank you for spotting this. Therefore, following reviewer's suggestion we will change the title of Section 4.1 from \"Adapting to Rao-Blackwellized score function increments\"to  \"Adapting to score function increments\" and remove or de-emphasise any claim about Rao-Blackwellization being significantly better.      \n \nRegarding \"detailed comparison between the proposed methods\" is clearer than \"detailed results\". Thank you, we will add this change\nand resolve also the grammatical errors. \n\nRegarding the comment about \"Numerical performance of the non-centred version of FisherMALA\" we include in the rebuttal pdf file the Table 1 that includes also the \"FisherMALA (non-centered)\" version as requested by the reviewer. Note that due to space limitations in the Table we include only 5 benchmarks, i.e. the four from the main paper plus the Ripley logistic regression example.  Recall, that as mentioned in the main paper, the burn-in period before the Fisher is computed is 500 iterations in all runs. From this table we can see that \"FisherMALA (non-centered)\" performs worse than the other FisherMALA variants, and only on Ripley dataset works equally well with the rest. Of course, if we increase the 500-iterations burn-in window before Fisher computation starts, so that we get closer to convergence, \"FisherMALA (non-centered)\" could improve.  However, the main proposed FisherMALA versions that learn from the increments are much more robust to this choice and learn efficiently even when we start the adaptation early in the transient phase. We believe this robustness can be important in practice.  \n\nFor many targets (especially the high-dimensional ones) the empirical means of the gradients can be very far away from zero in the transient phase (due to space limitation we do not include a plot here in the rebuttal pdf).\n\nRegarding \"RB vs no-RB versions running times\", we would like to point to an answer given above to reviewer bdsU, where we provided a Table of actual running times (experiments were done by using a v100 GPU) in the most expensive MNIST benchmark. Note that all MALA variants have similar running times (besides mMALA) and this is because the gradient computation of the log target dominates the running time and the quadratic cost needed for FisherMALA is somehow insignificant compared to the gradient computation. mMALA is slower because it requires computation of the Hessian plus its $O(m^3)$ matrix decomposition which adds an extra cost, while the slowest is HMC because  it requites multiple gradient evaluations per iteration. FisherMALA-no-RB is not included in the table but it is like FisherMALA (i.e. recall that FisherMALA in the main paper corresponds to the variant FisherMALA-with-RB as denoted in Table 7 in the Appendix).        "}, {"review_id": "7IOGuldCpB", "reviewer": "Reviewer_K8aZ", "summary": "The authors study a preconditioning matrix for the Langevin diffusion. It is given by maximizing the expected squared jumped distance. It turns out that the preconditioning is an inverse Fisher information covariance matrix for the target distribution. The authors apply the MALA scheme with this preconditioner to compute the MCMC algorithms. It is shown that in the Gaussian target distribution, this preconditioner performs as the Newton-type method. Numerical examples demonstrate the effectiveness of this method. The authors compared it with different methods, such as the position-dependent Riemannian manifold MALA sampler. \n", "questions": "The authors also discuss the Hamiltonian Monte-Carlo methods with Fisher information preconditioners. Can authors briefly discuss the motivation why the Hamiltonian Monte-Carlo with preconditioners will behave better than the overdamped Langevin dynamics? Would you explain any potential advantages for the Hamiltonian Monte-Carlo methods, in terms of convergence analysis? ", "limitations": "There are no limitations. ", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The authors present an important example of preconditioners, namely Fisher information matrix for Langevin dynamics. \nIt is motivated by the Gaussian target distribution or the expected squared jumped distance. The method is similar to Newton's method in the Gaussian case.  \n\n2. The numerical examples demonstrate the method's effectiveness for the Gaussian target distributions. \n\n", "weaknesses": "1. The computation of the inverse of Fisher-information may not be simple. The authors propose some modification methods to approximate it. Some computational complexity analysis is needed. \n\n2. There is no theory supporting the convergence speed for non-Gaussian settings. Suppose the target distribution is a mixed Gaussian or a distribution that is non-log concave. Can the authors provide some numerical examples for these non-Gaussian target distributions to test the algorithm's performance?  It would be curious to know how Fisher information preconditioners behave. \n\n3.  There is some related literature in this direction. One picks different choices of preconditioners motivated by the Gaussian target distributions.  \n\nAlfredo Garbuno-Inigo, et al. Interacting Langevin Diffusions: Gradient Structure And Ensemble Kalman Sampler, SIAM Journal on applied dynamical sysetm, 2019. \nYifei Wang, et.al. Projected Wasserstein gradient descent for high-dimensional Bayesian inference, SIAM/ASA Journal on Uncertainty Quantification, 2022.\nYifei Wang, et.al. Accelerated Information Gradient Flow, Journal of Scientific computing, 2021\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Regarding the iterative computation of the inverse Fisher matrix, we will add the complexity analysis in Appendix A.4 (where we prove Proposition 4) to fully explain the $O(d^2)$ cost per iteration.  To describe here the main part of this analysis note that the main adaptive computation is the update of the square root of the Fisher inverse in Equation (13). This is done in the following order. Firstly,  the vector $\\phi_n = R_{n-1}^\\top s_n$ is computed which is a matrix-vector multiplication with $O(d^2)$ cost. The next step is to compute the scalar $r_n$ in $O(d)$ (involving the dot product $\\phi_n^\\top \\phi_n$) and then the scaled vector $\\phi_n' = \\frac{r_n}{1 + \\phi_n^\\top \\phi_n} \\phi_n$ also an $O(d)$ operation. Then we need two additional $O(d^2)$ multiplication operations to obtain firstly  the vector $t_n = R_{n-1} \\phi_n$ and secondly the outer vector product $t_n (\\phi_n')^\\top$. Finally, the update is $R_n = R_{n-1} - t_n (\\phi_n')^\\top$ which requires a final $O(d^2)$ addition operation of two matrices which is typically cheaper than $O(d^2)$ multiplication.  Therefore, the overall complexity is $O(d^2)$. Following reviewer's suggestion, we will include this analysis to the Appendix A.4. \n\nRegarding the theoretical justification and convergence analysis mentioned by the reviewer,  we would like to point out that our current work focuses on the methodological and computational/algorithmic development. Theoretical analysis requires substantial more work, and we plan to address it in the future.\n\nRegarding points 2 and 3 we agree with the reviewer that the preconditioner we propose is best suited for log-concave or close to log-concave distributions. For multi-modal distributions, as the mixture example mentioned by the reviewer, our precondioner could be much less effective (especially in case we have to average across modes of possibly \"conflicting\" curvature). To deal with such cases, our method needs to be extended, by following e.g. the suggestions from the literature given by the reviewer.  We can discuss this in our final section 7.    \n\nIn the discussion (Section 7) we mentioned the potential application of our method to HMC, motivated also by the wide popularity of HMC. One common intuitive argument in the literature (see e.g. Neal 2011) in favour of HMC compared to MALA is that HMC often can avoid random walk behaviour due to the use of the momentum. However, an interesting question is whether this argument still holds when an effective preconditioning is used for both MALA and HMC, or in such case any benefit from using HMC is reduced. We don't know the answer! But we believe that this is an interesting question to investigate in the future both experimentally and theoretically. "}, {"review_id": "kxnCdtKEgr", "reviewer": "Reviewer_Zcpk", "summary": "The paper presents an “optimal” preconditioning for the Langevin diffusion by analytically maximizing the expected squared jumped distance. The authors have identified the optimal preconditioning as an inverse Fisher information covariance matrix. They apply this result to the Metropolis adjusted Langevin algorithm (MALA) and derive an efficient adaptive MCMC scheme that learns the preconditioning from the gradient history produced during the algorithm's execution. The authors show through experiments that their proposed algorithm can outperform other standard methods.", "questions": "The Fisher information preconditioner, which is the expectation of the negative Hessian of the log-density, is very reminiscent of the Newton–Langevin diffusion which uses the same but without the expectation; see S. Chewi, T. Le Gouic, C. Lu, T. Maunu, P. Rigollet, A. J. Stromme, \"Exponential ergodicity of mirror-Langevin diffusion,\" and Appendix A.4 of R. Li, M. Tao, S. S. Vempala, A. Wibisono's work, \"The mirror Langevin algorithm converges with vanishing bias.\" Please provide a comparison with this method, both in theory and practice.\n\nThe Newton–Langevin diffusion in particular enjoys convergence guarantees in continuous time that attest to its affine-invariant nature. Can anything be proven regarding convergence rates with this choice of preconditioner?\n\nI did not see any claim that $\\pi$ is the stationary distribution of Algorithm 1. Can the authors provide justification or proof for this?", "limitations": "Yes.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper introduces an innovative approach with the use of the inverse Fisher information covariance matrix as a preconditioner for MALA; to my knowledge this has not been previously proposed. The preconditioner is shown to maximize the expected squared jump distance, which could potentially improve the convergence of the algorithm.\n\nThe adaptive learning of the preconditioning from the history of gradients generated during the algorithm execution is a useful idea.\n\nThe experiments presented demonstrate the algorithm’s promise.\n", "weaknesses": "The paper lacks theoretical justification to support the empirical findings, which might raise questions about the universality and optimal applicability of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for bringing the interesting mirror Langevin diffusion work to our attention, and we will discuss it in the related work. By reading the papers of S. Chewi et al and Li et al, we can observe that for Gaussian target $\\pi(x)=\\mathcal{N}(x|\\mu, \\Sigma)$ the discrete version of Newton Langevin Diffusion (NLD) is precisely the Riemannian manifold MALA (mMALA). This is because in such case  the NLD proposal (e.g. by discretizing the Equation 33 in Appendix A.4 in Li et al) reduces to \n$$\nx_{n+1} = x_n + \\frac{\\sigma^2}{2} \\Sigma \\nabla \\log \\pi(x_n) + \n\\sigma \\sqrt{\\Sigma} \\epsilon_n \n$$\nwhich is the mMALA from Girolami and Galderhead (2011). We will clarify that mMALA in such case reduces to NLD. For the logistic regression case mMALA uses as a preconditioner the negative Hessian matrix (see also Section 7 in Girolami and Calderhead (2011)). \n$$\nA(\\theta) \n= Z^\\top \\Lambda(\\theta) Z + I_d\n$$\nwhere $\\Lambda(\\theta)$ is a diagonal matrix with entries $\\lambda_i = \\sigma(\\theta^\\top z_i) (1 - \\sigma(\\theta^\\top z_i))$. This is precisely the negative Hessian matrix of the log joint $- \\frac{1}{2} || \\theta ||^2 + \\sum_{i=1}^m y_i \\log  \\sigma(\\theta^\\top z_i) (1 - \\sigma(\\theta^\\top z_i))$ (see section 5.3 in our paper). We will clarify this by mentioning the connection with NLD. Therefore, for Bayesian logistic regression the mMALA that we have included in the comparison is very similar to NLD. Of course, NLD for this non-Gaussian example includes also an additional non-zero drift term that requires 3-order derivatives. While the numerical Newton iterative method, described in Appendix E.2 in S. Chewi et al, avoids 3-order derivatives it will still increase considerably the running time. For instance, \nnote that mMALA alone which requires in every iteration a Choleksy decomposition of the Hessian matrix (an operation of cubic cost) is already more than twice slower than FisherMALA and the other MALA schemes in the largest MNIST dataset. \n\nRegarding the theoretical justification and proving convergence rates with this choice of preconditioner, reviewer's comment  is valid since our work does not provide such results. We would like to point out that our current work focuses on the methodological and computational/algorithmic developments. Theoretical analysis requires substantial more work which is beyond the scope of the current paper, and we plan to address it in the future. \n"}]

Author rebuttal:
We thank all reviewer for the time and for their constructive comments.  Responses to comments of each individual reviewer are given after each corresponding review.

Here, in the general comments we would like to point out to all reviewers that after the submission we added a correction to Section 3 of the paper since we noticed  that the optimal preconditioner must be the minimum and not the maximum of the expected squared jumped distance (ESJD).  And indeed the inverse Fisher matrix that we propose in the paper is the minimum of the ESJD. Here, we provide the part of Section 3 that has been re-written.  Nothing else changes in the rest of the paper apart from the above change, ie.  that the optimal preconditioner which is the inverse Fisher, is the minimum and not the maximum of ESJD. For completeness here we provide the text that has been re-written in Section 3 and also the proof that the inverse Fisher is the minimizer of ESJD.      
   
**The main changes in Section 3 are the following**

We develop a method for selecting the preconditioning through the optimization of an objective function.  This method uses the observation that an effective preconditioning should make the target more well-conditioned or symmetrized which subsequently allows to use the largest possible discretization step size $\sigma^2$ in MALA.  An illustration of this is shown in Figure 1 in the attached pdf of this rebuttal
where our proposed inverse Fisher preconditiorer indeed gives the largest $\sigma^2$ when compared to other MALA samplers.   
 
Given the time discretized Langevin diffusion  
$$
x_{t + \delta} - x_t = \frac{\delta}{2} A \nabla \log \pi(x_t)  + \sqrt{A} (B_{t+ \delta} - B_t),  \ \ \text{where} \ \  B_{t+ \delta} - B_t \sim \mathcal{N}(0,\delta I).
$$
the expected squared jumped distance is $J(\delta, A) = E[||x_{t+\delta}  - x_t ||^2]$ where  $x_t \sim \pi$.   To control discretization error we can impose an upper bound constraint $J(\delta, A) \leq \epsilon$ for small $\epsilon > 0$. To obtain a preconditioning $A$ that "symmetrizes" the target as much as possible we would like to maximize the discretization step size $\delta$ subject to $J(\delta, A) \leq \epsilon$. Since $J(\delta, A)$ monotonically increases with $\delta$, the maximum $\delta^*$ satisfies $\min_A J(\delta^*, A) = \epsilon$. This means that the optimal preconditioning $A^*$ is obtained by minimizing the expected squared jumped distance, i.e.  by solving $\min_A J(\delta, A)$ under some global scale constraint on $A$, as stated next.  


**Proposition 3**

Suppose $A$ is a symmetric positive definite matrix satisfying 
$\text{tr}(A) = c$, with $c>0$ a constant.  Then the objective $J(\delta, A)$ is miminized for $A^*$ given by
$$
A^* = k \mathcal{I}^{-1}, \ \ k=\frac{c}{\sum_{i=1}^d \frac{1}{\mu_i} }, \ \ \ \mathcal{I} = E_{\pi(x)} 
\left[ \nabla \log \pi(x) 
\nabla \log \pi(x)^\top \right],
$$
where $\mu_i$s are the eigenvalues of $\mathcal{I}$ assumed to satisfy $0 < \mu_i < \infty$. 



**For completeness also we provide here in the rebuttal the new proof of Proposition 3** 

**Proof of Proposition 3** 

The expected squared jumped distance is written as 
$$
J(\delta, A) = \text{tr} \left(\frac{\delta^2}{4} A \mathcal{I} A +  \delta A \right)
=\frac{\delta^2}{4} 
\text{tr}(A \mathcal{I} A) + \delta c, 
$$
where we used 
$\text{tr}(A) = c$. Since $c$ is  a constant to minimize $J(\delta, A)$ is the same as minimizing $\text{tr}(A \mathcal{I} A)$ under the constraint that $A$ is symmetric positive definite matrix and $\text{tr}(A)=c$.  To deal with the equality constraint  we consider the Lagrangian 
$
\text{tr}(A \mathcal{I} A) - \lambda ( \text{tr}(A) - c). 
$
By taking derivatives wrt $A$ (using the matrix derivative identities  $\frac{\partial}{\partial X} \text{tr}(X B X) = X^\top B^\top + B^\top X^\top$ and $\frac{\partial}{\partial X} \text{tr}(X) = I_d$ for arbitrary $d \times d$ square matrices $X, B$) and setting to zero we obtain the linear equation
$
A^\top \mathcal{I} + \mathcal{I} A^\top =  \lambda I_d,
$
where we used that $\mathcal{I}$ is a symmetric matrix.  This is a set of linear equations
and given that each eigenvalue $\mu_i$ of $\mathcal{I}$ satisfies $0 < \mu_i < \infty$,  so that $\mathcal{I}$ is invertible,  there is an unique solution given by $A = (1/2) \lambda \mathcal{I}^{-1}$. The Lagrange multiplier $\lambda$  is chosen so that $\text{tr}(A)=c$ which leads to the optimal $A^*$ 
$$
A^* = \frac{c}{\sum_{i=1}^d \frac{1}{\mu_i}} \mathcal{I}^{-1}.  
$$
$A^*$ turned out to be symmetric and  positive definite as  desired.  For this $A^*$ the optimal loss value is $\text{tr}(A^* \mathcal{I} A^*) = \frac{c^2}{\sum_{i=1}^2  \frac{1}{\mu_i}}$, for which we further need to disambiguate whether this is the global minimum or maximum. We can do this by choosing a different matrix that satisfies the constraint $\text{tr}(A)=c$ and  compare its loss with the optimal loss  $\frac{c^2}{\sum_{i=1}^d  \frac{1}{\mu_i}}$ . For example, one such matrix is $A = \frac{c}{d} I_d$, which has loss value $\frac{c^2  (\sum_{i=1}^d \mu_i)}{d^2}$. Then by using the Cauchy-Schwarz inequality $d^2 = (\sum_{i=1}^d \frac{\sqrt{\mu_i}}{\sqrt{\mu_i}})^2 \leq (\sum_{i=1}^d \mu_i ) (\sum_{i=1}^d \frac{1}{\mu_i})$ 
we obtain $\frac{c^2  (\sum_{i=1}^d \mu_i)}{d^2} \geq  \frac{c^2  (\sum_{i=1}^d \mu_i)}{  (\sum_{i=1}^d \mu_i ) (\sum_{i=1}^d \frac{1}{\mu_i})}$ = $ \frac{c^2}{\sum_{i=1}^d  \frac{1}{\mu_i}}$.
This shows that $A^*$ achieves the global minimum which completes the proof.  

Meta-review: Reviewers for this paper were quite positive. I tend to side with the authors regarding some of the reviewer’s concerns about a lack of theoretical justification; I see this paper as a methods/algorithms paper and I don’t think a thorough theoretical analysis is necessary for acceptance. I encourage the authors to incorporate their responses from the rebuttal into the paper, in addition to making an editing pass throughout the paper focusing on grammar and typos. Overall, I would like to congratulate the authors on their great work.
Decision: Accept (poster)