Backpropagation-Free Metropolis-Adjusted Langevin Algorithm

TL;DR

This paper introduces the first backpropagation-free gradient-based MCMC algorithms using forward-mode automatic differentiation, reducing computational costs and improving performance in Bayesian inference tasks.

Contribution

It presents four novel backpropagation-free MALA algorithms that incorporate forward-mode AD and Hessian information, expanding the toolkit for scalable Bayesian inference.

Findings

Forward MALA is computationally cheaper than traditional MALA.

Forward-mode samplers can outperform standard MALA depending on the model.

The proposed algorithms are effective on hierarchical models and Bayesian neural networks.

Abstract

Recent work on backpropagation-free learning has shown that it is possible to use forward-mode automatic differentiation (AD) to perform optimization on differentiable models. Forward-mode AD requires sampling a tangent vector for each forward pass of a model. The result is the model evaluation with the directional derivative along the tangent. In this paper, we illustrate how the sampling of this tangent vector can be incorporated into the proposal mechanism for the Metropolis-Adjusted Langevin Algorithm (MALA). As such, we are the first to introduce a backpropagation-free gradient-based Markov chain Monte Carlo (MCMC) algorithm. We also extend to a novel backpropagation-free position-specific preconditioned forward-mode MALA that leverages Hessian information. Overall, we propose four new algorithms: Forward MALA; Line Forward MALA; Pre-conditioned Forward MALA, and Pre-conditioned Line Forward MALA. We highlight the reduced computational cost of the forward-mode samplers and show that forward-mode is competitive with the original MALA, while even outperforming it depending on the probabilistic model. We include Bayesian inference results on a range of probabilistic models, including hierarchical distributions and Bayesian neural networks.