arXiv CS open access

On the Design Space of Continuous Diffusion Language Models

Oscar Davis et al. · arxiv_cs

arXiv CS · Papers

deep-learning machine-learning neural-networks

Continuous diffusion language models (CDLMs) promise to bring the benefits of continuous flow modelling to language, including accelerated sampling and controllable generation. The field has grown rapidly, and the capabilities of CDLMs have begun to rival those of discrete diffusion. Yet, unlike for

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

On the Design Space of Continuous Diffusion Language Models Oscar Davis1,∗ , Daan Roos2,∗ , Alice Bizeul, Amitis Shidani, Victor Turrisi, Louis Béthune, Pierre Ablin, Marco Cuturi

arXiv:2610.13834v1 [cs.LG] 9 Oct 2026

Apple, 1 University of Oxford, 2 University of Amsterdam Continuous diffusion language models (CDLMs) promise to bring the benefits of continuous flow modelling to language, including accelerated sampling and controllable generation. The field has grown rapidly, and the capabilities of CDLMs have begun to rival those of discrete diffusion. Yet, unlike for continuous modalities such as images or videos, our grasp of the levers that matter when training CDLMs remains limited. In this work, we focus on two of these levers, namely selfconditioning and embedding geometry, and shed light on some of the mechanisms behind them, which in turn suggest cheaper and more effective designs. We begin by showing that self-conditioning corrects a train/inference mismatch by exposing the model to inference-like states that are otherwise rare during training. These gains can be recovered without an extra forward pass by a two-level interpolant, coined DoubleCat, which improves TinyGSM accuracy more than threefold. Self-conditioning also aids in recognising tokens dominated by large-norm Gaussians—an issue known as token identifiability. To address it directly, we introduce discrete reveals to DoubleCat, yielding SemiCat, which achieves discrete diffusion-level performance on TinyGSM, reaching 49.94% with standard training hyperparameters and a tuned sampler. At the 1.7B scale, SemiCat outperforms discrete diffusion on four of six likelihood-based evaluations, and surpasses one-hot flows on unconditional generation. On embedded spaces, we find that cross-entropy drives embeddings towards a geometry whose decoding error rate closely matches the one-hot one. We briefly explore a range of further options, such as hybrid corruption processes and noise schedules, including asynchronous ones. We believe this exploration can help the community better understand and advance the design space of CDLMs. Correspondence: Oscar Davis: oscar.davis@cs.ox.ac.uk, Marco Cuturi: m_cuturi@apple.com Contributions: ∗ Work done while at Apple Date: October 12, 2026

Introduction

Continuous diffusion language models (CDLMs) hold several promises for text generation. Ideally, they should port to the world of language modelling all of the key properties of the continuous toolbox, such as fast and controllable generation, consistency models (Song et al., 2023; Geng et al., 2025; Boffi et al., 2025b), or stochastic optimal control-based steering (Potaptchik et al., 2026a; Holderrieth et al., 2026; Potaptchik et al., 2025). Yet, some of the seminal advances that have managed to shrink the performance gap between continuous and discrete diffusion methods have relied instead on adaptations that deviate significantly from the continuous diffusion literature. One such technique is self-conditioning (Chen et al., 2023), which augments inputs using previous outputs, costing an additional forward pass during training. Yoo et al. (2026) interpret self-conditioning as a fixed-point iteration, and Shabalin et al. (2026) understand it through the lens of multi-modality. Yet, it remains unclear which deficiencies of CDLMs it corrects, and whether these can be addressed more cheaply. Another innovative line of development lies in embedded methods (Deschenaux and Gulcehre, 2026; Chemseddine et al., 2026; Hu et al., 2026; Chen et al., 2026; Batzolis et al., 2026), which alleviate the computational and memory cost of one-hot methods (Roos et al., 2026; Lee et al., 2026; Potaptchik et al., 2026b). They rely on carefully chosen noise schedules. However, it remains unclear how such schedules interact with the learnt embeddings.

In this work, we show first that self-conditioning addresses a discrepancy between inference- and train-time trajectories (Section 3.1), and increases the density of otherwise under-sampled regions at train-time (consistent with Shabalin et al. 2026). This gain can be recovered by a simple two-level interpolant, called DoubleCat, which requires no additional forward pass of the model, and which more than triples the accuracy on TinyGSM (Liu et al., 2023) from 4.42% to 16.00%. Moreover, we find that self-conditioning helps in token identifiability (Pynadath et al., 2026) (Section 4.3): given the noisy state, the model cannot easily identify the expected endpoint of the trajectory, due to the norm mismatch between the noise and the token’s embedding, whereas self-conditioning is supported on the convex hull formed by the embeddings. Replacing that signal with a better representation of the state recovers a substantial part of the performance gain at no additional training cost. We address the token identifiability issue in DoubleCat (Section 3.2). By introducing discrete stochastic reveals akin to CANDI (Pynadath et al., 2026), the model is able to condition correctly on parts of its context, helping identifiability. This method, called SemiCat (semi-continuous categorical flow), further boosts the performance up to 32.90%, reaching the discrete diffusion baseline of about 33.74%. Using standard hyperparameters on the task, and a confidence-based sampler, SemiCat reaches 49.94% without any selfconditioning. At the 1.7B scale, SemiCat outperforms discrete diffusion models on 4 out of 6 likelihood-based benchmarks and matches them on a fifth one; it also surpasses one-hot flows in unconditional generation quality (Section 5). We also cover embedded methods. Because of architectural decisions, one-hot flows effectively operate on an embedded space with an anisotropic noise defined by the embeddings’ geometry (Section 4.1), on which we can directly train. We demonstrate that embeddings may not be trained with cross-entropy naively, as it maximises their spread. This itself induces a one-hot-like geometry on the embedded space when considering its corresponding decoding error rate (Section 4.2). This may explain why carefully crafted schedules are needed. We argue that training should update the geometry of all embeddings jointly. Lastly, we briefly explore other avenues such as asynchronous schedules, dynamic time allocation, hybrid methods, and acceleration, thus covering a large part of the design space of CDLMs. We believe that our results on our benchmarks show the promise of continuous and hybrid methods, and unveil important avenues left to explore to understand continuous diffusion for discrete data.

Background

Let V ⊂ N be the finite vocabulary of tokens, and x = (x1 , . . . , xL ) ∈ V L a sequence of length L ∈ N⋆ . Our goal is to model some data distribution, pdata , given independent and identically distributed samples therefrom. Model parameters are denoted by θ. A uniform distribution over I is denoted by U(I). We defer to Appendix D.1 a brief presentation of discrete diffusion models. Flow-based generative modelling.

A common option for the modelling of pdata is through the dynamic transport of measure (Lipman et al., 2023; Liu et al., 2022; Albergo et al., 2023). Given a simple starting distribution, p0 , and a target distribution p1 ≡ pdata , the aim is to create a probability path between the two, (pt )t∈[0,1] , via an ODE: ẋt = vt (xt ), with x0 ∼ p0 . Let I : [0, 1]×Rd ×Rd −→ Rd be an interpolant of the form It (X0 , X1 ) = αt X0 +βt X1 , for some differentiable α and β, X0 ∼ p0 , X1 ∼ p1 , and such that I0 (X0 , X1 ) = X0 and I1 (X0 , X1 ) = X1 . By choosing vt (x) = E[I˙t | It = x], then Law(xt ) = Law(It ) for any t, making xt available simulation-free through It . It thus reduces the problem of generative modelling to the regression: L(θ) = EX0 ,X1 ,t∼U (0,1) vtθ (It ) − I˙t

2 2

(2.1)

Sampling is done by integrating v θ . When βt ≡ 1 − αt , we write It (X0 , X1 ) = (1 − γt )X0 + γt X1 . Flow-based language modelling.

Flow modelling has also been adapted to the setting where pdata is a discrete distribution in V L (Eijkelboom et al., 2025; Roos et al., 2026; Lee et al., 2026; Potaptchik et al., 2026b).

Double

It` X0`

πtθ (xt ) π0θ (x0 )

Cat

X0`

Yt` Y0`

Cat

Rocks (a) One-hot flow, train-time.

It`

Rocks (b) One-hot flow, inference.

Rocks (c) DoubleCat, train time.

Figure 1 (a) One-hot flows are always exposed to nearly clean samples for t ≈ 1, and only track the clean target, X1ℓ . (b) At inference, they track instead the dynamic predictions of their denoiser, π θ , which always lies on the interior of the simplex. (c) DoubleCat addresses this issue at train-time by simply adding an additional corruption process on the output space of its denoiser, the simplex.

One choice of modelling is that of one-hot flows, where the final distribution is one-hot encoded into R|V | |It =x]−x (i.e., X1ℓ = exℓ ). Using the endpoint parametrisation of the generating vector field, vt (x) = γ̇t E[X11−γ , it t is possible to train using a cross-entropy loss instead: L(θ) = −EX0 ,X1 ,t

L X

D E log X1ℓ , πtθ,ℓ (It ) ,

(2.2)

ℓ=1

where π θ : [0, 1] × RL×|V | −→ RL×|V | is a denoising network. In this work, we also consider the possibility of evolving the flow on an embedded space, which is first embedded through a trainable codebook (Deschenaux and Gulcehre, 2026; Chemseddine et al., 2026; Hu et al., 2026; Chen et al., 2026). We denote by E ⊆ Rd the embedded space, with d ≤ |V |, and by emb : V −→ E the embedding function. For one-hot values ev (e.g., X1ℓ ), we extend the mapping as emb(ev ) = emb(v); similarly, the embedding of a sequence is the element-wise embedding. Let W be the matrix such that Wv = emb(v). Time schedule. The choice of time schedule is crucial for one-hot flows (Lee et al., 2026; Potaptchik et al., 2026b). Following Davis et al. (2026); Potaptchik et al. (2026b), we refer to and use by default for one-hot flows a mixture schedule, γt = λγtDER + (1 − λ)t, where γ DER is the decoding error rate (DER) schedule (Pynadath et al., 2026; Sahoo et al., 2025; Lee et al., 2026), with λ = 0.75. γ DER makes uniform progress in the DER (Lee et al., 2026, eq. (25)): L 1X Pe (t) := P(argmax Itℓ ̸= argmax X1ℓ ). (2.3) L ℓ=1

Hybrid methods.

Pynadath et al. (2026) introduce a hybrid flow, CANDI, that continuously evolves each state of the sequence until some randomly sampled unmasking time, τ ℓ , per position. Formally, for each ℓ, we let i.i.d. τ ℓ ∼ U (0, 1); for t < τ ℓ , Itℓ (X0 , X1 ) = (1 − γt )X0ℓ + γt X1ℓ , and Itℓ (X0 , X1 ) = X1ℓ otherwise. The clean token is directly revealed as a one-hot encoding at τ ℓ .1 Agarwal et al. (2026) introduce a similar construction via Flow Maps (Boffi et al., 2025b,a): FMLM+ adds to one-hot flows time t = 0 random reveals of the sequence’s tokens at positions C, and learns the flow map from X0 to X1 , i.e., X0,1 (X0 ) = X1 , at positions {1, . . . , L} \ C. The sampler goes back and forth between times 0 and 1, similar to γ-sampling from Kim et al. (2024), committing a subset of tokens based on the model’s confidence at each round. 1 Note the reversal in time convention to match the flow matching one, and note that we here ignore the difference in interpolants incurred by typical diffusion schedules.

Plaid. Introduced in Gulrajani and Hashimoto (2023) and recently revisited in Yang et al. (2026), Plaid is a Variational Diffusion Model (VDM; Kingma et al., 2021) for language modelling, trained by minimising the negative evidence lower bound (NELBO): E D PL − log pθ (x) ≤ LNELBO (θ) := KL (q(X0 | x) ∥ p(X0 )) − EI1 |x ℓ=1 log exℓ , π1θ,ℓ (I1 ) h i (2.4) 2 + 21 Et,It |x SNR′ (t) ∥eθ (It , t) − emb(x)∥ ,

where eθ : E L × [0, 1] −→ E L is given by eℓθ : (x, t) 7−→ πtθ,ℓ (x)W for each ℓ, p ≡ p0 (Gaussian prior), and SNR′ being the derivative of the signal-to-noise ratio, SNR : [0, 1) −→ R, t 7−→ βt2 /αt2 . The schedules are learnt, following VDM. Self-conditioning.

Self-conditioning has been shown to be extremely beneficial for CDLMs (Chen et al., 2023; Hu et al., 2026; Chen et al., 2026; Batzolis et al., 2026). It extends the denoiser to accept an additional input, π̃ θ : [0, 1] × Rd × Rd −→ Rd , and feeds the model its outputs from the previous inference step when generating samples. At train time, it is provided to the model with a probability psc : let C ∼ Ber(psc ), then x̂1 = C · sg(π̃tθ (It , 0)) is the self-conditioning signal, where sg is the stop-gradient operator, training then uses π̃tθ (It , x̂1 ). Concretely, the cross-entropy loss becomes: L(θ) = −EX0 ,X1 ,t,x̂1

L X

D E log X1ℓ , π̃tθ,ℓ (It , x̂1 ) .

(2.5)

ℓ=1

Validation

Enhancing one-hot flows

In this section, we identify two issues that one-hot flows suffer from. First, they lack trajectory ambiguity: they are trained on clean sequences at t ≈ 1, which mismatches the statistics of their inference trajectories. Second, they are unable to identify the tokens that they are generating, which we attribute to the high norm Gaussian noise. We show two simple interventions, DoubleCat and SemiCat, respectively, both of which drastically improve the resulting performance.

H(πtθ (·))

Inference

6 4 2 0

0.25

0.5

0.75

t Figure 2 Entropy of a one-hot flow’s denoiser, π θ ,

along validation set and inference trajectories.

Benchmark.

We mostly rely on TinyGSM (Liu et al., 2023): the model is trained on generating Python solutions to grade school mathematical problems, and is then evaluated on GSM8K (Cobbe et al., 2021). To understand failure modes of trained models, we measure the percentage of programs that run, the percentage that have name errors (referencing non-existent variables), and the final accuracy. We use the settings of Kim et al. (2026). Finally, retained methods are also evaluated on the standard OpenWebText (OWT) (Gokaslan and Cohen, 2019) unconditional generation task on which we report the generative perplexity (Gen. PPL) as measured by GPT-2 (Radford et al., 2019) alongside the entropy of generated samples, reported as Pareto fronts. See Appendix E.1 for more detail. As baselines, we provide masked discrete diffusion (MDM) and uniform-state diffusion (USDM), for discrete diffusion models, and FMLM+ (Agarwal et al., 2026) and CANDI (Pynadath et al., 2026) for hybrid methods. We retrain all baselines under the same settings, and report those numbers.

3.1

Trajectory ambiguity

For a one-hot flow, near t ≈ 1, the model is only trained on sequences in a very small radius around the final endpoint: It is a highly concentrated Gaussian with covariance (1 − γt )2 I and mean γt X1 . Moreover, at inference time, the model follows the learnt denoiser, π θ , which has non-zero entropy for any t (Figure 2), and therefore never points towards a clean vertex of the simplex. This creates a mismatch between inference and

Accuracy (%, ↑)

Method

Best (%)

T = 1 T = 0.5 T = 0.2 top-k=1 Runs (↑) Name err. (↓) Discrete diffusion baselines MDM USDM

8.79 12.43

21.68 29.34

27.14 33.06

28.28 33.74

88.63 91.96

11.22 7.81

One-hot based Base one-hot flow with self-cond. DoubleCat Base one-hot sampling

3.28 6.44 9.70 11.45

4.42 6.75 13.50 15.85

3.89 8.34 14.18 16.00

3.64 6.97 14.78 15.92

50.91 61.79 63.61 69.67

48.07 37.60 36.16 30.02

Hybrid methods, one-hot CANDI (Pynadath et al., 2026) 7.20 FMLM+ (Agarwal et al., 2026) 16.83 SemiCat 17.51 with standard hyperparameters† 21.08 + InfoNoise (Raya et al., 2026) 21.46 + Algorithm 2 (ρ/N = 0.75, λ = 0.25, ηmax = 0.25) –

9.55 18.88 29.34 35.41 36.24 –

9.63 18.27 32.22 37.68 41.09 49.73

8.72 18.50 32.90 38.67 38.06 49.94

69.22 90.52 89.23 90.75 86.58 95.36

29.26 9.02 10.39 9.10 13.04 4.38

Table 1 TinyGSM results at 256 NFEs. Accuracy is reported per sampler; the last two columns report the best value

attained over those samplers. T denotes the temperature applied to π θ . The models are trained with batch size 256 for 200k steps. † refers to standard hyperparameters for the TinyGSM benchmark in the literature (including batch size of 512, 250k steps). See Appendix E.1 for more details.

training, where the model, instead of following a prediction for X1ℓ , follows the path defined by the denoiser (πtθ (xt ))t∈[0,1] on the interior of the simplex. On the other hand, self-conditioning provides a signal with the same support as the denoiser, which perturbs the Gaussian inputs towards inference-like regions (Shabalin et al., 2026). Double continuous corruption.

We introduce DoubleCat, a doubly categorical flow, with a corruption process on the simplex, tracked by a normal Gaussian one. Formally, we define Ytℓ = (1 − κt )Y0ℓ + κt X1ℓ Itdc,ℓ (X0 , X1 ) = (1 − γt )X0ℓ + γt Ytℓ ,

(3.1)

for some Y0ℓ prior, and κ a differentiable schedule, with κ0 = 0 and κ1 = 1. Equivalently, it can be written as Itdc,ℓ (X0 , X1 ) = (1 − γt )X0ℓ + γt (1 − κt )Y0ℓ + γt κt X1ℓ .

(3.2)

Note that this forms a valid probability path between p0 and p1 . This formulation comes from the observation that the inference trajectory (Iˆtℓ )t∈[0,1] (induced by the denoiser-driven ODE) can be written as Iˆtℓ = (1 − γt )X0ℓ + γt Ŷtℓ ,

(3.3)

for (Ŷtℓ )t∈[0,1] a trajectory over the simplex corresponding to a weighted average of the predictions of the chosen denoiser; the derivation is given in Appendix B. Ŷtℓ being a convex combination of predictions, by the concavity of the entropy, H(Ŷtℓ ) is at least the weighted average entropy of the denoiser along the path, which is positive at all t (Figure 2). For DoubleCat, if we choose Y0ℓ to have mean |V1 | , then the entropy evolves roughly from H( |V1 | ) = log |V | to H(Y1ℓ ) = 0, thus remaining positive until late, like Ŷtℓ . DoubleCat therefore trains on inference-like states. Pure one-hot flow training is recovered for κt ≡ 1, which, in contrast, puts no mass on simplex components with non-zero entropy, since H(X1ℓ ) = 0. For inference, we use the same denoiser for both parts of the flow as an approximation. We ablate the choice of prior in Table 6, and find the uniform √ distribution over the simplex to work best. We also ablate the schedule κ in Table 7, and retain κ : t 7−→ t. Figure 1 illustrates our intuition. 5

Results. The results on TinyGSM are available in Table 1. DoubleCat’s performance improves by more than 3× over one-hot flows. Discarding the Yt term at inference—i.e., sampling as in a usual one-hot flow—further improves the performance: up to over 4× (with T = 0.2), with an accuracy reaching 16.00%, name error percentages lower by 18 points, and valid programs almost 19 points higher. This suggests that DoubleCat imitates the dynamics at inference of a one-hot flow, which tracks the denoiser’s trajectory, (πtθ (xt ))t∈[0,1] , and that the simplex noise, Y0ℓ , biases sampling.

We also show that this mechanism, while being substantially cheaper computationally, outperforms plain self-conditioning, in Table 8. Nonetheless, self-conditioning further improves DoubleCat (up to 20.77% at T = 0.5); we further explore this in Section 4.3. Takeaway. Self-conditioning increases the density of highly likely inference regions during training, diminishing the train–test mismatch. DoubleCat recovers a large part of the benefit by introducing a double-level flow, imitating self-conditioning. Discrete vs. continuous.

A natural question is why this discrepancy affects discrete data generation, whereas it is far less of an issue in the typical continuous case (e.g., images). One possible explanation that we provide for this is the degeneracy of the target denoiser, Dt⋆ (x) = E[X1 | It = x]. Its Jacobian is proportional to the posterior covariance, γt ∇xℓ Dt⋆,ℓ (x) = cov(X1ℓ | It = x), (3.4) (1 − γt )2 2 ⋆,ℓ γt 1 − D (x) for one-hot data. For (continuous) log-concave densities the trace of which is (1−γ 2 t t) (e.g., Gaussian), this Jacobian remains bounded by 1/γt . For discrete data, by contrast, the covariance does not shrink between vertices. Call a state x uncommitted at position ℓ if its posterior is spread over several tokens, i.e., x ∈ Uηt,ℓ := {x : 1 − Dt⋆,ℓ (x)

≥ η} for a fixed

η ∈ (0, 1). On Uηt,ℓ , the trace of the Jacobian is at least ηγt /(1 − γt )2 , which diverges as t −→ 1. Near t ≈ 1, a small perturbation may change the target prediction drastically. Uncommitted states are also rarely seen in training: under a uniform prior, we show that the expected fraction of uncommitted positions is bounded as L 1X 2Pe (t) P(It ∈ Uηt,ℓ ) ≤ −−−→ 0. (3.5) t−→1 L η ℓ=1

At inference, however, H(Ŷtℓ ) > 0, and trajectories travel through this sparsely trained interior, where the model must extrapolate. A smooth extrapolation of a smooth target is nearly correct, but that of a target varying this sharply is not: any misplaced boundary flips the predicted token. These claims are detailed in Appendix B. Gaps in training coverage seem to be more benign for continuous data, but costly for discrete data, which DoubleCat partially addresses by placing training mass in those regions.

3.2

Token identifiability

Despite the improvement offered by DoubleCat, the model still heavily suffers from name errors (30.02%). We attribute this to token identifiability (Pynadath et al., 2026): when referencing another variable, the model struggles to identify current tokens because of the plarge noise in xt , and therefore cannot condition on it precisely: the Gaussian noise has an expected norm |V |, but the signal has a norm of 1. To mitigate this issue, we add a revealing mechanism akin to the one in CANDI: at random times τ ℓ , a token is revealed as a one-hot input. This way, the model can cleanly condition on revealed inputs. The resulting method is SemiCat. Formally, the SemiCat interpolants are given by: Itsc,ℓ (X0 , X1 ) = Itdc,ℓ (X0 , X1 )1(t < τ ℓ ) + 1(t ≥ τ ℓ )X1ℓ , 6

(3.6)

Gen. PPL

One-hot

DoubleCat

SemiCat

0 50 100 NFE = 64 1

NFE = 128 4

NFE = 512 4

NFE = 1024 4

Entropy Figure 3 Gen. PPL/entropy Pareto frontiers as measured by GPT-2 for one-hot flows trained on OpenWebText with

sequence lengths of 1024 tokens, using the GPT-2 tokeniser.

where 1 is the indicator function. An optimal model carries over tokens for which t ≥ τ ℓ , and so we parametrise it as such (Sahoo et al., 2024). The training therefore reduces to the following cross-entropy loss: L(θ) = −EX0 ,X1 ,t,τ

L X

D E 1(t < τ ℓ ) log X1ℓ , πtθ,ℓ (Itsc ) ,

(3.7)

ℓ=1

where τ = (τ 1 , . . . , τ L ). Following DoubleCat, we ignore the (Yt )t∈[0,1] terms at inference time, thus amounting to CANDI-like sampling: each still-continuous position is revealed between times tn and tn+1 with −tn θ probability tn+1 1−tn , committing a token drawn from π , while the remaining positions follow the one-hot flow, as summarised in Algorithm 1. In Table 1, SemiCat produces only 10.39% runs with name errors and programs that run 89.23% of the time, and reaches the discrete diffusion level with an accuracy of 32.90% without self-conditioning, on TinyGSM. We note that revealing inputs à la CANDI is not sufficient on its own. On OWT, SemiCat also dominates the Gen. PPL/Entropy Pareto frontier of one-hot flows at all NFEs (Figure 3). For better comparison, we also train SemiCat with more standard hyperparameters (including batch size 512 and 250k training steps, see Appendix E.1 and Table 5). We reach 38.67% with a naive sampler. Adaptive noise allocation (InfoNoise, Raya et al., 2026) brings our performance further up to 41.09%, and a better tuned confidence-based sampler allows SemiCat to attain 49.94%. Takeaway. Large Gaussian noise makes the tokens hard to identify from the current state, xt . Introducing reveals like CANDI considerably helps in this regard. While SemiCat obtains high performance, its stochastic reveals and piecewise continuous trajectories make it a harder candidate to accelerate using typical continuous techniques. We discuss in Appendix G.2 the possibility of designing asynchronous schedules (i.e., schedules per token position, ℓ) to alleviate this issue. Besides, we briefly explore more general hybrid corruption processes in Appendix E.5, and show that SemiCat can obtain competitive performance compared to FMLM+ when distilling our model using D-MMD (Hoogeboom et al., 2026; Salimans et al., 2024) in Appendix E.6, in the few-step regime.

Designing embeddings

One-hot flows are costly, especially given that modern LM vocabulary sizes can exceed 200k tokens, and incur a large memory and computational penalty, requiring the materialisation and manipulation of an L × |V | tensor per data point in a batch. It is therefore natural to depart from this representation and work on smaller embedded spaces. We show that their DER can become problematic with pure cross-entropy training, and that self-conditioning aids in token identifiability.

Accuracy (%, ↑)

Method

Best (%)

T = 1 T = 0.5 T = 0.2 top-k=1 Runs (↑) Name err. (↓) Ambient sphere d = 16 d = 128 d = 768 Plaid, d = 16 self-cond. without self-cond. Plaid, d = 128 self-cond. without self-cond. Plaid, d = 768 self-cond. without self-cond. Plaid variants, d = 16, without self-cond. with posterior DoubleCat stochastic codebook with posterior as self-cond. self-cond. at inference

6.97 7.88 8.64

9.40 9.25 12.89

8.95 9.63 13.12

8.87 10.24 12.28

58.83 56.48 66.72

40.71 42.76 32.75

19.71 8.19

28.43 10.84

11.30 11.98

2.81 11.14

84.38 72.25

9.17 26.91

17.89 6.75

22.44 9.55

16.22 9.48

10.46 9.78

82.94 62.77

12.66 36.54

17.97 4.62

24.56 5.23

12.96 8.79

5.00 8.04

78.32 57.70

3.56 41.77

12.89 9.55 10.84 16.60

15.85 16.15 16.53 18.04

13.87 18.35 17.97 19.86

14.18 18.42 18.35 20.24

74.07 68.31 67.10 76.12

25.32 31.08 32.07 23.12

Table 2 TinyGSM results at 256 NFE. Plaid learns its log-SNR endpoints; the variants all pin them. “Without self-

cond.” refers to checkpoints trained without self-conditioning. T denotes the temperature applied to π θ . The models are trained with batch size 256 for 200k steps. See Appendix E.1 for more details.

4.1

Your one-hot flow is secretly embedded

The open-sourced one-hot flow works for language modelling (Roos et al., 2026; Lee et al., 2026) use the adapted DiT architecture (Peebles and Xie, 2023) from Lou et al. (2024), with a simple linear projection for the input, which we denote by W ∈ R|V |×d . This means that the first layer of the architecture after the projection receives an input distributed as ∀ℓ ∈ {1, . . . , L}, (It W )ℓ ∼ N γt (X1ℓ )⊤ W, (1 − γt )2 W ⊤ W . (4.1) This effectively shows that one-hot flows operate on an embedded space defined by their projection matrix, but which also modifies the injected noise, giving it a covariance structure defined by W ⊤ W . Note that Pynadath et al. (2026) make a similar observation and use it to accelerate inference. This naturally holds for DoubleCat and SemiCat too. We discuss how to train SemiCat directly on the embedded space in Appendix F.1, and attain the same end performance but with a more unstable training due to a Cholesky decomposition.

4.2

Your embedded flow is secretly one-hot

Considering embedded methods, we know that DoubleCat and SemiCat on the embedded space from Section 4.1 obtain high performance, and Plaid (Gulrajani and Hashimoto, 2023) does too. In contrast, naive embedded training underperforms: denoted as “ambient sphere”, we train (with default cross-entropy) a simple flow interpolating between a Gaussian and the final unit-normalised embedding, i.e., ∥emb(v)∥ = 1 for any v ∈ V . Mathematically, we write Itℓ (X0 , X1 ) = (1 − t)X0ℓ + temb(X1ℓ ).

(4.2)

In Table 2, the ambient sphere underperforms Plaid with self-conditioning; besides, we note that the ambient sphere also largely underperforms on OWT, with a very high Gen. PPL (beyond 200) at matched data entropy. The axis that helps us understand these results is the training objective. The simple cross entropy objective of the ambient sphere is pathological in two ways: (i) The model only updates the embeddings of the tokens 8

Sphere

Plaid

One-hot

Plaid

SemiCat

USDM

MDM

Gen. PPL

Pe (t)

0.5

50 100

0 0

0.25

0.5

0.75

Entropy

(a) Decoding error rates, Pe .

(b) Gen. PPL/Entropy Pareto fronts on OWT.

Figure 4 DER of embedded geometries, and the quality/diversity trade-off across methods on OWT. (Left) The DER

of a trained sphere model on OWT and that of Plaid both with d = 768, against the known one-hot DER. (Right) Gen. PPL/Entropy Pareto frontiers on OWT, with Plaid for d = 16.

involved in the batch, thus not updating the entire geometry of E. Plaid updates all embeddings through its MSE diffusion term, and SemiCat and DoubleCat do so through the covariance structure of their embedded noise (Section 4.1). (ii) Cross-entropy ignores the geometry of the embeddings, and spreads tokens apart on the sphere approximately recovering the one-hot geometry in the DER sense (Pynadath et al., 2026; Sahoo et al., 2025), i.e., Pe (t) tends towards that of the one-hot case. To see this, observe that the DER for embeddings can be generalised to L

Pe (t) :=

1X 2 P(argminv Itℓ − emb(v) ̸= argmax X1ℓ ). L

(4.3)

ℓ=1

For emb : v 7−→ ev , this recovers the usual case of (2.3). Proposition 1 exhibits that the DER (4.3) depends solely on the interactions emb(j) − emb(i) and ⟨emb(i), emb(j)⟩ for any i, j ∈ V . In Appendix F.3, we show informally that, under a Bayes-optimal denoiser, the cross-entropy pushes only the true token’s embedding away from the posterior-mean embedding (i.e., away from the tokens it is confused with). Plaid’s MSE term contains an opposite, attracting component, which on its own would collapse embeddings to zero. Plaid thus balances the two, but pure cross-entropy training only repels, and we recover the one-hot DER regime. Indeed, Appendix F.2 shows that unit-norm embeddings maximally spread in dimension d have inner product terms that tend to 0 as d increases, and in turn the DER tends to the one-hot case. Empirically, in Figure 4a, we measure through Monte Carlo estimates the DER, Pe , across time for a trained ambient sphere and a Plaid model (d = 768) on OWT. The measured and one-hot DER curves are indistinguishable except for a small gap at t ≈ 0.85. Remark that Plaid’s embeddings never induce Pe (t) = 0 even at t = 1 due to the learnt VDM schedule endpoints (see Appendix D.2). At d = 768, the ambient sphere’s spreading of tokens makes them more identifiable, yielding lower name errors and higher run rate when compared with Plaid without self-conditioning. Takeaway. Cross-entropy training spreads tokens apart, inducing the same DER as a one-hot flow. This may explain why prior work requires more advanced schedules (Deschenaux and Gulcehre, 2026; Hu et al., 2026; Chen et al., 2026); an alternative is updating the geometry of all embeddings jointly.

4.3

Self-conditioning on embedded spaces

In Table 2, the performance of the Plaid runs without self-conditioning monotonically worsens with d. This substantiates our findings: in higher dimensions, the discrepancy between the train- and inference- time statistics, Itℓ and xℓt , is further widened. For this reason, we turn to adapting DoubleCat. Naively porting DoubleCat to Plaid, by simply tracking Y0ℓ = v⊤ W until the true This happens because the model is given an input noisier than expected by the statistics of the ELBO, and applying the large weight SNR′ to a non-zero loss makes the Embedded DoubleCat.

Y1ℓ = emb(X1ℓ ), for v ∼ U(∆|V | ), breaks.

model diverge. To solve this, we attempted two techniques that provide more stability. Note that we freeze the endpoints of the learnt VDM schedule for additional stability. The first one avoids having non-zero terms near t ≈ 1: we instead flow the model towards its token-wise Bayesian posterior expectation at that point, Itℓ = xℓt , and time t, assuming a uniform prior over the tokens and unit-normalised embeddings. That is to say, we set ! ℓ β emb(v), x t t W, (4.4) Dtℓ := E emb(X1ℓ ) | Itℓ = xℓt = softmaxv∈V αt2 and construct the input instead as I˜tℓ = αt X0ℓ + βt Dtℓ . Note that we recover similar statistics, since Y0ℓ is ⊤ deterministically |V1 | W and Y1ℓ = emb(X1ℓ ). As shown in Table 2, this modestly improves over the base 11.98%, reaching 15.85%. The second one is denoted by “stochastic codebook” in Table 2, where we defined embeddings to be Gaussians, regularised to 1/d variance per coordinate and unit-norm mean. The goal is to increase the density of undertrained regions by this added stochasticity. This is most effective at low temperatures, but does not recover the self-conditioning performance. This resembles the noising tricks of Hu et al. (2026); Meshchaninov et al. (2026). Self-conditioning and token identifiability.

We hypothesise that self-conditioning also enhances token identifiability. To show that, instead of removing self-conditioning outright, we attempt to replace it with a cheaper proxy at train time, which is rendered possible with small d. While the token identifiability issue is less severe in small dimension, the mismatch in norms between the Gaussian and the embedding of norm one is still present. In contrast, the self-conditioning signal, x̂1 = eθ (xt , 0, t) ≡ EX1 ∼πtθ (·,0,t) [emb(X1 ) | It = xt ] ,

(4.5)

lies exactly on the convex hull formed by the embeddings. Instead of self-conditioning, we feed the model the Bayesian posterior expectation (i.e., (4.4)), x̂1 = Dt , at train time. At inference, we try both the posterior (x̂1 = Dt , where Dt = (Dt1 , . . . , DtL )) and true self-conditioning (x̂1 = πtθk−1 (xtk−1 )). Since Dt is deterministic given It = xt , it adds no information to the input, and gains must come from a better representation of the state. This approach attains 18.35% (vs. 11.98% without self-conditioning), and 20.24% with true self-conditioning at inference—at no additional forward pass cost at train-time. We support this claim with the following idealised analysis, proven in Appendix F.6: we assume a perfect denoiser, self-conditioning computed from It as at train-time, and unit-norm embeddings; we use the Plaid path, Itℓ (X0 , X1 ) = αt X0ℓ + βt emb(X1ℓ ). Let R be a function of It , i.e., R = f (It ) ∈ Rd . Define DB (R) to be the largest possible distortion between g(emb(X1ℓ )) and g(R) over all B-Lipschitz functions g : Rd −→ Rd : DB (R) :=

sup Lip(g)≤B

E g(R) − g(emb(X1ℓ ))

(4.6)

This definition quantifies the worst-case error gap for Lipschitz maps evaluated on representations of It . Since 2 the sup is attained by g : x 7−→ Bx, we have DB (R) = B 2 E R − emb(X1ℓ ) . Define the posterior mean embedding as Dt⋆,ℓ = E[emb(X1ℓ ) | It ]. For any R, we know that DB Dt⋆,ℓ ≤ DB (R) . (4.7) That is to say, Dt⋆,ℓ is optimal. Furthermore, we compare Itℓ (normalised by βt to account for the norm mismatch) and Dt⋆,ℓ , and prove that d ⋆,ℓ ℓ ℓ 2 − MMSEt , (4.8) DB (It /βt ) − DB (Dt ) = B SNR(t) where MMSEℓt is the minimal MSE at time t and position ℓ. (Appendix F.6 treats the more general case of an imperfect denoiser.) We note that the best attainable MMSE over codebooks of dimension d is non-increasing 10

Gen. PPL

One-hot

SemiCat

0 50 NFE = 256 1

NFE = 512 4

NFE = 1024 4

NFE = 2048 4

Entropy Figure 5 Gen. PPL as measured by Qwen 2.5 7B for models trained on Nemotron with sequence lengths of 2048 tokens,

using the TikToken tokeniser.

Method

ARC-E

BoolQ

OBQA

PIQA

RACE

SIQA

Chance

24.7

50.4

26.6

51.6

24.2

32.2

LLaDA-8B-Base (Nie et al., 2025) AR† (autoregressive) AR⋆ (autoregressive)

– 72.7 73.1

– 71.9 63.6

– 40.4 40.2

74.4 78.1 76.9

– 36.2 36.3

– 41.9 43.5

Duo† (uniform diffusion) MDLM† (masked diffusion)

53.4 50.5

59.6 62.8

33.0 32.0

62.7 62.2

35.0 34.7

39.0 39.2

One-hot flow SemiCat (ours)

42.6 61.8

37.9 64.8

26.8 30.2

56.4 70.3

28.2 35.3

36.4 40.4

Table 3 Zero-shot accuracy (%) on likelihood-based benchmarks. ⋆ Models from Davis et al. (2026). † Models from Sahoo

et al. (2026). Both with 1.7B parameters, trained on Nemotron for 2.1T tokens. SemiCat and one-hot flow likelihoods are estimated by the cross-entropy loss of the continuation (see equations (2.2) and (3.7)). Best non-AR results at 1.7B parameters per column are in bold.

in d, while the distortion gap grows at least linearly in d. This mechanism also provides another possible explanation for the growing discrepancies of Plaid with d in Table 2, but it does not cover the aforementioned train/inference mismatch nor inference-time self-conditioning. Takeaway. Self-conditioning holds a more identifiable representation of the state, on the convex hull formed by the embeddings, which aids in token identifiability. Cross-method Pareto frontiers are available in Figure 7, and for 1024 NFE in Figure 4b.

Scaling SemiCat

We train a larger SemiCat model following the protocol of Davis et al. (2026); Sahoo et al. (2026) closely: the architecture is the same DiT (Peebles and Xie, 2023; Lou et al., 2024) with 1.7B parameters, trained on the Nemotron Pretraining dataset (Basant et al., 2025) for 2.1T tokens using the TikToken tokeniser (OpenAI, 2023), with sequences of length 2048. We use 256× H100 GPUs for the pre-training run with a batch size of 10 per GPU. We also train a one-hot flow baseline on the same settings (Davis et al., 2026). See Appendix E.1 for more details. Likelihood-based benchmarks. We evaluate our model over the same multiple-choice question answering (MCQA) benchmarks as in Davis et al. (2026); Sahoo et al. (2026): ARC-easy (Clark et al., 2018), BoolQ (Clark et al., 2019), OBQA (Mihaylov et al., 2018), PIQA (Bisk et al., 2019), RACE (Lai et al., 2017) and SIQA (Sap et al., 2019). Instead of using an ELBO, we find the plain cross-entropy loss to be a better proxy for the likelihood for both the one-hot flow and SemiCat.

The results are available in Table 3. SemiCat outperforms our one-hot flow baseline on every MCQA task,

and outperforms the discrete diffusion baselines on ARC-E by 8.4 points, on BoolQ by 2.0, PIQA by 7.6, and on SIQA by 1.2. On OBQA, our results are behind by 2.8 points, and we slightly exceed Duo and MDLM on RACE (+0.3). Unconditional generation.

Following the setting of Davis et al. (2026), SemiCat outperforms one-hot flows on unconditional generation, as measured by the Gen. PPL/Entropy Pareto frontiers provided in Figure 5. We generate sequences of 2048 tokens, and sweep over NFEs and temperature. The Gen. PPL is computed using Qwen 2.5 7B (Yang et al., 2025). At all NFEs, SemiCat dominates the one-hot flow baseline.

Related Work

Continuous diffusion for language modelling.

Gaussian interpolation on one-hot encodings was introduced from a variational perspective in VFM (Eijkelboom et al., 2025), and has been shown to be amenable to selfdistillation and steering by CFM/FMLM/DFM (Roos et al., 2026; Lee et al., 2026; Potaptchik et al., 2026b). Yoo et al. (2026) analyse self-conditioning through a fixed-point iteration lens; Shabalin et al. (2026) do so through that of multi-modality. Embedded space design. Since those advances, S-FLM and von Mises-Fisher flows (Deschenaux and Gulcehre, 2026; Chemseddine et al., 2026) were the first to show reasonable performance on TinyGSM for continuous methods, by flowing on the sphere but by taking different probability path constructions. In parallel, geometrically simpler approaches have been developed. LangFlow (Chen et al., 2026) use a simple linear projection, and ELF (Hu et al., 2026) attempt learnt embeddings but have found using a pre-trained encoder to be much more effective. Meshchaninov et al. (2026) show how to leverage a pretrained AR model to obtain strong embeddings.

We provide additional references in Appendix A.

Conclusion

In this work, we show that self-conditioning aids in two important aspects of CDLMs: trajectory ambiguity (increasing the density of inference-like points at train time) and token identifiability (the ability to identify the expected endpoint of the trajectory given its current state). For the former, DoubleCat can recover a large part of those benefits; for the latter, providing a better representation of the state, such as feeding the analytically available Bayesian posterior, helps substantially at a lower cost. We show that one-hot methods in fact evolve over an embedded space with anisotropic noise, while embedded methods trained naively recover one-hot properties in terms of their DER. SemiCat, a hybrid method, can attain discretediffusion-level performance on our benchmarks. At scale, SemiCat outperforms discrete diffusion on four out of six of our likelihood tasks, and surpasses one-hot flows on unconditional generation. We leave to future work accelerating these methods with the continuous toolbox, and understanding how that toolbox interacts with our design choices.

References Agarwal, M., Shah, S., Lee, C., Yoo, J., Huang, J., Hong, S., Raghunathan, A., Kim, J., and Boffi, N. M. (2026). Posterior refinement: Fast language generation via any-order flow maps. Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. (2023). Stochastic interpolants: A unifying framework for flows and diffusions. Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and van den Berg, R. (2023). Structured denoising diffusion models in discrete state-spaces. Bansal, P. and Sanghavi, S. (2026). Token time continuous diffusion for language modeling.

Basant, A., Khairnar, A., Paithankar, A., Khattar, A., Renduchintala, A., Malte, A., Bercovich, A., Hazare, A., Rico, A., Ficek, A., et al. (2025). Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model. Batzolis, G., Girolami, M., and Ambrogioni, L. (2026). Cobit: Language modeling with bitstream diffusion. Bisk, Y., Zellers, R., Bras, R. L., Gao, J., and Choi, Y. (2019). Piqa: Reasoning about physical commonsense in natural language. Boffi, N. M., Albergo, M. S., and Vanden-Eijnden, E. (2025a). Flow map matching with stochastic interpolants: A mathematical framework for consistency models. Boffi, N. M., Albergo, M. S., and Vanden-Eijnden, E. (2025b). How to build a consistency model: Learning flow maps via self-distillation. Brascamp, H. J. and Lieb, E. H. (1976). On extensions of the brunn-minkowski and prékopa-leindler theorems, including inequalities for log concave functions, and with an application to the diffusion equation. Journal of Functional Analysis, 22:366–389. Campbell, A., Benton, J., Bortoli, V. D., Rainforth, T., Deligiannidis, G., and Doucet, A. (2022). A continuous time framework for discrete denoising models. Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, T. (2024). Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. Chemseddine, J., Kornhardt, G., and Steidl, G. (2026). Spherical flows for sampling categorical data. Chen, T., Zhang, R., and Hinton, G. (2023). Analog bits: Generating discrete data using diffusion models with self-conditioning. Chen, Y., Liang, C., Sui, H., Guo, R., Cheng, C., You, J., and Liu, G. (2026). Langflow: Continuous diffusion rivals discrete in language modeling. Cheng, C., Li, J., Peng, J., and Liu, G. (2025). Categorical flow matching on statistical manifolds. Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. (2019). Boolq: Exploring the surprising difficulty of natural yes/no questions. Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. (2018). Think you have solved question answering? try arc, the ai2 reasoning challenge. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. (2021). Training verifiers to solve math word problems. Davis, O., Filippova, A., Ablin, P., Turrisi, V., Shidani, A., Cuturi, M., and Béthune, L. (2026). Scaling categorical flow maps. Davis, O., Kessler, S., Petrache, M., İsmail İlkan Ceylan, Bronstein, M., and Bose, A. J. (2024). Fisher flow matching for generative modeling over discrete data. Deschenaux, J. and Gulcehre, C. (2026). Language modeling with hyperspherical flows. Dieleman, S., Sartran, L., Roshannai, A., Savinov, N., Ganin, Y., Richemond, P. H., Doucet, A., Strudel, R., Dyer, C., Durkan, C., Hawthorne, C., Leblond, R., Grathwohl, W., and Adler, J. (2022). Continuous diffusion for categorical data. Eijkelboom, F., Bartosh, G., Naesseth, C. A., Welling, M., and van de Meent, J.-W. (2025). Variational flow matching for graph generation. Gat, I., Remez, T., Shaul, N., Kreuk, F., Chen, R. T. Q., Synnaeve, G., Adi, Y., and Lipman, Y. (2024). Discrete flow matching. Geng, Z., Deng, M., Bai, X., Kolter, J. Z., and He, K. (2025). Mean flows for one-step generative modeling. Geng, Z., Pokle, A., Luo, W., Lin, J., and Kolter, J. Z. (2024). Consistency models made easy. Gokaslan, A. and Cohen, V. (2019). Openwebtext corpus. http://Skylion007.github.io/OpenWebTextCorpus. Gourevitch, S., Janati, Y., Shariatian, D., Simsekli, U., Moulines, E., Xing, E. P., and Durmus, A. (2026). Uniform diffusion models revisited: Leave-one-out denoiser and absorbing state reformulation.

Graves, A., Srivastava, R. K., Atkinson, T., and Gomez, F. (2025). Bayesian flow networks. Gulrajani, I. and Hashimoto, T. B. (2023). Likelihood-based diffusion language models. Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. Holderrieth, P., Chen, D., Eyring, L., Shah, I., Anantharaman, G., He, Y., Akata, Z., Jaakkola, T., Boffi, N. M., and Simchowitz, M. (2026). Diamond maps: Efficient reward alignment via stochastic flow maps. Hoogeboom, E., Nielsen, D., Jaini, P., Forré, P., and Welling, M. (2021). Argmax flows and multinomial diffusion: Learning categorical distributions. Hoogeboom, E., Ruhe, D., Heek, J., Mensink, T., and Salimans, T. (2026). Beyond single tokens: Distilling discrete diffusion models via discrete mmd. Hu, K., Qiu, L., Lu, Y., Zhao, H., Li, T., Kim, Y., Andreas, J., and He, K. (2026). Elf: Embedded language flows. Jutras-Dubé, P., Pynadath, P., Lu, J., Gao, Y., and Zhang, R. (2026). Sticky jump diffusions: A unifying view of masked, continuous, and hybrid diffusion. Kim, D., Lai, C.-H., Liao, W.-H., Murata, N., Takida, Y., Uesaka, T., He, Y., Mitsufuji, Y., and Ermon, S. (2024). Consistency trajectory models: Learning probability flow ode trajectory of diffusion. Kim, J., Geuter, J., Alvarez-Melis, D., Kakade, S., and Chen, S. (2026). Stop training for the worst: Progressive unmasking accelerates masked diffusion training. Kingma, D. P., Salimans, T., Poole, B., and Ho, J. (2021). Variational diffusion models. Lai, G., Xie, Q., Liu, H., Yang, Y., and Hovy, E. (2017). Race: Large-scale reading comprehension dataset from examinations. Lee, C., Yoo, J., Agarwal, M., Shah, S., Huang, J., Raghunathan, A., Hong, S., Boffi, N. M., and Kim, J. (2026). Flow map language models: One-step language modeling via continuous denoising. Li, X. L., Thickstun, J., Gulrajani, I., Liang, P., and Hashimoto, T. B. (2022). Diffusion-lm improves controllable text generation. Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling. Liu, B., Bubeck, S., Eldan, R., Kulkarni, J., Li, Y., Nguyen, A., Ward, R., and Zhang, Y. (2023). Tinygsm: achieving >80% on gsm8k with small language models. Liu, X., Gong, C., and Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. Lou, A., Meng, C., and Ermon, S. (2024). Discrete diffusion modeling by estimating the ratios of the data distribution. Meshchaninov, V., Shabalin, A., Chimbulatov, E., Gushchin, N., Koziev, I., Korotin, A., and Vetrov, D. (2026). How to train your latent diffusion language model jointly with the latent space. Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. (2018). Can a suit of armor conduct electricity? a new dataset for open book question answering. Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. (2025). Large language diffusion models. OpenAI (2023). tiktoken: A fast BPE tokenizer for use with OpenAI models. Peebles, W. and Xie, S. (2023). Scalable diffusion models with transformers. Potaptchik, P., Lee, C.-K., and Albergo, M. S. (2025). Tilt matching for scalable sampling and fine-tuning. Potaptchik, P., Saravanan, A., Mammadov, A., Prat, A., Albergo, M. S., and Teh, Y. W. (2026a). Meta flow maps enable scalable reward alignment. Potaptchik, P., Yim, J., Saravanan, A., Holderrieth, P., Vanden-Eijnden, E., and Albergo, M. S. (2026b). Discrete flow maps. Prabhudesai, M. and Geng, Z. (2026). Latent thought flows with text compression. Technical report, Carnegie Mellon University.

Pynadath, P., Shi, J., and Zhang, R. (2026). Candi: Hybrid discrete-continuous diffusion models. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. (2019). Language models are unsupervised multitask learners. Raya, G., Nguyen, B., Batzolis, G., Takida, Y., Stancevic, D., Murata, N., Lai, C.-H., Mitsufuji, Y., and Ambrogioni, L. (2026). Noise scheduling as information-guided allocation in diffusion training. Roos, D., Davis, O., Eijkelboom, F., Bronstein, M., Welling, M., İsmail İlkan Ceylan, Ambrogioni, L., and van de Meent, J.-W. (2026). Categorical flow maps. Sahoo, S. S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J. T., Rush, A., and Kuleshov, V. (2024). Simple and effective masked diffusion language models. Sahoo, S. S., Deschenaux, J., Gokaslan, A., Wang, G., Chiu, J., and Kuleshov, V. (2025). The diffusion duality. Sahoo, S. S., Lemercier, J.-M., Yang, Z., Deschenaux, J., Liu, J., Thickstun, J., and Jukic, A. (2026). Scaling beyond masked diffusion language models. Salimans, T., Mensink, T., Heek, J., and Hoogeboom, E. (2024). Multistep distillation of diffusion models via moment matching. Sap, M., Rashkin, H., Chen, D., LeBras, R., and Choi, Y. (2019). Socialiqa: Commonsense reasoning about social interactions. Schiff, Y., Sahoo, S. S., Phung, H., Wang, G., Boshar, S., Dalla-torre, H., de Almeida, B. P., Rush, A., Pierrot, T., and Kuleshov, V. (2024). Simple guidance mechanisms for discrete diffusion models. Shabalin, A., Elistratov, S., Meshchaninov, V., Sadrtdinov, I., and Vetrov, D. (2026). Why gaussian diffusion models fail on discrete data and how to prevent it? Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. K. (2025). Simplified and generalized masked diffusion for discrete data. Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. (2015). Deep unsupervised learning using nonequilibrium thermodynamics. Song, J., Meng, C., and Ermon, S. (2022). Denoising diffusion implicit models. Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. (2023). Consistency models. Wang, G., Schiff, Y., Sahoo, S. S., and Kuleshov, V. (2026). Remasking discrete diffusion models with inference-time scaling. Wu, T., Fan, Z., Liu, X., Gong, Y., Shen, Y., Jiao, J., Zheng, H.-T., Li, J., Wei, Z., Guo, J., Duan, N., and Chen, W. (2023). Ar-diffusion: Auto-regressive diffusion model for text generation. Yang, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Li, C., Liu, D., Huang, F., Wei, H., Lin, H., et al. (2025). Qwen2.5 technical report. Yang, Z., Guo, W., Zhang, S., Sahoo, S. S., Chen, Y., Vahdat, A., Mardani, M., and Thickstun, J. (2026). Continuous diffusion scales competitively with discrete diffusion for language. Yoo, J., Kim, W., Eijkelboom, F., Lee, C., Boffi, N. M., Hong, S., and Kim, J. (2026). Self-conditioned flow map language models via fixed-point flows. Zheng, H., Gong, S., Zhang, R., Chen, T., Gu, J., Zhou, M., Jaitly, N., and Zhang, Y. (2025). Continuously augmented discrete diffusion model for categorical generative modeling. Zhou, C., Yang, C., Hu, Y., Wang, C., Zhang, C., Zhang, M., Mackey, L., Jaakkola, T., Bates, S., and Zhang, D. (2026). Coevolutionary continuous discrete diffusion: Make your diffusion language model a latent reasoner.

Extended related work

Earlier approaches.

Earlier approaches of continuous diffusion for language modelling include CDCD (Dieleman et al., 2022), Diffusion-LM (Li et al., 2022), Bayesian Flow Networks (Graves et al., 2025). Geometric constructions.

Earlier geometric work includes Davis et al. (2024); Cheng et al. (2025), which flowed both on Fisher-Rao simplex and on the positive orthant on the sphere, but that failed to scale to larger vocabulary sizes. Beyond token-wise embeddings.

Latent Thought Flow (Prabhudesai and Geng, 2026) explores the idea of compressing the sequence length. We believe that this general direction is also very promising, but the exploration thereof should be deferred to another dedicated work. Other works consider using an existing strong AR model for encoding, such as Meshchaninov et al. (2026), which uses the contextual embeddings of GPT-2 (Radford et al., 2019). While offering better performance, the approach is naturally more expensive computationally at train time. Hybrid methods.

Besides CANDI (Pynadath et al., 2026) and FMLM+ (Agarwal et al., 2026), we note also the existence of the literature of discrete diffusion with continuous augmentations, for instance via latents (Zheng et al., 2025; Zhou et al., 2026). Concurrently, Jutras-Dubé et al. (2026) formalised these hybrids, together with masked and continuous diffusion, as special cases of a single process on the embedding space. Tokens form discrete anchors that the reverse process can jump onto, and the forward process determines when and where those jumps happen. In their framework, the CANDI-style hybrid underlying SemiCat is recovered as the identity-kernel case. We see this as a principled way to design hybrid corruption processes, complementary to the variants we explore in Appendix E.5.

DoubleCat derivations

Interpolant form of the inference trajectory.

A similar result can also be found in Lee et al. (2026); Potaptchik et al. (2026b), where the authors provide an integral form of the two-time denoiser. Fix a position ℓ and write πt = πtθ,ℓ (Iˆt ) for the prediction of any denoiser along the inference trajectory, which solves πt − Iˆtℓ dIˆtℓ = γ̇t , dt 1 − γt

Iˆ0ℓ = X0ℓ .

(B.1)

i h Rt d ˆℓ It /(1 − γt ) = γ̇t πt /(1 − γt )2 . Integrating with γ0 = 0, and using 0 γ̇s /(1 − γs )2 ds = γt /(1 − γt ), Then dt we obtain, for t ∈ (0, 1), Iˆtℓ = (1 − γt )X0ℓ + γt Ŷtℓ ,

Ŷtℓ :=

Z t wt (s)πs ds, 0

wt (s) :=

1 − γt γ̇s . γt (1 − γs )2

(B.2)

Since wt ≥ 0 integrates to one on [0, t], Ŷtℓ is a convex combination of the predictions, and thus lies on the Rt simplex. Finally, by concavity of the entropy and Jensen’s inequality, H(Ŷtℓ ) ≥ 0 wt (s)H(πs )ds. Write at := γt /(1 − γt )2 and σt := (1 − γt )/γt , so that at σt2 = 1/γt . Since Itℓ (X0 , X1 ) = (1 − γt )X0ℓ + γt X1ℓ , the posterior reads ! 2 ∥x − γt x1 ∥ p(x1 | It = x) ∝ p(x1 ) exp − ∝ p̃(x1 ) exp (at ⟨x, x1 ⟩) , (B.3) 2(1 − γt )2 Degeneracy of the target denoiser.

where p̃(x1 ) := p(x1 ) exp(−γt2 ∥x1 ∥ /(2(1 − γt )2 )) does not depend on x. Hence, ∇x log p(x1 | It = x) = at (x1 − Dt⋆ (x)), and differentiating Dt⋆ (x) = E[X1 | It = x] under the expectation gives (B.4) ∇Dt⋆ (x) = at E X1 (X1 − Dt⋆ (x))⊤ | It = x = at cov(X1 | It = x). 16

This holds for any prior p. Consider the (ℓ, ℓ) block, ∇xℓ Dt⋆,ℓ (x) = at cov(X1ℓ | It = x). For one-hot data, X1ℓ ∈ {ev }v∈V and Dt⋆,ℓ (x) =: q is the posterior over tokens, so that cov(X1ℓ | It = x) = diag(q) − qq ⊤ , the 2 trace of which is 1 − ∥q∥ . In particular, on Uηt,ℓ , tr∇xℓ Dt⋆,ℓ (x) ≥ ηat , which blows up as t tends to 1. In contrast, for continuous densities, if p1 is log-concave, the posterior is p1 times a Gaussian of covariance σt2 I, and the Brascamp–Lieb variance inequality (Brascamp and Lieb, 1976) bounds the variance of any linear 2 function x 7−→ ⟨u, x⟩ by σt2 ∥u∥ , i.e., cov(X1 | It = x) ⪯ σt2 I. Therefore, for all x and t, 0 ⪯ ∇Dt⋆ (x) ⪯ at σt2 I =

1 I, γt

(B.5)

and so the Jacobian remains bounded as t tends to 1, since γ1 = 1. Finally, for the training mass, consider the token-wise posterior under a uniform prior, q = q(Itℓ ) (which 2 2 coincides with Dt⋆,ℓ for this choice of prior), as in (4.4), and let qmax := maxv qv . Since qmax ≤ ∥q∥ ≤ qmax , 2

1 − qmax ≤ 1 − ∥q∥ ≤ (1 − qmax )(1 + qmax ) ≤ 2(1 − qmax ).

(B.6)

Under a uniform prior, argmaxv qv = argmax Itℓ , so that E[1 − qmax ] = P(argmax Itℓ ̸= argmax X1ℓ ) =: Peℓ (t). 2 Hence, Peℓ (t) ≤ E[1 − q(Itℓ ) ] ≤ 2Peℓ (t), and Markov’s inequality gives P(It ∈ Uηt,ℓ ) ≤ 2Peℓ (t)/η, which averages over positions to 2Pe (t)/η.

Extended SemiCat presentation

The sampling algorithm for SemiCat is provided in Algorithm 1. It follows the sampling procedure of CANDI (Pynadath et al., 2026). Moreover, we propose Algorithm 2 for a more sophisticated sampling scheme, which mostly introduces DDIM-style (Song et al., 2022) renoising, and confidence-based “remasking” (return to continuous state) inspired by the procedure of ReMDM (Wang et al., 2026). Algorithm 1 SemiCat sampling. Require: Denoiser π θ , schedule γ, time grid 0 = t0 < · · · < tN = 1, temperature T . 1: xℓ ∼ p0 and r ℓ ← 0 for all ℓ ▷ rℓ = 1: revealed 2: for n = 0, . . . , N − 1 do 3: p ← πtθn (x) 4: for ℓ such that rℓ = 0 do γt −γtn 5: xℓ ← xℓ + n+1 pℓ − x ℓ ▷ Euler step towards π θ 1−γt n −tn 6: Draw M ℓ ∼ Ber tn+1 1−tn if M ℓ = 1 then Draw v ℓ ∼ Cat (pℓ )1/T xℓ ← evℓ 10: rℓ ← 1 11: end if 12: end for 13: end for L 14: return argmaxv xℓv ℓ=1 7: 8: 9:

Additional background

D.1

Discrete diffusion

Discrete diffusion models corrupt a sequence by acting directly on V L , and learn to invert that corruption. We keep the flow matching time convention of the main text: t = 1 carries the data and t = 0 the noise, so the signal level γ : [0, 1] −→ [0, 1] is increasing with γ0 = 0 and γ1 = 1, exactly as for the flow schedules of Section 2. Many works in the discrete diffusion literature use the opposite diffusion convention. 17

Algorithm 2 SemiCat tuned sampling. Require: Denoiser π θ , schedule γ, uniform time grid tn = n/N , DDIM-stochasticity ηmax , remask budget ρ, uniform share λ. 1: xℓ ∼ p0 and r ℓ ← 0 for all ℓ ▷ rℓ = 1: revealed 2: for n = 0, . . . , N − 1 do 3: for ℓ such that rℓ = 1 do ▷ Confidence-based remask 4: Draw Rℓ ∼ Ber min 1, Nρ λ + (1 − λ)(1 − cℓ ) 5: if Rℓ = 1 then 6: Draw ϵℓ ∼ p0 7: xℓ ← (1 − γtn )ϵℓ + γtn xℓ 8: rℓ ← 0 9: end if 10: end for 11: p ← πtθn (x) (1−tn+1 ) ▷ Paced reveal 12: q ← 1 − L|{ℓ:r ℓ =0}| if n < N − 1, else q ← 1 13: q ← clamp(q, 0, 1) r γ (1−γ 2 t tn+1 ) 14: η ← ηmax 1 − γtn (1−γ t ) n+1

for ℓ such ℓthat rℓℓ = 0 do p x −γ p x̂ℓ0 ← 1−γttn , draw ϵℓ ∼ p0 , set x̂ℓ0 ← 1 − η 2 x̂ℓ0 + ηϵℓ n 17: xℓ ← (1 − γtn+1 )x̂ℓ0 + γtn+1 pℓ 18: Draw M ℓ ∼ Ber(q) 19: if M ℓ = 1 then 20: v ℓ ← argmaxv′ pℓv′ 21: xℓ ← evℓ 22: cℓ ← pℓvℓ 23: rℓ ← 1 24: end if 25: end for 26: end for L 27: return argmaxv xℓv ℓ=1 15: 16:

▷ DDIM step towards π θ

▷ top-k=1 sampling ▷ Record confidence for remasking

The forward process.

The two families below are instances of D3PM (Austin et al., 2023), which builds on Sohl-Dickstein et al. (2015); Hoogeboom et al. (2021); Ho et al. (2020), and whose continuous-time limits are continuous-time Markov chains (Campbell et al., 2022). The corruption factorises over positions, and each position either keeps its token or is replaced by a draw from a fixed reference distribution qref on V : q(Xtℓ | xℓ ) = γt δxℓ + (1 − γt ) qref .

(D.1)

Masked diffusion models (MDMs) take qref = δm , for an absorbing mask token m ∈ / V , so that Xt ∈ (V ∪{m})L while uniform-state diffusion models (USDMs) take qref = U(V ), so that Xt ∈ V L . The two choices give different objectives and different samplers. Masked diffusion.

The denoiser takes the corrupted sequence and returns a distribution over V at every position, i.e., πtθ,ℓ (Xt ) ∈ R|V | . Two constraints make it exact: it puts no mass on m, and it copies positions that are already unmasked, since the forward process never changes them (Sahoo et al., 2024; Shi et al., 2025). The reverse process only unmasks at rate R̂tℓ (m, j) =

γ̇t 1[j = xℓ ], 1 − γt

R̂tθ,ℓ (m, j) =

γ̇t π θ,ℓ (Xt )j , 1 − γt t

(D.2)

P and since πtθ,ℓ (Xt ) sums to one, the total outflow j R̂tθ,ℓ (m, j) does not depend on θ. Under these constraints the continuous-time bound collapses to a masked cross-entropy,   Z 1 D E X γ̇t EXt ∼q(· | x)  − log pθ (x) ≤ − log X1ℓ , πtθ,ℓ (Xt )  dt, (D.3) 1 − γ t 0 ℓ ℓ : Xt =m

where X1ℓ = exℓ as in Section 2, and which, for the linear schedule γt = t, weighs the loss by 1/(1−t). Sampling reverses (D.1) by unmasking a subset of positions per step; a token, once revealed, is never revisited. MDMs are currently the most common setting of recent diffusion language models at scale (Nie et al., 2025). Uniform-state diffusion. With qref = U(V ) the corrupted sequence contains no marker of which positions are wrong, so the model must both locate and correct the errors. Nothing is frozen, so the sampler may revise any token at any time rather than committing once. For completeness, we give the closed forms which are those of the uniform kernel of Austin et al. (2023); Hoogeboom et al. (2021).

Let us write the marginal at position ℓ, and the same quantity with the data replaced by the denoiser’s prediction as 1 − γt 1 − γt , qtθ,ℓ (j) := γt πtθ,ℓ (Xt )j + . (D.4) qtℓ (j) := γt 1[j = xℓ ] + |V | |V | Because the uniform kernels compose Ka Kb = Kab , the transition from a cleaner time s to a noisier time t < s retains with probability γt /γs , and Bayes’ rule gives the exact reverse posterior as an element wise product γt 1 γt 1 ℓ ℓ ℓ q(Xs | Xt , x ) ∝ , (D.5) eXtℓ + 1 − ⊙ γs exℓ + (1 − γs ) γs γs |V | |V | with pθ (Xsℓ | Xtℓ ) obtained by substituting πtθ,ℓ (Xt ) for exℓ . This is the kernel a finite-step sampler uses. Its continuous-time limit is a jump process: writing νt := γ̇t /γt for the corruption rate, the noising process leaves each position at rate νt towards U(V ), and time reversal gives, for j ̸= i R̂tℓ (i, j) =

νt qtℓ (j) · , |V | qtℓ (i)

R̂tθ,ℓ (i, j) =

νt qtθ,ℓ (j) . · |V | qtθ,ℓ (i)

(D.6)

The prior term of the ELBO vanishes, since γ0 = 0 makes q(X0 | x) exactly the uniform prior, and the remainder is the integrated KL rate between the two jump processes (Campbell et al., 2022), which after dropping θ-independent terms would become  ! Z 1 L X θ,ℓ θ,ℓ ℓ X 1 q (j) q (j) q (j) t t  dt + const. − ℓ t ℓ log θ,ℓ (D.7) νt EXt ∼q(· | x)  − log pθ (x) ≤ θ,ℓ ℓ) ℓ) |V | 0 q (X ) q (X q (X t t ℓ t t t t ℓ=1 j̸=Xt

Unlike (D.2), the total outflow

θ,ℓ ℓ j̸=Xtℓ R̂t (Xt , j) does

depend on θ, so the first term does not drop and

the objective is not a cross-entropy: the model is scored on the ratios qtθ,ℓ (j)/qtθ,ℓ (Xtℓ ) rather than on the probability it assigns to the true token, and every position contributes at every t, not only the corrupted ones. See Schiff et al. (2024); Gourevitch et al. (2026) for practical estimators of (D.7), and Lou et al. (2024), which sidesteps it by learning the ratios pdata (y)/pdata (x) directly with a score-entropy loss. Relation to the continuous case. USDMs are the discrete family closest to the flows of Section 2: Sahoo et al. (2025) show that a Gaussian diffusion pushed through an argmax matches the marginals of a uniform state diffusion. Discrete flow matching (Gat et al., 2024; Campbell et al., 2024) reformulates both families as probability paths with a prescribed interpolant, in a formalism much closer to that of Section 2.

D.2

The VDM noise schedule

Plaid (Gulrajani and Hashimoto, 2023; Yang et al., 2026) is a variational diffusion model (Kingma et al., 2021) on token embeddings, so the weight in (2.4) is the derivative of a learnt signal-to-noise ratio. We recall here where it comes from, and why its endpoints are treated differently from its shape. Following Section 2, Plaid corrupts the embedded sequence as It = αt X0 + βt emb(X1 ), with X0 ∼ N (0, I) and emb applied position-wise, and sets β2 SNR(t) := t2 . (D.8) αt VDM takes the path variance-preserving, αt2 + βt2 = 1, so that βt2 = sigmoid(log SNR(t)) and the whole path is determined by log SNR alone which it parametrises with a monotone network in t. As shown in (2.4), the continuous time bound then splits into a prior term, a reconstruction term, and a diffusion term with the diffusion term being h i 1 2 (D.9) Ldiff (θ) = Et∼U (0,1), X0 SNR′ (t) ∥eθ (It , t) − emb(x)∥ , 2 with SNR′ > 0. The schedule enters (D.9) twice: through the weight SNR′ (t) and through It . The second dependence carries no extra freedom, since βt2 = sigmoid(log SNR(t)) and αt2 = 1−βt2 are functions of SNR(t) alone. Let us write r r 1 w Jw := X0 + emb(x) (D.10) 1+w 1+w for the corrupted sequence at SNR level w, we have It = JSNR(t) : two schedules that reach the same SNR give the same noisy input, whatever time they reach it at. The denoiser’s second argument behaves the same way as VDM conditions it on log SNR(t) rather than on t; we write eθ (·, w) for it below. Substituting w = SNR(t) so that dw = SNR′ (t) dt, Ldiff (θ) =

1 2

Z SNR(1)

EX0 ∥eθ (Jw , w) − emb(x)∥ dw,

(D.11)

SNR(0)

so, at a fixed denoiser, the bound depends on the schedule through its two endpoints SNR(0) and SNR(1) and nothing else. VDM uses both halves of this statement: the endpoints are learnt by the ELBO, and the shape in between is learnt by minimising the variance of the Monte Carlo estimate of (D.9). Plaid inherits both, which is what we mean by learnt log-SNR endpoints in Table 2.

Additional experimental results

E.1

Hyperparameters and experimental setup

Hyperparameters.

We provide all our hyperparameters in Table 4. Throughout, we use the standard diffusion transformer (DiT Peebles and Xie, 2023) adapted by Lou et al. (2024) for discrete diffusion. For one-hot flows, we also use a simple linear projection in R|V |×h , where h is the hidden dimension of the transformer, following Roos et al. (2026); Lee et al. (2026). Models are evaluated with the EMA weights, except for in the large Nemotron case. We evaluate using the final checkpoint. OWT and TinyGSM checkpoints are trained using 8× H100. Nemotron 1.7B checkpoints are trained with 256× H100. All models are trained with AdamW (Loshchilov and Hutter, 2019); the hyperparameters for which are provided in Table 4 too. L2 Weighting. For one-hot flows (including base one-hot, CANDI, DoubleCat, SemiCat), for training, we weight the loss as in Davis et al. (2026), following Potaptchik et al. (2026b); Geng et al. (2024), inversely proportional to the square r root of the L2 distance of the prediction to the target. Precisely, we divide the

loss by sg(∆) for ∆ =

πtθ,ℓ (It ) − X1ℓ + ϵ, with ϵ = 10−2 .

OWT setup.

We use the typical OWT experimental setup: 1024 long sequences, tokenised by the tokeniser of GPT-2 (Radford et al., 2019).

Architecture (DiT) Hidden size adaLN cond. dim Layers Attention heads FFN (SwiGLU) multiplier Dropout Sequence length L Vocabulary |V | Parameters, trunk Parameters, in + out proj. Parameters, total Optimisation (AdamW) Optimiser Weight decay Gradient clipping Schedule Warmup steps Peak / final LR Total steps Global batch size Precision EMA decay

OpenWebText

TinyGSM

Nemotron (1.7B)

768 128 12 12 4 0 1024 50 257 (GPT-2) 121M 77M 199M

768 128 12 12 4 0 512 151 680 (Qwen2) 121M 233M 354M

2304 128 18 18 3 0 2048 100 352 (TikToken) 1 281M 462M 1 743M

(β1 =0.9, β2 =0.999, ϵ=10−8 ) 0.1 1.0 Cosine decay w/ linear warmup 3 000 3×10−4 / 10−5 1 000 000 200 000 512 256 bfloat16 mixed 0.9999

(β1 =0.9, β2 =0.95, ϵ=10−8 ) 0.1 1.0 Cosine decay w/ linear warmup 2 000 3×10−4 / 3×10−5 400 457 2 560 bfloat16 mixed none

Table 4 Hyperparameters of our models on OpenWebText, TinyGSM and Nemotron.

TinyGSM.

The TinyGSM dataset (Liu et al., 2023) contains 11.8M grade school mathematics questions paired with Python code traces that generate the correct solution. Our default hyperparameters are provided in Table 4. Since those differ slightly from the standard settings in the literature, we also use the standard ones for our final SemiCat run in Table 5. See Table 5 for a comparison.

E.2

DoubleCat ablations

The ablations over the choice of prior and the choice of schedule for DoubleCat are available in Tables 6 and 7. We also provide the full ablation of DoubleCat and one-hot flow with self-conditioning on and off in Table 8. TinyGSM (ours) Global batch size Total steps Schedule Warmup steps Peak / final LR Weight decay Optimiser Gradient clipping Precision EMA decay Architecture

TinyGSM (standard)

256 512 200 000 250 000 Cosine decay w/ linear warmup Constant w/ linear warmup 3 000 2 500 3×10−4 / 10−5 3×10−4 / 3×10−4 0.1 0 AdamW (β1 =0.9, β2 =0.999, ϵ=10−8 ) 1.0 bfloat16 mixed 0.9999 As in Table 4

Table 5 A comparison of TinyGSM hyperparameters between ours and the standard ones found in the literature. We

run most of our experiments on our own, and report SemiCat’s numbers on the standard configuration too (Table 1).

Accuracy (%)

Best (%)

Y0ℓ

T =1

T = 0.5

T = 0.2

top-k=1

Runs (↑)

Name err. (↓)

Barycentre Uniform Uniform vertex

4.70 6.14 5.31

6.52 9.17 7.05

6.29 10.84 7.81

7.51 11.30 8.49

55.57 53.53 51.48

43.90 46.10 47.99

Table 6 Ablation on the simplex prior Y0ℓ for DoubleCat, at NFE 256. Accuracy is reported per sampler; the last

two columns report the best value attained over those samplers. “Barycentre” is δ 1 , “Uniform” is uniform over the |V |

simplex, U (∆|V | ), and “uniform vertex” is uniform over the vertices of the simplex (i.e., the one-hot encoding of U(V )). The inner schedule is set to κt ≡ t2 . We retain the uniform distribution which consistently outperforms the other choices on accuracy.

Accuracy (%) κt √ t t t2 t3

Best (%)

T =1

T = 0.5

T = 0.2

top-k=1

Runs (↑)

Name err. (↓)

9.70 6.97 6.14 5.00

13.50 9.55 9.17 7.35

14.18 10.61 10.84 7.05

14.78 10.84 11.30 8.79

63.61 55.19 53.53 51.18

36.16 44.43 46.10 48.45

Table 7 Ablation on the simplex schedule κt for DoubleCat, at NFE 256, with Y0ℓ ∼ U(∆|V | ) for all runs. Accuracy is

reported per sampler; the last two columns report the best value attained over those samplers.

Accuracy (%, ↑)

Best (%)

Method

Self-cond.

T =1

T = 0.5

T = 0.2

top-k=1

Runs (↑)

Name err. (↓)

One-hot flow

Off On

3.28 6.44

4.42 6.75

3.89 8.34

3.64 6.97

50.91 61.79

48.07 37.60

DoubleCat

Off On

11.45 16.83

15.85 20.77

16.00 20.70

15.92 19.79

69.67 75.21

30.02 24.34

Table 8 Self-conditioning effect on one-hot flows and DoubleCat for schedule κt ≡

√ t.

Gen. PPL

d = 16

d = 128

d = 768

0 50 100

NFE = 64 0

NFE = 128 3

NFE = 512 3

NFE = 1024 3

Entropy Figure 6 Gen. PPL as measured by GPT-2 for models trained on OpenWebText with sequence lengths of 1024 tokens,

using the GPT-2 tokeniser. All Plaid runs for different embedding dimensionalities, d.

Gen. PPL

Plaid

SemiCat

USDM

MDM

0 50 100 150

NFE = 64 1

NFE = 128 4

NFE = 512

NFE = 1024 4

Entropy Figure 7 Gen. PPL as measured by GPT-2 for models trained on OpenWebText with sequence lengths of 1024 tokens,

using the GPT-2 tokeniser. Plaid is for d = 16.

E.3

OpenWebText Pareto Frontiers

We provide additional OWT Pareto frontiers in Figures 6 and 7.

E.4

Corruption as data augmentation

We show, here, that corruption seen as data augmentation, instead of a proper change in the interpolant, lowers the performance of the resulting model on the unconditional generation task of OWT by substantially diminishing its diversity (entropy). To be precise, for U ℓ ∼ U(V ) i.i.d., we use interpolants instead pointing towards X̃1ℓ = ζt X1ℓ + (1 − ζt )U ℓ , (E.1) with various choices of ζt ∼ Ber(1 − ρt ), constant and scheduled. Results are provided in Table 9. Any score comparable in Gen. PPL and entropy is found for very low constant choices of ρt ; lower (better) Gen. PPL scores imply entropy collapse.

E.5

General hybrid corruption processes

We also explore the design of hybrid methods based on continuous flows. Inspired by AUDM (Gourevitch et al., 2026), we introduce α-Cat, a generalisation of CANDI (Pynadath et al., 2026), where the model is instead exposed to the following interpolants: letting τ ℓ ∼ U(0, 1) i.i.d., as before, and U ℓ | X1ℓ = xℓ ∼ Cat(α |V1 | + (1 − α)exℓ ) i.i.d., we define the interpolants to be, for each ℓ, for τ ℓ > t, Itℓ (X0 , X1 ) = (1 − γt )X0ℓ + γt eU ℓ ,

(E.2)

Itℓ (X0 , X1 ) = X1ℓ .

(E.3)

and, for τ ℓ ≤ t, The model is thus effectively exposed to various levels α of additional discrete corruption, recovering CANDI at α = 0. Simply put, at each position, the model flows towards either a random vertex with probability α or to the true one otherwise, until its reveal time τ ℓ . We report the TinyGSM results in Table 10. As we can see, with increasing α, we recover better performance until α = 0.75 as it recovers the effect of SemiCat. Nonetheless, for α = 1, the model learns a leave-one-out denoiser; that is to say, it learns E[X1ℓ | It−ℓ = x−ℓ t ], 23

Gen. PPL (↓) (entropy) Flip probability ρt

NFE = 50

NFE = 100

Baseline, 1M steps, NFE = 128 One-hot flow DoubleCat SemiCat

NFE = 200

85.64 (5.45) 66.02 (5.38) 117.00 (5.62)

no L2 weighting ρt ≡ 0.01 ρt ≡ 0.1 ρt ≡ 0.25 ρt ≡ 0.5 ρt = t ρt = γ t

127.70 (5.44) 91.58 (5.09) 68.97 (4.72) 48.94 (4.13) 22.88 (3.12) 26.80 (2.86)

118.74 (5.52) 81.71 (5.19) 61.64 (4.83) 44.23 (4.16) 22.07 (3.30) 22.41 (2.86)

115.59 (5.51) 79.62 (5.12) 65.67 (4.88) 47.68 (4.18) 23.13 (3.18) 24.29 (2.99)

with L2 weighting ρt ≡ 0.01 ρt ≡ 0.1 ρt ≡ 0.25 ρt ≡ 0.5 ρt = t ρt = γ t

100.44 (5.42) 75.00 (5.14) 57.64 (4.84) 44.43 (4.28) 27.99 (3.37) 24.60 (3.00)

86.72 (5.43) 65.62 (5.19) 53.68 (4.89) 39.28 (4.30) 28.55 (3.46) 21.63 (2.98)

85.33 (5.45) 63.50 (5.18) 51.50 (4.89) 41.99 (4.31) 29.26 (3.36) 23.82 (3.13)

Table 9 OWT Gen. PPL (GPT-2), token entropy in parentheses, at 200k steps, temperature 1. The corruption flips

each X1ℓ independently, with probability ρt , to a token drawn uniformly from V ; it enters the interpolant only, leaving the loss target clean. Trained with γ being the usual mixture schedule. Accuracy (%, ↑)

Best (%)

T =1

T = 0.5

T = 0.2

top-k=1

Runs (↑)

Name err. (↓)

0 0.25 0.5 0.75 1

7.20 13.50 19.26 15.47 14.56

9.55 18.35 23.88 25.25 23.86

9.63 17.89 25.17 29.57 27.74

8.72 17.97 25.93 29.26 27.37

69.22 79.38 84.91 87.11 86.73

29.26 19.86 14.71 12.81 12.81

Table 10 TinyGSM accuracy at NFE = 256, for α-Cat at corruption levels α ∈ {0, 0.25, 0.5, 0.75, 1}.

which ignores position ℓ to predict token ℓ, like discrete diffusion methods (Gourevitch et al., 2026). This is natural, since I(Itℓ ; X1ℓ ) = 0 for t < τ ℓ and α = 1, where I denotes mutual information. Variants with α > 0 thus partially defy the purpose of hybrid methods, as they rely on discretely revealed positions.

E.6

Acceleration via D-MMD

We apply D-MMD distillation (Hoogeboom et al., 2026; Salimans et al., 2024) to our pre-trained SemiCat model. We tune the temperature of the teacher T , and attempt some simple techniques such as sampling boundary pairs (s, t) = (0, 1) with a probability p, and a gap curriculum updated every n steps, both inspired by Lee et al. (2026). We find that the resulting model is competitive with FMLM+ (Agarwal et al., 2026), and tends to have higher accuracy in the 1-step and 32-step case, while FMLM+ retains a consistent and balanced accuracy for different numbers of steps. Our results are obtained with simple naive sampling, and we hypothesise that more advanced samplers could further improve the performance (e.g., confidence-based).

Accuracy (%, ↑) NFE 1

FMLM+ (Agarwal et al., 2026)

0.2

1.2

5.7

12.4

17.8

19.9

SemiCat teacher

–

5.5

12.1

20.1

T =1 T = 0.5 T = 0.3 T = 0.2 T = 0.05 + boundary pairs 3% + boundary pairs 25% + gap curriculum 2k + gap curriculum 2k, 3% + gap curriculum 5k, 3%

0.2 0.7 0.5 0.6 1.1 1.1 0.6 1.1 0.6 0.7

0.5 1.2 0.8 1.8 1.8 2.5 1.1 2.4 2.0 1.5

1.6 2.8 3.9 4.8 3.6 5.6 1.9 5.4 3.9 3.7

5.4 9.8 9.2 10.9 9.0 9.6 4.9 11.7 6.7 7.8

13.0 16.5 16.3 16.3 17.1 15.9 10.0 19.1 13.8 12.4

19.0 24.0 22.7 22.5 24.2 20.0 15.5 22.1 20.0 20.4

Student

Table 11 TinyGSM accuracy for D-MMD distillation of SemiCat, at teacher temperature T . Students train for 20k

steps at constant learning rate 2 × 10−5 , with early stopping on GSM8K test set accuracy at 32 NFE.

Embedded spaces

F.1

Reproducing SemiCat’s statistics on an embedded space

Let us formally write out the interpolation incurred by SemiCat. Given i.i.d. times τ ℓ ∼ U(0, 1), at each position, we have, for 0 ≤ t < τ ℓ , Itℓ (X0 , X1 ) = (1 − γt )X0ℓ + γt Ytℓ , (F.1) and, for t ≥ τ ℓ , Itℓ (X0 , X1 ) = X1ℓ .

(F.2)

Thus, conditioning on the same times, we have either, for unrevealed times, W ⊤ Itℓ (X0 , X1 ) ∼ N γt W ⊤ Ytℓ , (1 − γt )2 W ⊤ W .

(F.3)

Ytℓ = (1 − κt )Y0ℓ + κt X1ℓ ,

(F.4)

Moreover, since for a Dirichlet prior, the mean conditional on X1ℓ is (1 − κt )W ⊤

1 + κt W ⊤ X1ℓ , |V |

(F.5)

the first term corresponding to the barycentre of the embeddings. For times after reveal, W ⊤ Itℓ (X0 , X1 ) = W ⊤ X1ℓ = emb(X1ℓ ). The only difficulty to obtain the same statistics directly in the embedded space is to obtain the same covariance structure. To do so efficiently, we can compute W ⊤ W , which we can then decompose by Cholesky as W ⊤ W = LL⊤ , and thus apply it to a Gaussian draw in lower dimension, LZ, for Z ∼ N (0, I), thus obtaining a covariance of LL⊤ = W ⊤ W . Our trained models obtain the same final performance as normal SemiCat training, but it is less stable due to the positive definiteness requirement of Cholesky.

F.2

Embedding-dependent schedules

The decoding error rate for one-hot encoded data has the well-known form from Pynadath et al. (2026); Sahoo et al. (2025). We give here the building block needed to generalise it to arbitrary embeddings: the pairwise nearest-neighbour comparison. The full DER then follows from the joint behaviour of these |V | − 1 comparisons, which we do not treat here as they are not independent.

Proposition 1 (Embedding-dependent pairwise decoding scores). Given an embedding matrix W ∈ R|V |×d , and the interpolant Itℓ = (1 − γt )X0ℓ + γt emb(X1ℓ ) between a standard Gaussian and the token’s 2 embedding, define the nearest-neighbour decoding scores St (v) := emb(v) − Itℓ . Then, for any pair of tokens i ̸= j, writing ∆ij := emb(j) − emb(i), 2 St (i) − St (j) ∼ N mij (t), 4(1 − γt )2 ∥∆ij ∥ , (F.6) and hence, since decoding is an argmin over squared distances, token i is preferred over token j with probability mij (t) , (F.7) ∀0 ≤ t ≤ 1, P (St (i) < St (j)) = 1 − Φ 2(1 − γt ) ∥∆ij ∥ where Φ is the c.d.f. of the standard normal distribution and the mean score gap is 2

mij (t) := 2γt ∆ij , emb(X1ℓ ) + ∥emb(i)∥ − ∥emb(j)∥ .

(F.8)

Proof. Let ℓ ∈ {1, . . . , L}. We have that Itℓ ∼ N γt emb(X1ℓ ), (1 − γt )2 I .

(F.9)

For all v ∈ V , define the scores St (v) = emb(v) − Itℓ

= ∥emb(v)∥ + Itℓ

− 2 emb(v), Itℓ .

(F.10)

For the nearest neighbour decoding to yield the true final token, X1ℓ , let us compare the pairwise errors, for indices i ̸= j: 2 2 St (i) − St (j) = ∥emb(i)∥ − ∥emb(j)∥ − 2 emb(i), Itℓ + 2 emb(j), Itℓ . (F.11) Expanding the dot product with embedding i (or identically for j), by linearity, emb(j), Itℓ = (1 − γt ) emb(j), X0ℓ + γt emb(j), emb(X1ℓ ) ,

(F.12)

we can then rewrite the difference of scores, for ∆ij := emb(j) − emb(i), as 2

St (i) − St (j) = 2γt ∆ij , emb(X1ℓ ) + ∥emb(i)∥ − ∥emb(j)∥ + 2(1 − γt ) ∆ij , X0ℓ . 2 Since X0ℓ ∼ N (0, I), the last term follows N 0, 4(1 − γt )2 ∥∆ij ∥ . Hence, 2 St (i) − St (j) ∼ N mij (t), 4(1 − γt )2 ∥∆ij ∥ , where ∀0 ≤ t ≤ 1,

(F.13)

(F.14) 2

mij (t) := 2γt ∆ij , emb(X1ℓ ) + ∥emb(i)∥ − ∥emb(j)∥ .

So the probability of the event of interest can be expressed as mij (t) −mij (t) P (St (i) < St (j)) = Φ =1−Φ . 2(1 − γt ) ∥∆ij ∥ 2(1 − γt ) ∥∆ij ∥

(F.15)

(F.16)

Recovering the one-hot DER. Informally, consider unit-norm embeddings with pairwise inner products ρij := ⟨emb(i), emb(j)⟩, i ̸= j. If cross-entropy spreads the embeddings roughly uniformly over the sphere, then, since two independent √ uniform points on the sphere have an inner product of mean 0 and variance 1/d, the ρij are of order 1/ d, and thus vanish as d grows. Taking i = X1ℓ to be the true token, the norm terms 2 of Proposition 1 cancel, ∥∆ij ∥ = 2(1 − ρij ) and mij (t) = −2γt (1 − ρij ), so that ! r γt 1 − ρij γt P (St (i) < St (j)) = Φ −−−−−→ Φ √ , (F.17) ρij −→0 1 − γt 2 2(1 − γt )

the one-hot value. In the same limit, the comparisons also recover the correlation structure of the one-hot case, up to O(ρij ) corrections, so the full DER tends to the one-hot expression of Pynadath et al. (2026); Sahoo et al. (2025). An embedding spread out by cross-entropy therefore approximately inherits the one-hot DER.

F.3

Cross-entropy, mean squared error and the embedding geometry

To gain some intuition, we compare, informally and approximately, how the cross-entropy and the MSE terms of Plaid (Gulrajani and Hashimoto, 2023) move the input embeddings W ∈ R|V |×d , with rows emb(v). Fix a position ℓ with true token i = xℓ , write σt = 1 − γt and Itℓ = σt X0ℓ + γt emb(i). We assume the denoiser is Bayes-optimal, so that its prediction is the posterior 2 πtθ,ℓ,v (It ) ∝ p(v) exp − Itℓ − γt emb(v) /2σt2 , (F.18) P for some prior p over tokens. Denote by µ := πtθ,ℓ W = v πtθ,ℓ,v emb(v) the expected embedding, and observe P that µ − emb(i) = v̸=i πtθ,ℓ,v (emb(v) − emb(i)) points towards the tokens carrying the error mass. The cross-entropy sees the embeddings only through the interpolant, through emb(i), and recall that X1ℓ is ℓ the one-hot random variable for the token at position ℓ (i.e., X1ℓ,x = 1 and X1ℓ,v = 0 for v ̸= xℓ ). Hence, omitting the arguments for the denoiser at time t, πtθ , for conciseness, ∇emb(i) − log πtθ,ℓ,i = γt ∇Itℓ − log πtθ,ℓ,i (F.19) Itℓ − γt emb(i) 2σt2

= γt ∇Itℓ + log

X v

I ℓ − γt emb(v) p(v) exp − t 2σt2

!! (F.20) !

γt = 2 σt

Itℓ − γt emb(i) −

πtθ,ℓ,v · Itℓ − γt emb(v)

(F.21)

X θ,ℓ,v γt X ℓ,v ℓ = 2 X1 It − γt emb(v) − πt · Itℓ − γt emb(v) σt v v γt X ℓ,v θ,ℓ,v ℓ = 2 X1 − πt It − γt emb(v) σt v ! X ℓ,v X ℓ,v γt θ,ℓ,v θ,ℓ,v X1 − πt = 2 Itℓ −γt X1 − πt emb(v) σt v v | {z }

! (F.22) (F.23)

(F.24)

∇emb(i) − log πtθ,ℓ,i

γ2 = t2 (µ − emb(i)) . σt

(F.25)

For (F.21), we used the Bayes optimality of the denoiser. Writing av := ∇Itℓ av = (Itℓ − γt emb(v))/σt2 , ∇Itℓ log

p(v)e−av = −

X v

Itℓ − γt emb(v)

/2σt2 , so that

1 X θ,ℓ,v p(v)e−av ℓ P ∇ π · I − γ emb(v) , ℓ av = − t I t t 2 −a t u σt v u p(u)e

(F.26)

with the last equality coming from the definition of the Bayes optimal denoiser. The wrong embeddings, emb(v) for v ̸= i, receive no gradient from the cross-entropy through the input, since they do not enter Itℓ . On the other hand, for the MSE term, 12 wt (πtθ,ℓ − X1ℓ )W ∇emb(v) 12 wt (πtθ,ℓ − X1ℓ )W

, with weight wt ≥ 0 for all t,

= wt πtθ,ℓ,v − X1ℓ,v (µ − emb(i)) + wt ∇emb(v) πtθ,ℓ W (µ − emb(i)) . 27

(F.27)

0.5

Gen. PPL

Val. loss

cos (Wc , Wx )

1.0 7 6 5

0.0 0

100

150

100

104 103

150

100

150

Step (×103 ) Figure 8 Statistics across training steps, for a Plaid model trained on OWT. (Left) The evolution of cosine similarity

between Wc (the part of the input matrix for self-conditioning) and Wx (the part that multiplies the input). (Middle) The validation loss. (Right) The Gen. PPL.

The cross-entropy gradient and the first term of the MSE gradient act on emb(i) along the same direction, µ − emb(i), but with opposite signs. A descent step on the cross-entropy moves the true token away from the predicted mean—i.e., away from the tokens it is confused with—weighted by how much it is confused with them, and without preference for any other structure, thus spreading the embeddings apart. A descent step on the first term of the MSE instead moves emb(i) towards µ with weight 1 − πtθ,ℓ,i , and translates every confused token emb(v) with weight πtθ,ℓ,v along the common direction emb(i) − µ, from the predicted mean to the true token, which brings emb(v) towards emb(i) when the confusion concentrates on v; on its own, the MSE is trivially minimised by the collapse W = 0. The second term has no definite direction. Again, the above is only informal : a trained network merely approximates the posterior.

F.4

Plaid self-conditioning incorporation

With a self-conditioning input signal x̂1 , alongside the usual input xt , the Plaid architecture forwards the two inputs to the rest of the diffusion transformer as Wc ℓ h0 = x̂1 xt . (F.28) Wx It thus doubles the size of the effective input matrix (i.e., from d × h to 2d × h).

F.5

Embedded self-conditioning: beyond Euler steps

Another observation which is consistent with our understanding from Section 4.3 is in Figure 8. Observe that one can split the input projection matrix in two parts: the one that multiplies the usual input, Wx , and the one that multiplies the self-conditioning, Wc (see Appendix F.4 for additional detail). We can then measure the cosine similarity between their vectorised forms. For the first 80k steps, the model converges to representations that are highly aligned (with a cosine similarity of 0.9), and the input projection thus amounts approximately to a linear combination followed by a common projection thereon. After that, the loss, performance and the embedding alignment of the model plateau for roughly 30k iterations, and then continue improving anew when the cosine similarity drops down to 0.6, alluding to the fact that the model learns a more complex function than a mere Euler-like step.

F.6

Self-conditioning and token identifiability

We use, here, the path Itℓ = αt X0ℓ + βt emb(X1ℓ ) (as in Appendix D.2), with X0 ∼ N (0, I) independent of X1 , and with X̂1ℓ being It -measurable (which holds at train time). Recall the distortion definition: DB (R) :=

sup Lip(g)≤B

E g(R) − g(emb(X1ℓ ))

(F.29)

where R is any square integrable representation of It . For any B-Lipschitz function g : Rd −→ Rd , we have that 2 2 E g(R) − g(emb(X1ℓ )) ≤ B 2 E R − emb(X1ℓ ) . (F.30) 28

The linear function g : x 7−→ Bx attains the RHS bound, and thus DB (R) = B 2 E R − emb(X1ℓ ) minimisation by Dt⋆,ℓ , observe that D E D h iE E R − Dt⋆,ℓ , Dt⋆,ℓ − emb(X1ℓ ) = E R − Dt⋆,ℓ , E Dt⋆,ℓ − emb(X1ℓ ) | It = 0,

. For the

(F.31)

since the RHS of the inner product is zero, by definition. The first equality is true by the tower property of the expectation, and R and Dt⋆,ℓ are both It -measurable. Therefore, we can add and subtract a Dt⋆,ℓ term to show that 2 2 (F.32) DB (R) = B 2 E Dt⋆,ℓ − emb(X1ℓ ) + B 2 E R − Dt⋆,ℓ , which implies minimisation when R = Dt⋆,ℓ almost surely. For the comparison, define ϵℓt := E X̂1ℓ − Dt⋆,ℓ

and MMSEℓt := E emb(X1ℓ ) − Dt⋆,ℓ

(F.33) 2

(F.34)

Therefore, DB (Itℓ ) = B 2 E αt X0ℓ − (1 − βt )emb(X1ℓ )

i h 2 . = B 2 dαt2 + (1 − βt )2 E emb(X1ℓ )

(F.35)

On the other hand, using (F.32), DB (X̂1ℓ ) = B 2 E emb(X1ℓ ) − Dt⋆,ℓ It follows that

+ B 2 E X̂1ℓ − Dt⋆,ℓ

= B 2 MMSEℓt + B 2 ϵℓt .

DB (Itℓ ) − DB (X̂1ℓ ) 2 = dαt2 + (1 − βt )2 E emb(X1ℓ ) − MMSEℓt − ϵℓt . B2

(F.36)

(F.37)

For the normalisation, Itℓ /βt , remark that Itℓ /βt − emb(X1ℓ ) =

αt ℓ X + emb(X1ℓ ) − emb(X1ℓ ) = SNR(t)−1/2 X0ℓ . βt 0

Whence: DB (Itℓ /βt ) − DB (X̂1ℓ ) = B 2 Note that MMSEℓt ≤ E emb(X1ℓ )

d − MMSEℓt − ϵℓt . SNR(t)

(F.38)

(F.39)

= 1 for unit-normalised embeddings.

Dependence on dimension.

We now show that, for unit-normalised embeddings, the MMSE does not depend on d explicitly, but only on SNR(t), the data distribution and the geometry of the codebook through its Gram matrix G ∈ R|V |×|V | , Guv := ⟨emb(u), emb(v)⟩ for u, v ∈ V , i.e., G = W W ⊤ . L m Let x1 = (x11 , . . . , xL 1 ) ∈ V . Since the X0 are independent standard Gaussians, the likelihood of It given X1 = x1 is ! L L 2 Y Y )∥ ∥Itm − βt emb(xm βt 1 m p(It | X1 = x1 ) ∝ exp − ∝ exp S , Svm := 2 ⟨Itm , emb(v)⟩ , m x 2 1 2α α t t m=1 m=1 (F.40) 2 where the proportionality constants do not depend on x1 , as ∥emb(v)∥ = 1 for all v ∈ V . Therefore, by Bayes’ rule, the posterior qvℓ := P(X1ℓ = v | It ) is a function of S = (Svm )m,v and of the law of X1 only, and P Dt⋆,ℓ = v∈V qvℓ emb(v). Conditionally on X1 = x1 ,

Svm = SNR(t) Gxm + SNR(t)1/2 ⟨X0m , emb(v)⟩ , 1 v

(F.41)

Linear

Entropic

Mixture

DER

1 log V = 11.93

S(γ(t))

γ(t)

0.5

0.25

0.5

0.75

0.25

0.5

0.75

t Figure 9 The entropic schedule against the base DER, mixture and linear schedules. (Left) The schedules themselves. (Right) The expected remaining entropy, S(γt ) = EX ℓ [H(X1ℓ | Itℓ )]. 0

where, for each m, (⟨X0m , emb(v)⟩)v∈V is a centred Gaussian vector with covariance G, independent across m. Hence, the joint law of (X1 , S), and thus that of (X1 , q ℓ ), depends only on the law of X1 , SNR(t) and G. Finally, writing δv ∈ R|V | for the indicator vector of v, 2

MMSEℓt = E

δX1ℓ − q

v∈V

ℓ

emb(v)

δX1ℓ − q

ℓ

⊤

G δX1ℓ − q

ℓ

(F.42)

which depends on the same quantities only. Consequently, two codebooks of dimension d and d′ sharing a Gram matrix attain the same MMSEℓt at the same SNR(t), whereas the distortion of the normalised input, B 2 d/SNR(t), scales linearly with the dimension. That is to say, at a fixed level of difficulty of token identification, the advantage of the posterior mean, d DB (Itℓ /βt ) − DB (X̂1ℓ ) = B 2 − MMSEℓt − ϵℓt , (F.43) SNR(t) grows linearly in d for a fixed approximation error ϵℓt : the noise lying in the extra directions is a nuisance which Dt⋆,ℓ averages out, but which Itℓ carries into the network. Observe that the set of Gram matrices of unit-normalised codebooks realisable in dimension d grows with d ′ (any codebook in Rd embeds into Rd for d′ > d by zero-padding, preserving G), so that the lowest MMSEℓt attainable by such a codebook is non-increasing in d, whereas the distortion of the normalised input increases linearly in d.

Further schedules

G.1

Entropic schedule

We attempt, here, to understand why the mixture schedule performs better than the raw DER schedule (Lee et al., 2026; Potaptchik et al., 2026b; Davis et al., 2026). For the usual one-hot flow case with a uniform prior, the true token-wise Bayesian posterior reads γt Itℓ ∀0 ≤ t ≤ 1, Dtℓ = softmax . (G.1) (1 − γt )2 Let S : [0, 1] −→ R+ be the expected remaining entropy; that is to say, ∀0 ≤ t ≤ 1, S(γt ) := EX0ℓ H(X1ℓ | Itℓ ) . 30

(G.2)

Gen. PPL

One-hot

One-hot (ent. sched.)

DoubleCat

0 50 100

NFE = 64 1

NFE = 128 4

NFE = 512 4

NFE = 1024 4

Entropy Figure 10 Gen. PPL as measured by GPT-2 for models trained on OpenWebText with sequence lengths of 1024 tokens,

using the GPT-2 tokeniser.

Instead of linearising directly the schedule in the probability of argmax Itℓ ̸= argmax X1ℓ as the DER schedule does, we propose to choose a schedule to advance linearly in S. Concretely, we are looking for a solution, γ : [0, 1] −→ [0, 1], to the following: ∀0 ≤ t ≤ 1,

S(γt ) = (1 − t) log |V |.

(G.3)

The solution has no closed form, so we instead build it with a look up table, which incurs a constant cost once before training, exactly as the DER schedule. Let us call the solution the “entropic” schedule. We plot the DER, mixture and entropic schedules in Figure 9, alongside the resulting S, for each choice of γ. We remark that the mixture schedule in fact closely follows the entropic schedule both in γ and S ◦ γ. Furthermore, we also plot the Gen. PPL/Entropy Pareto fronts of a one-hot flow with the mixture schedule, the entropic schedule and DoubleCat in Figure 10. While it closely matches the curve of the mixture schedule, it exhibits an outlier near 3 nats in entropy, especially at low NFE; however, we report that the resulting performance on TinyGSM is the same as that of the mixture schedule. We believe that this entropic schedule and the mixture one allude to the fact that making uniform progress in the DER probability Pe might be the wrong criterion.

G.2

Asynchronous schedules

While appealing, discrete reveals incur two properties that render the usage of usual continuous diffusion tools difficult: the path is now instead piecewise continuous, with discontinuity points that are stochastic. For example, the formalism from Flow Map Matching (Boffi et al., 2025a,b) then does not port trivially, because of the deterministic and continuous requirements on the flow map. To return to a continuous and (conditionally on the endpoints) deterministic process, we note that the revealing procedure can be seen as a limiting case of asynchronous schedules with added randomness. Instead of having a single time for the evolution of the interpolant, we can have a distinct time for each position, tℓ . In principle, a schedule that depends on the current state or the model’s denoising distribution is appealing, as parts of the sequence may benefit from being committed earlier than others. Nonetheless, training such a schedule either requires two forward passes from the denoising model, π θ , or training uniformly on the [0, 1]L hypercube to produce a model that is agnostic to the specific choice of schedule; but the latter exposes the model to sequences that have uninformative statistics for inference. We thus decide to solve this tension by resolving different tokens at different speeds, following Wu et al. (2023). Specifically, local times tℓ are derived from a global time tg , i.e., tℓ = F (tg , ℓ), which is monotonic in ℓ and approximately mean-preserving across positions, so that averaging tℓ over ℓ recovers tg . The results are also available in Table 12, with inner schedule κt = t. Asynchronous times improve accuracy, valid programs and name errors in three of the four matched settings, the exception being the linear schedule with InfoNoise. InfoNoise raises accuracy almost uniformly, but for the mixture schedule it lowers the fraction of valid programs, so that the best accuracy and the fewest name errors are reached by different configurations. All settings remain well below SemiCat (Table 1); we hypothesise that crafting a better time schedule is crucial here.

` = 32 ` = 64

tg = 0.1 tg = 0.3

` = 96 ` = 128

1.0

0.8

local time t`

`=1 ` = 16

0.6 0.4 0.2

tg = 0.5 tg = 0.7

tg = 0.9

0.6 0.4 0.2

0.0

0.0 0.0

0.2

0.4

0.6

global time t

0.8

1.0

100

120

position `

Figure 11 Local per-position times tℓ as a function of global time tg at concentration k = 1 and sequence length L = 128.

(Left) Each curve traces one position as global time advances. (Right) Each curve is the spatial profile at one global time, monotone decreasing in ℓ: the sequence is revealed left to right. The dashed horizontals mark the per-position mean, which equals tg at every setting.

Implementation. We follow Wu et al. (2023) and the specific parametrisation by Bansal and Sanghavi (2026). We introduce a global time tg that we sample as before: according to some density p(t) that is often to be taken uniform, but in the case of InfoNoise can be adapted dynamically during training. During inference we step uniformly in tg . For each position ℓ = 1, . . . , L we introduce a local time tℓ = F (tg , ℓ) that is a function of the global time and the position. Each position ℓ gets a fixed rank rℓ ∈ (0, 1), given by rℓ = 1 − (ℓ − 0.5)/L. The rank is then pushed through the quantile function of a Beta distribution whose mean is tg : tℓ = Qt (rℓ ), where Qt is the quantile function of Beta(k/(1 − t), k/t) where k specifies its concentration: as k −→ ∞ we recover a shared time tℓ = tg . As k −→ 0, Beta tends to a Bernoulli with mean tg , so position ℓ switches from noise to data exactly when tg crosses (ℓ − 0.5)/L: a deterministic left-to-right reveal. In our experiments we use k = 1; the resulting local times are shown in Figure 11.

G.3

Adaptive allocation of training density

Like Deschenaux and Gulcehre (2026); Batzolis et al. (2026) we adapt InfoNoise, the information-guided time allocation of Raya et al. (2026), to the cross-entropy objective used in our model. The adaptation changes only the distribution π(t) from which training times are sampled; the loss weighting itself remains fixed. R1 R1 Writing L̄(t) for the expected loss at time t, the objective thus changes from 0 p(t)L̄(t) dt to 0 π(t)L̄(t) dt, where p is the base (uniform) density. For a fixed reference batch Bref of size |Bref | = 256 (a single batch), we periodically estimate the cross-entropy across a grid of corruption times {ti }M i=1 , Ĥ(ti ) =

X 1 LCE (x, ti ). |Bref |

(G.4)

x∈Bref

We measure at M = 32 levels: 24 at equal mass of the current density, plus 8 spaced uniformly over the interior of [0, 1]. Cross-entropy upper-bounds the conditional entropy of the clean token given its corrupted input, with equality for the optimal denoiser, so we use Ĥ(t) directly as a proxy for the conditional-entropy profile. This differs from applications of InfoNoise to continuous data, where the information profile must first be estimated indirectly from a denoising-error statistic. The measured values are projected onto a monotone profile and interpolated with a cubic spline, yielding a

Accuracy (%, ↑) InfoNoise

Best (%)

Schedule

Asynchronous

T =1

T = 0.5

T = 0.2

top-k=1

Runs (↑)

Name err. (↓)

Linear

Off On

6.82 7.88

9.40 10.99

11.14 12.28

9.40 10.99

63.61 66.26

36.01 33.43

Mixture

Off On

10.39 12.21

14.78 15.47

13.65 16.00

13.87 14.63

68.16 73.39

31.54 26.38

Linear

Off On

9.63 11.37

17.44 18.20

18.42 17.82

15.69 17.44

67.78 66.11

31.69 33.43

Mixture

Off On

10.84 15.54

16.45 19.11

14.48 19.41

13.80 18.04

62.85 68.23

36.85 31.31

Off

Table 12 TinyGSM accuracy at NFE = 256. For linear and mixture outer schedules γt and inner schedule κt = t, we

compare a scalar time shared across all positions against an asynchronous left-to-right time per token position, each with and without dynamic allocation via InfoNoise. All settings are DoubleCat sampled as a one-hot flow. Better of each matched pair in each column is shown in bold.

e smooth estimate H(t). We then construct an information-rate profile from its derivative, r(t) =

e dH(t) , dt

(G.5)

and allocate training times in proportion to this rate. Thus, corruption regions over which the conditional entropy changes rapidly receive more training samples, while regions in which little information is gained or lost receive proportionally fewer. To avoid starving regions assigned very low probability, we constrain the resulting density before normalisation. Denoting the unconstrained allocation by π e(t) ∝ r(t), we use π new (t) = normalise [clip (e π (t), πmin , πmax )] ,

(G.6)

with πmin = 0.1 and πmax = 10. The lower threshold ensures that every corruption level continues to receive training signal, which is particularly important for an online allocation because a region that is never revisited cannot subsequently develop an estimated information rate that would cause it to regain probability mass. The upper threshold prevents the allocation from becoming excessively concentrated. We represent the sampling distribution on a fixed grid of 256 points. After each update, the density is numerically integrated on this grid using the trapezoidal rule to obtain a cumulative distribution function F new . Rather than replacing the current allocation immediately, we damp successive updates in CDF space, F (n) = (1 − λ)F (n−1) + λF new ,

λ = 0.5.

(G.7)

Averaging normalised cumulative distributions rather than raw loss or entropy values makes the update insensitive to changes in the overall scale of the loss during training and smooths changes in where probability mass is allocated. Using a fixed grid likewise provides a common coordinate system across successive measurements. We recompute the entropy profile and update the allocation every ndynamic = 2500 training iterations. Between updates, corruption times are sampled from the current tabulated CDF by inverse-transform sampling, u ∼ U(0, 1),

t = F (n)

−1

(u),

(G.8)

where the inverse is obtained by linear interpolation between neighbouring points. The resulting samples determine only which corruption times are presented during subsequent training iterations; all other components of the objective are unchanged.

On the Design Space of Continuous Diffusion Language Models

Analyze & explore related research in Conceptio

Related documents