Optimal Preconditioning and Fisher Adaptive Langevin Sampling

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.

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. The adaptive learning of the preconditioning from the history of gradients generated during the algorithm execution is a useful idea. The experiments presented demonstrate the algorithm’s promise.

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.

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. The 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? I did not see any claim that $\pi$ is the stationary distribution of Algorithm 1. Can the authors provide justification or proof for this?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Yes.

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.

Strengths

1. The authors present an important example of preconditioners, namely Fisher information matrix for Langevin dynamics. It 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. 2. The numerical examples demonstrate the method's effectiveness for the Gaussian target distributions.

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. 2. 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. 3. There is some related literature in this direction. One picks different choices of preconditioners motivated by the Gaussian target distributions. Alfredo Garbuno-Inigo, et al. Interacting Langevin Diffusions: Gradient Structure And Ensemble Kalman Sampler, SIAM Journal on applied dynamical sysetm, 2019. Yifei Wang, et.al. Projected Wasserstein gradient descent for high-dimensional Bayesian inference, SIAM/ASA Journal on Uncertainty Quantification, 2022. Yifei Wang, et.al. Accelerated Information Gradient Flow, Journal of Scientific computing, 2021

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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

There are no limitations.

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.

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. The 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". There are numerous grammatical errors and awkward phrasings, and the document would benefit from another read through with this in mind.

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? The 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations are addressed in weaknesses/questions. There is little view for negative societal impact in this type of work.

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. A nice feature is that the paper is very readable which will enhance the uptake of the idea. The numerical experimentation is good but not exhaustive, nonetheless convincing to this reader.

Strengths

Reasonable overview of literature. Careful presentation. Readable style.

Weaknesses

Limited experimentation in high dimensional settings. Limited ablation study.

Questions

What were wall clock times for all the different computations? In particular they should be reported for the larger models. You give many results in the table for Bayesian logistic regression. How did you parameterize the various methods? How does the performance of the method presented in the paper scale up in a parallel computing setting?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

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.

Response

Thank you for the response. I am satisfied with the rebuttal and I would like to raise my score to a 6.

Reply to authors

Do you have any comments and discussions on the references?