-
Conditioning diffusion models by explicit forward-backward bridging
Authors:
Adrien Corenflos,
Zheng Zhao,
Simo Särkkä,
Jens Sjölund,
Thomas B. Schön
Abstract:
Given an unconditional diffusion model $π(x, y)$, using it to perform conditional simulation $π(x \mid y)$ is still largely an open question and is typically achieved by learning conditional drifts to the denoising SDE after the fact. In this work, we express conditional simulation as an inference problem on an augmented space corresponding to a partial SDE bridge. This perspective allows us to im…
▽ More
Given an unconditional diffusion model $π(x, y)$, using it to perform conditional simulation $π(x \mid y)$ is still largely an open question and is typically achieved by learning conditional drifts to the denoising SDE after the fact. In this work, we express conditional simulation as an inference problem on an augmented space corresponding to a partial SDE bridge. This perspective allows us to implement efficient and principled particle Gibbs and pseudo-marginal samplers marginally targeting the conditional distribution $π(x \mid y)$. Contrary to existing methodology, our methods do not introduce any additional approximation to the unconditional diffusion model aside from the Monte Carlo error. We showcase the benefits and drawbacks of our approach on a series of synthetic and real data examples.
△ Less
Submitted 22 May, 2024;
originally announced May 2024.
-
BlackJAX: Composable Bayesian inference in JAX
Authors:
Alberto Cabezas,
Adrien Corenflos,
Junpeng Lao,
Rémi Louf,
Antoine Carnec,
Kaustubh Chaudhari,
Reuben Cohn-Gordon,
Jeremie Coullon,
Wei Deng,
Sam Duffield,
Gerardo Durán-Martín,
Marcin Elantkowski,
Dan Foreman-Mackey,
Michele Gregori,
Carlos Iguaran,
Ravin Kumar,
Martin Lysy,
Kevin Murphy,
Juan Camilo Orduz,
Karm Patel,
Xi Wang,
Rob Zinkov
Abstract:
BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by taking a functional approach to the algorithms' implementation. BlackJAX is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs. The library integrates well w…
▽ More
BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by taking a functional approach to the algorithms' implementation. BlackJAX is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs. The library integrates well with probabilistic programming languages by working directly with the (un-normalized) target log density function. BlackJAX is intended as a collection of low-level, composable implementations of basic statistical 'atoms' that can be combined to perform well-defined Bayesian inference, but also provides high-level routines for ease of use. It is designed for users who need cutting-edge methods, researchers who want to create complex sampling methods, and people who want to learn how these work.
△ Less
Submitted 22 February, 2024; v1 submitted 16 February, 2024;
originally announced February 2024.
-
Nesting Particle Filters for Experimental Design in Dynamical Systems
Authors:
Sahel Iqbal,
Adrien Corenflos,
Simo Särkkä,
Hany Abdulsamad
Abstract:
In this paper, we propose a novel approach to Bayesian experimental design for non-exchangeable data that formulates it as risk-sensitive policy optimization. We develop the Inside-Out SMC$^2$ algorithm, a nested sequential Monte Carlo technique to infer optimal designs, and embed it into a particle Markov chain Monte Carlo framework to perform gradient-based policy amortization. Our approach is d…
▽ More
In this paper, we propose a novel approach to Bayesian experimental design for non-exchangeable data that formulates it as risk-sensitive policy optimization. We develop the Inside-Out SMC$^2$ algorithm, a nested sequential Monte Carlo technique to infer optimal designs, and embed it into a particle Markov chain Monte Carlo framework to perform gradient-based policy amortization. Our approach is distinct from other amortized experimental design techniques, as it does not rely on contrastive estimators. Numerical validation on a set of dynamical systems showcases the efficacy of our method in comparison to other state-of-the-art strategies.
△ Less
Submitted 29 May, 2024; v1 submitted 12 February, 2024;
originally announced February 2024.
-
Risk-Sensitive Stochastic Optimal Control as Rao-Blackwellized Markovian Score Climbing
Authors:
Hany Abdulsamad,
Sahel Iqbal,
Adrien Corenflos,
Simo Särkkä
Abstract:
Stochastic optimal control of dynamical systems is a crucial challenge in sequential decision-making. Recently, control-as-inference approaches have had considerable success, providing a viable risk-sensitive framework to address the exploration-exploitation dilemma. Nonetheless, a majority of these techniques only invoke the inference-control duality to derive a modified risk objective that is th…
▽ More
Stochastic optimal control of dynamical systems is a crucial challenge in sequential decision-making. Recently, control-as-inference approaches have had considerable success, providing a viable risk-sensitive framework to address the exploration-exploitation dilemma. Nonetheless, a majority of these techniques only invoke the inference-control duality to derive a modified risk objective that is then addressed within a reinforcement learning framework. This paper introduces a novel perspective by framing risk-sensitive stochastic control as Markovian score climbing under samples drawn from a conditional particle filter. Our approach, while purely inference-centric, provides asymptotically unbiased estimates for gradient-based policy optimization with optimal importance weighting and no explicit value function learning. To validate our methodology, we apply it to the task of learning neural non-Gaussian feedback policies, showcasing its efficacy on numerical benchmarks of stochastic dynamical systems.
△ Less
Submitted 21 December, 2023;
originally announced December 2023.
-
Parallel-in-Time Probabilistic Numerical ODE Solvers
Authors:
Nathanael Bosch,
Adrien Corenflos,
Fatemeh Yaghoobi,
Filip Tronarp,
Philipp Hennig,
Simo Särkkä
Abstract:
Probabilistic numerical solvers for ordinary differential equations (ODEs) treat the numerical simulation of dynamical systems as problems of Bayesian state estimation. Aside from producing posterior distributions over ODE solutions and thereby quantifying the numerical approximation error of the method itself, one less-often noted advantage of this formalism is the algorithmic flexibility gained…
▽ More
Probabilistic numerical solvers for ordinary differential equations (ODEs) treat the numerical simulation of dynamical systems as problems of Bayesian state estimation. Aside from producing posterior distributions over ODE solutions and thereby quantifying the numerical approximation error of the method itself, one less-often noted advantage of this formalism is the algorithmic flexibility gained by formulating numerical simulation in the framework of Bayesian filtering and smoothing. In this paper, we leverage this flexibility and build on the time-parallel formulation of iterated extended Kalman smoothers to formulate a parallel-in-time probabilistic numerical ODE solver. Instead of simulating the dynamical system sequentially in time, as done by current probabilistic solvers, the proposed method processes all time steps in parallel and thereby reduces the span cost from linear to logarithmic in the number of time steps. We demonstrate the effectiveness of our approach on a variety of ODEs and compare it to a range of both classic and probabilistic numerical ODE solvers.
△ Less
Submitted 2 October, 2023;
originally announced October 2023.
-
Debiasing Piecewise Deterministic Markov Process samplers using couplings
Authors:
Adrien Corenflos,
Matthew Sutton,
Nicolas Chopin
Abstract:
Monte Carlo methods - such as Markov chain Monte Carlo (MCMC) and piecewise deterministic Markov process (PDMP) samplers - provide asymptotically exact estimators of expectations under a target distribution. There is growing interest in alternatives to this asymptotic regime, in particular in constructing estimators that are exact in the limit of an infinite amount of computing processors, rather…
▽ More
Monte Carlo methods - such as Markov chain Monte Carlo (MCMC) and piecewise deterministic Markov process (PDMP) samplers - provide asymptotically exact estimators of expectations under a target distribution. There is growing interest in alternatives to this asymptotic regime, in particular in constructing estimators that are exact in the limit of an infinite amount of computing processors, rather than in the limit of an infinite number of Markov iterations. In particular, Jacob et al. (2020) introduced coupled MCMC estimators to remove the non-asymptotic bias, resulting in MCMC estimators that can be embarrassingly parallelised. In this work, we extend the estimators of Jacob et al. (2020) to the continuous-time context and derive couplings for the bouncy, the boomerang and the coordinate samplers. Some preliminary empirical results are included that demonstrate the reasonable scaling of our method with the dimension of the target.
△ Less
Submitted 27 June, 2023;
originally announced June 2023.
-
Variational Gaussian filtering via Wasserstein gradient flows
Authors:
Adrien Corenflos,
Hany Abdulsamad
Abstract:
We present a novel approach to approximate Gaussian and mixture-of-Gaussians filtering. Our method relies on a variational approximation via a gradient-flow representation. The gradient flow is derived from a Kullback--Leibler discrepancy minimization on the space of probability distributions equipped with the Wasserstein metric. We outline the general method and show its competitiveness in poster…
▽ More
We present a novel approach to approximate Gaussian and mixture-of-Gaussians filtering. Our method relies on a variational approximation via a gradient-flow representation. The gradient flow is derived from a Kullback--Leibler discrepancy minimization on the space of probability distributions equipped with the Wasserstein metric. We outline the general method and show its competitiveness in posterior representation and parameter estimation on two state-space models for which Gaussian approximations typically fail: systems with multiplicative noise and multi-modal state distributions.
△ Less
Submitted 19 June, 2023; v1 submitted 11 March, 2023;
originally announced March 2023.
-
Auxiliary MCMC and particle Gibbs samplers for parallelisable inference in latent dynamical systems
Authors:
Adrien Corenflos,
Simo Särkkä
Abstract:
We introduce two new classes of exact Markov chain Monte Carlo (MCMC) samplers for inference in latent dynamical models. The first one, which we coin auxiliary Kalman samplers, relies on finding a linear Gaussian state-space model approximation around the running trajectory corresponding to the state of the Markov chain. The second, that we name auxiliary particle Gibbs samplers corresponds to der…
▽ More
We introduce two new classes of exact Markov chain Monte Carlo (MCMC) samplers for inference in latent dynamical models. The first one, which we coin auxiliary Kalman samplers, relies on finding a linear Gaussian state-space model approximation around the running trajectory corresponding to the state of the Markov chain. The second, that we name auxiliary particle Gibbs samplers corresponds to deriving good local proposals in an auxiliary Feynman--Kac model for use in particle Gibbs. Both samplers are controlled by augmenting the target distribution with auxiliary observations, resulting in an efficient Gibbs sampling routine. We discuss the relative statistical and computational performance of the samplers introduced, and show how to parallelise the auxiliary samplers along the time dimension. We illustrate the respective benefits and drawbacks of the resulting algorithms on classical examples from the particle filtering literature.
△ Less
Submitted 1 March, 2023;
originally announced March 2023.
-
Parallel square-root statistical linear regression for inference in nonlinear state space models
Authors:
Fatemeh Yaghoobi,
Adrien Corenflos,
Sakira Hassan,
Simo Särkkä
Abstract:
In this article, we introduce parallel-in-time methods for state and parameter estimation in general nonlinear non-Gaussian state-space models using the statistical linear regression and the iterated statistical posterior linearization paradigms. We also reformulate the proposed methods in a square-root form, resulting in improved numerical stability while preserving the parallelization capabiliti…
▽ More
In this article, we introduce parallel-in-time methods for state and parameter estimation in general nonlinear non-Gaussian state-space models using the statistical linear regression and the iterated statistical posterior linearization paradigms. We also reformulate the proposed methods in a square-root form, resulting in improved numerical stability while preserving the parallelization capabilities. We then leverage the fixed-point structure of our methods to perform likelihood-based parameter estimation in logarithmic time with respect to the number of observations. Finally, we demonstrate the practical performance of the methodology with numerical experiments run on a graphics processing unit (GPU).
△ Less
Submitted 5 April, 2023; v1 submitted 29 June, 2022;
originally announced July 2022.
-
De-Sequentialized Monte Carlo: a parallel-in-time particle smoother
Authors:
Adrien Corenflos,
Nicolas Chopin,
Simo Särkkä
Abstract:
Particle smoothers are SMC (Sequential Monte Carlo) algorithms designed to approximate the joint distribution of the states given observations from a state-space model. We propose dSMC (de-Sequentialized Monte Carlo), a new particle smoother that is able to process $T$ observations in $\mathcal{O}(\log T)$ time on parallel architecture. This compares favourably with standard particle smoothers, th…
▽ More
Particle smoothers are SMC (Sequential Monte Carlo) algorithms designed to approximate the joint distribution of the states given observations from a state-space model. We propose dSMC (de-Sequentialized Monte Carlo), a new particle smoother that is able to process $T$ observations in $\mathcal{O}(\log T)$ time on parallel architecture. This compares favourably with standard particle smoothers, the complexity of which is linear in $T$. We derive $\mathcal{L}_p$ convergence results for dSMC, with an explicit upper bound, polynomial in $T$. We then discuss how to reduce the variance of the smoothing estimates computed by dSMC by (i) designing good proposal distributions for sampling the particles at the initialization of the algorithm, as well as by (ii) using lazy resampling to increase the number of particles used in dSMC. Finally, we design a particle Gibbs sampler based on dSMC, which is able to perform parameter inference in a state-space model at a $\mathcal{O}(\log(T))$ cost on parallel hardware.
△ Less
Submitted 4 February, 2022;
originally announced February 2022.
-
Temporal Gaussian Process Regression in Logarithmic Time
Authors:
Adrien Corenflos,
Zheng Zhao,
Simo Särkkä
Abstract:
The aim of this article is to present a novel parallelization method for temporal Gaussian process (GP) regression problems. The method allows for solving GP regression problems in logarithmic O(log N) time, where N is the number of time steps. Our approach uses the state-space representation of GPs which in its original form allows for linear O(N) time GP regression by leveraging the Kalman filte…
▽ More
The aim of this article is to present a novel parallelization method for temporal Gaussian process (GP) regression problems. The method allows for solving GP regression problems in logarithmic O(log N) time, where N is the number of time steps. Our approach uses the state-space representation of GPs which in its original form allows for linear O(N) time GP regression by leveraging the Kalman filtering and smoothing methods. By using a recently proposed parallelization method for Bayesian filters and smoothers, we are able to reduce the linear computational complexity of the temporal GP regression problems into logarithmic span complexity. This ensures logarithmic time complexity when run on parallel hardware such as a graphics processing unit (GPU). We experimentally demonstrate the computational benefits on simulated and real datasets via our open-source implementation leveraging the GPflow framework.
△ Less
Submitted 17 May, 2021; v1 submitted 19 February, 2021;
originally announced February 2021.
-
Differentiable Particle Filtering via Entropy-Regularized Optimal Transport
Authors:
Adrien Corenflos,
James Thornton,
George Deligiannidis,
Arnaud Doucet
Abstract:
Particle Filtering (PF) methods are an established class of procedures for performing inference in non-linear state-space models. Resampling is a key ingredient of PF, necessary to obtain low variance likelihood and states estimates. However, traditional resampling methods result in PF-based loss functions being non-differentiable with respect to model and PF parameters. In a variational inference…
▽ More
Particle Filtering (PF) methods are an established class of procedures for performing inference in non-linear state-space models. Resampling is a key ingredient of PF, necessary to obtain low variance likelihood and states estimates. However, traditional resampling methods result in PF-based loss functions being non-differentiable with respect to model and PF parameters. In a variational inference context, resampling also yields high variance gradient estimates of the PF-based evidence lower bound. By leveraging optimal transport ideas, we introduce a principled differentiable particle filter and provide convergence results. We demonstrate this novel method on a variety of applications.
△ Less
Submitted 30 June, 2021; v1 submitted 15 February, 2021;
originally announced February 2021.
-
Parallel Iterated Extended and Sigma-point Kalman Smoothers
Authors:
Fatemeh Yaghoobi,
Adrien Corenflos,
Sakira Hassan,
Simo Särkkä
Abstract:
The problem of Bayesian filtering and smoothing in nonlinear models with additive noise is an active area of research. Classical Taylor series as well as more recent sigma-point based methods are two well-known strategies to deal with these problems. However, these methods are inherently sequential and do not in their standard formulation allow for parallelization in the time domain. In this paper…
▽ More
The problem of Bayesian filtering and smoothing in nonlinear models with additive noise is an active area of research. Classical Taylor series as well as more recent sigma-point based methods are two well-known strategies to deal with these problems. However, these methods are inherently sequential and do not in their standard formulation allow for parallelization in the time domain. In this paper, we present a set of parallel formulas that replace the existing sequential ones in order to achieve lower time (span) complexity. Our experimental results done with a graphics processing unit (GPU) illustrate the efficiency of the proposed methods over their sequential counterparts.
△ Less
Submitted 31 January, 2021;
originally announced February 2021.