Tensor-Train Joint Modeling for Few-Step Discrete Diffusion

2026-07-04

AI Summary

Existing discrete diffusion models suffer from parallelization bias under few-step generation due to the conditional independence assumption, making it challenging to balance efficiency and quality. This work proposes the first tensor decomposition–based joint modeling framework for discrete diffusion, explicitly capturing the conditional clean distribution via low-rank tensors (CPD/TTD) and introducing iterative marginal inference tailored to predefined position schedules. Theoretical analysis, grounded in Oseledets’ theorem, reveals that TTD inherently favors structured dependencies among neighboring tokens. Requiring only lightweight fine-tuning of a pretrained MDM, the method achieves substantial improvements in generation quality with extremely few sampling steps, while incurring significantly lower training costs compared to training from scratch.

Abstract

Discrete diffusion promises orders-of-magnitude faster generation than autoregressive (AR) models for sequential discrete data, yet its full potential of few-step generation has remained out of reach due to a fundamental structural limitation. The conditional-independence assumption underlying current discrete diffusion models introduces a systematic parallelization bias that compounds with the number of tokens unmasked per step, becoming severe in the few-step regime that fast generation requires. We address this with 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. The framework supports both Canonical Polyadic (CPD) and Tensor-Train (TTD) decompositions, and we identify a structural bias of TTD toward dependencies between nearby tokens, formalized through Oseledets' theorem relating TT-rank to unfolding-matrix rank, which is well-suited to sequential data such as natural language and line notations for molecular data. To enable efficient generation, we present an iterative marginal inference procedure with specialization for predetermined position schedules. Our framework integrates into pretrained MDMs through lightweight fine-tuning, yielding substantial improvements in few-step generation at a fraction of the cost of training from scratch.

Problem

• discrete diffusion
• few-step generation
• conditional independence
• parallelization bias
• sequential discrete data

Innovation

• Tensor-Train Decomposition
• Discrete Diffusion
• Joint Distribution Modeling
• Few-Step Generation
• Sequential Data