Skip to content
Preprint

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

Jul 2026 · 0 citations · 72 references
Mathematics Computer Science

TL;DR

This work introduces a novel and efficient solver for the underlying HJB equation based on the functional tensor train (FTT) format, which leverages latent low-rank structures to efficiently approximate high-dimensional functions, enabling both model compression and rapid computation.

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.

View source

Similar papers

Preprint Jul 2026

Tensor-Train Joint Modeling for Few-Step Discrete Diffusion

The first framework for explicit joint distribution modeling in discrete diffusion via tensor decomposition, which represents the conditional clean distribution as a low-rank tensor with controllable expressivity, is presented.

Byoungkwon Kim, Minhyuk Sung · 0 citations
Jul 2026

Variational Inference and Density Estimation with Non-negative Tensor Train

An efficient numerical approach for compressing a high-dimensional discrete distribution function into a non-negative tensor train (NTT) format and observing that the proposed NTT fitting procedure exhibits drastically faster convergence than an alternative multiplicative update method that has been previously proposed is observed.

Xun Tang, R. Dwaraknath, Lexing Ying · 0 citations
Conference Jul 2026

DATT: Dimension-Augmented Tensor-Train Decomposition for Neural Network Compression

Traditional model reduction techniques for neural network compression usually leverage existing matrix/tensor decomposition methods to explore the redundancy among model parameters, which may overlook the local structures of the models. In this work, we propose a novel low-rank approximation algorithm, called Dimension-Augmented Tensor-Train (DATT) decomposition, which integrates dimension slicing with the Tensor-Train decomposition. DATT allows different partitions to be represented using distinct, non-shared spanning sets. This enables the decomposition to adapt more flexibly to localized structure, potentially achieving better compression. With the localized low-rank structures, DATT provides a more compact representation of neural network weights. We evaluated the performance of DATT with the other state-of-the-art low-rank approximation-based model compression algorithms, and demonstrated that DATT outperforms them in terms of parameter reduction and model accuracy. In particular, for compressing ResNet-20 on CIFAR-10, DATT achieves a 77.00% reduction in parameters with only a 0.37% drop in top-1 accuracy.

Yu-Chuan Tai, Cheng-Yu Sie, Che-Rung Lee · 0 citations
Preprint Aug 2026

Accelerated Learning of High Dimensional Functions with a Tensor-Featured Training Network

This work presents a method to accelerate the optimization of learning high dimensional functions using deep neural network (DNN), and studies the effect of adding features which distill pretrained DNN into TNs using a discretize and decompose strategy.

Karl Pierce, Y. Khoo, Haizhao Yang · 0 citations
Preprint Jul 2026

Tensor-Based Reduced-Order Modeling for Optimization-Based Inverse Problems

We develop a tensor reduced-order modeling (TROM) framework for optimization-based inverse problems governed by parameter-dependent dynamical systems. The approach approximates the parameter-to-observation map directly in tensor-train format, using either TT-SVD or TT-Cross compression, and integrates the resulting representation into a regularized nonlinear least-squares formulation. Beyond accelerating forward evaluations, the low-rank tensor structure is used to reformulate the inverse problem in reduced coordinates, assemble the Gauss--Newton quantities without forming the full observation-space Jacobian, and perform TROM-based objective minimization over the discrete parameter grid. This tensor optimization step can be used either as a stand-alone approximate minimization procedure or as a data-informed initialization for a subsequent Gauss--Newton solve. The method is studied for two inverse problems: an inverse heat-transfer problem in a heterogeneous medium, where the unknown parameters describe the locations of multiple low-conductivity inclusions, and a FitzHugh--Nagumo parameter-estimation problem with a highly nonconvex optimization landscape. Numerical experiments assess the effects of ROM approximation error, measurement noise, regularization, initialization, spatial discretization, and increasing parameter dimension. The results show that TROM can reproduce the behavior of full-order inversion at a substantially reduced online cost. The experiments also demonstrate that reduced-coordinate inversion, tensor-based optimization, and appropriate regularization improve robustness in higher-dimensional, noisy, and strongly nonconvex regimes.

S. Islam, Andreas Mang, Maxim A. Olshanskii · 1 citation
Preprint Aug 2026

Regularized High-Dimensional Additive Tensor Autoregressive Model

High-dimensional time series has diverse applications in econometrics and finance. Recent models for capturing temporal dependence have employed a bilinear representation for matrix time series, or the Tucker-decomposition based representation in case of tensor time series. A Tucker-decomposition based temporal effect is difficult to interpret on many occasions, along with its computational complexity due to the non-convex nature of the underlying optimization problem. Moreover, the existing tensor models have not sufficiently explored the possibilities of imposing any lower-dimensional pattern on the transition matrices. In this work, we propose a regularized additive tensor autoregressive model with additive interaction of row-wise, column-wise and tube-wise temporal dependence, that offers more interpretability, less computational burden due to its convex nature and estimation of the underlying low rank plus sparse pattern of its transition matrices. We address the issue of identifiability of the various components in our model and subsequently develop a scalable alternating block minimization algorithm for estimating the parameters. We provide a finite sample error bound under high-dimensional scaling for the model parameters. Finally, the efficacy of the proposed model is demonstrated on synthetic and real data.

D. Ghosh, Nilanjana Chakraborty, S. Roy · 0 citations