# Generative Modelling with Tensor Train approximations of Hamilton--Jacobi--Bellman equations

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 47 references that are in the map.
- First posted on arXiv 23 Feb 2024.

## Builds on (most cited first, 16 of 16)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Fast Sampling of Diffusion Models with Exponential Integrator](https://mlsanity.com/paper/2204.13902/) (2022)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [A Variational Perspective on Diffusion-Based Generative Models and Score Matching](https://mlsanity.com/paper/2106.02808/) (2021)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space](https://mlsanity.com/paper/2005.05409/) (2020)
- [Wasserstein Control of Mirror Langevin Monte Carlo](https://mlsanity.com/paper/2002.04363/) (2020)
- [Actor-Critic Method for High Dimensional Static Hamilton--Jacobi--Bellman Partial Differential Equations based on Neural Networks](https://mlsanity.com/paper/2102.11379/) (2021)
- [Numerically Solving Parametric Families of High-Dimensional Kolmogorov Partial Differential Equations via Deep Learning](https://mlsanity.com/paper/2011.04602/) (2020)
- [Ensemble Kalman Methods: A Mean Field Perspective](https://mlsanity.com/paper/2209.11371/) (2022)
- [Improving the convergence of reversible samplers](https://mlsanity.com/paper/1601.08118/) (2016)
- [Approximating optimal feedback controllers of finite horizon control problems using hierarchical tensor formats](https://mlsanity.com/paper/2104.06108/) (2021)
- [Dynamical low-rank approximations of solutions to the Hamilton-Jacobi-Bellman equation](https://mlsanity.com/paper/2111.14540/) (2021)

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

- [Importance Corrected Neural JKO Sampling](https://mlsanity.com/paper/2407.20444/) (2024)
- [Approximation and learning with compositional tensor trains](https://mlsanity.com/paper/2512.18059/) (2025)
- [Sampling from Boltzmann densities with physics informed low-rank formats](https://mlsanity.com/paper/2412.07637/) (2024)

## Abstract (arXiv)

Sampling from probability densities is a common challenge in fields such as Uncertainty Quantification (UQ) and Generative Modelling (GM). In GM in particular, the use of reverse-time diffusion processes depending on the log-densities of Ornstein-Uhlenbeck forward processes are a popular sampling tool. In Berner et al. [2022] the authors point out that these log-densities can be obtained by solution of a textitHamilton-Jacobi-Bellman (HJB) equation known from stochastic optimal control. While this HJB equation is usually treated with indirect methods such as policy iteration and unsupervised training of black-box architectures like Neural Networks, we propose instead to solve the HJB equation by direct time integration, using compressed polynomials represented in the Tensor Train (TT) format for spatial discretization. Crucially, this method is sample-free, agnostic to normalization constants and can avoid the curse of dimensionality due to the TT compression. We provide a complete derivation of the HJB equation's action on Tensor Train polynomials and demonstrate the performance of the proposed time-step-, rank- and degree-adaptive integration method on a nonlinear sampling task in 20 dimensions.