LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling

In plain words

Most language models write one token after another, from left to right. A different family of models starts from random noise over a whole passage and cleans it up step by step, so every step can look at the entire passage at once. Until now, these "diffusion" models for text have worked on discrete symbols, and models that work on continuous vectors have lagged behind them.

LangFlow keeps the vectors. Each word is turned into a vector, the vectors are mixed with noise, and a network learns to guess which word sits under the noise. The noise level is measured on a logarithmic scale, and the paper picks the levels to sample from by checking where the guesses actually get better. The model is also allowed to see its own previous guess during each step, which the paper finds helps the continuous version a lot. Generation is a smooth, fixed path from noise to vectors, followed by a choice of the most likely word at each position.

The main claim is that this kind of model now reaches roughly the same quality as the best discrete diffusion models on standard benchmarks. The trade-off is that generating text takes many steps.

TL;DR

• LangFlow is a continuous diffusion language model. It denoises token embeddings, trains with token cross-entropy, and samples with a deterministic Euler ODE solver §3.1 §3.2 alg. 2.
• Three design choices carry the result: a log noise-to-signal time variable, a learned Gumbel noise schedule, and self-conditioning that is switched on for the evaluation §4.1 §4.2 §5.
• Author-reported: LM1B PPL 30.0 (best of the DLMs in their table) and OpenWebText PPL 24.6 (third of the DLMs). Gen. PPL is 92.2 on LM1B and 36.5 on OpenWebText tab. 2.
• Use it when you want a continuous latent space that is editable and compatible with flow-based few-step methods. Expect many sampling steps, and note the training code is not yet released (see Paper vs code).

Problem

Autoregressive models generate one token at a time and cannot use the right context while writing the left §1. Discrete diffusion models for text work on symbols and corrupt them with masks or random replacements; they perform well but give up a continuous latent space, which limits controllable and few-step generation §1. Continuous diffusion over the probability simplex is sparse and hard to score §1. Embedding-space diffusion avoids both problems, but before this paper it lacked (a) a training objective that is simple to implement, (b) an ODE-based likelihood estimate instead of stochastic bounds, and (c) tested design choices §1 §2. The paper gives a training objective with a Bregman-divergence view §3.1, an ODE likelihood bound §3.2 [thm. 3.1], and two design choices §4 §5.

The paper's own diagnosis is that a uniform noise schedule over the usual time variable spends more than half its steps at noise levels where the loss is almost zero §4.1. It also reports that self-conditioning helps continuous diffusion much more than discrete diffusion, which earlier evaluations had hidden because they turned it off §4.2.

Method

Overview. Tokens are mapped to vectors with a learned embedding table E (V x D). The noisy vector z_gamma is produced on a variance-preserving path indexed by gamma, the log noise-to-signal ratio §3.1. A network reads z_gamma, gamma and an optional self-conditioning input, and outputs a probability for each vocabulary item at each position §3.1. Training uses only the token cross-entropy. Sampling integrates an ODE from pure noise at high gamma to low gamma, and each step predicts a clean embedding by taking the probability-weighted average of the embedding table §3.1 alg. 2. The final tokens are the argmax of the network output at the lowest gamma alg. 2.

1. Embed tokens. z = E[x]. Each row of E is rescaled to unit norm and multiplied by sqrt(D), so that its variance matches the noise. This is described in the architecture notes app. "Model architecture". The code does the same at langflow/model.py model.py.

2. Noise on the gamma path. With the variance-preserving constraint, the coefficients are alpha_gamma^2 = sigmoid(-gamma) and sigma_gamma^2 = sigmoid(gamma), and z_gamma = alpha_gamma z + sigma_gamma eps with eps standard normal §3.1. A time variable t can be mapped to gamma; gamma is the one the network is conditioned on §3.1 §4.1.

3. Denoiser. The network outputs token probabilities x_hat (L x V) rather than embeddings. The continuous denoiser used by the sampler is z_hat = E^T x_hat §3.1. This is the expectation of the embedding over the predicted distribution.

4. Training loss. For each gamma, the loss is the mean token cross-entropy over positions. Choosing f(p) = p log p in the Bregman divergence turns the flow-matching objective into this cross-entropy §3.1. gamma is drawn from the learned Gumbel proposal described in step 5 §4.1 alg. 1.

5. Learned Gumbel noise schedule. The entropy profile H_gamma is modelled as H_inf exp(-exp(-(gamma - P_mu) / P_beta)), a Gumbel shape. The three parameters H_inf, P_mu and P_beta are learned. They are fitted by a separate loss (sg(CE) - H_gamma)^2, where sg denotes stop-gradient, so the main model is not affected §4.1 alg. 1. Samples of gamma come from the inverse CDF of the Gumbel distribution, and gamma is clipped to the 1e-5 quantiles a and b of that distribution §4.1.

6. Self-conditioning. With probability 0.25 during training, a first pass is run with zero self-conditioning input. Its embedding prediction, with gradients stopped, is fed back as a second input §5 alg. 1. During sampling it is always on: it starts at zero and then uses the previous step's embedding prediction §4.2 alg. 2.

7. Backbone. A DiT-style transformer with rotary positions, adaptive layer norm conditioned on gamma, 12 layers, hidden size 768 and 12 heads, about 130M parameters §5 app. "Model architecture". The self-conditioning input enters as an extra linear projection of the concatenation of z_gamma and the previous prediction, with zero initialisation app. "Model architecture" [code: langflow/model.py#L329-L332] model.py.

8. Bias skip connection. The logits get an added term alpha/sigma^2 times the dot product of z_gamma with the embedding table. The paper ramps its weight from 0 to 1 over the first 5000 training steps app. "Model architecture". At inference the weight is 1 model.py.

9. Sampling. Integrate the ODE with an Euler method on a grid of gamma values set by the quantiles of the Gumbel proposal alg. 2 app. "Numerical Solver". The update is written in terms of z / sigma; see Algorithm 2 for the exact expression.

10. Likelihood bound. An upper bound on the negative log-likelihood is obtained by integrating along the probability-flow ODE, with a reconstruction term at the clean end, an integral of the divergence of the denoiser, and a Gaussian prior term at the noisy end [thm. 3.1]. For evaluation the ODE is solved with a 128-step Heun-2 method, and the divergence is estimated with a Hutchinson trace estimator app. "Additional Evaluation Details".

Training pseudocode, following Algorithm 1 alg. 1. Variable names are ours:

# One training step  [alg. 1]
# Input: tokens (B, L); embedding table (V, D); p_self_cond = 0.25  [§5]
# Output: total loss (scalar)

def training_step(tokens, embedding, p_self_cond):
    # Clean embeddings of the tokens  [§3.1] [app. "Model architecture"]
    clean = normalized_embedding(embedding)[tokens]  # (B, L, D)

    # Draw a noise level from the learned Gumbel proposal  [§4.1] [alg. 1]
    # (the paper uses a low-discrepancy sampler for the uniform draw)
    quantile = clip(uniform(0, 1), 1e-5, 1 - 1e-5)  # (B,)
    gamma = stop_gradient(
        P_mu - P_beta * log(-log(quantile)))  # (B,)
    alpha = sqrt(sigmoid(-gamma))  # (B,)
    sigma = sqrt(sigmoid(gamma))  # (B,)

    # Add noise on the gamma path  [§3.1] [alg. 1]
    noisy = alpha[:, None, None] * clean + (
        sigma[:, None, None] * normal(0, 1, clean.shape))  # (B, L, D)

    # Self-conditioning input, zero unless drawn  [§5] [alg. 1]
    self_cond = zeros_like(noisy)
    if bernoulli(p_self_cond):
        # First pass without self-conditioning; no gradient through it
        # [alg. 1]
        probs = model(noisy, gamma, self_cond)  # (B, L, V)
        self_cond = stop_gradient(
            embedding_expectation(probs, embedding))  # (B, L, D)

    # Main pass: predicted token probabilities  [§3.1] [alg. 1]
    probs = model(noisy, gamma, self_cond)  # (B, L, V)

    # Token cross-entropy averaged over positions  [§3.1]
    ce_loss = mean(-log(probs[batch, position, tokens]))

    # Fit the noise schedule to the observed loss  [§4.1] [alg. 1]
    entropy = H_inf * exp(-exp(-(gamma - P_mu) / P_beta))  # (B,)
    schedule_loss = mean((stop_gradient(ce_loss) - entropy) ** 2)

    # One optimizer step on both terms  [alg. 1]
    return ce_loss + schedule_loss

Sampling pseudocode, following Algorithm 2 alg. 2:

# Euler sampler on the gamma grid  [alg. 2]
# Input: model, embedding (V, D), N steps
# Output: token ids (B, L)

def sample(model, embedding, num_steps, batch, length):
    # Gamma grid from the quantiles of the Gumbel proposal  [alg. 2]
    quantiles = [clip(1 - k / num_steps, 1e-5, 1 - 1e-5)
                 for k in range(num_steps + 1)]
    gammas = [P_mu - P_beta * log(-log(q)) for q in quantiles]

    # Start from noise at the highest gamma  [alg. 2]
    # (sigma(g) = sqrt(sigmoid(g)), alpha(g) = sqrt(sigmoid(-g)))
    z = normal(0, sigma(gammas[0]), (batch, length, D))
    z_hat = zeros_like(z)

    for k in range(num_steps):
        # Predict probabilities with the previous step's embedding
        # prediction as self-conditioning  [alg. 2] [§4.2]
        probs = model(z, gammas[k], z_hat)
        z_hat = embedding_expectation(probs, embedding)

        # Euler update in z / sigma  [alg. 2] [app. "Numerical Solver"]
        z = sigma(gammas[k + 1]) * (
            z / sigma(gammas[k])
            + (alpha(gammas[k + 1]) / sigma(gammas[k + 1])
               - alpha(gammas[k]) / sigma(gammas[k])) * z_hat)

    # Final tokens are the argmax at the last gamma  [alg. 2]
    probs = model(z, gammas[num_steps], z_hat)
    return argmax(probs, axis=-1)

Key equations, in the paper's notation:

\alpha_\gamma^2 = \operatorname{sigmoid}(-\gamma), \qquad \sigma_\gamma^2 = \operatorname{sigmoid}(\gamma), \qquad \gamma_t = \log\left(\sigma_t^2 / \alpha_t^2\right) §3.1

H_\gamma = H_{+\infty} \exp\left(-\exp\left(-\frac{\gamma - P_\mu}{P_\beta}\right)\right) §4.1

\log p(\mathbf{x}) \ge \mathbb{E}_{\mathbf{z}} \Bigg[ \frac{LD}{2} - \frac{\|\mathbf{z}_b\|^2}{2\sigma_b^2} + \sum_{i=1}^{L} \log \hat{x}_\theta^{(i, x^{(i)})}(\mathbf{z}_a, a) - \int_a^b \frac{\alpha_\gamma}{2} \nabla \cdot \hat{\mathbf{z}}_\theta(\mathbf{z}_\gamma, \gamma)\, \mathrm{d}\gamma \Bigg] [thm. 3.1]

The bound is written with z_a drawn from N(alpha_a E^T x, sigma_a^2 I), the reverse ODE trajectory z_gamma for a <= gamma <= b, and the divergence taken with respect to z_gamma [thm. 3.1].

Hyperparameters

Values from the paper's text and appendix. "Not stated" means the paper does not give the value in its active text. Repository defaults are listed where they exist.

Results

Author-reported. Setup: LM1B and OpenWebText, 130M DiT-style models trained for 1M steps. Baselines were retrained on LM1B and the OWT ones use the Duo checkpoints §5. Gen. PPL is measured with GPT-2 Large over 1024 samples §5. PPL for DLMs is an upper bound §5.

Language modelling, from tab. 2:

Zero-shot perplexity, models trained on OWT for 1M steps, DLM values are upper bounds tab. 3:

The paper's claim that LangFlow beats the autoregressive Transformer on 4 of 7 benchmarks checks out against this table: PTB, Lambada, PubMed and Arxiv tab. 3.

Self-conditioning ablation on LM1B, from tab. 1:

The LangFlow LM1B Gen. PPL in tab. 1 (81.5) does not match the LangFlow LM1B Gen. PPL in tab. 2 (92.2). The paper's NFE table gives 92.24 at 128 steps, which agrees with tab. 2. See Pitfalls.

Sampling-budget sweep for LangFlow on LM1B (no distillation), from the appendix table "NFE comparison on LM1B":

Data entropy is 4.32 on LM1B (same table). For OWT, the appendix table "NFE comparison on OWT" gives LangFlow at 1024 steps Gen. PPL 36.53 and entropy 5.25, against a data entropy of 5.44. Duo gives 77.69 and 5.55 at the same step count, and MDLM gives 104.85 and 5.63. At 64 steps LangFlow gives Gen. PPL 80.34.

Ablations reported in the text, as cited: Gumbel scheduling lowers LangFlow's Gen. PPL from about 1000 to 154.2 §4.1. Self-conditioning is what closes the PPL gap to discrete models §4.2 tab. 1.

Paper vs code

Repository: https://github.com/nealchen2003/LangFlow at commit a712b08570c56c787d6ef24f8e906a2dcdf768f5. It is linked from the abstract and its README lists the paper's authors. The repository has no LICENSE file, so its code is not quoted here; only paths and lines are cited.

• Training code is not in the repository. The README lists "Training code" and "All trainable checkpoints" as "after paper acceptance" README.md. Only sampling, likelihood evaluation and generative-perplexity scripts exist. Training losses, the scheduler update and the self-conditioning training pass cannot be checked against code, and must be taken from the paper's Algorithm 1 alg. 1.

• Gumbel values. The paper's active text gives no numbers for P_mu, P_beta or H_inf; they are learned §4.1. The repository's defaults are 4.723, 0.852 and 7.02 config.py config.json. Those are initial values for learnable parameters in the code, and the repo does not say whether they are the trained values in the released checkpoint.

• Sampler update. The code's Euler step is written as an EDM-style interpolation model.py. Algebra by hand, at gamma_t = 0 and gamma_s = 1 and generally: the code's coefficient on the current state is w * alpha_s / alpha_t, where w = sigma_s / sigma_t. Algorithm 2 alg. 2 gives w. These differ unless alpha_s = alpha_t. I did not run the code to confirm, so treat this as a discrepancy to check, not a verified bug.

• Number of network calls. The code builds num_steps gamma values, makes num_steps - 1 Euler steps and one final call, so it makes num_steps network calls model.py. Algorithm 2 indexes gamma from k = 0 to N and calls the network in a loop of N steps plus one final call alg. 2, which is N + 1 calls. Quote "128 steps" as the paper does, but count network calls from the code if you need a cost.

• Sample entropy. The paper says entropy counts each token that reaches at least 1/128 (LM1B) or 1/1024 (OWT) of the sequence [app. "Additional Evaluation Details"]. The script counts all unique tokens with no threshold gen_ppl.py.

• ODE solver default. The evaluation config sets Heun-2 with 128 steps config.yaml, which matches the paper. The NLL script falls back to Euler if the config lacks the key eval_ppl.py. Make sure the config is loaded.

• Trace estimator. The divergence probe is a Gaussian draw, one per ODE evaluation eval_ppl.py. Results vary with the random seed.

• Tokenizer. The paper uses gpt2-large for OWT §5. The inference script loads the gpt2 tokenizer inference.py. Both are GPT-2 byte-pair encodings, but I did not check their vocabularies against each other.

• Consistent with the paper. The embedding normalisation [code: langflow/model.py#L225-L228], the bias skip weight at inference [code: langflow/model.py#L472-L477], the zero-initialised self-conditioning projection model.py, and the gamma clipping at 1e-5 [code: langflow/model.py#L289-L294] all agree with the paper's text and appendix. The bias skip coefficient alpha/sigma^2 equals the code's exp((softplus(-gamma) - gamma)/2) when both are squared; I checked this by hand.

Pitfalls

• Two LangFlow LM1B Gen. PPL values in the paper. tab. 1 gives 81.5 and tab. 2 gives 92.2 for the same model and dataset. The appendix NFE table ("NFE comparison on LM1B") supports 92.2. Use 92.2 unless you can confirm the other value.

• Plaid is in the main table. tab. 2 lists Plaid with the best LM1B Gen. PPL (77.3), which the text does not discuss. The paper's text says LangFlow ranks second on LM1B Gen. PPL §5. Plaid's appendix says it uses a different architecture and a 16-dimensional embedding app. "Plaid Baseline".

• Likelihood is an ODE bound with a random probe. PPL in this paper is an upper bound from [thm. 3.1], computed with 128-step Heun-2 and a Hutchinson trace estimate app. "Additional Evaluation Details". Repeated runs differ slightly. The paper reports the upper bound for each DLM baseline §5, so compare with care across methods with different bounds.

• Self-conditioning changes the metric. Earlier DLM work reported PPL without self-conditioning §4.2. Compare LangFlow's PPL only with baselines evaluated under the same protocol. The paper's tab. 1 shows MDLM's PPL gets worse with self-conditioning, while LangFlow's gets better tab. 1.

• Gumbel schedule is learned and clipped. Clip gamma to the 1e-5 quantiles §4.1. From the repository defaults (not the paper's trained values) those quantiles are about 2.64 and 14.53 (my arithmetic from [code: langflow/config.py#L38-L41]). At sampling time the grid is the N-quantile grid alg. 2.

• Stop-gradients keep the schedule separate. The schedule loss uses the stopped cross-entropy, so the scheduler does not change the main network §4.1 alg. 1. If you implement the two losses in one graph without the stop-gradient, the main model will be pushed toward the schedule fit.

• Self-conditioning in training needs a second forward pass. With probability 0.25 a training step first runs the network once without gradient, then runs it again with that prediction as input alg. 1. The extra pass is therefore taken on about a quarter of the steps.

• Euler sampling is many calls. Sampling runs for about as many network calls as the step count: 128 on LM1B and 1024 on OWT in the paper's Gen. PPL setup §5. At 16 steps on LM1B, Gen. PPL rises to 179.6 (appendix NFE table on LM1B).

• Low entropy is not the same as repetition. LangFlow's sample entropy is lower than the discrete baselines' on OWT. The paper attributes this to a few frequent content words rather than local repetition, and it leaves the entropy gap for future work app. "Quantitative Sample Quality" [conclusion, Limitations].

• Training budget is large. About 292 GPU hours per GPU for LangFlow, against about 375 for Plaid (appendix table "Model and training setup"). The OWT run uses 32 A100s app. "Training Details".

• Plaid's MSE objective collapses embeddings. The paper argues that regressing the clean embedding with mean squared error collapses the embedding table, with nearest-neighbour distances near zero for Plaid [app. "Embedding Collapse under Mean Square Error"]. The paper uses cross-entropy for this reason §3.1.

Lineage

• Builds on: Diffusion-LM for embedding-space diffusion over word vectors §1 §2; flow matching §2; the discrete diffusion training recipe of MDLM and Duo for architecture and data §5; Variational Diffusion Models for the log noise-to-signal time variable §3.1.
• Compared against: AR Transformer, D3PM, DiffusionBert, SEDD (absorbing and uniform), MDLM, UDLM, Duo, Diffusion-LM, Plaid and FLM §5 tab. 2 tab. 3.
• Superseded by: not known from content in this wiki.