Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling
This work addresses the challenges of solving high-dimensional Hamilton-Jacobi-Bellman (HJB) equations within diffusion models, which typically suffer from prolonged training times and high sensitivity to hyperparameters. The paper introduces, for the first time, a functional tensor train (FTT) low-rank structure into this setting, integrating it with backward stochastic differential equations (BSDEs) and a reverse-time iterative algorithm. This combination enables efficient approximation of high-dimensional density functions and facilitates rapid score-based sampling. The proposed method substantially enhances sampling efficiency and stability for complex high-dimensional distributions, significantly reducing training time and diminishing reliance on careful hyperparameter tuning across multiple benchmarks.