Tensor-Train Joint Modeling for Few-Step Discrete Diffusion

Technical Summary: Tensor-Train Joint Modeling for Few-Step Masked Diffusion

Problem Statement
Discrete diffusion models, particularly Masked Diffusion Models (MDMs), offer a promising alternative to autoregressive (AR) models for generating sequential discrete data (e.g., text, molecules) by enabling parallel token generation. However, their potential for few-step generation remains unrealized due to a fundamental structural limitation: the conditional-independence assumption. Current MDMs model the conditional clean distribution $p_\theta(x|x_t)$ as a product of independent marginals ($\prod p_\theta(x_i|x_t)$). This assumption introduces a systematic "parallelization bias" that compounds with the number of tokens unmasked per step. In few-step regimes, where many tokens are unmasked simultaneously, this bias severely degrades sample quality.

Explicitly modeling the joint distribution $p_\theta(x|x_t)$ to resolve this bias is theoretically desirable but computationally intractable, as it requires representing an $N$-dimensional tensor with $V^N$ entries (where $V$ is vocabulary size and $N$ is sequence length). Existing approaches attempt to sidestep this by using auxiliary AR models or latent variables, but these either inherit external inference costs or rely on implicit correlations rather than explicit joint parameterization.

Methodology
The authors propose a framework that explicitly parameterizes the joint distribution $p_\theta(x|x_t)$ within the discrete diffusion model itself using tensor decomposition. The core methodology involves two main components:

1. Tensor Decomposition for Joint Modeling:
   The framework represents the conditional distribution as a low-rank tensor, supporting two specific decompositions:

   • Canonical Polyadic Decomposition (CPD): Approximates the tensor as a sum of rank-1 tensors. While it generalizes standard MDMs (rank-1 case), it treats all token positions symmetrically.
   • Tensor-Train Decomposition (TTD): Approximates the tensor as a sequence of cores (matrices) connected in a chain. The authors identify a crucial structural bias in TTD: the TT-rank at any split point corresponds to the rank of the unfolding matrix of the joint distribution. For sequential data dominated by local dependencies (e.g., natural language, molecular line notations), these unfolding ranks remain small, allowing TTD to capture dependencies between nearby tokens efficiently with a low rank. This is theoretically grounded by Oseledets' theorem.

2. Efficient Sampling via Iterative Marginal Inference:
   Direct sampling from the joint distribution is intractable. The authors present a chain-rule-based sampling procedure that performs iterative marginal inference.

   • General Procedure: The method samples tokens one by one (or in batches) by conditioning on previously sampled positions. It utilizes caching and parallel prefix sums to compute marginal probabilities efficiently without evaluating the full joint distribution.
   • Predetermined Schedules: For fixed unmasking schedules, the method optimizes inference by contracting unconditioned cores and using an auxiliary head to predict contracted cores directly, avoiding redundant computations.

3. Integration and Fine-Tuning:
   The framework is designed to integrate into pretrained MDMs via lightweight fine-tuning. Architectural modifications involve replacing the final output head with a tensor-decomposed structure. To preserve the pretrained model's marginal predictions, weights are initialized as copies of the original head (with small noise to break symmetry), allowing the model to learn joint dependencies with minimal cost compared to training from scratch.

Key Contributions

• First Explicit Joint Modeling in Discrete Diffusion: This work introduces the first framework to parameterize the conditional clean distribution of discrete diffusion via tensor decomposition, strictly generalizing standard MDMs as the rank-one case.
• Structural Bias of TTD for Sequential Data: The authors formalize the suitability of Tensor-Train Decomposition for sequential data. They demonstrate that TTD's structure aligns with the locality of dependencies in natural language and molecular notations, outperforming CPD in these domains.
• Efficient Sampling Algorithm: A novel iterative marginal inference procedure is presented, enabling tractable sampling from the joint distribution with minimal overhead through caching and parallel prefix sums.
• Lightweight Fine-Tuning: A recipe for integrating this framework into pretrained models is provided, achieving substantial few-step improvements at a fraction of the cost of training from scratch.

Experimental Results
The framework was evaluated on text generation (OpenWebText, LM1B) and molecule generation (QM9) using base models MDLM and VADD.

• Text Generation: TTD-based fine-tuning significantly reduced generative perplexity compared to vanilla models, particularly in the few-step regime. For instance, on OpenWebText, TTD reduced 8-step perplexity by 32.7% over vanilla VADD. CPD showed modest or no improvement over baselines.
• Molecule Generation: On QM9, the TTD approach consistently achieved higher validity, uniqueness, and novelty scores compared to baselines, with the largest gains observed in left-to-right generation where local dependencies are critical.
• Efficiency: Despite the theoretical complexity, the sampling overhead was minimal. On OpenWebText, the TTD-enhanced VADD was only 1.7% slower than the original VADD at 128 timesteps.

Significance and Claims
The paper claims to be the first to successfully introduce tensor decomposition for explicit joint probability modeling within discrete diffusion models. It argues that by moving beyond the conditional-independence assumption through a tractable low-rank parameterization, the systematic parallelization bias can be mitigated. The authors emphasize that TTD is particularly well-suited for sequential data due to its theoretical alignment with local token dependencies, offering a practical path to high-quality, few-step generation without relying on external autoregressive models. The work demonstrates that even low-rank approximations are sufficient to capture the necessary joint structure for significant performance gains in few-step regimes.