# Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 30 of its 72 references that are in the map.
- First posted on arXiv 7 Jul 2026.

## Placed in

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

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

- [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)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations](https://mlsanity.com/paper/1706.04702/) (2017)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems](https://mlsanity.com/paper/2307.08423/) (2023)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [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)
- [VarGrad: A Low-Variance Gradient Estimator for Variational Inference](https://mlsanity.com/paper/2010.10436/) (2020)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Score-Based Diffusion meets Annealed Importance Sampling](https://mlsanity.com/paper/2208.07698/) (2022)
- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [Underdamped Diffusion Bridges with Applications to Sampling](https://mlsanity.com/paper/2503.01006/) (2025)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Interpolating between BSDEs and PINNs: deep learning for elliptic and parabolic boundary value problems](https://mlsanity.com/paper/2112.03749/) (2021)
- [Dynamical Measure Transport and Neural PDE Solvers for Sampling](https://mlsanity.com/paper/2407.07873/) (2024)
- [Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators](https://mlsanity.com/paper/2412.00088/) (2024)
- [From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training](https://mlsanity.com/paper/2501.06148/) (2025)
- [Solving high-dimensional parabolic PDEs using the tensor train format](https://mlsanity.com/paper/2102.11830/) (2021)
- [FP64 is All You Need: Rethinking Failure Modes in Physics-Informed Neural Networks](https://mlsanity.com/paper/2505.10949/) (2025)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [Diffusion-PINN Sampler](https://mlsanity.com/paper/2410.15336/) (2024)
- [From continuous-time formulations to discretization schemes: tensor trains and robust regression for BSDEs and parabolic PDEs](https://mlsanity.com/paper/2307.15496/) (2023)

## Abstract (arXiv)

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.