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

Robert Gruhlke, Julius Berner, David Sommer, Lorenz Richter
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:36905-36941, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-gruhlke26a, title = {Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling}, author = {Gruhlke, Robert and Berner, Julius and Sommer, David and Richter, Lorenz}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {36905--36941}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/gruhlke26a/gruhlke26a.pdf}, url = {https://proceedings.mlr.press/v306/gruhlke26a.html}, 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.} }
Endnote
%0 Conference Paper %T Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling %A Robert Gruhlke %A Julius Berner %A David Sommer %A Lorenz Richter %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-gruhlke26a %I PMLR %P 36905--36941 %U https://proceedings.mlr.press/v306/gruhlke26a.html %V 306 %X 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.
APA
Gruhlke, R., Berner, J., Sommer, D. & Richter, L.. (2026). Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:36905-36941 Available from https://proceedings.mlr.press/v306/gruhlke26a.html.

Related Material