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

• URL: http://arxiv.org/abs/2607.06841v1
• Date: Tue, 07 Jul 2026 22:24:10 GMT
• Title: Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling
• Abstract summary: Diffusion models offer a powerful framework for sampling from complex probability densities.<n>A common approach involves solving for the time-reversed differential equation (SDE)<n>We introduce a novel and efficient solver for the underlying HJB equation based on the functional tensor train (FTT) format.
• Score: 17.153317772943968
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Parallel Diffusion Solver via Residual Dirichlet Policy Optimization [88.7827307535107]
  Diffusion models (DMs) have achieved state-of-the-art generative performance but suffer from high sampling latency due to their sequential denoising nature.<n>Existing solver-based acceleration methods often face significant image quality degradation under a low-dimensional budget.<n>We propose the Ensemble Parallel Direction solver (dubbed as EPD-EPr), a novel ODE solver that mitigates these errors by incorporating multiple gradient parallel evaluations in each step.
  arXiv  Detail & Related papers  (2025-12-28T05:48:55Z)
• Sublinear iterations can suffice even for DDPMs [20.394390032998036]
  We show that a "shifted composition rule" algorithm enjoys favorable discretization properties under appropriate assumptions.<n>This is the first sublinear complexity bound for pure DDPM sampling.<n>We also provide experimental validation of the advantages of our method, showing that it performs well in practice with pre-trained image synthesis models.
  arXiv  Detail & Related papers  (2025-11-06T22:15:19Z)
• Diffusion Models for Solving Inverse Problems via Posterior Sampling with Piecewise Guidance [52.705112811734566]
  A novel diffusion-based framework is introduced for solving inverse problems using a piecewise guidance scheme.<n>The proposed method is problem-agnostic and readily adaptable to a variety of inverse problems.<n>The framework achieves a reduction in inference time of (25%) for inpainting with both random and center masks, and (23%) and (24%) for (4times) and (8times) super-resolution tasks.
  arXiv  Detail & Related papers  (2025-07-22T19:35:14Z)
• An Adaptive Random Fourier Features approach Applied to Learning Stochastic Differential Equations [0.8947831206263182]
  This study considers Ito diffusion processes and a likelihood-based loss function derived from the Euler-Maruyama integration introduced in citerichDiet20 and citedridi 2021stochasticdynamicalsystems.<n>Across all cases, the ARFF-based approach matches or surpasses the performance of conventional Adam-based optimization in both loss minimization and convergence speed.
  arXiv  Detail & Related papers  (2025-07-21T09:52:33Z)
• A hybrid FEM-PINN method for time-dependent partial differential equations [9.631238071993282]
  We present a hybrid numerical method for solving evolution differential equations (PDEs) by merging the time finite element method with deep neural networks. The advantages of such a hybrid formulation are twofold: statistical errors are avoided for the integral in the time direction, and the neural network's output can be regarded as a set of reduced spatial basis functions.
  arXiv  Detail & Related papers  (2024-09-04T15:28:25Z)
• Dynamical Measure Transport and Neural PDE Solvers for Sampling [77.38204731939273]
  We tackle the task of sampling from a probability density as transporting a tractable density function to the target. We employ physics-informed neural networks (PINNs) to approximate the respective partial differential equations (PDEs) solutions. PINNs allow for simulation- and discretization-free optimization and can be trained very efficiently.
  arXiv  Detail & Related papers  (2024-07-10T17:39:50Z)
• On the Trajectory Regularity of ODE-based Diffusion Sampling [79.17334230868693]
  Diffusion-based generative models use differential equations to establish a smooth connection between a complex data distribution and a tractable prior distribution. In this paper, we identify several intriguing trajectory properties in the ODE-based sampling process of diffusion models.
  arXiv  Detail & Related papers  (2024-05-18T15:59:41Z)
• Gaussian Mixture Solvers for Diffusion Models [84.83349474361204]
  We introduce a novel class of SDE-based solvers called GMS for diffusion models. Our solver outperforms numerous SDE-based solvers in terms of sample quality in image generation and stroke-based synthesis.
  arXiv  Detail & Related papers  (2023-11-02T02:05:38Z)
• A Deep Unrolling Model with Hybrid Optimization Structure for Hyperspectral Image Deconvolution [50.13564338607482]
  We propose a novel optimization framework for the hyperspectral deconvolution problem, called DeepMix.<n>It consists of three distinct modules, namely, a data consistency module, a module that enforces the effect of the handcrafted regularizers, and a denoising module.<n>This work proposes a context aware denoising module designed to sustain the advancements achieved by the cooperative efforts of the other modules.
  arXiv  Detail & Related papers  (2023-06-10T08:25:16Z)
• Reflected Diffusion Models [93.26107023470979]
  We present Reflected Diffusion Models, which reverse a reflected differential equation evolving on the support of the data. Our approach learns the score function through a generalized score matching loss and extends key components of standard diffusion models.
  arXiv  Detail & Related papers  (2023-04-10T17:54:38Z)
• Blending Neural Operators and Relaxation Methods in PDE Numerical Solvers [3.2712166248850685]
  HINTS is a hybrid, iterative, numerical, and transferable solver for partial differential equations. It balances the convergence behavior across the spectrum of eigenmodes by utilizing the spectral bias of DeepONet. It is flexible with regards to discretizations, computational domain, and boundary conditions.
  arXiv  Detail & Related papers  (2022-08-28T19:07:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.