1 Introduction↩︎

Discrete diffusion has emerged as a compelling alternative to AR models for generating sequential discrete data such as text and proteins [1]–[4]. Its key promise is parallel token generation, producing multiple tokens per step rather than one. Most notably, Masked Diffusion Models (MDMs) [3], [4] have scaled to large language models [5], [6] with competitive quality and faster generation.

In practice, however, this speed potential remains unrealized, as pushing toward few-step generation systematically degrades sample quality. The root cause is structural; to remain tractable, discrete diffusion models represent the conditional clean distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) under the conditional-independence assumption \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t) = \prod_{i=1}^{N} p_\theta(x^i | \boldsymbol{x}_t)\). This assumption introduces a systematic parallelization bias [7] when multiple tokens are unmasked simultaneously. The bias compounds with the number of tokens unmasked per step, becoming especially severe in the few-step regime that fast generation requires.

Resolving this bias requires modeling the joint distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) explicitly. However, naïve joint modeling is infeasible since the joint over \(N\) tokens with vocabulary size \(V\) forms an \(N\)-dimensional tensor with \(V^N\) entries. Some existing approaches therefore sidestep direct joint modeling, introducing latent factors that induce token correlations [8], [9] or borrowing dependence from auxiliary autoregressive models [10], [11]. Compared to AR-auxiliary approaches, our method models the joint distribution explicitly as a tractable parametric family within the discrete diffusion model itself, without requiring any external pretrained model. Also, compared to latent-variable approaches, our explicit parameterization is orthogonal to their implicit one, and the two can be combined for further gains. Concurrent work [12] addresses the same factorization barrier by augmenting the factorized neural output with a tractable probabilistic circuit prior, requiring manual structural design for sequential locality, whereas Tensor-Train decomposition provides the same tractability with a theoretically grounded locality bias.

In this work, we introduce a framework that explicitly parameterizes the joint distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) via tensor decomposition [13]. Compared to prior work, our framework provides explicit joint modeling within the diffusion model itself, without relying on auxiliary models. Our contributions are four-fold.

• First explicit joint modeling via tensor decomposition. To our knowledge, this is the first work to parameterize the conditional clean distribution of discrete diffusion via tensor decomposition, strictly generalizing standard MDMs as the rank-one case.

• TTD for sequential structure. We develop the framework for both Canonical Polyadic Decomposition (CPD) [14] and Tensor-Train Decomposition (TTD) [15], identifying a structural bias of TTD toward dependencies between nearby tokens, formalized through Oseledets’ theorem relating TT-rank to unfolding-matrix rank.

• Efficient sampling via marginal inference. Direct sampling from a joint of size \(V^N\) is intractable. We present a chain-rule procedure built on iterative marginal inference via caching and parallel prefix sums, with specialization for predetermined position schedules.

• Lightweight fine-tuning of pretrained MDMs. Our framework integrates into pretrained MDMs through small architectural modifications initialized to preserve the marginal predictions, yielding substantial few-step improvements at a fraction of the cost of training from scratch.

In experiments on text generation (OpenWebText [16], LM1B [17]) and molecule generation (QM9 [18], [19]), our framework delivers substantial few-step improvements across multiple base models, including MDLM [3] and VADD [9]. As predicted by our theory, TTD consistently outperforms CPD on these sequential domains, with the largest gains concentrated in the few-step regime. For instance, TTD-based fine-tuning of VADD reduces 8-step generative perplexity on OpenWebText by 32.7% over vanilla VADD, while being only 1.7% slower at 128 timesteps.

2 Related Work↩︎

Several lines of work attempt to address the parallelization bias arising from the conditional-independence assumption in discrete diffusion. The central challenge is the exponential size of the joint distribution: an \(N\)-token sequence over vocabulary \(V\) admits \(V^N\) possible configurations, making direct modeling intractable.

Auxiliary-model approaches distill dependency information from separately trained autoregressive models via copula-based reweighting [10] or energy-based reweighting [11]. While effective, these methods inherit the inference cost of an external model and become inapplicable in domains lacking strong AR models, such as generation of line notations for molecular data. Our framework, in contrast, requires no additional model.

Latent-variable approaches introduce latent factors that induce token correlations through marginalization, either via a continuous mixing variable [8] or a VAE-style latent vector [9]; our explicit modeling is orthogonal to these implicit approaches and can be combined with them for further gains.

Sampling strategies mitigate the bias without modifying the model, selecting positions via low-discrepancy schedules [20] or adaptively bounding the number of simultaneously unmasked tokens through entropy criteria [21]; these strategies integrate naturally within our framework.

Concurrent work [12] couples the factorized neural output with a tractable probabilistic circuit prior, addressing the same factorization barrier with a different parameterization. Although probabilistic circuits offer general-purpose tractable inference, they require manual structural design for sequential locality, whereas Tensor-Train decomposition provides the same tractability with a theoretically grounded locality bias (4.2), making it particularly suited to natural language and line notations for molecular data without further structural engineering.

The parallelization bias has also been studied in the context of speculative decoding for autoregressive language models, where [22] were the first to apply Canonical Polyadic Decomposition for explicit token-wise dependency modeling. We extend this idea to discrete diffusion for the first time, with two main technical contributions beyond their work: (1) identifying Tensor-Train Decomposition as particularly suited to sequential structure, with a theoretically grounded inductive bias toward dependencies between nearby tokens; and (2) presenting an efficient sampling procedure that introduces minimal overhead when sampling from the joint distribution.

3 Background & Overview↩︎

3.1 Background: Masked Diffusion Models (MDM)↩︎

A Masked Diffusion Model (MDM) [3] is a discrete diffusion model whose forward process is an absorbing process. The absorbing token is commonly referred to as the mask token \(m\) and thus lends the model its name. The forward process maps the data distribution \(p(\boldsymbol{x})\) to the prior distribution, which is a Dirac delta concentrated at the fully masked sequence \(\mathbf{m} = [m, m, \ldots, m]\), as \(t\) progresses from \(0\) to \(1\). Formally, \[q(x_t^i | \boldsymbol{x}) = t\,\delta_{m}(x_t^i) + (1-t)\,\delta_{x^i}(x_t^i),\] where \(x_t^i \in [V] = \{1, 2, \ldots, V\}\) denotes the \(i\)-th token at time \(t \in [0,1]\), \(\boldsymbol{x} = [x^1, x^2, \ldots, x^N]\) is the clean sequence, \(\boldsymbol{x}_t = [x_t^1, \ldots, x_t^N]\) is its partially masked version, and \(\delta_a\) denotes the point mass at \(a\).

The model is trained to predict the original values of the masked tokens in a partially masked sequence, typically under an assumption of conditional independence across tokens. Namely, \[p_\theta(\boldsymbol{x} | \boldsymbol{x}_t) = \prod_{i=1}^{N} p_\theta(x^i | \boldsymbol{x}_t). \label{eq:mdm-factorization}\tag{1}\] Given a partially masked sequence \(\boldsymbol{x}_t\), for every position \(i\) and every \(x^i \in [V]\), the model predicts the marginal probability \(p_\theta(x^i | \boldsymbol{x}_t)\) that the \(i\)-th token of the clean sequence equals \(x^i\). Under the factorization, the negative log-likelihood (NLL) loss \[\mathcal{L}_{\text{NLL}} = \mathbb{E}_{t \sim U(0,1),\; \boldsymbol{x} \sim p(\boldsymbol{x}),\; \boldsymbol{x}_t \sim q(\boldsymbol{x}_t | \boldsymbol{x})} \!\left[-\log p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\right] \label{eq:nll}\tag{2}\] reduces to the sum of per-position marginal NLLs: \[\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.

3.2 Method Overview↩︎

Despite the promising progress of masked diffusion models, a key limitation lies in the conditional-independence assumption (1 ) imposed on \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\), which introduces a systematic parallelization bias [7] that becomes especially severe in the few-step regime. Explicitly modeling the joint distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\), however, is fundamentally a problem of scale. In continuous diffusion, the joint admits a tractable parameterization as a multivariate Gaussian. In the discrete setting, by contrast, the joint corresponds to an \(N\)-dimensional tensor of size \(V^N\), which rapidly becomes intractable as the vocabulary size \(V\) and the sequence length \(N\) grow.

To address this challenge, we propose to model \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) explicitly via tensor decomposition. In 4, we investigate two decompositions particularly suited to this purpose—Canonical Polyadic Decomposition (CPD) [14] and Tensor-Train Decomposition (TTD) [15]—both supporting efficient training and sampling with only \(O(VN)\) memory. We describe how each parameterizes the joint distribution and how generation is carried out under it. Our analysis identifies a structural limitation of CPD that motivates the use of TTD, whose form is naturally aligned with dependencies between nearby tokens, a property especially well suited to sequential data such as natural language [23] and line notations for molecular data. Note that either decomposition reduces to the standard marginal probability formulation (1 ) in the rank-one special case. To our knowledge, this is the first work introducing tensor decomposition to joint probability modeling in discrete diffusion.

4 Joint Probability Modeling via Tensor Decomposition↩︎

In this section, we describe how the joint distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) can be modeled using two tensor decompositions: Canonical Polyadic Decomposition (CPD; 4.1) and Tensor-Train Decomposition (TTD; 4.2). We discuss the advantages of TTD over CPD in capturing dependencies between nearby tokens, a property crucial for many practical domains such as text and line notations for molecular data. Building on the formulations, 4.3 describes the sampling procedure under either decomposition.

4.1 Modeling via Canonical Polyadic Decomposition (CPD)↩︎

Given an order-\(N\) tensor \(\boldsymbol{A}\), Canonical Polyadic Decomposition (CPD) [14] approximates \(\boldsymbol{A}\) by a sum of \(r\) rank-1 tensors, each formed as the outer product of \(N\) vectors: \[\boldsymbol{A}(\nu_1, \ldots, \nu_N) \approx \sum_{\alpha=1}^{r} \prod_{i=1}^{N} \boldsymbol{u}_{i, \alpha}(\nu_i),\] where each \(\boldsymbol{u}_{i, \alpha} \in \mathbb{R}^{V}\) is a one-dimensional factor vector, and \(r\) is referred to as the CP-rank, which we may write as \(r_{\text{CPD}}\). Building on this, we model the joint distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) over a sequence of \(N\) tokens via CPD, with an additional rank weight (as in the autoregressive setting of [22]): \[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), \label{eq:cpd}\tag{3}\] where \(w^{(\alpha)}_\theta(\boldsymbol{x}_t) \geq 0\) is the weight produced for the \(\alpha\)-th rank-1 component, and \(p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t) \geq 0\) is the corresponding output for token \(x^i\) at position \(i\) in that component. For normalization, we enforce \[\sum_{\alpha=1}^{r} w^{(\alpha)}_\theta(\boldsymbol{x}_t) = 1, \qquad \sum_{x^i=1}^{V} p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t) = 1,\] which together ensure that \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) is a valid probability distribution while maintaining expressibility (proof in Appendix 7). These constraints further allow \(p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)\) to be interpreted as the marginal probability of \(x^i\) at position \(i\) within the \(\alpha\)-th component.

Note that when \(r=1\), the formulation reduces to the standard MDM, which predicts only per-position marginal probabilities. Moreover, any lower-rank representation with \(r' < r\) can be recovered by zeroing the surplus weights, so the family of joint distributions expressible at rank \(r\) strictly contains those expressible at any lower rank.

In practice, \(r\) must be kept small for computational and memory efficiency. To maximize the benefit of a low-rank approximation, the key is to model inter-token dependencies in a manner aligned with the structure of the data modality. For natural language, it is well known that nearby tokens are the most strongly correlated [23]. To capture such local dependencies effectively, we next introduce Tensor-Train Decomposition (TTD) and discuss its advantages over CPD.

4.2 Modeling via Tensor-Train Decomposition (TTD)↩︎

Tensor-Train Decomposition (TTD) [15] approximates an order-\(N\) tensor \(\boldsymbol{A}\) by a sequence of cores \(\boldsymbol{G}_1, \ldots, \boldsymbol{G}_N\), with each entry obtained as a product of the corresponding matrix slices: \[\boldsymbol{A}(\nu_1, \ldots, \nu_N) \approx \boldsymbol{G}_1(\nu_1)\, \boldsymbol{G}_2(\nu_2) \cdots \boldsymbol{G}_N(\nu_N),\] where each core \(\boldsymbol{G}_i \in \mathbb{R}^{V \times r_{i-1} \times r_i}\) has matrix slices \(\boldsymbol{G}_i(\nu_i) := \boldsymbol{G}_i(\nu_i, :, :) \in \mathbb{R}^{r_{i-1} \times r_i}\), with boundary ranks \(r_0 = r_N = 1\) so that the right-hand side is a scalar. The tuple of internal ranks \((r_1, r_2, \ldots, r_{N-1})\) is referred to as the TT-rank of the decomposition.

We adopt this form to model the joint distribution \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\), identifying each tensor mode with a sequence position. Each core \(\boldsymbol{G}_{i,\theta}(x^i | \boldsymbol{x}_t) \in \mathbb{R}_{\geq 0}^{r_{i-1} \times r_i}\) is conditioned on the masked input. For practical implementation, we let every core share the common shape \(V \times r \times r\), i.e., \(r_0 = r_1 = \cdots = r_N = r\) (which we may also write as \(r_{\text{TTD}}\)), and contract the boundary dimensions with all-ones vectors \(\boldsymbol{1}_r \in \mathbb{R}^{r}\) to recover a scalar: \[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. \label{eq:ttd}\tag{4}\] To ensure that \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) is a valid probability distribution, we impose the following normalization constraint on each core, which also maintains expressibility (proof in Appendix 7): \[\sum_{x^i, k} \boldsymbol{G}_{i,\theta}(x^i, j, k | \boldsymbol{x}_t) = 1, \quad \forall i, j.\] As with CPD, when \(r=1\), the formulation reduces to the standard MDM, which predicts only per-position marginal probabilities. Moreover, any lower-rank representation with \(r' < r\) can be recovered by zeroing surplus columns of the cores, so the family of joint distributions expressible at rank \(r\) strictly contains those expressible at any lower rank.

A key advantage of TTD over CPD is its structural bias toward modeling dependencies between nearby tokens. To make this precise, we recall a result from [15] that connects each TT-rank \(r_i\) to the dependency structure of \(\boldsymbol{A}\).

For each split position \(i \in \{1, \ldots, N-1\}\), partition the \(N\) token positions into two contiguous blocks: a left block consisting of positions \(1, \ldots, i\), and a right block consisting of positions \(i+1, \ldots, N\). The unfolding matrix \(\boldsymbol{A}_i\) at this split is obtained by reshaping \(\boldsymbol{A}\) into a matrix whose rows are indexed by the joint values of the left block and whose columns by those of the right: \[\boldsymbol{A}_i \in \mathbb{R}^{V^i \times V^{N-i}}, \qquad \boldsymbol{A}_i\!\bigl([\nu_1, \ldots, \nu_i],\,[\nu_{i+1}, \ldots, \nu_N]\bigr) = \boldsymbol{A}(\nu_1, \ldots, \nu_N).\] The rank of \(\boldsymbol{A}_i\) captures how much information must flow across the cut at position \(i\): a small \(\operatorname{rank}(\boldsymbol{A}_i)\) means that the joint values of the left block can be summarized into a low-dimensional vector that fully determines the distribution over the right block, indicating that the dependence between the two sides is limited.

[15] establishes that the minimum TT-rank required at level \(i\) is exactly the rank of the corresponding unfolding matrix: \[r_i = \operatorname{rank}(\boldsymbol{A}_i).\] Intuitively, this is because any TT representation of \(\boldsymbol{A}\) already factorizes \(\boldsymbol{A}_i\) into a product of a \(V^i \times r_i\) matrix (built from the first \(i\) cores) and a \(r_i \times V^{N-i}\) matrix (built from the remaining cores), so \(r_i \geq \operatorname{rank}(\boldsymbol{A}_i)\); the reverse direction is shown constructively in the same theorem.

The crucial implication is that each \(r_i\) depends only on the dependence structure across the \(i\)-th cut, not on the overall complexity of \(\boldsymbol{A}\). For sequential data dominated by local dependencies, the unfolding ranks \(\operatorname{rank}(\boldsymbol{A}_i)\) can remain small at every cut and not scale with the sequence length, so a small TT-rank \(r\) can suffice for an accurate approximation. CPD, by contrast, treats all positions symmetrically and offers no analogous reduction when locality is the dominant structure of the data.

As a concrete illustration of how extreme this asymmetry can be, Appendix 8 constructs a specific distribution over \(2N\) tokens for which \(r_\text{CPD} = 2^N\) but \(r_\text{TTD} = 2\).

4.3 Efficient Sampling via Iterative Marginal Inference↩︎

The tensor-decomposed forms (CPD and TTD) of \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) admit efficient evaluation at any \(\boldsymbol{x}\), so training proceeds identically to a standard MDM with the same NLL loss (2 ). Sampling, however, is nontrivial: naively sampling from \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)\) would require evaluating it at all \(V^N\) possible sequences, an exponential cost in the sequence length \(N\). For efficient sampling, we present a general procedure based on the chain rule, followed by a special case that exploits additional structure.

The chain rule of probability decomposes the joint into a sequence of conditionals. For any ordering \((i_1, \ldots, i_K)\) of the \(K\) positions, \[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).\] We can therefore sample the subsequence one position at a time by conditioning on previously sampled positions and repeatedly applying marginal inference. Specifically, given a set \(S \subseteq \{1, \ldots, N\}\) of conditioned positions with values \(\boldsymbol{x}_S := \{x^j\}_{j \in S}\), we compute the marginal \(p(x^i | \boldsymbol{x}_S, \boldsymbol{x}_t)\) at every unconditioned position \(i \notin S\). Sampling the full sequence then proceeds by initializing \(S = \emptyset\) and iteratively (i) performing marginal inference, (ii) selecting a position \(i \notin S\), (iii) sampling \(x^i \sim p(x^i | \boldsymbol{x}_S, \boldsymbol{x}_t)\), and (iv) adding \(i\) to \(S\). The position selection in step (ii) can follow any logit-based strategy (e.g., top probability-based [24] or entropy-based [6] selection). The number of tokens sampled in a single step can be either fixed or determined by an adaptive sampler such as EB-Sampler [21]. We now describe how to perform marginal inference under each decomposition.

4.3.0.1 CPD.

For each position \(j\) and rank component \(\alpha\), define the conditioning factor \[\bar{p}^{(\alpha, j)}_\theta = \begin{cases} p^{(\alpha, j)}_\theta(x^j | \boldsymbol{x}_t) & \text{if } j \in S, \\ 1 & \text{otherwise}. \end{cases}\] Substituting into the CPD form (3 ), for any \(i \notin S\), \[p(x^i | \boldsymbol{x}_S, \boldsymbol{x}_t) = \frac{p(x^i, \boldsymbol{x}_S | \boldsymbol{x}_t)}{p(\boldsymbol{x}_S | \boldsymbol{x}_t)} \propto p(x^i, \boldsymbol{x}_S | \boldsymbol{x}_t) = \sum_{\alpha=1}^{r} w^{(\alpha)}_\theta(\boldsymbol{x}_t)\, p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t) \prod_{j=1}^{N} \bar{p}^{(\alpha, j)}_\theta.\] Since \(\bar{p}^{(\alpha, i)}_\theta = 1\) for \(i \notin S\), the product \(\prod_{j=1}^{N} \bar{p}^{(\alpha, j)}_\theta\) does not depend on \(i\) and can therefore be precomputed once and reused across all positions, enabling parallel evaluation of \(p(x^i, \boldsymbol{x}_S | \boldsymbol{x}_t)\) for every \(i \notin S\). The marginal \(p(x^i | \boldsymbol{x}_S, \boldsymbol{x}_t)\) is then obtained by normalizing \(p(x^i, \boldsymbol{x}_S | \boldsymbol{x}_t)\) over the vocabulary. (Division by \(p(\boldsymbol{x}_S | \boldsymbol{x}_t)\) is also possible but unnecessary and less numerically stable.)

4.3.0.2 TTD.

Analogously, for each position \(j\), define the contracted core \[\bar{\boldsymbol{G}}_j = \begin{cases} \boldsymbol{G}_{j,\theta}(x^j | \boldsymbol{x}_t) & \text{if } j \in S, \\ \displaystyle \sum_{x^j = 1}^{V} \boldsymbol{G}_{j,\theta}(x^j | \boldsymbol{x}_t) & \text{otherwise}. \end{cases}\] When \(j \notin S\), \(\bar{\boldsymbol{G}}_j\) does not depend on the chain-rule conditioning \(\boldsymbol{x}_S\) and can be cached across queries. For any \(i \notin S\), \[p(x^i, \boldsymbol{x}_S | \boldsymbol{x}_t) = \tfrac{1}{r}\, \boldsymbol{1}_r^{\top}\, \bar{\boldsymbol{G}}_1 \cdots \bar{\boldsymbol{G}}_{i-1}\, \boldsymbol{G}_{i,\theta}(x^i | \boldsymbol{x}_t)\, \bar{\boldsymbol{G}}_{i+1} \cdots \bar{\boldsymbol{G}}_N\, \boldsymbol{1}_r.\] The left-cumulative products \(\bar{\boldsymbol{G}}_1 \cdots \bar{\boldsymbol{G}}_{i-1}\) and right-cumulative products \(\bar{\boldsymbol{G}}_{i+1} \cdots \bar{\boldsymbol{G}}_N\) for all \(i\) can be obtained via parallel prefix sums. Normalizing \(p(x^i, \boldsymbol{x}_S | \boldsymbol{x}_t)\) over the vocabulary then yields \(p(x^i | \boldsymbol{x}_S, \boldsymbol{x}_t)\), again avoiding the explicit computation of \(p(\boldsymbol{x}_S | \boldsymbol{x}_t)\).

4.3.1 Special Case: Predetermined Sampling Positions↩︎

While the method above samples one position at a time, in many settings—random ordering, fixed orderings such as left-to-right — the positions to be sampled in the current step are determined before any computation, allowing them to be sampled jointly. Let \(T \subseteq \{1, \ldots, N\}\) denote the masked positions to be sampled jointly in the same step, with values \(\boldsymbol{x}_T := \{x^j\}_{j \in T}\). For CPD: \[p(\boldsymbol{x}_T | \boldsymbol{x}_t) = \sum_{\alpha=1}^{r} w^{(\alpha)}_\theta(\boldsymbol{x}_t) \prod_{j \in T} p^{(\alpha, j)}_\theta(x^j | \boldsymbol{x}_t)\] Sampling therefore requires evaluating \(p^{(\alpha, j)}_\theta(x^j | \boldsymbol{x}_t)\) only at positions \(j \in T\).

With TTD, positions \(j \notin T\) enter \(p(\boldsymbol{x}_T | \boldsymbol{x}_t)\) only through the contracted cores \(\bar{\boldsymbol{G}}_j\). By absorbing each \(\bar{\boldsymbol{G}}_j\) into a neighboring core, the remaining cores form a TT representation over only the \(|T|\) positions in \(T\). On this we can run a highly optimized version of the general TTD method, by caching cumulative products and computing \(p(x^i | \boldsymbol{x}_S, \boldsymbol{x}_t)\) for only the next immediate value of \(i\).

The reduction above assumes access to \(\bar{\boldsymbol{G}}_j\) for each \(j \notin T\). Computing it from the definition by outputting \(\boldsymbol{G}_{j,\theta}(x^j | \boldsymbol{x}_t)\) for all \(V\) values of \(x^j\) and summing is computationally wasteful. We sidestep this overhead by introducing an additional head \(\bar{\boldsymbol{G}}_{j,\theta}(\boldsymbol{x}_t)\) that directly produces the contracted core, trained to be consistent with \(\boldsymbol{G}_{j,\theta}(x^j | \boldsymbol{x}_t)\): \[\mathcal{L}_{\text{C}} = \frac{1}{N} \sum_{j=1}^{N} \Bigl\lVert \bar{\boldsymbol{G}}_{j,\theta}(\boldsymbol{x}_t) - \text{stopgrad}(\sum_{x^j} \boldsymbol{G}_{j,\theta}(x^j | \boldsymbol{x}_t)) \Bigr\rVert_F^2 \label{eqn:l95c}\tag{5}\] then optimize \[\mathcal{L} = \mathcal{L}_{\text{NLL}} + \mathcal{L}_{\text{C}}.\]

4.4 Extension to Latent-Variable Approaches↩︎

Our framework is flexible: beyond a base MDM, it can also be combined with latent-variable approaches. We illustrate this with Variational Autoencoding Discrete Diffusion (VADD) [9], a few-step MDM framework based on a latent variable model. Integrating our explicit modeling of \(p_{\boldsymbol{\theta}}(\boldsymbol{x} | \boldsymbol{x}_t)\) into VADD yields the training loss \[\widehat{\mathcal{L}}_{\lambda}(\boldsymbol{x}_0;\boldsymbol{\theta},\boldsymbol{\phi}) = \int_{0}^{1} \mathbb{E}_{q(\boldsymbol{x}_t|\boldsymbol{x}_0)}\mathbb{E}_{r_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x}_0, \boldsymbol{x}_t)}\frac{-\alpha_t'}{1-\alpha_t} \left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_0|\boldsymbol{x}_t,\boldsymbol{z})-\lambda\log\left(\frac{r_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x}_0, \boldsymbol{x}_t)}{p(\boldsymbol{z})}\right)\right] \mathrm{d} t,\] with the sampling procedure \[\boldsymbol{z} \sim p(\boldsymbol{z}), \qquad \boldsymbol{x}_{t_{i-1}} \sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_{i-1}}|\boldsymbol{x}_{t_i},\boldsymbol{z}),\] where our tensor decomposition parameterizes both \(p_{\boldsymbol{\theta}}(\boldsymbol{x}_0|\boldsymbol{x}_t,\boldsymbol{z})\) and \(p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_{i-1}}|\boldsymbol{x}_{t_i},\boldsymbol{z})\). Results for this combination are reported in 5.1.

4.5 Fine-Tuning Pretrained MDMs and Implementation Details↩︎

Our tensor-decomposition framework can be integrated into pretrained MDMs via fine-tuning. We describe the necessary architectural modifications and other implementation details below.

4.5.0.1 Architecture.

A standard MDM produces \(NV\) logits, but tensor-decomposed forms require \(EVN\), where \(E = r\) (CPD; plus \(r\) rank weights) or \(E = r^2\) (TTD). Our 1-layer method replaces the final head \(\boldsymbol{W} \in \mathbb{R}^{H \times V}\) with \(\boldsymbol{W}' \in \mathbb{R}^{H \times EV}\), which is costly for large \(V\). Following [22], the 2-layer method instead inserts \(\widetilde{\boldsymbol{W}} \in \mathbb{R}^{H \times EH}\) before the final head, then applies \(\boldsymbol{W}\) to each of the \(E\) outputs, which significantly reduces the parameter count since \(H \ll V\) in language models. Schematics of different architectures are shown in 2.

4.5.0.2 Initialization.

We initialize the added weights to preserve pretrained marginal predictions: \(\boldsymbol{W}'\) as \(E\) copies of \(\boldsymbol{W}\) (1-layer), or \(\widetilde{\boldsymbol{W}}\) as \(E\) stacked identity matrices with \(\boldsymbol{W}\) frozen (2-layer). Small Gaussian noise is added to break rank-symmetry.

4.5.0.3 Masked-token-only modeling.

Let \(\mathcal{M} := \{i : x_t^i = m\}\) denote the set of masked token indices. During the forward process, unmasked tokens remain unchanged (i.e., \(x_t^i = x^i\) for \(i \notin \mathcal{M}\)), so the likelihood factorizes as \(p_\theta(\boldsymbol{x} | \boldsymbol{x}_t) = p_\theta(\boldsymbol{x}_\mathcal{M} | \boldsymbol{x}_t) \prod_{i \notin \mathcal{M}} \delta_{x_t^i}(x^i)\). We therefore restrict network forward passes to predicting \(p_\theta(\boldsymbol{x}_\mathcal{M} | \boldsymbol{x}_t)\) over masked tokens only for efficiency. The NLL loss is normalized by \(|\mathcal{M}|\) accordingly.

5 Experimental Results↩︎

5.1 Text Generation↩︎

We demonstrate the application of our methods on both MDLM [3] and VADD [9]. For MDLM [3], we use the official checkpoint trained on OpenWebText (OWT) [16] as the base model, and for VADD [9] we train a base model from scratch for both OWT and LM1B [17]. These models are augmented with the 2-layer architecture [22] before fine-tuning, with \(r_{\text{CPD}}=16, r_{\text{TTD}}=4\). For training and sampling with VADD and other baselines [8], [11], we use their official implementation, while fixing the numerical issue in categorical sampling reported by [25], using 64-bit precision for logits. Additional details on training and sampling can be found in App. 10.1.

We sample 1024 samples from each using random ordering, then measure the generative perplexity with gpt2-large [26], and entropy of the generated sequences as a measure of diversity. The sample length is 1024 for OWT [16] and 128 for LM1B [17].

Results on OWT [16] show that TTD significantly improves over base MDLM [3], reducing generative perplexity by 24.22% and 14.88% in the challenging 8- and 16-step regimes, while CPD yields modest to no improvement (¿tbl:tab:owt-genppl?, Top). A similar trend holds for VADD+TTD on the same dataset (¿tbl:tab:owt-genppl?, Bottom) and on LM1B [17], where large perplexity gaps persist across all sampling steps (¿tbl:tab:lm1b-genppl?). We conjecture that the gain of TTD on VADD stems from two directions of joint probability modeling: VADD captures high-level semantics via a latent variable, while TTD models low-level dependencies between neighboring tokens through lightweight fine-tuning with the objective in 4.3.1.

We further analyze the proposed method along two axes: rank size impact and runtime overhead. We begin by ablating the rank size using MDLM [3] trained on OWT [16], by varying the rank size from \(r=2\) to \(r=4\). The results are summarized in ¿tbl:tab:owt-genppl-rank?. As shown, generative perplexity largely improves as rank increases, validating that greater expressibility of the decomposition helps in capturing the joint distribution.

r0.55

Due to the large \(V\) in language models, naïvely applying our methods can impose a significant overhead. However, the techniques discussed in 4.3.1 substantially reduce this cost. To demonstrate this, we measure sampling time across methods at 128 timesteps, and report the results in 1. Notably, on VADD+TTD, the best-performing approach on OWT [16] by generative perplexity, our method is only 1.7% slower compared to the original VADD implementation, demonstrating practical efficiency. Moreover, we stress that the overhead stems from the size of output tensors rather than model size, and thus does not scale with model depth; we expect it to be negligible for larger diffusion language models.

5.2 Molecule Generation↩︎

We evaluate on the molecule generation task using QM9 [18], [19], with MDLM [3], parameterized using DiT [27], as the base model. Given the smaller vocabulary size (\(V=40\)) relative to text generation, a 1-layer architecture suffices, with \(r_{\text{CPD}}=64, r_{\text{TTD}}=8\). Training and sampling details are described in App. 10.2.

For the DiT architecture, we modify the rotary positional encoding to apply the rotary matrices to values as well. We stress that this is crucial for any MDM-like sampling methods. Without such a modification, the first sampling step collapses entirely for both the baseline MDLM and our methods due to the uniformity of the inputs in every position. The necessity of such a modification is proved in App. 9.

We evaluate each method by sampling 1024 molecules with random ordering across varying numbers of sampling steps, then compute validity, and uniqueness and novelty among valid samples. These metrics are summarized in 3. We find that our methods bring significant improvements despite small increases in model size. Additional experiments, including sampling strategy comparisons and rank size ablations, are provided in App. 11.1. We observe that, as predicted from our discussion in 4, TTD outperforms CPD of the same output size, and that this gap is the largest in left-to-right generation, where the structural bias of TTD can be exploited to the fullest extent.

6 Conclusion↩︎

We presented the first framework for explicit joint distribution modeling in discrete diffusion, parameterizing the conditional clean distribution \(p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{x}_t)\) as a low-rank tensor via CPD or TTD and thus generalizing standard MDMs. Our analysis identifies a structural bias of TTD toward dependencies between nearby tokens, making it particularly well-suited to sequential data. Through an efficient sampling procedure and a lightweight fine-tuning recipe, the framework integrates with any base MDM or latent-variable method and substantially improves few-step generation on both text and molecule benchmarks.

Societal Impacts↩︎

This work targets efficiency improvements for discrete diffusion generation. On the positive side, faster sampling can reduce the computational cost and energy use of sequential generation, and may accelerate beneficial downstream applications such as molecular design and biological sequence modeling. As with any advance in generative model efficiency, there is also potential for misuse (e.g., faster generation of misleading textual content), though such risks are inherent to the broader generative modeling paradigm rather than introduced by our framework.

Limitations and Future Work↩︎

A fundamental challenge in discrete diffusion modeling is that the joint distribution \(p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{x}_t)\) over \(N\) tokens corresponds to an \(N\)-dimensional tensor with \(V^N\) entries. For arbitrary target distributions this tensor can be of prohibitively high rank, so any tractable parameterization, including ours, must rely on a low-rank approximation, and arbitrarily increasing the rank in CPD or TTD is infeasible in practice. Nonetheless, our experiments show that even low ranks yield substantial few-step improvements across both text and molecule generation, indicating that low-rank structure suffices in the regime where parallelization bias is most severe. As future work, we plan to explore richer tensor formats, such as Tucker decomposition, which may capture dependency structures beyond the locality bias of TTD and offer alternative trade-offs between rank and expressivity.

Appendix↩︎

7 Normalization Constraints↩︎

In 4, we introduced normalization constraints for both decompositions. Here, we prove that (i) these constraints ensure \(p_\theta(\boldsymbol{x} \vert \boldsymbol{x}_t)\) is a valid probability mass function (i.e., sums to 1 over \(V^N\)), and (ii) they do not reduce expressibility: any normalized tensor representable by the given decomposition is also representable under the constraints. We first prove the claim for Canonical Polyadic Decomposition (CPD), then for Tensor Train Decomposition (TTD).

7.1 Canonical Polyadic Decomposition (CPD)↩︎

For CPD, recall that the constraints are given as: \[\sum_{\alpha=1}^{r} w^{(\alpha)}_\theta(\boldsymbol{x}_t) = 1, \qquad \sum_{x^i=1}^{V} p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t) = 1.\] We check (i) by:

\[\begin{gather} \sum \limits_{\boldsymbol{x}} p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)= \sum \limits_{\boldsymbol{x}} \sum_{\alpha=1}^{r} w^{(\alpha)}_\theta(\boldsymbol{x}_t) \prod_{i=1}^{N} p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)= \\ \sum \limits_{\alpha=1}^r w^{(\alpha)}_\theta(\boldsymbol{x}_t) \prod \limits_{i=1}^N \sum \limits_{x^i=1}^V p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)=\sum \limits_{\alpha=1}^r w^{(\alpha)}_\theta(\boldsymbol{x}_t) \prod \limits_{i=1}^N 1 = 1. \end{gather}\] We prove (ii) by starting with a representation of a normalized tensor, then converting it to a constraint-satisfying representation of the same tensor. Given \(w^{(\alpha)}_\theta(\boldsymbol{x}_t)\) and \(p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)\), for every \(\alpha\), we divide each \(p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)\) by \(\sum \limits_{x^i=1}^V p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)\), and multiply \(w^{(\alpha)}_\theta(\boldsymbol{x}_t)\) by the same amount, keeping the tensor unchanged. Then we have \(\sum \limits_{x^i=1}^V p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)=1\). And by \[1 = \sum \limits_{\boldsymbol{x}} p(\boldsymbol{x})= \sum \limits_{\alpha=1}^r w^{(\alpha)}_\theta(\boldsymbol{x}_t) \prod \limits_{i=1}^N \sum \limits_{x^i=1}^V p^{(\alpha, i)}_\theta(x^i | \boldsymbol{x}_t)=\sum \limits_{\alpha=1}^r w^{(\alpha)}_\theta(\boldsymbol{x}_t)\] we get \(\sum \limits_{\alpha=1}^r w^{(\alpha)}_\theta(\boldsymbol{x}_t)=1\). It represents the same tensor, and satisfies the constraints.

7.2 Tensor-Train Decomposition (TTD)↩︎

For TTD, recall that the constraint is given as: \[\sum_{x^i, k} \boldsymbol{G}_{i,\theta}(x^i, j, k | \boldsymbol{x}_t) = 1, \quad \forall i,j.\] Similar to our proof for CPD, we begin by verifying (i): \[\begin{gather} \sum \limits_{\boldsymbol{x}} p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)=\sum \limits_{\boldsymbol{x}} \tfrac{1}{r} \boldsymbol{1}_r^T \boldsymbol{G}_{1,\theta}(x^1 | \boldsymbol{x}_t) \boldsymbol{G}_{2,\theta}(x^2 | \boldsymbol{x}_t) ... \boldsymbol{G}_{N,\theta}(x^N | \boldsymbol{x}_t) \boldsymbol{1}_r \\ =\tfrac{1}{r} \boldsymbol{1}_r^T \sum \limits_{x^1=1}^V \boldsymbol{G}_{1,\theta}(x^1 | \boldsymbol{x}_t) ... \sum \limits_{x^N=1}^V \boldsymbol{G}_{N,\theta}(x^N | \boldsymbol{x}_t) \boldsymbol{1}_r. \end{gather}\] Since by constraint \(\sum \limits_{x^i=1}^V \boldsymbol{G}_{i,\theta}(x^i | \boldsymbol{x}_t) \boldsymbol{1}_r = \boldsymbol{1}_r\), the summations collapse: \[\sum \limits_{\boldsymbol{x}} p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)=\tfrac{1}{r} \boldsymbol{1}_r^T \sum \limits_{x^1=1}^V \boldsymbol{G}_{1,\theta}(x^1| \boldsymbol{x}_t) ... \sum \limits_{x^{N-1}=1}^V \boldsymbol{G}_{N-1,\theta}(x^{N-1} | \boldsymbol{x}_t) \boldsymbol{1}_r = ... = \tfrac{1}{r} \boldsymbol{1}_r^T \boldsymbol{1}_r = 1\] We prove (ii) in the same way as CPD. Given the cores for a normalized tensor, for position \(i\), for every \(j\), we divide each \(\boldsymbol{G}_{i,\theta}(x^i,j,k | \boldsymbol{x}_t)\) by \(\sum \limits_{x^i, k} \boldsymbol{G}_{i,\theta}(x^i,j,k | \boldsymbol{x}_t)\), and multiply the same term to each \(\boldsymbol{G}_{i-1,\theta}(x^{i-1},l,j | \boldsymbol{x}_t)\), to keep the tensor unchanged. After this, each \(\boldsymbol{G}_{i,\theta}(x^i | \boldsymbol{x}_t)\) satisfies the constraints at the cost of disrupting \(\boldsymbol{G}_{i-1,\theta}(x^{i-1} | \boldsymbol{x}_t)\). This is applied for \(i=N,\dots,2\) in decreasing order, leaving all but \(\boldsymbol{G}_{1,\theta}(x^1 | \boldsymbol{x}_t)\) well-constrained. For \(\boldsymbol{G}_{1,\theta}(x^1 | \boldsymbol{x}_t)\), we set \(\boldsymbol{G}_{1,\theta}(x^1, l, j | \boldsymbol{x}_t)=\tfrac{1}{r} \sum \limits_{l} \boldsymbol{G}_{1,\theta}(x^1, l, j | \boldsymbol{x}_t)\), i.e. we run in-place column-wise mean for each matrix. This preserves the tensor, since \(\boldsymbol{1}_r^T \boldsymbol{G}_{1,\theta}(x^1 | \boldsymbol{x}_t)\) remains the same. Also it gives the final constraint, because \[1=\sum \limits_{\boldsymbol{x}} p_\theta(\boldsymbol{x} | \boldsymbol{x}_t)=\tfrac{1}{r} \boldsymbol{1}_r^T \sum \limits_{x^1=1}^V \boldsymbol{G}_{1,\theta}(x^1| \boldsymbol{x}_t) ... \sum \limits_{x^N=1}^V \boldsymbol{G}_{N,\theta}(x^N| \boldsymbol{x}_t) \boldsymbol{1}_r=\tfrac{1}{r} \boldsymbol{1}_r^T \sum \limits_{x^1=1}^V \boldsymbol{G}_{1,\theta}(x^1| \boldsymbol{x}_t) \boldsymbol{1}_r\] is equivalent to \(\sum \limits_{x^1, j, k} \boldsymbol{G}_{1,\theta}(x^1,j,k | \boldsymbol{x}_t)=r\), thus after running the column-wise mean \(\sum \limits_{x^1, k} \boldsymbol{G}_{1,\theta}(x^1, j, k | \boldsymbol{x}_t)=1\) for every \(j\). It satisfies all constraints.

8 An Example Demonstrating TTD’s Advantage↩︎

As claimed in 4.2, we provide a simple distribution with neighbor-only dependence that can be modeled with a low TT-rank but not with a low CP-rank. Consider a tensor \[\boldsymbol{P}=\underbrace{I_2 \otimes I_2 \otimes ... \otimes I_2}_{N-\text{times}} / 2^N,\] representing a PMF defined over \(2N\)-sequences consisting of binary tokens where the paired neighbors share identical values (e.g., \(1100111100\)). We show that \(\boldsymbol{P}\) has \(r_{\text{CPD}}=2^N\), but \(r_{\text{TTD}}=2\).

We first show that \(r_{\text{CPD}} \leq 2^N\). Note that \(\boldsymbol{P}\) can be written as a sum of scaled Dirac deltas, with each indicating a possible sequence, while being rank-1 tensors. There are \(2^N\) of these, so the rank is at most \(2^N\).

Next, we show that \(r_{\text{CPD}} \geq 2^N\) by using the technique of matrix flattening. Sending the odd indices to the row, and even indices to the column, and we obtain a matrix of size \(2^N \times 2^N\). From the definition of \(\boldsymbol{P}\), this becomes the scaled identity matrix. This matrix has rank \(2^N\), and since the tensor rank is always greater than or equal to the matrix rank of any of its flattenings, we get \(r_{\text{CPD}} \geq 2^N\).

On the other hand \(r_{\text{TTD}} \leq 2\): We construct an explicit representation with \(r_{\text{TTD}}=2\) : \[\boldsymbol{G}_{2k-1}(0)=\begin{bmatrix}1/2 & 0 \\ 1/2 & 0\end{bmatrix}, \boldsymbol{G}_{2k-1}(1)=\begin{bmatrix}0 & 1/2 \\ 0 & 1/2\end{bmatrix}\] \[\boldsymbol{G}_{2k}(0)=\begin{bmatrix}1/2 & 1/2 \\ 0 & 0\end{bmatrix}, \boldsymbol{G}_{2k}(1)=\begin{bmatrix}0 & 0 \\ 1/2 & 1/2\end{bmatrix}\] then \[\boldsymbol{P}(\nu_1, \ldots, \nu_{2N})=\tfrac{1}{r} \boldsymbol{1}_r^T \boldsymbol{G}_1(\nu_1) \boldsymbol{G}_2(\nu_2) ... \boldsymbol{G}_{2N-1}(\nu_{2N-1}) \boldsymbol{G}_{2N}(\nu_{2N}) \boldsymbol{1}_r\]

\(r_{\text{TTD}} \geq 2\): \(r_{\text{TTD}}=1\) is equivalent to being a rank-1 tensor, which is not the case since \(r_{\text{CPD}} \geq 2^N\).

Thus \(r_{\text{CPD}}=2^N, r_{\text{TTD}}=2\). \(\square\)

9 A Pitfall of RoPE in Masked Diffusion↩︎

Rotary Position Embedding (RoPE) [28] 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\] where \(R^d_{\Theta, m}\) and \(R^d_{\Theta, n}\) are appropriately defined rotary matrices. (Here we write position indices in subscripts to match [28].)

This method, however, 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 neither changes 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 \[\begin{align} o_m=\sum \limits_{n=1}^N a_{m,n}v_n=v\sum \limits_{n=1}^N a_{m,n}=v \end{align}\] also becomes uniform regardless of the attention weights \(a\). All the other components, such as normalizing layers or dense layers, are applied independently to each position, maintaining this unintended uniformity. Indeed, we confirm that under this formulation, the outputted logits of the model are identical for every position in the first step.

Although this issue is less pronounced in settings such as conditional generation (where the first step is not fully masked), long sequence generation (where the impact of an incorrectly chosen first token is negligible), or datasets where true marginal probabilities are similar across positions, we observed noticeable performance degradation in our QM9 [18], [19] experiments, where the model performs unconditional generation of short sequences. This motivated a fix: applying RoPE rotary matrices to values in addition to queries and keys in attention. Note that we retain the previous behavior in 5.1 for compatibility with the pretrained model.

10 Experiment Details↩︎

10.1 Text Generation↩︎

For OWT [16] and LM1B [17], we matched the training time for each method. Each resulted in a total training time of approximately 10 hours on four NVIDIA A6000 GPUs.

For MDLM [3] we fine-tuned the base model with global batch size 64, constant learning rate \(1 \times 10^{-5}\) with Adam for 100k steps, and for CPD and TTD we used global batch size 4 for 150k steps.

Di4C [8] model was trained for 13k steps with the base model as the teacher model, with other hyperparameters set to default values in the released code of global batch size 4 and constant learning rate \(3 \times 10^{-5}\) with warmup for 2500 steps with AdamW [29] optimizer, alongside exponential moving average (EMA) with decay rate of 0.9999.

For EDLM [11], we used the checkpoints from MDLM with the AR-based method, and used the default sampling parameter of sampling size \(k=2\) but widened the sampling window to \(w=1\) to follow the setting outlined in their paper.

For DCD [10], we followed the released code and used sedd-medium [2] as the diffusion model and gpt2 [26] as the AR model with a chunk size of 1 to simulate fully random generation.

For VADD [9], we trained the base model with the default values in the released code of global batch size 512 and constant learning rate \(3 \times 10^{-4}\) with warmup for 2500 steps with AdamW and EMA with decay rate 0.9999 for 5k steps on latent dimension of 512. VADD and our methods were fine-tuned from the base model in the same setting, with VADD for 1k steps and our methods for 400 steps. For sampling we followed their default setting using fixed latent.

For measuring the sampling times, we generated 64 sequences with 128 timesteps on batch size of 1 for every method and computed the sampling time of a single sequence. Sampling time measurements were conducted on a single NVIDIA PRO6000 GPU.

10.2 Molecule Generation↩︎

The base MDLM [3] was trained with global batch size 8192, constant learning rate \(1 \times 10^{-4}\) with Adam [30] for 50k steps. For MDLM and our methods, we fine-tuned the base model under the same settings for 10k steps.

Di4C [8] model was trained for 10k steps with our base model as the teacher model, with global batch size 512 and other hyperparameters set to default values in the released code of constant learning rate \(3 \times 10^{-5}\) with warmup for 2500 steps with AdamW [29] and exponential moving average (EMA) with decay rate 0.9999.

VADD [9] model was trained for 60k steps from scratch with global batch size 2048 and other hyperparameters set to default values in the released code of constant learning rate \(3 \times 10^{-4}\) with warmup for 2500 steps with AdamW and EMA with decay rate 0.9999 on latent dimension of 512. For sampling of VADD we also followed their default setting using fixed latent.

Experiments were conducted using four GPUs of either NVIDIA A6000 or PRO6000.

11 Additional Results↩︎

11.1 Molecule Generation↩︎

In 4 where we generate molecules with different ordering strategies, we observe that the gap between TTD and CPD is most clearly pronounced in left-to-right order generation. This agrees with our theory that TTD outperforms CPD due to its ability to model the dependence between neighboring tokens. Orthogonal to the ordering strategy, we also adapt EB-Sampler [21], which adaptively determines the number of unmasked tokens to validate whether our models can be paired with advanced samplers. As shown in 5, we observe that our methods outperform the baseline even with a different choice of sampler.

From 6, we can see that the performance of TTD monotonically increases with the number of ranks used, as expected.

References↩︎