Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling
Robert Gruhlke, Julius Berner, David Sommer, Lorenz Richter
stat.ML, cs.LG
Submitted: 2026-07-07
Journal ref: ICML 2026
Code: https://github.com/robertgruhlke/TTD
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion models offer a powerful framework for sampling from complex probability densities by learning to reverse a noising process.
Terminology
Abstract
Diffusion models offer a powerful framework for sampling from complex probability densities by learning to reverse a noising process. A common approach involves solving for the time-reversed stochastic differential equation (SDE), which requires the score function of the evolving sample distribution. The logarithm of this distribution's density is governed by a Hamilton-Jacobi-Bellman (HJB) type partial differential equation (PDE). However, current methods for solving this PDE, such as PINNs or trajectory-based techniques, often suffer from long training times and significant sensitivity to hyperparameter tuning. In this work, we introduce a novel and efficient solver for the underlying HJB equation based on the functional tensor train (FTT) format. The FTT representation leverages latent low-rank structures to efficiently approximate high-dimensional functions, enabling both model compression and rapid computation. By integrating this efficient representation with a backward-in-time iterative scheme derived from backward stochastic differential equations (BSDEs), we develop a fast, robust and accurate sampling method. Our approach overcomes primary bottlenecks of existing techniques, enabling high-fidelity sampling from challenging target distributions with improved efficiency.
Sources
- NETS: A Non-Equilibrium Transport Sampler
- Approximation Theory of Tree Tensor Networks: Tensorized Univariate Functions -- Part I
- Approximation Theory of Tree Tensor Networks: Tensorized Univariate Functions -- Part II
- Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching
- Approximation Theory of Tree Tensor Networks: Tensorized Multivariate Functions
- Sequential Controlled Langevin Diffusions
- Approximation by tree tensor networks in high dimensions: Sobolev and compositional functions
- MCMC for multi-modal distributions
- Flow Annealed Importance Sampling Bootstrap
- Dynamical Measure Transport and Neural PDE Solvers for Sampling
- FP64 is All You Need: Rethinking Failure Modes in Physics-Informed Neural Networks
- Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems
- Diffusion-PINN Sampler
Related papers
- Behavior of prediction performance metrics with rare events
- Optimal Estimation of Generic Dynamics by Path-Dependent Neural Jump ODEs
- A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
- Scalable Posterior Uncertainty for Flexible Density-Based Clustering
- A Transport Geometry of Benign Overfitting
- Efficient quadratic entropy with distance sketches