Abstract
The conditional-independence assumption underlying current discrete diffusion models introduces a systematic parallelization bias in the few-step regime. 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.
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 which is well-suited to sequential data.
Problem
Masked diffusion models (MDM) represent the conditional clean distribution with marginal probabilities. $$ p_\theta(\boldsymbol{x} | \boldsymbol{x}_t) = \prod_{i=1}^{N} p_\theta(x^i | \boldsymbol{x}_t) $$ This ignores any interdependence between tokens, causing large degradations in few-step regime.
Background
Masked Diffusion Models (MDM): a discrete diffusion model whose forward process is an absorbing process. $$ q(x_t^i | \boldsymbol{x}) = t\,\delta_{m}(x_t^i) + (1-t)\,\delta_{x^i}(x_t^i) $$ It is trained to predict the clean sequence: $$ \mathcal{L}_{\text{NLL}} = \mathbb{E} \left[-\log p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\right] $$ Which with the independence assumption becomes: $$ \mathcal{L}_{\text{marginal}} = \mathbb{E}\!\left[\sum_{i=1}^{N} -\log p_\theta(x^i | \boldsymbol{x}_t)\right] $$ Generation proceeds by initializing a fully masked sequence and iteratively unmasking tokens until a clean, mask-free sequence is obtained. In few-step generation, it must generate multiple tokens in a single step, which causes the joint dependence error.
Joint Probability Modeling
Key Idea: We consider the true distribution \(p(\boldsymbol{x} | \boldsymbol{x}_t)\) as an \(N\)-dimensional tensor of size \(V^N\) and model it with tensor decomposition methods.
Canonical Polyadic Decomposition (CPD): approximation by sum of \(r\) rank-1 tensors. $$ p_\theta(\boldsymbol{x} | \boldsymbol{x}_t) = \sum_{\alpha=1}^{r} w^{(\alpha)}_\theta(\boldsymbol{x}_t) \prod_{i=1}^{N} p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t) $$
Tensor-Train Decomposition (TTD): approximation by sequence of \(N\) cores of size \(V \times r_{i-1} \times r_i\) with each entry being the product of the corresponding matrix slices. $$ p_\theta(\boldsymbol{x} | \boldsymbol{x}_t) = \tfrac{1}{r} \boldsymbol{1}_r^{\top}\, \boldsymbol{G}_{1,\theta}(x^1 | \boldsymbol{x}_t)\, \boldsymbol{G}_{2,\theta}(x^2 | \boldsymbol{x}_t) \cdots \boldsymbol{G}_{N,\theta}(x^N | \boldsymbol{x}_t)\, \boldsymbol{1}_r $$
Why Tensor Trains?
For tensor \(\boldsymbol{A}\), take its flattening: $$ \boldsymbol{A}_i\!\bigl([\nu_1, \ldots, \nu_i],\,[\nu_{i+1}, \ldots, \nu_N]\bigr) = \boldsymbol{A}(\nu_1, \ldots, \nu_N) $$ Then the minimum TT-rank required at position \(i\) to represent \(\boldsymbol{A}\) is \(\text{rank}(\boldsymbol{A}_i)\), which can be interpreted as the amount of dependence across the position. For sequential domains where local dependencies dominate, \(r_i = \text{rank}(\boldsymbol{A}_i)\) can remain small, allowing rank-efficient approximation.
Training & Sampling
Like MDM, it is trained to predict the clean sequence with \(\mathcal{L}_{\text{NLL}}\), but with \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) computed directly from the tensor decomposition formulation.
For efficient sampling, we use iterative marginal inference. For any ordering \(i_1, i_2, \ldots, i_K\): $$ p(x^{i_1}, \ldots, x^{i_K} | \boldsymbol{x}_t) = \prod_{j=1}^{K} p(x^{i_j} | x^{i_1}, \ldots, x^{i_{j-1}}, \boldsymbol{x}_t) $$ which can be sampled efficiently.
Results
OpenWebText unconditional generation
| Method | Timesteps | Generative Perplexity (↓) | Entropy |
|---|---|---|---|
| MDLM | 8 / 16 / 32 / 64 / 128 | 840.57 / 347.10 / 195.09 / 146.09 / 122.08 | 7.77 / 7.67 / 7.60 / 7.58 / 7.56 |
| EDLM | 8 / 16 / 32 / 64 / 128 | 891.18 / 367.48 / 212.02 / 153.79 / 129.38 | 7.83 / 7.73 / 7.66 / 7.62 / 7.60 |
| Di4C | 8 / 16 / 32 / 64 / 128 | 832.52 / 360.89 / 207.50 / 152.63 / 127.31 | 7.79 / 7.70 / 7.63 / 7.60 / 7.57 |
| MDLM+CPD (Ours) | 8 / 16 / 32 / 64 / 128 | 848.87 / 358.25 / 205.46 / 148.10 / 126.85 | 7.81 / 7.71 / 7.65 / 7.60 / 7.60 |
| MDLM+TTD (Ours) | 8 / 16 / 32 / 64 / 128 | 636.98 / 295.43 / 183.43 / 140.97 / 119.20 | 7.75 / 7.66 / 7.60 / 7.57 / 7.54 |
QM9 molecule generation
Both CPD and TTD work well for short sequence generation (QM9), but only TTD scales for long sequences (OpenWebText), due to its advantages in modeling local dependencies.
Pitfall of RoPE in Masked Diffusion (Appendix)
Rotary Position Embedding (RoPE) incorporates positional information into the attention layer by "rotating" queries and keys by $$ q_m=R^d_{\Theta, m} W_q x_m, k_n=R^d_{\Theta, n} W_k x_n $$ We have found, however, that this method should be used with care on MDMs. In the very first step of MDM generation, the input is a fully-masked sequence. Since default RoPE changes neither the token embeddings nor the values in attention layers, it loses its ability to distinguish between positions.
If the initial \(x_m=x\) is uniform across position, \(v_m=W_v x_m=W_v x=v\) is also uniform, and thus the output of attention layer $$ o_m=\sum \limits_{n=1}^N a_{m,n}v_n=v\sum \limits_{n=1}^N a_{m,n}=v $$ also becomes uniform regardless of the attention weights \(a\). Indeed, we confirm that under this formulation, the outputted logits of the model are identical for every position in the first step.
This greatly affects the quality of all MDM-like sampling, including our baseline MDLM and our proposed methods. As a fix, we applied RoPE rotary matrices to values in addition to queries and keys in attention.
BibTeX
@article{kim2026tensortrainjointmodeling,
title={Tensor-Train Joint Modeling for Few-Step Discrete Diffusion},
author={Byoungkwon Kim and Minhyuk Sung},
journal={arXiv preprint arXiv:2607.03788},
year={2026},
}