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

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.

Generating realistic data requires sampling from complex probability distributions. A particularly challenging setting is when no data from the target distribution is available, and one must rely on the distribution's mathematical description alone. Diffusion models can handle this by learning to reverse a noising process, which requires solving a mathematical equation tracking how probability evolves over time. Existing solvers for this equation are slow and sensitive to hyperparameter choices. We introduce a new solver based on a compact representation called a functional tensor train, which exploits hidden low-dimensional structure to efficiently approximate high-dimensional functions. Combined with an algorithm that works backwards through time, this yields a method that is faster, more stable, and easier to use than existing approaches. Our method enables accurate sampling from challenging distributions that were previously difficult to handle - a useful advance for scientific computing and generative machine learning.