{"id":"35db8442-38f9-49f9-b0da-10f0e4850f40","arxiv_id":"2601.22450","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":1,"one_line_summary":"The masked diffusion objective decomposes into signal and implicit-regularizer terms, and restricting mask sampling to a signal-rich window improves language model pretraining and fine-tuning at scales up to 8B parameters.","lead":"Masked diffusion language models can be split into a signal part that teaches features and a noise part that acts as a regularizer, which the authors show prevents grokking on a synthetic parity task. The paper then tunes the mask probability range to boost perplexity and downstream accuracy, with measured gains of up to 8.8 points on an 8-billion parameter model.","discovery_kind":"first_principles","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The Signal/Noise theory is proved only for a uniform-attention MLP; the F.1 ablation shows sufficiency, not that learned attention is irrelevant, leaving the LLaDA-8B transfer theoretically unsupported.","rationale":"The reader's weakest assumption correctly identifies the uniform-attention reduction as the load-bearing gap. The theory's theorems are stated for the reduced MLP, and the F.1 ablation supports only the claim that attention is unnecessary for the no-grokking phenomenon in one parity setting, not that the learned-attention dynamics obey the same energy landscape. This is the single most consequential weakness because it connects the theoretical contribution to both the parity experiments and the LLaDA-8B language claims. If the reduction is unfaithful, the language results remain an empirical schedule heuristic rather than a consequence of the derived mechanism. The concern does not amount to a rejection: the paper has independent empirical support—the clear no-grokking result, the verification of the predicted optimal mask interval on parity, and the F.1 reduced-model ablation—so the appropriate verdict remains conditional on resolving this gap. I agree with the reader's assessment rather than escalating to rejection.","tokens_in":22843,"tokens_out":14948,"duration_ms":137592,"concrete_test":"Train the (20,6) nanoGPT parity model under matched hyperparameters in two conditions: (i) full learned attention, and (ii) attention frozen to uniform, recording steps to reach 95% test accuracy and inspecting learned attention maps on signal versus noise masks. If the frozen-uniform model matches the learned-attention model in convergence speed and the learned attention is approximately uniform, the reduction is adequate for the parity claim; if learned attention converges faster or systematically attends to secret positions, the MLP reduction does not capture the mechanism operating in the actual architecture.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's central mechanism—the Signal/Noise decomposition (Theorem 4.3), the energy landscape (Theorem 4.4), and the derived optimal mask rate (Corollary 4.7)—is proved for a two-layer MLP on the aggregated input after fixing attention to uniform. The only justification is the empirical ablation in Appendix F.1, which shows that a uniform-attention MLP can learn parity without grokking. That establishes sufficiency, not necessity: the actual nanoGPT model, and certainly LLaDA-8B, trains learned attention jointly with the MLP layers, and the theory gives no account of how attention interacts with the noise-regime regularizer. In particular, Theorem 4.4's feature-learning collapse and the claim that P_S acts as a dynamic gain factor rely on a scalar readout v of the reduced MLP; a transformer's per-token readouts and attention weights are not governed by this landscape. If learned attention is load-bearing—for example, if it is needed to route visible secret bits or to implement the identifiability computation—then the claimed mechanism does not apply to the actual models, and the language gains reduce to an empirical schedule heuristic. The paper reports no attention analysis for the 8B models, so this gap is unresolved.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"This paper studies the generalization properties of masked diffusion language models (MDLMs) through the lens of the k-parity problem. The authors decompose the masked diffusion objective into a Signal regime (masks that leave the parity bit identifiable) and a Noise regime (masks that do not), and prove for a reduced two-layer MLP with uniform attention that the noise term acts as an implicit regularizer penalizing confident outputs on unidentifiable inputs. They derive an optimal masking rate t=1/(k+1) for signal maximization and a sample-complexity-optimal schedule. Experiments on (20,6) parity with a nanoGPT-style transformer show that masked diffusion avoids the grokking plateau, and the best interval U[0,0.246] roughly matches the computed signal-optimal interval for t0=0. The authors then transfer the idea to language modeling: a 50M-parameter ablation on WikiText locates the best masking interval near t in [0.45,0.55], and using this interval for LLaDA-8B pretraining and SFT yields downstream gains up to 8.8 points on ARC-Easy and 5.8 points on GPQA, with the caveat that generative math reasoning favors wider intervals extending to t=1.","tokens_in":23071,"tokens_out":15966,"duration_ms":135164,"significance":"If the decomposition is taken at face value for the reduced model, the paper provides a clean explanation for the absence of grokking under masked diffusion and a falsifiable prediction for the optimal mask rate, which is qualitatively confirmed by the parity experiment. The core loss decomposition (Theorem 4.3) is derived rather than fitted, and the 8B pretraining and SFT gains are substantial and practically useful. The paper is honest about several limitations but overstates the theoretical scope: the theory applies to a uniform-attention MLP, while the language-model window is selected empirically. The contribution is therefore a mechanism-inspired schedule with partial theoretical backing, and the manuscript would benefit from a more precise separation of what is proven, what is empirically established, and what is conjecture.","major_comments":[{"comment":"The theoretical claims (Theorems 4.3, 4.4, Corollaries 4.5-4.7) are all proved for the aggregated-input two-layer MLP after fixing attention to uniform. The only justification for this reduction is the ablation in Appendix F.1, which demonstrates sufficiency (uniform attention can learn parity without grokking) but not necessity. The nanoGPT experiments in Section 4.5 and the LLaDA-8B experiments in Section 5 train learned attention jointly, and the paper contains no analysis of how attention interacts with the noise-regime regularizer. The statement in the introduction that the noise regime 'prove[s]' the mechanism for masked diffusion generally therefore overstates the theoretical scope; the language-model results should be presented as an empirically motivated schedule, with the theory applying strictly to the reduced MLP.","section":"Section 4.1 and Appendix F.1"},{"comment":"The sentence 'Since E(W) is proportional to P_S^2' is not a consequence of the theorem as stated. With Sigma(W) = P_S E_S[|Mm|/(2t) hh^T] + P_N E_N[|Mm|/(2t) hh^T] and c(W) = P_S E_S[|Mm|/(2t) f* h], the energy is E(W) = P_S^2 a^T (P_S A + P_N B)^dagger a, which is not proportional to P_S^2 in general. This invalidates the 'dynamic gain factor' interpretation as written, although the qualitative statement that the signal term scales c(W) remains true. Please correct the proportionality claim and re-derive any subsequent conclusions that rely on it.","section":"Section 4.3, after Theorem 4.4"},{"comment":"The optimal language window t in [0.45,0.55] is selected from the 50M WikiText ablation (Figure 3), not derived from the parity theory. The parity theory gives t=1/(k+1) for a known k, which has no direct analogue for natural language. The 8B results therefore validate an empirical schedule heuristic, not the theoretical prediction. This should be stated explicitly, and the abstract's phrasing 'leveraging our theoretical insights' should be softened accordingly.","section":"Sections 5.1 and 5.2"},{"comment":"The generative reasoning results show the opposite of the signal-optimal prescription: shifting the sampling interval toward t=1 improves GSM8K and MATH, and the [0.5,1.0] interval is best. The paper's post hoc explanation is plausible, but it means the proposed 'signal-optimal window' is task-dependent, not a universal principle. The conclusion and abstract should present the discriminative/generative distinction as part of the empirical finding rather than as a corollary of the parity theory.","section":"Section 5.2.2 and Table 3"},{"comment":"The training objective uses a 1/t factor on the cross-entropy loss. Under U[0,1], samples with t near 0 receive extremely large weights; restricting to [0.45,0.55] removes this instability. The 50M and 8B improvements could be due in part to this reweighting artifact rather than to the signal-noise mechanism. To support the theoretical attribution, add a control schedule that excludes low t but includes high t (e.g., U[0.1,1] or U[0.05,0.95]) and show that the gains are not explained by the 1/t weighting alone.","section":"Section 5, objective with 1/t weighting"}],"minor_comments":[{"comment":"'quesiton' is a typo for 'question'.","section":"Section 1"},{"comment":"'colume-wise' should be 'column-wise'.","section":"Definition 3.3"},{"comment":"The theorem statement refers to 'theorem 3.4'; it should cite Definition 3.4.","section":"Theorem 4.2"},{"comment":"The text refers to 'theorem 4.5' when it means 'Corollary 4.5'.","section":"After Corollary 4.5"},{"comment":"The figure caption says the baseline is a 'dashed blue line' while the text refers to a 'dashed black line'; make the references consistent.","section":"Figure 3 and Section 5.1"},{"comment":"The table includes a 'Base' column but the caption says 'Evaluation of SFT models'; clarify whether 'Base' is the pretrained model without SFT.","section":"Table 3"},{"comment":"The choice t in [0.45,0.55] is said to be suggested by Figure 3, but Figure 3's best intervals are wider ([0.4,0.5] and [0.5,0.6]); justify the narrower choice.","section":"Section 5.1"},{"comment":"'Experiments was conducted' should be 'Experiments were conducted'.","section":"Appendix F.5"},{"comment":"The main text defines Sigma(W) without the explicit P_S/P_N weighting that appears in the proof; align the two definitions.","section":"Theorem 4.4 statement vs Appendix D.1"}],"recommendation":"major_revision","confidential_remarks":"The paper is likely to be citable and the 8B results are valuable, but the theoretical scope is narrower than the title and abstract suggest. The uniform-attention reduction and the empirical window search are the main gaps. If the authors add the proposed control for the 1/t confound and correct the proportionality claim, I would support acceptance; in its current form the overclaim prevents acceptance."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"First thing to know: the k-parity decomposition is the real deal, and the no-grokking observation is clean. The language-model story is a useful heuristic, not a proven consequence of the theory.\n\nWhat's actually new: the paper gives a Signal/Noise decomposition of the masked diffusion loss for a two-layer MLP with uniform attention, and proves that the noise term behaves as an implicit regularizer. That is a genuine result. The optimal mask rate t = 1/(k+1) is derived, not fitted, and the parity experiment confirms both the prediction and the absence of grokking. The F.1 ablation shows that a uniform-attention MLP still learns parity without grokking, which supports the reduction. This is a real contribution to the diffusion-language-model theory literature, where prior work simplified the objective but did not analyze the regularizer.\n\nThe main soft spot is the bridge to actual transformers. The theory holds for the reduced MLP; F.1 establishes sufficiency, not necessity. The nanoGPT and LLaDA-8B models train learned attention, and the paper gives no account of how attention interacts with the noise-regime regularizer. So the 8B transfer is theoretically unsupported. That doesn't make it wrong, but it makes it a heuristic.\n\nThe language mask window [0.45,0.55] is chosen empirically at 50M scale, then rationalized via the theory. That's an honest engineering recipe, not a derivation. Table 3 complicates the story further: on GSM8K and MATH the mid-range window underperforms, and shifting toward high masking helps. So the abstract's blanket claim of 'superior results' is too strong. The 8B numbers are single runs with no error bars.\n\nNet: the parity result deserves a serious referee. I'd ask for multiple seeds or error bars at 8B, a tone-down of the abstract, and a clearer statement that the theory covers the reduced MLP. Those are revision-level issues, not fatal flaws.\n\nI'd take it to reading group and would cite the parity decomposition in my own work. Send it to review.","headline":"Parity theory is solid and worth refereeing; the language-scale claims need tempering.","tokens_in":23635,"tokens_out":2559,"would_cite":true,"duration_ms":21390,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":["68T07","68Q32"],"pacs":[],"model":"deepseek-v4-flash","headline":"The masked diffusion loss splits into signal and noise, and tuning the mask schedule to the signal-rich window speeds up parity learning and 8B-scale language training.","keywords":["masked diffusion language models","k-parity","grokking","implicit regularization","mask schedule","signal-noise decomposition","language model pretraining","generalization"],"falsifier":"Train a masked-diffusion transformer on $k$-parity with learnable attention and compare its generalization curve and optimal mask rate with the uniform-attention MLP; a grokking plateau in the learnable-attention run, or an optimal rate far from $1/(k+1)$, would show the attention reduction hides the real mechanism. A language-scale test would hold compute fixed and swap the mask schedule to, say, high-noise intervals during pretraining; if a high-noise schedule matches or beats the mid-range window on knowledge-heavy benchmarks, the claim that $t\\approx0.5$ is the general signal-optimal window would be weakened.","tokens_in":22578,"feed_emoji":"🎛️","tokens_out":11867,"duration_ms":88983,"temperature":0.7,"pith_summary":"The paper claims that the masked diffusion training objective contains an implicit regularizer that explains why these models generalize without the long plateau seen in standard training. Studying the $k$-parity problem, where transformers normally 'grok' suddenly after a long chance-level plateau, the authors decompose the diffusion loss into a Signal Regime that drives feature learning and a Noise Regime that penalizes confident predictions on inputs where the target is information-theoretically unidentifiable. This decomposition predicts an optimal mask rate $t=1/(k+1)$ for maximizing signal. The paper then shows that restricting language-model mask sampling to the empirically located signal-rich window near $t=0.5$ improves perplexity and downstream accuracy, with gains up to 8.8 absolute points on ARC-Easy at the 8B scale. A sympathetic reader would take the central claim to be that the mask schedule is a tunable implicit regularizer, not just a training detail.","feed_headline":"Masked diffusion skips grokking; a tuned mask schedule adds 8.8 points","feed_subtitle":"Signal/noise split of the diffusion loss explains fast parity learning and improves 8B-scale pretraining.","key_machinery":"The load-bearing object is the Signal/Noise decomposition of the masked diffusion objective, paired with the reduction of the transformer to a two-layer MLP acting on the aggregated embedding $\\tilde z$ after attention is fixed to uniform. In that reduced model, Theorem 4.3 rewrites the loss as a signal term that fits $f^*(\\tilde z)$ on identifiable examples plus a noise term that drives outputs to zero on unidentifiable examples. Theorem 4.4 turns gradient flow on this loss into ascent on the energy $E(W)=c(W)^\\top \\Sigma(W)^\\dagger c(W)$, where $c$ is the correlation with the target and $\\Sigma$ is the feature covariance; the signal probability $P_S$ multiplies this energy and therefore sets the effective learning-rate scale for the hidden weights. Maximizing $P_S$ over a uniform mask distribution yields the single-point optimum $t_0=t_1=1/(k+1)$, which is the theoretical anchor for the language-model schedule. The same decomposition also produces the sample-complexity bound of Theorem 4.2, which the paper uses to derive a separate sample-complexity-optimal schedule.","core_discovery":"On the authors' own terms, the discovery is that the masked diffusion loss for $k$-parity separates exactly (up to constants) into $$L_{\\mathrm{eff}}(\\$\\theta$)=P_S\\,\\mathbb{E}_{\\tilde z\\mid R_S}\\!\\left[\\frac{|M_m|}{2t}(f_\\$\\theta$(\\tilde z)-f^*(\\tilde z))^2\\right]+P_N\\,\\mathbb{E}_{\\tilde z\\mid R_N}\\!\\left[\\frac{|M_m|}{2t}f_\\$\\theta$(\\tilde z)^2\\right],$$ with $P_S=(k+1)\\mathbb{E}_t[t(1-t)^k]$ and $P_N=1-P_S$. Signal examples leave exactly one member of the secret set unmasked and therefore determine the target; noise examples either hide too much or hide too little and make the target unidentifiable. The noise term acts as an implicit regularizer that keeps the model from memorizing, and the landscape analysis in Theorem 4.4 shows that feature learning is driven by an energy $E(W)=c(W)^\\top \\Sigma(W)^\\dagger c(W)$ whose ascent speed is set by $P_S$. From this, the paper derives $t^*=1/(k+1)$ as the signal-optimal mask rate and verifies on $(20,6)$-parity that masked diffusion reaches train-and-test accuracy together, without a grokking plateau, and fastest near the predicted range. At language scale, the paper reports that restricting mask sampling to $[0.45,0.55]$ improves pretraining and fine-tuning results for 50M and 8B models relative to uniform $[0,1]$ sampling.","pith_inferences":["The paper does not say this, but the optimal window should depend on the effective dependency structure of the task; longer-range dependencies would shift the signal maximum, so a fixed $[0.45,0.55]$ window is unlikely to be optimal for every benchmark.","An extension the authors leave implicit is a curriculum: pretrain on the signal-rich window first, then anneal the schedule toward higher mask rates for reasoning tasks, which their range ablations suggest would combine the knowledge and reasoning gains.","The uniform-attention reduction implies a sharper test than the paper runs: an MLP and a transformer with identical schedules should show the same optimal mask rate, and the appendix only checks that uniform attention is trainable, not that the rate is identical.","The same decomposition should apply to any reconstruction loss over independently masked coordinates, so the schedule-tuning idea could be tested in masked autoencoders outside language modeling."],"forward_implications":["Masked diffusion training on $k$-parity should reach train and test accuracy at the same time, with no grokking plateau, whenever the mask schedule keeps $P_S$ away from zero.","For a known dependency order $k$, the signal-optimal schedule is concentrated at $t=1/(k+1)$; for $k=6$ this predicts the fast range observed near $t\\approx0.246$.","The standard uniform schedule $t\\sim U[0,1]$ spends a large fraction of its compute on low-signal regimes, so restricting training to the empirically located signal-rich window should improve convergence speed and final perplexity.","At 8B scale, the restricted schedule improves zero-shot HellaSwag by 4.6 absolute points and ARC-Easy by 8.8 absolute points after 15,000 pretraining steps relative to uniform sampling.","For supervised fine-tuning, mid-range masking helps discriminative benchmarks (GPQA 0.402 vs 0.344 uniform), while shifting toward high-noise intervals improves generative reasoning (GSM8K 0.785 on $[0.5,1.0]$ vs 0.768 on $[0,1]$)."],"supporting_citations":[{"why":"Supplies the grokking phenomenon—prolonged chance-level plateau followed by sudden generalization—that the paper claims masked diffusion eliminates.","marker":"Power et al. (2022)"},{"why":"Establishes k-parity as a canonical task where transformers and MLPs exhibit grokking, giving the paper its testbed and comparison.","marker":"Barak et al. (2022)"},{"why":"Provides the lazy-readout and energy-landscape analysis, including the claim that feature learning can collapse without regularization, which Theorem 4.4 adapts.","marker":"Tian (2025)"},{"why":"Shows the masked diffusion objective can be rewritten as a weighted average of masked language modeling losses, the structural simplification the paper's decomposition builds on.","marker":"Sahoo et al. (2024)"},{"why":"Establishes the connection between absorbing discrete diffusion and conditional distributions of clean data, used to justify analyzing the objective rather than sampling dynamics.","marker":"Ou et al. (2025)"},{"why":"Documents that masked diffusion models maintain performance without weight decay, the empirical generalization puzzle the paper's implicit-regularizer explanation addresses.","marker":"Ni et al. (2025b)"},{"why":"Supplies the LLaDA-8B architecture, evaluation protocol, and base-model checkpoints used for the large-scale pretraining and SFT experiments.","marker":"Nie et al. (2025b)"},{"why":"Provides the WikiText dataset used in the 50M-parameter interval ablation that locates the signal-rich window around $t=0.5$.","marker":"Merity et al. (2017)"}],"fun_headline_variants":["Masked diffusion skips grokking with a tunable implicit regularizer","Tune mask rate to skip grokking and gain 8.8% on 8B","k-parity reveals masked diffusion's fast generalization trick","Signal-noise split of masked diffusion loss explains no grokking"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The theory assumes attention can be fixed to uniform without changing the learning dynamics, and all theorems are proved for the resulting two-layer MLP; if attention itself is doing load-bearing work in real masked diffusion transformers, the mechanism would not transfer to the large models.","fun_headline_variants_meta":{"raw":{"variants":["Masked diffusion skips grokking with a tunable implicit regularizer","Tune mask rate to skip grokking and gain 8.8% on 8B","k-parity reveals masked diffusion's fast generalization trick","Signal-noise split of masked diffusion loss explains no grokking"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000868,"raw_usage":{"total_tokens":3852,"prompt_tokens":1127,"completion_tokens":2725,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":743,"completion_tokens_details":{"reasoning_tokens":2645}},"tokens_in":743,"tokens_out":2725,"duration_ms":18206,"temperature":1.0,"reasoning_tokens":2645,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T15:38:18.885694+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train a masked-diffusion transformer on $k$-parity with learnable attention and compare its generalization curve and optimal mask rate with the uniform-attention MLP; a grokking plateau in the learnable-attention run, or an optimal rate far from $1/(k+1)$, would show the attention reduction hides the real mechanism. A language-scale test would hold compute fixed and swap the mask schedule to, say, high-noise intervals during pretraining; if a high-noise schedule matches or beats the mid-range window on knowledge-heavy benchmarks, the claim that $t\\approx0.5$ is the general signal-optimal window would be weakened.","supporting_citations":[{"cited_title":"S., Arriola, M., Gokaslan, A., Marroquin, E","cited_arxiv_id":null,"evidence_quote":"Shows the masked diffusion objective can be rewritten as a weighted average of masked language modeling losses, the structural simplification the paper's decomposition builds on."},{"cited_title":"Pointer sentinel mixture models","cited_arxiv_id":null,"evidence_quote":"Provides the WikiText dataset used in the 50M-parameter interval ablation that locates the signal-rich window around $t=0.5$."}],"review_version":2}