Any-Order GPT as Masked Diffusion Model: Decoupling Formulation and Architecture

Summary

The paper argues that comparisons between autoregressive (AR) language models and masked diffusion models (MDMs) are systematically confounded by convention: AR is run in causal decoder-only transformers while MDMs are run in full-attention encoder-only transformers, so any judgment about which paradigm is "better" mixes the loss (fixed-order vs. any-order) with the attention mechanism. To disentangle these, the authors build AO-GPT, a decoder-only implementation of the any-order autoregressive (AO-AR) objective — leveraging that AO-AR and the ELBO of masked diffusion are equivalent — and use it to answer two questions: within a shared decoder-only architecture, how does the fixed-order AR objective compare with the AO-AR/MDM objective; and within the AO-AR objective, how do encoder-only and decoder-only implementations differ. The conceptual move — treating AO-AR and MDM as a single formulation and instantiating it inside a σ-GPT-style decoder with AdaLN target-position injection and EMA — is a genuinely useful clarification of an increasingly muddled literature. The empirical program yields three noteworthy observations: language's L2R bias means uniform-permutation training converges slowly, and a 10% L2R mixture repairs much of the gap; encoder-only MDMs express $n \cdot 2^{n-1}$ order-invariant conditionals whereas decoders must cover $\approx e \cdot n!$ order-dependent ones, and averaging over context-order permutations at evaluation time closes the perplexity gap; and, with KV-caching plus a Bernoulli-then-sample decoding trick, decoder-only AO-GPT is roughly 25× faster than a comparable SEDD baseline at similar generation perplexity. The main conceptual limitation is that the "controlled" comparisons still slide multiple design changes into one ablation — AO-GPT‡ bundles adaLN, EMA, and 10% L2R mixing against a σ-GPT baseline lacking all three — and the two most novel claims (context-order ensembling as a remedy for the encoder/decoder gap, and the "cross-order benefit" of L2R data) sit uncomfortably with appendix ablations showing the opposite trend. Scale is capped at 350M and evaluation is restricted to language-model perplexity plus wall-clock time, with no downstream, infilling, or per-seed variance results.

Strengths

• Conceptual contribution. The framing — that AR vs MDM comparisons must decouple the objective (fixed vs any order) from the attention pattern (causal vs full) — is timely and clearly argued. Instantiating AO-AR inside a decoder-only stack via σ-GPT is a clean way to isolate the two axes.
• Counting argument for conditionals. The combinatorial derivation that encoder-only AO-AR expresses $n \cdot 2^{n-1}$ order-invariant conditionals while decoder-only AO-AR expresses $\approx e \cdot n!$ order-dependent ones is a compact, memorable characterization of the modeling-space asymmetry and directly motivates the context-order ensembling diagnostic.
• Concrete efficiency story. Combining KV-cache with the Bernoulli-then-sample identity (Lemma 5.1) yields an $O(n)$ decoder-only generation loop with a specialized parallel-decoding attention mask, and Figure 5's linear-vs-quadratic scaling is a clean empirical corroboration of the complexity argument.
• Serious architectural ablations. The AdaLN target-position injection and EMA choices are ablated systematically (Figures 7–9), and the AdaLN win over vanilla / per-layer PE variants is consistent across L2R and any-order objectives.
• Honest limitations passage. The conclusion openly flags the scale ceiling and the omission of infilling / downstream evaluation, which sets up the natural next work rather than overclaiming.

Weaknesses

• Confounded "controlled" comparisons. Figure 2's identical-architecture claim ignores that AO-GPT carries σ-GPT-style target-position injection that vanilla GPT-2 does not, and Table 2/4/5's AO-GPT‡ vs σ-GPT rows conflate three simultaneous changes (adaLN + EMA + 10% L2R mixing). The orthogonality claim for adaLN and EMA can only be read off Figure 9; using Table 2 to "further corroborate" it slides the L2R-mixing confound into the argument.
• Direct tension between Finding 3.2 and Figure 11. Finding 3.2 announces a cross-order benefit of 10% L2R data on any-order performance, while the Appendix caption for Figure 11 says the opposite ("increasing the left-to-right ratio … hurts the any-order loss, revealing a trade-off"). A non-monotone sweet spot at 10% is plausible, but the mechanism is never stated and the two framings read as contradictory.
• Finding 6 is scoped too broadly. As written it asserts decoder-only AO-AR falls short of encoder-only in general, but Table 4's Left-to-Right block shows AO-GPT‡ actually matches or beats SEDD/RADD on WikiText2, PTB, WikiText103, and 1BW. The finding is really about the any-order evaluation regime and should be restricted accordingly.
• Overclaims about the encoder side. The claim that encoder-only AO-AR "cannot easily specialize" to a fixed-order AR is not defended: restricting the masking distribution to prefix-only trivially produces the L2R chain. Similarly, "uniquely positions [decoder-only AO-AR] to interpolate between … paradigms" overstates a difference that is really one of degree.
• Joint parallel decoding under-discussed. The parallel-decoding attention mask is exactly the standard MDM independence approximation; the paper's statement that this "introduces no training/inference mismatch" holds only per-marginal and understates the joint-independence issue, which is central to any multi-token unmasking step.
• Efficiency framing. The 25× headline is real but is presented as a "trade-off between generation quality and practical inference speed" — yet under appropriate annealing Table 1 shows comparable perplexity at matched step counts. The trade-off is only present without annealing, and the framing should say so.
• Practical-vs-diagnostic status of context-order ensembling. Finding 6's closure of the encoder/decoder gap relies on M-permutation averaging that Remark 5.2 concedes is not deployable at inference. Both Finding 6 and the conclusion currently read as if the fix is available to a user of AO-GPT.
• Scale is capped at 350M. The AR-vs-AO-AR convergence penalty (Finding 1) and the ensemble-size dependence of Finding 6 are exactly the kind of effects that plausibly change with capacity. Even a 3-point scaling curve would materially strengthen the paper.
• No downstream / infilling evaluation. The strongest a priori case for AO-AR over L2R AR is infilling, editing, and constrained generation. Restricting evaluation to zero-shot perplexity and wall-clock speed sidesteps the main practical motivation.
• No variance / seed reporting. All headline numbers are single points, and Table 1's no-annealing SEDD row is non-monotone in steps (86.4 → 87.6 from 256 to 512 steps), which is exactly where seed noise would matter.
• Notation drift. Finding 7 references "Theorem 5.1" when only Lemma 5.1 is defined; the reverse-process factorization is written over $i = 0, \ldots, n-1$ while $L_{\mathrm{AO\text{-}AR}}$ elsewhere sums over $i = 1, \ldots, n$; and "OA-GPT" appears in Appendix A.3.2 and Figures 12–14 for a model consistently called "AO-GPT" elsewhere.

Reproducibility & code

• Encoder-only baselines are numerical citations, not runnable code. The released repository ships AO-GPT training, evaluation, ensembling, and generation but contains no SEDD or RADD implementation, checkpoint pointer, or timing pipeline. All encoder-only rows in Tables 1, 4, 5 and Figures 4–5, plus the 25× speedup ratio, therefore depend on an external SEDD/RADD reproduction that is not specified (version, hardware, Gumbel-noise precision, batch size, step count).
• Block-wise fixed-permutation experiment has no code. Only identity and uniform-random permutations are supported by the shipped sample_random_orders* methods; no config_order/train_AOGPT_blockwise*.py exists. Reproducing Figure 3(a)'s block-wise curve requires the reader to hand-craft the exact 1024-length permutation.
• AR (identity-order) baseline config is broken. The shipped config_order/train_AOGPT_AR.py does not define an order list, and the launcher then calls torch.tensor(None). The "identical decoder-only architecture" comparison in Figure 2 does not run out of the box.
• Positional-encoding ablation ships only the winner. Of the four variants in Figure 7 (Vanilla PE, Same-PE-every-layer, Different-PE-every-layer, AdaLN), only the AdaLN model class is in the repo. Reproducing the ablation requires re-implementing three model variants from Figure 6.
• EMA ablation ships only the winning decay. Only ema_rate = 1 - 1e-4 (0.9999) is configured; the {0.99, 0.999, No-EMA} baselines in Figure 8 require manual training-script edits, and the "No EMA" convention (decay 0 vs. skipping the EMA object) is undocumented.
• Speed script does not match Figure 5. speed_compare.py times only the first decoding step at batch 8, not the full 64–1024-step generation at batch 32 reported in Figure 5, and no SEDD-side timing exists. The 25× ratio cannot be produced from the shipped code.
• σ-GPT training driver missing. model_sigmaGPT.py is provided for evaluation only; no training script or config for the σ-GPT† baseline is shipped, though σ-GPT† is a headline baseline in Tables 2/4/5 and Figures 2/9.
• Any-Order Monte-Carlo count undocumented. eval_ppl_sigmagpt_random.py defaults to 10 MC samples but reproduced entries in the repo use 1; the paper does not state which value produced Tables 2/4/5, though the choice affects both the AO-GPT and σ-GPT rows.

Recommended Changes

Essential

• Reconcile Finding 3.2 with Figure 11. Add explicit any-order-loss numbers at 0%, 10%, 30%, 50%, 90% L2R mixing and, if the effect is non-monotone with a 10% sweet spot, state that mechanism. As written the "cross-order benefit" and the "trade-off" framings are contradictory.
• Scope Finding 6 to the any-order regime. Rephrase to say decoder-only AO-AR trails encoder-only on any-order evaluation, and either explain why the L2R rows of Table 4 do not undercut the finding or drop the general phrasing.
• De-confound the AO-GPT‡ vs σ-GPT comparison. Add an AO-GPT row that carries adaLN + EMA but not the 10% L2R data (matching Figure 9 rather than the "fully enhanced" recipe), and remove Table 2 from the corroboration of the adaLN×EMA orthogonality claim.
• Add per-seed variance for headline tables and the 25× speedup. Report at least 3-seed error bars in Tables 1/2/4/5 and in Figure 5, and state batch/sample counts per cell. In particular, address the non-monotone SEDD-Medium generation-perplexity row at 256/512 steps in Table 1.
• Ship a runnable SEDD/RADD comparison harness or a versioned pointer. Either add an adapter that reproduces the encoder-only rows from an upstream repo (with pinned commits, hardware, and precision), or document the exact configuration used to produce each SEDD/RADD number so the comparisons can be repeated.
• Fix the AR (identity-order) baseline config and add block-wise / PE-ablation configs. Set order = list(range(1024)) in config_order/train_AOGPT_AR.py, ship one config per block-wise permutation used in Figure 3(a), and ship model files or configs for the Vanilla-PE, Same-PE-per-layer, and Different-PE-per-layer variants in Figure 7.
• Match the shipped speed script to Figure 5. Extend speed_compare.py to time the full multi-step generation at batch 32, add a SEDD-side timing script under the same protocol, and record the exact hardware / CUDA / PyTorch versions behind the 25× headline.

Suggested

• Add a downstream or infilling result. Even a single controlled infilling / cloze benchmark comparing AO-GPT, σ-GPT, GPT-2, and SEDD would substantiate the practical value of the any-order flexibility that Section 4 defends.
• Add a scaling curve. A 3-point (Small / Medium / one larger size) sweep of the AR-vs-AO-AR loss gap and of the ensemble-size sensitivity would directly address the scale caveat in the conclusion.
• Restate the $O(1)$-per-step claim with the small-context assumption. Move the footnote-4 assumption to the point where the per-step cost is claimed, so readers do not compare $O(1)$-per-step with $O(n)$-per-step of encoders on different accounting bases.
• Tone down the encoder-side overclaims. Replace "cannot easily specialize to a strict, fixed-order autoregressive model" with a more careful statement about the practical inconvenience of a prefix-only masking schedule, and replace "uniquely positions" with "is well-suited."
• Clarify the parallel-decoding independence approximation. Explicitly state that the "no training/inference mismatch" claim is per-marginal and that co-sampled positions are drawn under a factorized approximation of the joint.
• Reframe the 25× as speed-at-comparable-quality when annealing is used. Update the framing sentence so readers do not read Finding 8.1 as a quality-for-speed trade, given that Table 1 with (top-p 0.95, temp 0.7) shows comparable perplexities at matched step counts.
• Fix the notation drift. Rename "Theorem 5.1" → "Lemma 5.1"; pick 0-based or 1-based indexing consistently between the reverse-process product and the $L_{\mathrm{AO\text{-}AR}}$ sum; and rename all "OA-GPT" occurrences to "AO-GPT".
• Document the Any-Order Monte-Carlo count. State how many permutation samples were used for the AO-GPT/σ-GPT Any-Order rows in Tables 2/4/5 and, ideally, report the MC standard error.
• Ship a σ-GPT training config. Provide the exact recipe used to train the σ-GPT† baseline (LR, batch, EMA, order distribution) so the baseline is fully self-contained in the repo.