PaperlayerView exact revision on arXiv

Source-linked AI summary

On Trajectory-Aware Training for Masked Diffusion Language Models

Manuel Madeira, Amitis Shidani, Alice Bizeul, Victor Turrisi, Louis Béthune, Bhavika Devnani, Dan Busbridge, Pierre Ablin, João Monteiro

arXiv:2609.37974v1cs.LGcs.AIcs.CL

TL;DR

MDMs train on random masks, while inference follows policy-shaped trajectories and carries no information from earlier steps. PUMBA trains across consecutive policy-induced steps, passes information between them, and jointly optimizes their losses; it improves the performance–NFE frontier in full-canvas and block diffusion generation.

  • Problem

    MDM training uses randomly masked sequences, whereas inference follows policy-shaped trajectories, and independent training examples do not teach the model to use information from prior steps.

    Source: §1 · p. 1

  • Method

    PUMBA trains on consecutive policy-induced trajectory steps, passes information between steps, and jointly optimizes their losses with backpropagation through time.

    Source: §1 · p. 1, §3.2 · p. 5

  • Results

    PUMBA improves the performance–NFE frontier in both LLaDA-8B formulations, using up to 22% fewer NFEs than an SFT run twice as long in full-canvas generation and up to 26% fewer than an equally long MDM control in block diffusion.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The Post-SFT recipe pushes a model beyond what additional SFT allows, and its reduced denoising steps amortize the short extra training stage.

    Source: §6 · p. 10

    Sustaining closer inference alignment across model scales and decoding policies beyond the confidence-based ones studied remains open.

    Source: §6 · p. 10

Abstract

from arXiv · show

Masked diffusion models (MDMs) generate text by unmasking several tokens per step, but they are trained and sampled under different conditions. The model is trained on randomly masked sequences, whereas inference follows a trajectory shaped by the model's own predictions. Additionally, each step has no access to what the previous one computed. Recent methods narrow these limitations from separate angles, leaving open how these choices interact. We introduce PUMBA, a unified framework for trajectory-aware training that trains the denoiser on consecutive steps of policy-induced trajectories, passes information between steps, and optimizes them jointly by backpropagation through time. A controlled study of this design space shows that i) exact train--inference alignment fails due to local overfitting, whereas a looser alignment still brings training masks closer to those seen at inference; ii) passing continuous information outperforms discrete gradient estimators through the commitment at each step; and iii) performance improves as backpropagation through time spans more steps, which we support theoretically. Combined, these components match the best checkpoint of a same-size autoregressive model. Building on these findings, we scale PUMBA to supervised fine-tuning of LLaDA-8B, where it improves the trade-off between performance and number of function evaluations (NFEs) in both full-canvas and block diffusion generation. At matched performance, it needs up to 22% fewer NFEs than standard fine-tuning with twice the budget in full-canvas generation, and up to 26% fewer than standard fine-tuning for the same number of steps in block diffusion.

1 Introduction

MDMs train on independently masked sequences, while inference follows policy-shaped trajectories and does not pass information between steps. PUMBA unifies trajectory-based training, cross-step information passing, and joint optimization across consecutive steps.

1 Introduction

  • MDM training uses random masks, unlike inference trajectories shaped by model predictions, and independently drawn training sequences do not teach the model to use information from prior steps.Recent methods address these gaps separately through cross-step continuous information, policy-induced trajectories, or joint training across consecutive steps.Source: §1 · p. 1
  • PUMBA trains the denoiser on consecutive policy-induced trajectory steps, passing information between steps and jointly optimizing their losses with BPTT.The framework spans the trajectory-building policy, the information carried between steps, and the BPTT window; PUMA is one PU instantiation that randomizes and anneals tokens revealed per step.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Across decoding budgets, the three axes compound generative gains, matching a same-size autoregressive model trained on the same data while decoding multiple tokens per step.Performance improves as the BPTT window grows, a trend the authors support theoretically.Source: §1 · p. 2
  • Up to 22% fewer NFEs than an SFT run twice as long in full-canvas generation, and up to 26% fewer than an equally long MDM control in block diffusion.The LLaDA-8B results improve the performance–NFE frontier in both formulations.Source: §1 · p. 2

2 Background

The background describes masked diffusion training as denoising independently masked sequences, then explains iterative inference policies that choose which masked positions to reveal. Confidence-based policies and block decoding trade denoiser calls against generation quality.

2 Background

  • The notation treats a sequence as tokens from a vocabulary containing a mask token and defines masked positions as those containing that token.The denoiser’s masking ratio provides its inferred noise level, so it needs no separate time input.Source: §2 · pp. 2–3
  • Each MDM training example independently masks clean-sequence tokens with probability t, then trains a denoiser to predict clean tokens at masked positions using weighted cross-entropy.The loss includes a 1/t factor and training draws t uniformly from [0, 1].Source: §2 · pp. 2–3
  • At inference, generation starts fully masked and iteratively reveals selected positions, filling them with tokens sampled from the denoiser’s predicted distributions.Revealing multiple positions at once is faster but samples from a product of marginals rather than their joint distribution.Source: §2 · p. 3
  • The top-u policy reveals the u most-confident masked positions, while Fast-dLLM reveals positions above a confidence threshold or falls back to the single most confident position.The paper uses top-u in controlled experiments and Fast-dLLM for LLaDA, paired with left-to-right block diffusion to reduce denoiser calls while preserving quality.Source: §2 · p. 3

3 Elucidating the Design Space of PUMBA

The design-space study tests policy-induced training trajectories, inter-step information transfer, and multi-step optimization in a controlled TinyGSM-to-GSM8K setting. It finds that looser alignment avoids local overfitting, continuous carries outperform discrete bridges, and longer BPTT windows improve accuracy.

3 Elucidating the Design Space of PUMBA

  • The study maps PUMBA’s design space by training a 125M-parameter bidirectional transformer on TinyGSM and evaluating zero-shot on GSM8K.The variants share architecture, data, batch size, and learning rate.Source: §3 · p. 3

3.1 Training–Inference Alignment via Progressive Unmasking

  • At small inference-like u, progressive unmasking locally overfits; training with larger u nevertheless reduces mask discrepancy versus MDM at u=2.At u=2, sequence completion loss falls after batch entry and rises again before exit; the PUMA recipe uses larger effective training u to retain some alignment without that local overfitting.Source: §3.1 · pp. 4–5, Figure 2 · p. 4

3.2 Propagating Information Across Denoising Steps

  • A continuous carry outperforms discrete gradient bridges, while BPTT trains the carry for subsequent predictions and its theoretical gap bound shrinks as the window grows.Each update backpropagates across W consecutive trajectory steps; the bound’s gap vanishes when the window covers the whole trajectory under the stated assumptions.Source: §3.2 · pp. 5–6

3.3 Composing the Design Choices

  • The three design choices each add accuracy, and together reach the best same-size AR checkpoint while revealing more than one token per step.Accuracy keeps increasing as the BPTT window grows; at matched denoiser passes, longer MDM training remains substantially below BPTT.Source: §3.3 · pp. 6–7, Table 2 · p. 6

4 A PUMBA Recipe for Large-Scale SFT

PUMBA applies progressive unmasking and hidden-state carry with BPTT to LLaDA-8B SFT, improving the performance–NFE trade-off in both full-canvas and block diffusion. The reveal count must balance alignment with inference and the diversity needed for learning.

4 A PUMBA Recipe for Large-Scale SFT

  • PUMBA trains on inference-policy trajectories at a larger u than inference, passes a hidden-state carry between steps, and trains it with BPTT over window W.The two-stage setup applies Post-SFT after standard SFT to distinguish PUMBA’s effect from learning the data distribution during SFT.Source: §4 · p. 7

4.1 Full-Canvas Generation

  • In full-canvas generation, Post-SFT MDM stays near its starting checkpoint, while 3-epoch SFT improves only slightly and only at higher NFEs.The authors interpret this as evidence that MDM training has nearly saturated on this data.Source: §4.1 · p. 7

4.2 Block Diffusion

  • At matched score, PUMBA needs 26% fewer NFEs than the Post-SFT MDM control in block diffusion.Its best configuration, u=16, W=2, gains 2.2 points in 4000 steps, more than the 50k SFT steps from 3 to 6 epochs.Source: Figure 4 · p. 8, §4.2 · p. 8

4.1 Full-Canvas Generation

  • At matched score, PUMBA needs up to 22% fewer NFEs than the 3-epoch SFT in full-canvas generation.The best configuration, u=32, W=8, gains 2.6 points; doubling the SFT budget yields less than one-third that gain.Source: §4.1 · p. 7, Figure 4 · p. 8

4.2 Block Diffusion

  • In block diffusion, PU and carry with W=1 bring modest gains, while wider BPTT windows yield clear improvement.Because blocks contain 32 positions, u also bounds the available window; for example, u=16 implies W ≤2.Source: §4.2 · p. 8

4.3 Failure Modes at Small Training u

  • Smaller training u reduces sample diversity and can cause local overfitting and training collapse; the best u balances alignment with inference against learning from enough distinct samples.For full-canvas u=1, a run visits 0.6% of the Post-SFT MDM control’s samples, versus 26% at u=128; block diffusion reaches 100% by u=32.Source: §4.3 · p. 8, Figure 5 · p. 9

5 Related Work

Related work addresses masked diffusion generation, train–inference alignment, and information transfer across denoising steps. PUMBA combines trajectory-aligned supervised post-training with continuous state propagation and multi-step optimization.

Masked diffusion language models

  • Masked diffusion scales to large language models through autoregressive conversion or training from scratch, and block diffusion combines autoregressive generation across blocks with parallel denoising within them [Arriola et al., 2025].The paper studies post-training for the states induced by these decoders.Source: §5 · p. 9

Structured inference

  • Inference policies choose reveal order and parallelism, inducing trajectories whose intermediate states differ from conventional random corruptions.Policies include confidence, probability-margin, and entropy criteria, along with planners and learned policies.Source: §5 · p. 9

Train–inference alignment

  • Prior alignment methods include progressive teacher-forced unmasking (PUMA) [Kim et al., 2026], planner-loss reweighting, reward optimization, on-policy distillation, and step distillation.This work investigates alignment within supervised instruction fine-tuning, without task-specific reward optimization or distillation.Source: §5 · p. 9

State propagation across denoising steps

  • Prior state-propagation approaches include self-conditioning, probability-weighted embeddings, persistent memory, and Relay’s differentiable representations with two-step truncated BPTT.PUMBA combines trajectory-aligned supervised post-training with continuous state propagation and studies component contributions and the alignment–learnability trade-off.Source: §5 · p. 9

6 Conclusion

The paper concludes that trajectory-aware training improves masked diffusion training by matching generation trajectories, propagating hidden states, and optimizing across steps. Its findings emphasize balancing trajectory alignment with sample diversity.

B.1 PUMBA training procedure

  • PUMBA trains on consecutive policy-induced trajectory steps, passes a hidden-state carry between steps, and jointly optimizes step losses with BPTT.Its training procedure advances a trajectory for up to W steps before updating parameters and stopping gradients across the window boundary.Source: §6 · p. 10, §B.1 · p. 16

6 Conclusion

  • PUMBA improves the performance–NFE frontier for LLaDA-8B under both full-canvas and block diffusion, and the short Post-SFT stage can be amortized by reduced denoising steps.The paper reports that a small number of training steps suffices to realize PUMBA’s gains.Source: §6 · p. 10
  • The paper identifies local overfitting at small reveal counts and leaves open how to sustain closer inference alignment across model scales and decoding policies beyond confidence-based ones.It also notes that richer carries and BPTT-specific methods could improve the trade-off among training cost, memory, and credit-assignment horizon.Source: §6 · p. 10

A Extended background

  • Masked diffusion training uses random masked-sequence corruptions, while generation iteratively unmasks several positions according to a policy that acts on model predictions.The inference policy determines both the reveal order and the degree of parallelism.Source: A Extended background · p. 15

B.2 Progressive unmasking: the two constructions

  • The paper uses stage-indexed progressive unmasking for TinyGSM but a confidence-threshold construction for LLaDA-8B, whose variable response lengths make a single stage schedule a poor match to decoding.The LLaDA-8B rule starts fully masked and commits the u most confident positions plus other positions above τ=0.9.Source: §B.2 · pp. 16–17

B.3 TinyGSM; B.4 LLaDA-8B

  • The controlled TinyGSM setting trains on TinyGSM and evaluates problem-solving accuracy on 1,319 GSM8K test problems.The LLaDA-8B setting uses Dolci-Instruct-SFT and a 4096-token canvas; prompt truncation removes 16.2% of prompt tokens, concentrated in the 2.5% of samples that do not fit.Source: §B.3 · p. 17, §B.4 · p. 18

B.4.1 SFT stage

  • Full-canvas SFT attends bidirectionally over the canvas, while block diffusion uses 32-token blocks and a vectorized clean/noised response layout.In full-canvas SFT, training on all padding tokens caused collapse to EOS or very short responses, so only the first EOS after the response is trained on.Source: §B.4.1 · pp. 19–20

B.4.2 Post-SFT stage

  • Post-SFT uses 4000 updates, and its carry is the last normalized hidden state per position, added to token embeddings through a newly initialized LayerNorm.The carry begins at zero and is propagated across trajectory steps; under block diffusion it is reset at each block during evaluation.Source: §B.4.2 · p. 20

B.4.3 Evaluation

Evaluation assesses instruction following, mathematics, and code across decoding thresholds, with compute also reported. The reported score averages three benchmark measures.

B.4.3 Evaluation

  • The evaluation averages IFEval prompt-level strict accuracy, GSM8K exact match, and MBPP pass@1.The benchmarks assess verifiable instruction following, mathematical reasoning, and program synthesis, respectively.Source: §B.4.3 · p. 20
  • Models are decoded with Fast-dLLM in left-to-right blocks of 32 tokens at T=0, sweeping τ∈{0.7, 0.8, 0.9, 1}.The generation budgets are 1024 tokens for GSM8K and MBPP and 1280 for IFEval; τ=1 recovers one-token-per-step decoding.Source: §B.4.3 · p. 21

B.5 Compute

  • The reported runs used about 85,000 H100 and 7,300 B200 GPU-hours, including about 24,000 H100 GPU-hours for evaluation.TinyGSM accounts for about 36,000 H100 GPU-hours and LLaDA-8B for about 49,000 H100 and 7,300 B200 GPU-hours.Source: §B.5 · p. 21

C.1 Accuracy and pass@k

These appendices compare diffusion variants on GSM8K and detail experiments on student forcing, mask discrepancy, and local overfitting. They report accuracy and pass@k outcomes, training setups, and methods for measuring training–inference mask differences.

C.1 Accuracy and pass@k

  • Carry with W=8 matches AR’s best GSM8K performance, while PU, carry, and BPTT improve pass@k over MDM; AR leads at higher temperatures, but higher-W BPTT can outperform AR at lower temperatures.Figure 7 reports accuracy across decoding budgets, and Figure 9 reports pass@k from 100 samples per problem at u=2 and u=3.Source: §C.1 · p. 22, Figure 7 · p. 22, Figure 9 · p. 24

C.2 Carry and gradient estimators

  • Carry trained with BPTT at W=2 and W=4 separates from the carry without BPTT and the two-step objective without a carry; no tested discrete gradient bridge exceeds the latter.The ordering holds across evaluated decoding budgets, with gaps narrowing at u=16 as all variants degrade.Source: §C.2 · p. 25, Figure 10 · p. 25, Figure 11 · p. 26

C.3 From teacher forcing to student forcing

  • No scheduled-sampling schedule improves on the PU baseline of 39.0%; accuracy reaches that baseline only as the average ground-truth-token ratio approaches 1.Below an average ratio of about 0.9, every run falls below the MDM baseline of 34.3%; at 0.24, accuracy is 13.7%. Detaching predicted tokens changes accuracy by less than one point.Source: §C.3 · p. 26, Figure 12 · p. 27

C.4 Estimating the mask discrepancy

  • Mask discrepancy is estimated by comparing similarities among training masks and between training masks and the deterministic inference mask, averaged over 50 validation problems.The re-binned MDM control differs from MDM by at most 0.0015 at u=2, so the reported PU–MDM gap is not attributed to binning.Source: §C.4 · p. 28

D Local overfitting at small training u

  • The TinyGSM study varies u while holding each trajectory in the batch for a fixed number of steps, then follows held-out sequences as they enter and leave training.Runs last 3,000 updates; K is selected to approximate L/u, with L about 460 maskable positions.Source: D Local overfitting at small training u · p. 29, §D.1 · p. 29

D.1 Setup

  • Inserted validation sequences are tracked through consecutive PU updates, with completion loss measured on the answer’s second half and compared against validation loss.Twelve sequences are inserted at three update times; four are never inserted as controls.Source: §D.1 · p. 29

D.2 Results

The results show how reveal count and sequence reuse affect completion and validation loss. Ablations indicate that changing masks does not remove the failure, while spreading visits makes the overfitting extend across the sequence pool.

D.2 Results

  • With MDM masks, inserted sequences are fitted even more deeply; with spread visits, the model memorizes the whole pool and validation loss exceeds PU at every tested u.The MDM-mask ablation has relative completion loss 0.02 at u=2 and remains 0.52 300 updates after the sequence leaves; the spread-visits run has training loss 0.03 at u=2.Source: §D.2 · pp. 29–30, Figure 14 · p. 31
  • PU beats MDM when sequences stay in the batch only a few steps, but trails it for u ≤8; at u=2, training loss falls to 0.07 while validation completion loss reaches 6.92.Figure 13 also shows the trajectory fitting its answer within a quarter of the sequence’s stay at u=2.Source: §D.2 · p. 29, Figure 13 · p. 30
  • Starting all trajectories in stages 0–31 raises validation loss at u≤4, while removing the 1/t weighting or halving the batch deepens the training-loss drop at u=2.Removing 1/t changes validation loss by at most 0.16; halving the batch raises validation loss at every u.Source: §D.2 · p. 30, Table 10 · p. 31

D.3 Why the loss rises again before the sequence leaves the batch

The sequence’s trajectory can fit its answer before its batch stay ends, after which later updates erase some of that improvement. Position-level tracking links the rise in loss to tokens becoming revealed under confidence ordering.

D.3 Why the loss rises again before the sequence leaves the batch

  • At u=2, trajectory loss falls below 0.1 nats 55 updates after insertion, before 95% of the answer is revealed; subsequent updates on other sequences erase what it taught the model.At u=4, fitting occurs 112 updates into a 115-update stay; at u=8, the trajectory has not fit the answer when it leaves.Source: §D.3 · p. 31
  • Under confidence order, revealed tokens’ loss rises by 1.10 ± 0.07 nats in the 50 updates after reveal, whereas random order shows no reveal-linked change.The 1,179 tokens revealed between stages 60 and 170 improve by 0.62 ± 0.11 nats while masked; under random order, masked and revealed tokens degrade at about the same rate after the trajectory fits its answer.Source: §D.3 · pp. 31–32, Figure 15 · p. 32

E Theory for the carry and the BPTT window

This section examines how carry and BPTT window length relate to likelihood and gradient properties in theory, and to performance and NFE trade-offs in diffusion SFT experiments. The theoretical guarantees concern an idealized aligned setting, while the block-diffusion results compare configurations with and without BPTT.

E Theory for the carry and the BPTT window

  • The theory derives a likelihood identity and bounds truncated-gradient bias under a contractive carry; under its PL conditions, the excess-KL bound shrinks with window length and vanishes at W=J.The likelihood identity connects the teacher-forced loss on revealed positions to the sampler’s likelihood; the truncation-bias result assumes regular parameters and a contractive carry.Source: E Theory for the carry and the BPTT… · pp. 34–35
  • The theoretical guarantee is idealized: it assumes training reveals positions as inference does, bounds excess KL rather than trained-model accuracy, and may not cover the smaller-u sampler evaluated.Training used a much larger effective u than inference on TinyGSM (about 133 to 67 versus 2), and Algorithm 1 updates after each window rather than at the fixed point analyzed by the theorem.Source: E Theory for the carry and the BPTT… · pp. 35–36
  • Without a carry, committed tokens provide no cross-step gradient credit; with a carry, a later loss can reach earlier steps only within its BPTT window.The theorem concerns exact gradients: surrogate estimators may help but do not match the carry, and teacher-forced carry gradients train steps to pass useful information forward rather than correct their own committed errors.Source: E Theory for the carry and the BPTT… · p. 36

F.1 Full-canvas SFT schedule

  • At 1.5 epochs, the best full-canvas SFT peak learning rate was 5 × 10^-6, scoring 60.9 versus 60.5 at 10^-5.Doubling the SFT schedule to 3 epochs raised the score at τ=1 from 60.9 to 61.6. At τ=1, IFEval rose from 23.3 to 68.0, GSM8K from 60.3 to 75.9, and MBPP from 36.6 to 38.8.Source: §F.1 · p. 37, Figure 17 · p. 38

F.2 Full-canvas BPTT

  • Full-canvas BPTT with carry beat the 61.2 Post-SFT MDM control in every tested configuration; the best score occurred at u=32, and the best window improved on PU at every u.The best-window gain over PU grew from +0.1 at u=16 to +1.4 at u=128, mostly because PU weakened as u increased.Source: §F.2 · p. 38, Figure 18 · p. 39, Table 11 · p. 39

F.3 Block diffusion SFT

  • At τ=1, block-diffusion SFT rose from 57.2 to 59.8 at 3 epochs and 61.4 at 6, showing it had not saturated at 1.5; full canvas gained 0.7 points from 1.5 to 3 epochs.Source: §F.3 · p. 40

F.4 Block diffusion u sweep and BPTT

  • Best-window block-diffusion BPTT improved over PU without carry by 1.0–1.5 points at u=4, 8, and 16, while carry-only did not.For u=1, trajectories averaged 16.0 steps, versus 160.2 on the full canvas; u=1 and u=2 still trailed other settings in the block-diffusion score–NFE trade-off.Source: §F.4 · p. 42, Table 12 · p. 42

F.5 The limit case of u equal to the block size in block diffusion

  • At u = 32, a training step commits the whole 32-token block, so each trajectory is one step and BPTT cannot be applied.This limit case trains only on fully masked blocks after a clean prefix; later within-block reveals receive no Post-SFT training signal and rely on pre-training and SFT.Source: §F.5 · p. 44
  • 63.1 at τ = 1 for PU at u = 32 exceeds 62.0 for Post-SFT MDM, and its score–NFE curve lies above PU at every other u.The authors attribute the gain to the short 4,000-step Post-SFT horizon: it improves the first reveal while pre-training and SFT support block completion; they expect longer Post-SFT at u = 32 to collapse.Source: Figure 24 · p. 44, §F.5 · p. 44
  • PU at u = 16 with BPTT at W = 2 dominates PU at u = 32 in score–NFE trade-off, although PU alone at u = 16 scores 62.1 against 63.1.The best reveal count without BPTT is not necessarily best with it; the experiments do not isolate why u = 32 suits PU alone.Source: Figure 24 · p. 44, §F.5 · p. 44

F.6 Model collapse as a function of training u

Small training u increases persistent sample overlap and slows early noise-level mixing, while confidence collapse speeds sequence turnover. Decreasing u also couples trajectory granularity to the balance of supervision.

F.6 Model collapse as a function of training u

  • Confidence collapse dampens this overlap by committing several high-confidence positions, so sequences complete sooner and are replaced faster.Source: §F.6 · p. 45
  • Small training u yields substantially higher sequence overlap persisting across large step gaps, while larger u shows faster decay.The comparison is measured by intersection over union between sequences visited at steps separated by Δ.Source: §F.6 · p. 44, Figure 25 · p. 45
  • At small training u, within-batch noise-level mixing progresses more slowly, though the difference affects only the first few hundred training steps.For u=1, the progression toward steady state takes longer than for u=32 and u=128.Source: §F.6 · p. 45
  • As u decreases, the ratio of supervised to committed positions grows, increasingly weighting gradients toward predictions at positions the step does not commit.Thus u jointly controls trajectory granularity and the balance of supervision over committed and uncommitted positions.Source: §F.6 · p. 45

F.6.1 Mitigating the gap at small train u

The section tests interventions for small-u PU’s learning and diversity problems, finding partial mitigation from combined changes and reduced sequence reuse with block diffusion.

F.6.1 Mitigating the gap at small train u

  • A pool of 10B training sequences advances non-overlapping batches, increasing distinct sequences visited but raising memory requirements.The pool stores sequences and masks, and its feasible size is further constrained by carry and BPTT tensors.Source: §F.6.1 · p. 45
  • At u=1, random initial masking ratios maintain higher within-batch standard deviation than PU from the first training step.Figure 26 tracks the standard deviation of unmasking steps across training.Source: §F.6.1 · p. 46, Figure 26 · p. 46
  • The adapted objective balances losses on tokens revealed now against tokens still masked, using λ=32/L to match the u=32 run’s relative weighting.The authors offer this as a proof of concept and do not thoroughly ablate λ.Source: §F.6.1 · p. 46
  • PU+ substantially reduces but does not close the small-u performance gap, while adding memory and tuning costs.The simpler alternative is choosing a training u that balances train–inference alignment against healthy training dynamics.Source: §F.6.1 · p. 46
  • Small-u PU is impractical despite targeting the typical inference regime of about 8 tokens per step, and its adjustments are harder to combine with carry or BPTT.The passage attributes this constraint to three effects.Source: §F.6.1 · p. 46
  • Block diffusion reuses each sequence L/b times less at the same u, consistent with less exposure to lost sample diversity and no u fully collapsing.With blocks of b tokens, a sequence stays in the batch for at most ⌈b/u⌉ steps rather than about L/u.Source: §F.6.1 · pp. 46–47

F.7 Relevance of confidence thresholding at training time

Training-time confidence thresholding changes both the trajectories and compute costs of PU, with its largest score gains at small training u.

F.7 Relevance of confidence thresholding at training time

  • At u=1, the threshold raises the estimated share of Post-SFT MDM samples visited from 0.3% to 0.6%; at u=8, the estimates are 2.3% and 2.5%.The difference is within uncertainty bands from u=8 onward.Source: §F.7 · p. 47, Figure 28 · p. 48
  • The threshold brings training trajectories closer to confidence-threshold decoding and commits already-confident positions sooner, shortening trajectories and increasing distinct samples visited.Training commits positions above τtrain and, if fewer than u qualify, fills up to u with the most confident remaining positions.Source: §F.7 · p. 47
  • 13.2–13.8 points at u=4 and 2.4–5.4 points at u=8: the threshold’s score gains across decode thresholds; at u≤2, both variants collapse below 6 points.At u∈{16, 32, 64} the difference ranges from −0.3 to +2.0 points, and at u=128 from −0.8 to +0.3; comparisons match u and τ, not NFEs.Source: §F.7 · p. 47, Figure 29 · p. 48
  • With the threshold, u∈{32, 64} runs lie above both the 1.5-epoch SFT checkpoint and Post-SFT MDM control over broad NFE ranges; without it, no u clearly dominates both.Figure 31 plots score against NFEs for runs without the training-time threshold.Source: §F.7 · p. 47, Figure 31 · p. 50

F.8 Comparison with published LLaDA results

The section cautions that published LLaDA-8B scores are not protocol-matched to the paper’s measurements. Prompt format shifts the measured base scores, while the Post-SFT starting checkpoints score within the range of published instruct-model figures.

F.8 Comparison with published LLaDA results

  • The chat prompt format scores 10.1 points lower on GSM8K, 3.2 lower on MBPP, and 2.4 higher on IFEval than the plain format, so absolute scores should not be compared directly with published ones.The prompt format is fixed across the paper’s runs, so this shift does not affect their internal comparisons.Source: §F.8 · p. 51, Table 13 · p. 51
  • Post-SFT starting checkpoints score 75.9 and 77.0 on GSM8K, 38.8 and 41.6 on MBPP, and 68.0 and 65.4 on IFEval, within the range of published LLaDA-8B-Instruct figures.The section also notes that published GSM8K figures for LLaDA-8B-Instruct disagree with one another.Source: §F.8 · p. 51, Table 13 · p. 51

G Relation to the closest work

PUMBA combines trajectory alignment, recurrent conditioning, and temporal credit assignment, building on related methods while distinguishing its training setup and evaluation scope.

G Relation to the closest work

  • PUMBA’s stated contribution is to study policy-aligned training states, information passed between them, and the span of temporal credit assignment together.The study also examines conditions under which closer alignment remains learnable.Source: G Relation to the closest work · p. 52

PUMA: from trajectory construction to temporal credit assignment

  • PUMA [Kim et al. (2026)] supplies teacher-forced trajectories, while PUMBA adds predecessor model representations and BPTT to connect later losses to earlier computations.PUMA’s coarse-to-fine schedule suggests reveal order need not match inference granularity; PUMBA traces small-count degradation to local overfitting from sample reuse.Source: G Relation to the closest work · p. 52

Loopholing: from self-conditioned context to trajectory-trained state

  • Loopholing [Jo et al. (2026)] carries continuous information across discrete sampling; PUMBA instead caches the preceding rollout step’s hidden-state representation.Caching removes Loopholing’s auxiliary context-generation pass but adds storage and BPTT memory and computation costs.Source: G Relation to the closest work · p. 52

RCD and DiffusionGemma: distribution-derived carries

MetaState: memory architecture and trajectory alignment

  • MetaState [Xia et al. (2026)] uses fixed-size working memory with a frozen backbone, whereas PUMBA studies full-model instruction post-training and explicitly evaluates instruction following.The paper does not compare architectures directly and frames combining more effective or compact memories with PUMBA as future work.Source: G Relation to the closest work · p. 53

Relay: shared mechanism, complementary evidence and analysis

  • Relay [Rozonoyer et al. (2026)] also combines model-selected rollouts, differentiable state, and truncated BPTT; PUMBA varies trajectory granularity and unroll horizon at 8B scale.Relay’s large-model study establishes a two-step unroll within block diffusion, while PUMBA tests both full-canvas and block decoding.Source: G Relation to the closest work · p. 53

Related papers

Loading 2609.37974v1…