# HJ-sampler: A Bayesian sampler for inverse problems of a stochastic process by leveraging Hamilton-Jacobi PDEs and score-based generative models

arXiv: https://arxiv.org/abs/2409.09614 · HTML: https://mlsanity.com/paper/2409.09614/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 56 of its 109 references that are in the map.
- First posted on arXiv 15 Sep 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 56)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Solving Inverse Problems in Medical Imaging with Score-Based Generative Models](https://mlsanity.com/paper/2111.08005/) (2021)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [Uncertainty Quantification in Scientific Machine Learning: Methods, Metrics, and Comparisons](https://mlsanity.com/paper/2201.07766/) (2022)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise](https://mlsanity.com/paper/1906.02355/) (2019)
- [Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization](https://mlsanity.com/paper/2310.02679/) (2023)
- [Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance](https://mlsanity.com/paper/2212.06359/) (2022)
- [Protein Conformation Generation via Force-Guided SE(3) Diffusion Models](https://mlsanity.com/paper/2403.14088/) (2024)
- [Learning Functional Priors and Posteriors from Data and Physics](https://mlsanity.com/paper/2106.05863/) (2021)
- [Correcting model misspecification in physics-informed neural networks (PINNs)](https://mlsanity.com/paper/2310.10776/) (2023)
- [L-HYDRA: Multi-Head Physics-Informed Neural Networks](https://mlsanity.com/paper/2301.02152/) (2023)
- [A mean-field games laboratory for generative modeling](https://mlsanity.com/paper/2304.13534/) (2023)
- [Uncertainty quantification for noisy inputs-outputs in physics-informed neural networks and neural operators](https://mlsanity.com/paper/2311.11262/) (2023)
- [Deep Relaxation: partial differential equations for optimizing deep neural networks](https://mlsanity.com/paper/1704.04932/) (2017)
- [Algorithms for Overcoming the Curse of Dimensionality for Certain Hamilton-Jacobi Equations Arising in Control Theory and Elsewhere](https://mlsanity.com/paper/1605.01799/) (2016)

## Built on by (most cited first, 2 of 2)

- [From PINNs to PIKANs: Recent Advances in Physics-Informed Machine Learning](https://mlsanity.com/paper/2410.13228/) (2024)
- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)

## Abstract (arXiv)

The interplay between stochastic processes and optimal control has been extensively explored in the literature. With the recent surge in the use of diffusion models, stochastic processes have increasingly been applied to sample generation. This paper builds on the log transform, known as the Cole-Hopf transform in Brownian motion contexts, and extends it within a more abstract framework that includes a linear operator. Within this framework, we found that the well-known relationship between the Cole-Hopf transform and optimal transport is a particular instance where the linear operator acts as the infinitesimal generator of a stochastic process. We also introduce a novel scenario where the linear operator is the adjoint of the generator, linking to Bayesian inference under specific initial and terminal conditions. Leveraging this theoretical foundation, we develop a new algorithm, named the HJ-sampler, for Bayesian inference for the inverse problem of a stochastic differential equation with given terminal observations. The HJ-sampler involves two stages: (1) solving the viscous Hamilton-Jacobi partial differential equations, and (2) sampling from the associated stochastic optimal control problem. Our proposed algorithm naturally allows for flexibility in selecting the numerical solver for viscous HJ PDEs. We introduce two variants of the solver: the Riccati-HJ-sampler, based on the Riccati method, and the SGM-HJ-sampler, which utilizes diffusion models. We demonstrate the effectiveness and flexibility of the proposed methods by applying them to solve Bayesian inverse problems involving various stochastic processes and prior distributions, including applications that address model misspecifications and quantifying model uncertainty.