REVIEW 4 major objections 6 minor 14 references
Learning the generation order of tokens—when to insert and when to unmask—improves insertion-based masked diffusion models for molecules and graphs, the paper argues.
Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →
2026-08-02 21:53 UTC pith:AUKVBKOK
load-bearing objection Solid formal machinery; the molecule-quality claim rests on unmatched nucleus-p comparisons and the abstract overstates results. the 4 major comments →
Insertion Based Sequence Generation with Learnable Order Dynamics
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
The central claim, stated in the paper's own terms, is that 'learning the target generation order can improve insertion-based diffusion models without giving up tractable training.' Concretely, the authors define per-position target rates through two event-time distributions — an insertion time T_in^i and an unmasking time T_um^i — parameterized as Kumaraswamy CDFs whose shape parameters are predicted by a transformer from the clean sequence. The generator matches these target rates via a projected rate-matching objective; gradients for the order parameters come from a REINFORCE leave-one-out estimator. The paper reports that this learnable order improves validity and quality of generated mo
What carries the argument
The load-bearing object is a per-position pair of CDFs F_in(t)=1-(1-t^a)^{b_in} and F_um(t)=1-(1-t^a)^{b_um} (Kumaraswamy distributions with shared shape a and learned multipliers b_in, b_um). These yield closed-form hazard rates λ_in = b_in · a t^{a-1}/(1-t^a) and λ_um = b_um · a t^{a-1}/(1-t^a); under the time-change τ=-log(1-t^a), each position's two events form an exponential race with rates b_in and b_um and the constraint insertion precedes unmasking. Because the hazard rates can be expressed in closed form, the likelihood p_φ_t(z|z1) is computable without numerical integration, so both the generator θ and the order parameters φ can be trained in a single backward pass per sample by pr
Load-bearing premise
The order dynamics assume that each position's insertion time and unmasking time are independent draws from position-specific Kumaraswamy distributions, and that unmasking times are a truncation of a base distribution onto [T_in, 1]; this family cannot express joint order preferences across positions, and the truncation makes the unmasking hazard insensitive to the realized insertion time.
What would settle it
Run the method on a task whose optimal order is inherently joint — e.g., star graphs where all tokens on the first arm must be completed before any token on the second arm — while keeping all other components identical. If exact-match accuracy does not exceed a fixed-schedule baseline (or a randomly reordered baseline), the independent per-position order family is the limiting factor rather than the learning procedure.
If this is right
- On the hard star-graph traversal task, learning the insertion order (with fixed unmasking rate) increases exact-match accuracy from 6.0% (FlexMDM) to 88.1%.
- On de novo SAFE molecule generation, validity improves from 98.9% to 99.9% and quality from 39.5 to 62.1, with the gains concentrated at smaller nucleus-sampling thresholds where the model generates more greedily.
- The unmasking hazard b_um should be fixed at 1 during training: trainable b_um destabilizes optimization, but fixing it still leaves room to learn insertion order, which can be elicited at decoding time via confidence-based position selection.
- The training procedure requires no simulation of the full trajectory — one forward/backward pass per sample with two MC samples suffices for both parameter sets.
- The framework reduces to FlexMDM when the order parameters are fixed to a=b=1, so the method is a strict generalization that preserves tractable training.
Where Pith is reading between the lines
- The order family is per-position independent, so it cannot represent joint preferences such as 'complete all of branch A before touching branch B'. If real-world generation requires such joint order structure, the current gains might come from the auxiliary conditioning signal or the evaluation regime rather than from order learning per se.
- The framework treats 'order' as a race of independent exponential clocks. A natural testable extension is to allow position-dependent shape parameter a or a hierarchical (Plackett–Luce) order distribution; the exponential-race view suggests a direct connection to top-k order statistics that could yield analytic order probabilities beyond the current equal-a case.
- The quality improvements concentrate at lower nucleus-sampling p values, which suggests learned schedules act partly as a prior that tightens the effective action space; one could test this by measuring decoder entropy or the rate of redundant insertions under the learned schedule.
- Because the target rates are computed from the clean sequence at training time but the generator sees only partial sequences, the method implicitly learns to predict 'what would be inserted next' from context; this could enable application to fragment-constrained generation, where conditioning on scaffolds is natural.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. This paper proposes LFlexMDM, an extension of FlexMDM's variable-length masked discrete flow matching in which per-position insertion and unmasking rates are learned from the clean target sequence. The target order is parameterized by Kumaraswamy CDFs (Eq. 9), and joint training of the target rates and the generator is performed by projected rate matching augmented with a REINFORCE leave-one-out gradient estimator (Eq. 12). Experiments on star-graph traversal and SAFE molecule generation report gains in exact match and validity/quality. The theoretical derivation (Props. 1, 3-5) is coherent, but the empirical evidence for the headline claim is weakened by evaluation confounds and missing comparisons.
Significance. If the central claim were supported, it would be a useful step: instead of hand-designing schedules in insertion-based discrete flow matching, one could learn data-dependent schedules while retaining a tractable training objective. The paper's main formal contributions — the per-position race construction in Prop. 1, the projected-rate-matching upper bound in Prop. 5, and the closed-form Kumaraswamy likelihoods in App. D.5.3 — appear internally consistent, and the authors provide code. However, the current experiments do not cleanly separate the effect of learned order from the effect of more aggressive nucleus truncation, increased auxiliary model capacity, or confidence-based decoding; the abstract's performance numbers are not reproducible from the tables.
major comments (4)
- [Abstract / Table 2 / Table 7] The abstract's 'up to 17.5%' improvement cannot be located in the reported tables. At matched nucleus p=1.0, LFlexMDM quality is 39.6±2.3 (medium), 43.8±0.8 (small), 43.1±1.8 (xtiny) versus FlexMDM's 39.5±1.6; the large 54.7–62.1 values occur only at p=0.2–0.5. The '17.5%' figure must be defined with row/column references, or removed. The abstract also claims fragment-constrained molecule generation (6.7%) but no fragment-constrained experiments appear in the paper.
- [Table 2 and Figure 7] The de novo quality comparison is not matched. FlexMDM is reported only for p=1.0, while LFlexMDM's best rows use p=0.2/0.5, which reduce diversity (e.g., 0.922→0.830–0.876) and can mechanically increase the quality metric. Figure 7's left panel is said to use 'respective best values,' so it too compares FlexMDM at p=1.0 with LFlexMDM at lower p. Report FlexMDM at p=0.2/0.5 (with and without confidence) before claiming a quality improvement from learned order.
- [5.1.1, Fig. 5] The deployed configuration's learned order correlates only weakly with the known optimal order: -0.22 (bum=1, no confidence) and -0.60 (bum=1, confidence), while the -0.91 correlation is obtained with the trainable-bum configuration that was abandoned for instability. This weakens the mechanistic claim that the gains come from learning the target order. Either provide evidence from the deployed configuration that its generation order matches a useful order, or soften the interpretation.
- [D.5.2 / Eq. 11] Because ϕ and θ jointly minimize the same Bregman loss, the learned target rates can be self-confirming; the paper itself notes degenerate solutions and adds a schedule regularizer. The comparison to FlexMDM does not control for the extra auxiliary transformer (clean-sequence conditioning), which may provide the benefit even if order dynamics are not learned. Please include an ablation in which the same LFlexMDM architecture has bin frozen (or the auxiliary network is ablated), and report the learned bin statistics (e.g., distribution of b values) to show they are data-dependent and not collapsed by the regularizer.
minor comments (6)
- [Tables 4 and 5] 'LNFlexMDM' appears to be a typo for 'LFlexMDM' in the model rows.
- [D.7] The phrase 'with a constrant Tin < Tum' contains a typo: should be 'constraint'.
- [Section 6] The sentence contains 'simultaneous simultaneous'; remove the duplicate.
- [Figure 7] The caption and text do not state whether the error bars/± values are standard errors across seeds or standard deviations; please clarify.
- [Algorithm 2] The confidence function c(Kθ, λθ) used in SELECT is not specified. Describe the exact confidence score used in the experiments.
- [Prop. 6 / D.5.4] The proof sketch for the order probability is terse. Include the exponential-race derivation or cite a source for the Plackett–Luce form.
Circularity Check
No significant circularity: the projection/rate-matching derivation is self-contained and externally benchmarked.
full rationale
The paper's derivation chain is self-contained. Proposition 1 defines target conditional rates from per-position insertion/unmasking time CDFs and proves they generate the desired marginals via the Kolmogorov Forward Equation (Appendix D.1). The projected rate matching loss (Equation 11) is derived from Bregman divergence upper bounds (Proposition 5 and Appendix D.4), and the REINFORCE leave-one-out gradient estimator is derived explicitly (Appendix D.5). The learned rates F^{φ,i}_* are a parameterization of the target schedule, optimized jointly with the generator by the same matching objective; this is a standard training setup, not a fitted parameter renamed as a prediction. The empirical evaluation is grounded in external benchmarks: SAFE/ZINC molecule validity/quality/diversity, star-graph exact match, and comparison against FlexMDM and other external baselines. The paper itself flags the need for schedule regularization to prevent degenerate endpoint-concentrated schedules (Appendix D.5.2), which is a training stability limitation, not a circular step. The reader's noted issues around unmatched nucleus truncation thresholds (FlexMDM only at p=1.0 vs LFlexMDM at p=0.2–0.5) and the weak correlation of the deployed order with the locally optimal order are experimental/confound concerns, not circularity: no claimed prediction reduces by construction to its own input, and no load-bearing step relies on a self-citation chain.
Axiom & Free-Parameter Ledger
free parameters (5)
- Kumaraswamy exponent a (shared a_in = a_um) =
not stated in paper
- Target insertion-rate multipliers b^{φ,i}_in(z1) =
learned (auxiliary transformer φ output)
- Frozen unmasking multiplier b_um =
1 (fixed)
- Schedule-regularization grid and cutoffs =
t_ε=0.01, δ=0.01, grid {t_k} with weights w_k
- Evaluation nucleus p and confidence selection =
p ∈ {0.2, 0.5, 0.9, 1.0}; conf on/off
axioms (8)
- domain assumption Per-position independence of target conditional paths and generator rates (Section 2, 'Positional Independence Assumption'; Eq. 86)
- domain assumption Two events per position: [D]→[M]→value, with T_in < T_um almost surely (Eqs. 6–8)
- ad hoc to paper Unmasking time = truncation of a base distribution onto [T_in, 1] (Eq. 35; Appendix D.7)
- ad hoc to paper Kumaraswamy CDF family for F_in and F_um with a_in = a_um = a (Eqs. 9–10)
- standard math Regular CTMC representation: any regular CTMC rate decomposes as Poisson arrivals × embedded transition kernel (Section 3.1)
- standard math Projected rate matching upper-bounds terminal KL (Prop. 5, Eqs. 31–33)
- standard math REINFORCE leave-one-out gradient estimator with n=2 is unbiased (Eq. 12, Appendix D.5)
- domain assumption Generator network θ is expressive enough to profitably match the learned target rates
Cite this review
Pith. "Pith review of Insertion Based Sequence Generation with Learnable Order Dynamics." pith.science (2026). https://pith.science/paper/AUKVBKOK
@misc{pith2026260218695,
author = {Pith},
title = {Pith review of: Insertion Based Sequence Generation with Learnable Order Dynamics},
year = {2026},
howpublished = {\url{https://pith.science/paper/AUKVBKOK}},
note = {Machine review of arXiv:2602.18695}
}
read the original abstract
Existing insertion-based masked diffusion models that generate sequences by interleaving token insertion with unmasking use fixed schedules that are not dependent on the data. For structured sequences like graphs and molecules, learning data-dependent generation orders can improve generation quality by reducing uncertainty over the action space. We propose LoFlexMDM, an insertion-based masked diffusion model with learnable order dynamics that learns data-dependent insertion and unmasking rates. We generalize the discrete flow matching framework to work with variable-length sequences, propose a tractable schedule parameterization and a training objective for joint training of the generator and the target order dynamics. On De Novo and fragment-constrained molecule generation, LoFlexMDM improves sample quality over FlexMDM by up to 17.5% and 6.7%, respectively. These results show that learning the target generation order can improve insertion-based diffusion models without giving up tractable training. We open source the code at https://github.com/dhruvdcoder/LoFlexMDM.
Figures
Reference graph
Works this paper leans on
-
[1]
Pick the indexing setC= (X 0, X1)
-
[2]
π(x0|x1) =p 0(x0) for some easy to samplep 0
Use p1 =p data and pick the coupling distribution π(x0|x1), typically independent i.e. π(x0|x1) =p 0(x0) for some easy to samplep 0
-
[3]
Construct a collection of CTMCs {X c t }c∈C with rates Rt(x, y|C)such that itgenerates pt|C, which satisfies p1|C(x|x1, x0) =δ x1 (x)andp 0|C(x|x0, x0) =δ x0 (x)
-
[4]
The marginal pt(x) = EC[pt|C(x|C)] satisfies p1(x) =p data(x) and Proposition 2 ensures that the CTMC with the rateR t(·,·) = EC[Rt(·,·|C)]generatesthe marginalsp t(x) = EC[pt|C(x|C)]
-
[5]
Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M
URL https://openreview.net/forum? id=tcvMzR2NrP. Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. Simplified and generalized masked diffusion for discrete data.Advances in neural information processing systems, 37:103131–103167, 2024. Uria, B., Murray, I., and Larochelle, H. A deep and tractable density estimator. InInternational Confer- ence on Ma...
2024
-
[7]
Zheng, H., Gong, S., Zhang, R., Chen, T., Gu, J., Zhou, M., Jaitly, N., and Zhang, Y
URL https://openreview.net/forum? id=EY6pXIDi3G. Zheng, H., Gong, S., Zhang, R., Chen, T., Gu, J., Zhou, M., Jaitly, N., and Zhang, Y . Continuously Augmented Dis- crete Diffusion model for Categorical Generative Model- ing, October 2025. URL http://arxiv.org/abs/ 2510.01329. 11 Insertion Based Sequence Generation with Learnable Order Dynamics Appendix fo...
arXiv 2025
-
[12]
Next we look at the optimization objective for learning the ratesR t
Construct an optimization problem to learn parametricR θ t (·,·). Next we look at the optimization objective for learning the ratesR t. 14 Insertion Based Sequence Generation with Learnable Order Dynamics B.2. Discrete Flow Matching We want to learn the rates Rt(x, y)using a parametric model Rθ t (x, y). Since we can easily write down the conditional rate...
2025
-
[13]
If the event times T i are independent and have exponential distribution with parameter bi, then the probability of an orderπis given by: P(π) = LY i=1 bπ(i) PL j=i bπ(j)
-
[14]
Then ¯Nu is a homogeneous Poisson process with warped time u= R t 0 β(s)ds
Let Nt be a nonhomogeneous Poisson process with hazard rate of the form λ(t) =bβ(t) . Then ¯Nu is a homogeneous Poisson process with warped time u= R t 0 β(s)ds. Let U1 = inf{u: ¯Nu = 1} , be the first arrival time of the homogeneous Poisson process. D.6. Sampling The generator network predicts learnable hazard ratesλ θ,i in,t(xt)for insertion at gap posi...
2024
-
[2019]
URL https://openreview.net/forum? id=r1lgTGL5DE. Kumaraswamy, P. A generalized probability density function for double-bounded random processes. Journal of Hydrology, 46(1):79–88, 1980. ISSN 0022-1694. doi: 10.1016/0022-1694(80)90036-0. URL https://www.sciencedirect.com/ science/article/pii/0022169480900360. Lee, S., Kreis, K., Veccham, S. P., Liu, M., Re...
arXiv 1980
-
[2020]
9 Insertion Based Sequence Generation with Learnable Order Dynamics Hoogeboom, E
URL https://openreview.net/forum? id=rygGQyrFvH. 9 Insertion Based Sequence Generation with Learnable Order Dynamics Hoogeboom, E. and Salimans, T. Blurring diffusion models. arXiv preprint arXiv:2209.05557, 2022. Hoogeboom, E., Gritsenko, A. A., Bastings, J., Poole, B., van den Berg, R., and Salimans, T. Autoregressive dif- fusion models, 2022. URL https...
Pith/arXiv arXiv 2022
-
[2022]
Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, T
URL https://openreview.net/forum? id=DmT862YAieY. Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, T. Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Pro- tein Co-Design. InProceedings of the 41st Interna- tional Conference on Machine Learning, pp. 5453–5512. PMLR, July 2024. URL https://proceedin...
Pith/arXiv arXiv 2024
-
[2023]
Wang, Z., Shi, J., Heess, N., Gretton, A., and Titsias, M
URL https://openreview.net/forum? id=X6dEqXIsEW. Wang, Z., Shi, J., Heess, N., Gretton, A., and Titsias, M. Learning-Order Autoregressive Models with Applica- tion to Molecular Graph Generation. InForty-Second International Conference on Machine Learning, June
-
[2025]
Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y
URL https://openreview.net/forum? id=RuP17cJtZo. Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y . The curious case of neural text degeneration. In International Conference on Learning Representations,
This paper was first reviewed by deepseek-v4-flash on August 2, 2026.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.