MotifRole-Diff: Risk-Optimal Role-Aware Corruption for Masked Molecular Graph Diffusion
Summary
MotifRole-Diff proposes a role-aware corruption schedule for masked discrete diffusion on molecular graphs, allocating masking rates based on denoising difficulty and graph-level perturbation impact, demonstrating improved validity and reduced FCD on QM9 and MOSES benchmarks.
View Cached Full Text
Cached at: 07/27/26, 07:39 AM
# Risk-Optimal Role-Aware Corruption for Masked Molecular Graph Diffusion
Source: [https://arxiv.org/html/2607.21634](https://arxiv.org/html/2607.21634)
###### Abstract
Masked discrete diffusion for molecular graph generation typically applies a uniform corruption schedule to all tokens in a lossless graph\-to\-sequence representation, implicitly treating structurally heterogeneous molecular components as equally difficult and equally important to reconstruct\. However, different molecular graph token roles exhibit substantial variation in denoising difficulty and their influence on the decoded molecule, motivating role\-specific corruption strategies\. We introduceMotifRole\-Diff, a role\-aware corruption process that allocates masking rates according to empirically measured denoising difficulty and graph\-level perturbation impact while preserving the model architecture, clean sequence space, and lossless molecular\-graph decoder\. We formulate schedule selection as the risk\-optimal allocation of a fixed masking budget across token roles\. Our theorem characterizes optimality for the modeled role\-weighted residual risk, while downstream generation performance is evaluated empirically\. Under matched architecture, training budget, and sampling compute,MotifRole\-Diffimproves validity on QM9 from 0\.905 to 0\.944 while reducing FCD from 1\.701 to 1\.609, and on MOSES improves validity from 0\.920 to 0\.938 while reducing FCD from 2\.125 to 1\.850\. Role\-wise diagnostics further show improved reconstruction across molecular graph token categories\. Together, these matched\-compute results indicate that structurally informed corruption is a more effective masking strategy than uniform schedules for serialized molecular graph diffusion\.
## Introduction
Generating valid, novel, and diverse molecular graphs is a fundamental problem in computational molecular discovery, and diffusion models have become one of the strongest paradigms for molecular graph generation\(Jo, Lee, and Hwang[2022a](https://arxiv.org/html/2607.21634#bib.bib21); Vignac et al\.[2023a](https://arxiv.org/html/2607.21634#bib.bib35); Chen et al\.[2023](https://arxiv.org/html/2607.21634#bib.bib6)\)\. Masked discrete language modeling \(MDLM\)\(Sahoo et al\.[2024](https://arxiv.org/html/2607.21634#bib.bib31)\), built on discrete denoising diffusion\(Austin et al\.[2021](https://arxiv.org/html/2607.21634#bib.bib1); Campbell et al\.[2022](https://arxiv.org/html/2607.21634#bib.bib4)\), provides a scalable formulation by serializing molecular graphs into token sequences, denoising them in parallel, and reconstructing the original graph through a lossless decoder, combining exact reconstruction with the efficiency of sequence modeling\.
This approach depends fundamentally on graph serialization\. Prior work has shown that serialization order influences sequence\-model efficiency\(Diamant et al\.[2023](https://arxiv.org/html/2607.21634#bib.bib9); Jang, Lee, and Ahn[2024](https://arxiv.org/html/2607.21634#bib.bib17)\), while hierarchical molecular generators demonstrate the importance of preserving chemically meaningful motifs\(Jin, Barzilay, and Jaakkola[2018a](https://arxiv.org/html/2607.21634#bib.bib18),[2020](https://arxiv.org/html/2607.21634#bib.bib20)\)\. Existing masked discrete diffusion models, however, apply an identical corruption schedule to every token position\(Austin et al\.[2021](https://arxiv.org/html/2607.21634#bib.bib1); Sahoo et al\.[2024](https://arxiv.org/html/2607.21634#bib.bib31)\), implicitly assuming that all serialized tokens are equally difficult to reconstruct and equally important for recovering the molecular graph\.
That assumption is often incorrect\. Serialized molecular graphs contain structurally distinct token roles: boundary markers, serialization syntax, motif\-interior atoms and bonds, and interface tokens connecting substructures\. Errors on these roles are not equally consequential, since some are reconstructed reliably with little effect on the decoded graph while others are harder to recover or disproportionately disrupt validity and connectivity, yet existing graph\-to\-sequence diffusion methods do not ask whether corruption should depend on these differences\. We call this property*role heterogeneity*, and hypothesize that masked molecular diffusion is therefore not only*order\-sensitive*but also*role\-sensitive*: uniform corruption is appropriate only when roles have comparable structural importance, and otherwise corruption should follow measured structural risk\.
To investigate this hypothesis, we introduceMotifRole\-Diff, a role\-sensitive framework for masked discrete molecular diffusion\. We extend the lossless SENT representation\(Chen, Krimmel, and Borgwardt[2026](https://arxiv.org/html/2607.21634#bib.bib7)\)with the motif\-aware serialization*mSENT*, exposing four structural token roles while preserving decoder compatibility\. We quantify each role’s reconstruction difficulty and graph\-level consequence using controlled perturbation experiments, then derive a role\-aware corruption process by allocating a fixed masking budget according to measured structural risk\. The resulting absorbing diffusion formulation generalizes standard MDLM and reduces to uniform corruption when all token roles possess identical structural criticality\.
Our hypothesis gives rise to four research questions\.\(RQ1\)Do serialized token roles differ significantly in reconstruction difficulty and graph\-level structural consequence?\(RQ2\)Can role\-aware corruption be derived as a principled constrained allocation of a fixed masking budget?\(RQ3\)Does empirically informed role\-aware corruption outperform uniform masking, and do the gains stem from informed structural allocation rather than mere non\-uniformity?\(RQ4\)Does motif\-aware serialization improve motif locality while preserving the lossless decoding guarantees of SENT?
Throughout, architecture, optimization, compute, and expected masking rate are held constant so that differences arise only from serialization and corruption design\. Our contributions are fourfold:
1. 1\.Characterizing role heterogeneity in serialized molecular diffusion \(RQ1\)\.We identify structural token roles as an overlooked property of graph serialization and show that they differ systematically in reconstruction difficulty and graph\-level structural consequence\.
2. 2\.A principled role\-aware corruption process \(RQ2\)\.We formulate corruption scheduling as constrained structural\-risk minimization under a fixed masking budget, yielding a role\-aware absorbing diffusion process that generalizes standard MDLM and recovers uniform corruption as a special case\.
3. 3\.Controlled evaluation of role\-aware corruption \(RQ3\)\.Under matched settings,MotifRole\-Diffimproves generation quality over uniform MDLM at equal compute; because the schedule is derived from measured criticalities rather than set arbitrarily, the gain reflects informed structural allocation rather than non\-uniform corruption alone\.
4. 4\.A motif\-aware role\-preserving serialization \(RQ4\)\.We introduce mSENT, a motif\-aware extension of SENT that preserves the original lossless decoder while improving motif locality and exposing four interpretable structural token roles\.
## Preliminaries and Problem Formulation
#### Notation\.
A molecular graphG=\(V,E,X,A\)G=\(V,E,X,A\)has atomsVV, bondsEE, featuresXX, and adjacencyAA\. A lossless serializationT\(⋅;π\)T\(\\cdot\\,;\\pi\)mapsGGunder a traversal orderπ\\pito a token sequencezπ=T\(G;π\)z^\{\\pi\}=T\(G;\\pi\)with decoderD\(zπ\)≅GD\(z^\{\\pi\}\)\\cong Gand\|zπ\|=O\(\|V\|\+\|E\|\)\|z^\{\\pi\}\|=O\(\|V\|\+\|E\|\)\. The motif\-aware tokenizer assigns atoms to motif blocksμ:V→\{1,…,K\}\\mu:V\\to\\\{1,\\dots,K\\\}and tags each token with a roleri∈ℛ=\{special,syntax,interior,interface\}r\_\{i\}\\in\\mathcal\{R\}=\\\{\\text\{special\},\\text\{syntax\},\\text\{interior\},\\text\{interface\}\\\}\(Table[1](https://arxiv.org/html/2607.21634#Sx2.T1)\);πr\\pi\_\{r\}denotes the fraction of role\-rrtokens\.
Table 1:Token roles induced by motif\-aware serialization\.RoleMeaningConsequence of errorSpecialcontrol /\[MASK\]tokensinvalid boundary or paddingSyntaxgrammar, reset, delimitersmalformed or undecodable graphInterioratoms/bonds within one motiflocal motif corruption; occasional scaffold / atom\-bond changeInterfacecross\-motif attachment tokensattachment ambiguity and possible connectivity error; highest single\-token perturbation impact
#### Masked discrete diffusion\.
An absorbing forward process replaces each token of a clean sequencez0z\_\{0\}by\[MASK\]with probabilityαt\\alpha\_\{t\}at diffusion timett, and a denoiserpθp\_\{\\theta\}reconstructs the masked tokens under the absorbing\-state ELBO,
ℒMDLM=𝔼z0,t,Mt\[w\(t\)∑i∈Mt−logpθ\(z0\(i\)∣zt,t\)\]\.\\mathcal\{L\}\_\{\\text\{MDLM\}\}=\\mathbb\{E\}\_\{z\_\{0\},t,M\_\{t\}\}\\\!\\Big\[w\(t\)\\\!\\\!\\sum\_\{i\\in M\_\{t\}\}\\\!\-\\log p\_\{\\theta\}\\big\(z\_\{0\}^\{\(i\)\}\\mid z\_\{t\},t\\big\)\\Big\]\.\(1\)The rateαt\\alpha\_\{t\}is token\-agnostic: every position is corrupted identically regardless of role, and a confidence\-based sampler reveals tokens in the reverse process\.
#### Role exposure and budget\.
For a role\-dependent schedule\{αt\(r\)\}\\\{\\alpha\_\{t\}^\{\(r\)\}\\\}, letρr=𝔼t\[αt\(r\)\]\\rho\_\{r\}=\\mathbb\{E\}\_\{t\}\[\\alpha\_\{t\}^\{\(r\)\}\]be the average exposure of rolerrandρ¯=∑rπrρr\\bar\{\\rho\}=\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}the total masking budget\. Uniform MDLM is the caseρr≡ρ¯\\rho\_\{r\}\\equiv\\bar\{\\rho\}\. Fixingρ¯\\bar\{\\rho\}equalizes the expected number of masked tokens, so schedules sharingρ¯\\bar\{\\rho\}are matched in budget and, under a shared backbone, in compute\.
#### Problem statement\.
Givenℛ\\mathcal\{R\}, the frequencies\{πr\}\\\{\\pi\_\{r\}\\\}, and a fixed budgetρ¯\\bar\{\\rho\}, we seek the role\-dependent schedule minimizing the expected graph\-level reconstruction error at equal budget,
minρ\\displaystyle\\min\_\{\\rho\}𝔼\[Δgraph\(ρ\)\]\\displaystyle\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\(\\rho\)\]\(2\)s\.t\.∑rπrρr=ρ¯,ρmin≤ρr≤ρmax\.\\displaystyle\\textstyle\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}=\\bar\{\\rho\},\\quad\\rho\_\{\\min\}\\leq\\rho\_\{r\}\\leq\\rho\_\{\\max\}\.This reallocates a fixed budget across roles rather than adding noise or capacity; the Theoretical Analysis makes𝔼\[Δgraph\]\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]concrete and derives the risk\-optimal allocation\.
## Role Sensitivity: A Diagnostic
On a lightweight*probe*denoiserpθ0p\_\{\\theta\_\{0\}\}trained with uniform masking, we measure for each rolerrthe difficultyDr=𝔼i:ri=r\[−logpθ0\(z0\(i\)∣zt,t\)\]D\_\{r\}=\\mathbb\{E\}\_\{i:r\_\{i\}=r\}\[\-\\log p\_\{\\theta\_\{0\}\}\(z\_\{0\}^\{\(i\)\}\\mid z\_\{t\},t\)\]and top\-1 errorErE\_\{r\}, witht∼𝒰\(0\.05,0\.95\)t\\sim\\mathcal\{U\}\(0\.05,0\.95\)\. Figure[1](https://arxiv.org/html/2607.21634#Sx4.F1)\(a\) reports a monotone, large ordering,interface≫interior\>syntax≫special\\text\{interface\}\\gg\\text\{interior\}\>\\text\{syntax\}\\gg\\text\{special\}\(probe NLLDr=0\.11/0\.23/0\.35D\_\{r\}=0\.11/0\.23/0\.35nats and top\-1 errorEr=4\.2/8\.8/12\.9%E\_\{r\}=4\.2/8\.8/12\.9\\%for syntax/interior/interface; special≈0\\approx 0\)\. The fully trained uniform MDLM reproduces the*same*role ordering at larger absolute magnitudes \(Table[4](https://arxiv.org/html/2607.21634#Sx7.T4)\), so it is the ordering, not the probe’s absolute scale, that we rely on\. This heterogeneity is the headroomMotifRole\-Diffexploits\.
## MotifRole\-Diff
Motivated by this role heterogeneity, we make the corruption schedule role\-dependent\. Although we evaluate it on molecular graphs, the framework applies to any losslessly serialized graph with identifiable token roles\. Figure[1](https://arxiv.org/html/2607.21634#Sx4.F1)previews the construction on real QM9 values\. LetΛ\(t\)=∫0tβ\(s\)𝑑s\\Lambda\(t\)=\\int\_\{0\}^\{t\}\\beta\(s\)\\,dsbe the base schedule andγr\>0\\gamma\_\{r\}\>0a per\-role exponent\. The absorbing process becomes role\-dependent,
αt\(r\)=1−exp\(−γrΛ\(t\)\),\\alpha\_\{t\}^\{\(r\)\}=1\-\\exp\\\!\\big\(\-\\gamma\_\{r\}\\,\\Lambda\(t\)\\big\),\(3\)soγr<1\\gamma\_\{r\}<1masks rolerrless \(visible longer, revealed earlier\) andγr\>1\\gamma\_\{r\}\>1masks it more\. Uniformγr≡1\\gamma\_\{r\}\\equiv 1recovers MDLM exactly\. The clean sequence, grammar, and decoder are unchanged, soDSENT\(z0\)≅GD\_\{\\text\{SENT\}\}\(z\_\{0\}\)\\cong Gfor any grammar\-validz0z\_\{0\}\(Lemma[1](https://arxiv.org/html/2607.21634#Thmlemma1)\)\.
00\.10\.20\.3051015DrD\_\{r\}\(NLL\)ErE\_\{r\}\(%\)SpecSynIntrIntf\(a\) measured difficultyγr⋆\\gamma\_\{r\}^\{\\star\}\(η=2\\eta\{=\}2\)γ=1\\gamma\{=\}1Intf0\.44Intr1\.07Syn2\.37\(b\) derived rate
Figure 1:Measured difficulty and the derived schedule, on real QM9 values\.\(a\)Measured role difficulty on the probe denoiser: probe NLLDrD\_\{r\}\(teal, left axis\) and top\-1 reconstruction errorErE\_\{r\}\(navy, right axis\)\. Both rank interface hardest and special trivial\.\(b\)Risk\-optimal ratesγr⋆\\gamma\_\{r\}^\{\\star\}from Eq\.[9](https://arxiv.org/html/2607.21634#Sx5.E9)atη=2\\eta\{=\}2\(Table[2](https://arxiv.org/html/2607.21634#Sx5.T2)\): the hardest role \(interface\) is*protected*\(γ<1\\gamma<1\) and easier roles corrupted more \(γ\>1\\gamma\>1\), at a fixed masking budget\. The lossless decoder is untouched\.#### Inverse\-exposure loss weighting\.
Protected \(low\-γ\\gamma\) roles are masked less often and would otherwise be under\-trained\. We reweight the per\-token loss by inverse exposure,
w\(t,ri\)=1Pr\(i∈Mt∣ri\)\+ϵ,w\(t,r\_\{i\}\)=\\frac\{1\}\{\\Pr\(i\\in M\_\{t\}\\mid r\_\{i\}\)\+\\epsilon\},\(4\)which is a no\-op atγr≡1\\gamma\_\{r\}\\equiv 1\(w≡1w\\equiv 1\)\.
#### Role\-aware sampling\.
Motif roles are not directly observed for masked tokens at generation time, but they need not be: the denoiser already induces a distribution over the clean token at each position, and the tokenizer supplies a training token–role priorP\(r∣v\)P\(r\\mid v\)\. We therefore estimate a*soft*role distributionqi\(r\)=∑vpθ\(xi=v∣zt\)P\(r∣v\)q\_\{i\}\(r\)=\\sum\_\{v\}p\_\{\\theta\}\(x\_\{i\}\{=\}v\\mid z\_\{t\}\)\\,P\(r\\mid v\)at every masked positioniiand carry the same predetermined role ratesγr\\gamma\_\{r\}used in training into the reverse process, giving each position an expected exponentγi=∑rqi\(r\)γr\\gamma\_\{i\}=\\sum\_\{r\}q\_\{i\}\(r\)\\,\\gamma\_\{r\}and a role\-aware reverse unmasking probabilitypiunmask=\(as\(i\)−at\(i\)\)/\(1−at\(i\)\)p^\{\\mathrm\{unmask\}\}\_\{i\}=\(a\_\{s\}^\{\(i\)\}\-a\_\{t\}^\{\(i\)\}\)/\(1\-a\_\{t\}^\{\(i\)\}\)withat\(i\)=exp\(−γiΛ\(t\)\)a\_\{t\}^\{\(i\)\}=\\exp\(\-\\gamma\_\{i\}\\Lambda\(t\)\)\. Positions are then revealed by the combined confidence\-and\-schedule scorelogmaxvpθ\(xi=v∣zt\)\+logpiunmask\\log\\max\_\{v\}p\_\{\\theta\}\(x\_\{i\}\{=\}v\\mid z\_\{t\}\)\+\\log p^\{\\mathrm\{unmask\}\}\_\{i\}\. Algorithm[1](https://arxiv.org/html/2607.21634#alg1)states the full procedure; theγr\\gamma\_\{r\}are fixed before sampling, so reverse sampling reuses the derived schedule end\-to\-end rather than deferring it to training alone\.
Algorithm 1Role\-aware sampling with predetermined roleγ\\gamma’s0:denoiser
pθp\_\{\\theta\}; length
LL; reverse steps
\{t→s\}\\\{t\\\!\\to\\\!s\\\}; role rates
γspecial,γsyntax,γinterior,γinterface\\gamma\_\{\\text\{special\}\},\\gamma\_\{\\text\{syntax\}\},\\gamma\_\{\\text\{interior\}\},\\gamma\_\{\\text\{interface\}\}; token–role prior
P\(r∣v\)P\(r\\mid v\); noise schedule
Λ\(⋅\)\\Lambda\(\\cdot\); reveal counts
ktk\_\{t\};\[MASK\],\[BOS\],\[EOS\]; clip floor
ϵ\\epsilon
0:generated graph
G^\\hat\{G\}
1:
z←\[\[BOS\],\[MASK\],…,\[MASK\],\[EOS\]\]z\\leftarrow\[\\texttt\{\[BOS\]\},\\texttt\{\[MASK\]\},\\dots,\\texttt\{\[MASK\]\},\\texttt\{\[EOS\]\}\]
2:foreach reverse step
t→st\\to sdo
3:
M←\{M\\leftarrow\\\{masked positions in
z\}∖\{BOS,EOS\}z\\\}\\setminus\\\{\\text\{BOS\},\\text\{EOS\}\\\}
4:if
M=∅M=\\emptysetthen
5:break
6:endif
7:
probs←pθ\(⋅∣z,t\)\\mathrm\{probs\}\\leftarrow p\_\{\\theta\}\(\\cdot\\mid z,t\)\{clean\-token probabilities\}
8:for all
i∈Mi\\in Mdo
9:
qi\(r\)←∑vprobs\[i\]\[v\]P\(r∣v\)q\_\{i\}\(r\)\\leftarrow\\sum\_\{v\}\\mathrm\{probs\}\[i\]\[v\]\\,P\(r\\mid v\)for each role
rr
10:
γi←∑rqi\(r\)γr\\gamma\_\{i\}\\leftarrow\\sum\_\{r\}q\_\{i\}\(r\)\\,\\gamma\_\{r\}
11:
at←e−γiΛ\(t\)a\_\{t\}\\leftarrow e^\{\-\\gamma\_\{i\}\\Lambda\(t\)\},
as←e−γiΛ\(s\)a\_\{s\}\\leftarrow e^\{\-\\gamma\_\{i\}\\Lambda\(s\)\}
12:
piunmask←clip\(\(as−at\)/\(1−at\),ϵ,1\)p^\{\\mathrm\{unmask\}\}\_\{i\}\\leftarrow\\mathrm\{clip\}\\\!\\big\(\(a\_\{s\}\-a\_\{t\}\)/\(1\-a\_\{t\}\),\\,\\epsilon,\\,1\\big\)
13:
confi←maxvprobs\[i\]\[v\]\\mathrm\{conf\}\_\{i\}\\leftarrow\\max\_\{v\}\\mathrm\{probs\}\[i\]\[v\]
14:
scorei←logconfi\+logpiunmask\\mathrm\{score\}\_\{i\}\\leftarrow\\log\\mathrm\{conf\}\_\{i\}\+\\log p^\{\\mathrm\{unmask\}\}\_\{i\}
15:endfor
16:
S←S\\leftarrowtop\-
ktk\_\{t\}positions of
MMranked by
scorei\\mathrm\{score\}\_\{i\}
17:for all
i∈Si\\in Sdo
18:
z\[i\]←argmaxvprobs\[i\]\[v\]z\[i\]\\leftarrow\\arg\\max\_\{v\}\\mathrm\{probs\}\[i\]\[v\]\(greedy\)orsample
∼probs\[i\]\\sim\\mathrm\{probs\}\[i\]
19:endfor
20:unselected positions remain\[MASK\]
21:endfor
22:
G^←decodeSENT/mSENT\(z\)\\hat\{G\}\\leftarrow\\mathrm\{decode\}\_\{\\text\{SENT/mSENT\}\}\(z\); sanitize
G^\\hat\{G\}with RDKit
23:return
G^\\hat\{G\}
#### Computational overhead\.
MotifRole\-Diffreparameterizes the forward corruption rather than adding a module: identical3\.323\.32M parameters, no extra forward passes, and only twoO\(L\)O\(L\)table operations per training step \(the per\-token exponentγri\\gamma\_\{r\_\{i\}\}and the inverse\-exposure weight\) and per reverse step, plus a one\-time schedule solve over\|ℛ\|=4\|\\mathcal\{R\}\|\{=\}4roles\. Wall\-clock therefore matches MDLM within run\-to\-run noise \(∼1\.25\\sim\\\!1\.25min/epoch;87\.487\.4vs\.91\.891\.8s per10001000samples on one TITAN RTX\)\. Full comparisons appear in Supplementary Table S3\.
## Theoretical Analysis
The theory formalizes one claim: if serialized roles differ in reconstruction difficulty and graph\-level impact, uniform masking is not risk\-optimal, and a role\-aware schedule should protect high\-criticality roles at an unchanged budget\. Serialization is lossless \(TπT\_\{\\pi\}a bijection, soH\(Zπ\)=H\(G\)H\(Z^\{\\pi\}\)=H\(G\)\), so schedules matter only because the denoiser is finite\-capacity and roles are not equally easy to reconstruct\.
###### Lemma 1\(Decoding preserved\)\.
MotifRole\-Diffchanges only the forward corruption kernelq\(zt∣z0,r\)q\(z\_\{t\}\\mid z\_\{0\},r\); the clean sequence space, grammar, and decoder are unchanged\. Hence, for any grammar\-validz0z\_\{0\},DSENT\(z0\)≅GD\_\{\\text\{SENT\}\}\(z\_\{0\}\)\\cong G\. Settingγr≡1\\gamma\_\{r\}\\equiv 1recovers uniform MDLM exactly \(proof in Supplementary Section S1\)\.
#### Role\-decomposed graph risk\.
LetIrI\_\{r\}be the expected graph\-level impact of a role\-rrreconstruction error,ϵr\\epsilon\_\{r\}its residual reconstruction error under a finite\-capacity denoiser, andρr=𝔼t\[αt\(r\)\]\\rho\_\{r\}=\\mathbb\{E\}\_\{t\}\[\\alpha\_\{t\}^\{\(r\)\}\]its average corruption exposure\. A role\-rrtoken harms the graph only when corrupted and then misreconstructed, so
𝔼\[Δgraph\]≤ℛ\(ρ\):=∑r∈ℛπrIrϵrρr\.\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]\\leq\\mathcal\{R\}\(\\rho\):=\\sum\_\{r\\in\\mathcal\{R\}\}\\pi\_\{r\}I\_\{r\}\\epsilon\_\{r\}\\rho\_\{r\}\.\(5\)
The bound separates role frequencyπr\\pi\_\{r\}, graph impactIrI\_\{r\}, and reconstruction errorϵr\\epsilon\_\{r\}, which we estimate by the measured probe NLLDrD\_\{r\}\.
#### Role criticality\.
Define the role criticality
cr:=Irϵr\.c\_\{r\}:=I\_\{r\}\\epsilon\_\{r\}\.\(6\)High\-criticality roles are both hard to reconstruct and damaging when wrong; the product, rather than a sum, keeps criticality low unless*both*factors are large, so a role extreme on only one axis does not dominate the allocation\. We use the plug\-in estimateCr=Norm\(D^r\)⋅Norm\(I^r\)C\_\{r\}=\\mathrm\{Norm\}\(\\widehat\{D\}\_\{r\}\)\\cdot\\mathrm\{Norm\}\(\\widehat\{I\}\_\{r\}\)from measured difficultyD^r\\widehat\{D\}\_\{r\}and impactI^r\\widehat\{I\}\_\{r\}\.
#### Corruption as constrained allocation\.
Uniform MDLM assigns every role the same exposureρ¯\\bar\{\\rho\}; we instead minimize the criticality\-weighted exposure at the same budget:
minρ∑rπrcrρrs\.t\.∑rπrρr=ρ¯,ρmin≤ρr≤ρmax\.\\min\_\{\\rho\}\\sum\_\{r\}\\pi\_\{r\}c\_\{r\}\\rho\_\{r\}\\quad\\text\{s\.t\.\}\\quad\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}=\\bar\{\\rho\},\\qquad\\rho\_\{\\min\}\\leq\\rho\_\{r\}\\leq\\rho\_\{\\max\}\.\(7\)
###### Theorem 1\(Risk\-optimal exposure allocation\)\.
Letcr:=Irϵr≥0c\_\{r\}:=I\_\{r\}\\epsilon\_\{r\}\\geq 0\. Over the fixed\-budget feasible set in Eq\.[7](https://arxiv.org/html/2607.21634#Sx5.E7), the minimizer has the water\-filling form
ρr⋆=\{ρmin,cr\>λ\(protect\),ρmax,cr<λ\(corrupt more\),∈\[ρmin,ρmax\],cr=λ,\\rho\_\{r\}^\{\\star\}=\\begin\{cases\}\\rho\_\{\\min\},&c\_\{r\}\>\\lambda\\quad\(\\text\{protect\}\),\\\\ \\rho\_\{\\max\},&c\_\{r\}<\\lambda\\quad\(\\text\{corrupt more\}\),\\\\ \\in\[\\rho\_\{\\min\},\\rho\_\{\\max\}\],&c\_\{r\}=\\lambda,\\end\{cases\}for a budget multiplierλ\\lambda\. Consequently: \(i\)ρr⋆\\rho\_\{r\}^\{\\star\}is non\-increasing incrc\_\{r\}; \(ii\) uniform MDLM is optimal iffcrc\_\{r\}is constant across roles; \(iii\) if criticalities are heterogeneous and the feasible box permits a nonzero exposure shift from a higher\- to a lower\-criticality role, thenℛ\(ρ⋆\)<ℛ\(ρuni\)\\mathcal\{R\}\(\\rho^\{\\star\}\)<\\mathcal\{R\}\(\\rho^\{\\mathrm\{uni\}\}\); and \(iv\) sinceρ\(γ\)=𝔼t\[1−exp\(−γΛ\(t\)\)\]\\rho\(\\gamma\)=\\mathbb\{E\}\_\{t\}\[1\-\\exp\(\-\\gamma\\Lambda\(t\)\)\]is strictly increasing inγ\\gamma, the optimal exponentsγr⋆=ρ−1\(ρr⋆\)\\gamma\_\{r\}^\{\\star\}=\\rho^\{\-1\}\(\\rho\_\{r\}^\{\\star\}\)inherit the same ordering\.
*Proof sketch\.*The objective is linear on the budget box, so KKT stationarity gives the water\-filling vertex and monotonicity ofρ\(γ\)\\rho\(\\gamma\)transfers the ordering to the exponents; full proof in Supplementary Section S1\.
###### Corollary 1\(Uniform null\)\.
Ifcrc\_\{r\}is constant the optimum is uniform, soMotifRole\-Diffcannot improve the modeled risk; conversely, any strict reduction in Eq\.[5](https://arxiv.org/html/2607.21634#Sx5.E5)certifies heterogeneous criticality\.
#### Smooth risk\-optimal schedule\.
The threshold rule of Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)is too sharp for training, so we add a convex penalty against large deviations from uniform MDLM:
minρ∑rπrcrρr\+1η∑rπrKL\(ρr∥ρ¯\)s\.t\.∑rπrρr=ρ¯\.\\min\_\{\\rho\}\\sum\_\{r\}\\pi\_\{r\}c\_\{r\}\\rho\_\{r\}\+\\frac\{1\}\{\\eta\}\\sum\_\{r\}\\pi\_\{r\}\\mathrm\{KL\}\(\\rho\_\{r\}\\\|\\bar\{\\rho\}\)\\quad\\text\{s\.t\.\}\\quad\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}=\\bar\{\\rho\}\.\(8\)
###### Proposition 1\(Smooth risk\-optimal schedule\)\.
The regularized objective in Eq\.[8](https://arxiv.org/html/2607.21634#Sx5.E8)is strictly convex on\(0,1\)\|ℛ\|\(0,1\)^\{\|\\mathcal\{R\}\|\}and has the unique minimizer
ρr⋆=σ\(σ−1\(ρ¯\)\+η\(λ−cr\)\),\\rho\_\{r\}^\{\\star\}=\\sigma\\\!\\left\(\\sigma^\{\-1\}\(\\bar\{\\rho\}\)\+\\eta\(\\lambda\-c\_\{r\}\)\\right\),\(9\)whereλ\\lambdaenforces the budget constraint andσ\\sigmais the logistic sigmoid\. Moreoverρr⋆\\rho\_\{r\}^\{\\star\}is strictly decreasing in the criticalitycrc\_\{r\}\.
Higher\-criticality roles thus receive lower exposure, andη\\etainterpolates from uniform MDLM \(η→0\\eta\\to 0\) to the threshold rule of Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)\(η→∞\\eta\\to\\infty\), setting the magnitude of reallocation, not its direction\.
Two further guarantees are deferred to Supplementary Section S1 \(Theorems S1 and S2, with all proofs\): the smooth allocation is stable to noisy criticality estimates \(maxr\|ρ^r⋆−ρr⋆\|=O\(ηδ\)\\max\_\{r\}\|\\widehat\{\\rho\}\_\{r\}^\{\\star\}\-\\rho\_\{r\}^\{\\star\}\|=O\(\\eta\\delta\)formaxr\|c^r−cr\|≤δ\\max\_\{r\}\|\\widehat\{c\}\_\{r\}\-c\_\{r\}\|\\leq\\delta\), and a Rademacher bound shows that at matched architecture, data size, and training budget, schedule comparisons are governed by the empirical role\-weighted riskℛ^\(ρ\)=∑rπrI^rϵ^r\(ρr\)\\widehat\{\\mathcal\{R\}\}\(\\rho\)=\\sum\_\{r\}\\pi\_\{r\}\\widehat\{I\}\_\{r\}\\widehat\{\\epsilon\}\_\{r\}\(\\rho\_\{r\}\)\.
#### Robustness toη\\eta\.
The only free hyperparameter isη\\eta; because Eq\.[9](https://arxiv.org/html/2607.21634#Sx5.E9)is monotone incrc\_\{r\}, Table[2](https://arxiv.org/html/2607.21634#Sx5.T2)confirms interface stays protected across datasets and everyη\>0\\eta\>0\.
Table 2:Sensitivity of the derivedγr\\gamma\_\{r\}toη\\eta\(Eq\.[9](https://arxiv.org/html/2607.21634#Sx5.E9)\), re\-derived from the measured criticalities without retraining\. For allη\>0\\eta\>0, the interface role is protected \(γr<1\\gamma\_\{r\}<1\), whereas the interior and syntax roles receive greater corruption \(γr\>1\\gamma\_\{r\}\>1\)\. Thus,η\\etacontrols the magnitude of the allocation rather than its direction\. The settingη=0\\eta=0corresponds to uniform MDLM; we useη=2\\eta=2\(highlighted\)\.QM9MOSESGuacaMolη\\etaInteriorInterfaceSyntaxInteriorInterfaceSyntaxInteriorInterfaceSyntax0\.00\.01\.001\.001\.001\.001\.001\.001\.001\.001\.000\.50\.51\.020\.811\.241\.080\.721\.051\.060\.781\.051\.01\.01\.030\.661\.541\.170\.521\.091\.130\.611\.112\.02\.01\.070\.442\.371\.340\.271\.171\.260\.371\.224\.04\.01\.130\.254\.001\.590\.251\.221\.490\.251\.408\.08\.01\.220\.254\.001\.790\.251\.051\.680\.251\.49
## Experimental Setup
#### Datasets\.
We evaluate on QM9\(Ramakrishnan et al\.[2014](https://arxiv.org/html/2607.21634#bib.bib30)\)\(molecules up to nine heavy atoms\) and MOSES\(Polykovskiy et al\.[2020](https://arxiv.org/html/2607.21634#bib.bib28)\)\(a larger, more diverse drug\-like set\), and use GuacaMol\(Brown et al\.[2019](https://arxiv.org/html/2607.21634#bib.bib3)\)as a third benchmark to test whether the derived schedule and role\-wise diagnostics generalize beyond the two generation datasets\. All are serialized losslessly by the motif\-aware tokenizer with the per\-token role labels that drive the schedule\.
#### Backbone and protocol\.
All conditions share one masked discrete\-diffusion backbone under matched architecture, optimizer, and training budget, so differences reflect the corruption and sampling schedule rather than capacity or training duration\.MotifRole\-Diffsamples with the role\-aware confidence sampler \(Algorithm[1](https://arxiv.org/html/2607.21634#alg1)\) using the inferred token\-role distribution and role\-specific rates, whereas uniform MDLM uses the standard confidence sampler with a single token\-agnostic schedule\. Unless noted, we generate5,0005\{,\}000samples per condition and decode with the same lossless decoder\.
#### Metrics\.
We report validity \(RDKit\-sanitizable fraction\), uniqueness, novelty, atom stability, QED, and role\-wise reconstruction NLL and top\-1 error, with distributional fidelity measured by Fréchet ChemNet Distance \(FCD; lower is better\) against a*held\-out*reference split\. Complete architecture, optimization, schedule\-construction, metric, hardware, and seed details are provided in Supplementary Section S2\.
## Results
We organize the empirical evaluation around the four research questions posed in the Introduction, treating uniform MDLM as the flat\-criticality special case ofMotifRole\-Diffthroughout\.
### Are Token Roles Heterogeneous? \(RQ1\)
Role\-aware corruption is justified only if the token roles are not interchangeable\. On a uniformly masked probe denoiser the roles separate sharply in reconstruction difficulty \(Table[4](https://arxiv.org/html/2607.21634#Sx7.T4)\): interface is the hardest role on every dataset \(QM9 NLL0\.4970\.497and top\-1 error0\.1660\.166, against0\.1200\.120and0\.0450\.045for syntax; MOSES interface NLL0\.7120\.712; GuacaMol interface NLL0\.9030\.903\), with interior consistently in between, and Figure[1](https://arxiv.org/html/2607.21634#Sx4.F1)\(a\) shows the same ordering on the probe\. The single\-token perturbation impact \(Table[3](https://arxiv.org/html/2607.21634#Sx7.T3)\) ranks the roles in the same direction: interface is the highest\-impact role \(Ir=0\.759I\_\{r\}=0\.759, 95% CI\[0\.742,0\.776\]\[0\.742,0\.776\]\), above interior \(0\.6200\.620\) and syntax \(0\.5660\.566\), because corrupting a cross\-motif attachment token most often produces an invalid or disconnected graph\. The two axes agree that interface is most critical, though difficulty separates the roles far more strongly \(roughly fourfold\) than impact \(∼1\.3×\\sim\\\!1\.3\\times\)\. Roles thus differ significantly along both axes on all three datasets, answering RQ1\.
Table 3:Single\-token perturbation impact by role \(QM9, 2430 trials/role, 95% CI\)\. Impact=P\(invalid or disconnected\)=P\(\\text\{invalid or disconnected\}\)\. The ordering is interface\>\>interior\>\>syntax: interface is the highest\-impact role, because corrupting a cross\-motif attachment token most often yields an invalid or disconnected graph\.RoleImpact95% CIInterface0\.759\\mathbf\{0\.759\}\[0\.742,0\.776\]\[0\.742,\\,0\.776\]Interior0\.61980\.6198\[0\.605,0\.632\]\[0\.605,\\,0\.632\]Syntax0\.5660\.566\[0\.552,0\.580\]\[0\.552,\\,0\.580\]
### Is the Schedule Derived, Not Tuned? \(RQ2\)
Given heterogeneous criticalities, does the risk\-optimal rule of Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)yield a usable schedule rather than a hand\-set one? Table[2](https://arxiv.org/html/2607.21634#Sx5.T2)instantiates Eq\.[7](https://arxiv.org/html/2607.21634#Sx5.E7): from measured difficultyDrD\_\{r\}and impactIrI\_\{r\}we formCr=Norm\(Dr\)Norm\(Ir\)C\_\{r\}=\\mathrm\{Norm\}\(D\_\{r\}\)\\,\\mathrm\{Norm\}\(I\_\{r\}\)and solve forγr\\gamma\_\{r\}\. On QM9, MOSES, and GuacaMol the top\-criticality role is interface, so at everyη\>0\\eta\>0the schedule protects it \(γinterface<1\\gamma\_\{\\text\{interface\}\}<1; e\.g\.0\.440\.44,0\.270\.27, and0\.370\.37respectively atη=2\\eta=2\) and corrupts the lower\-criticality interior and syntax roles more at fixed budget\. The allocation direction agrees across all three datasets, indicatingCrC\_\{r\}tracks a structural signal rather than dataset noise\. As established in RQ1, difficulty dominates the criticality ranking, so interface is protected because it is both hardest to reconstruct and highest\-impact\. The schedule is thus derived rather than tuned, answering RQ2\.
Table 4:Role\-wise reconstruction performance under uniform evaluation masking \(10001000graphs/checkpoint\)\. We report negative log\-likelihood \(NLL\) and top\-1 reconstruction error for each structural role on QM9, MOSES, and GuacaMol\.DatasetMethodRoleNLL↓\\downarrowTop\-1 err\.↓\\downarrowQM9Uniform MDLMinterior0\.2610\.095interface0\.4970\.166syntax0\.1200\.045MotifRole\-Diffinterior0\.2220\.085interface0\.3540\.132syntax0\.0970\.039MOSESUniform MDLMinterior0\.4060\.141interface0\.7120\.247syntax0\.5900\.201MotifRole\-Diffinterior0\.3900\.136interface0\.7080\.202syntax0\.2710\.105GuacamolUniform MDLMinterior0\.5280\.180interface0\.9030\.292syntax0\.4350\.166MotifRole\-Diffinterior0\.4800\.167interface0\.8990\.286syntax0\.3140\.118
### Does It Beat Uniform Masking? \(RQ3\)
Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)predicts that the heterogeneous criticality established in RQ1 makes uniform masking suboptimal\. Holding architecture, optimizer, training budget, and compute fixed isolates the corruption and sampling schedule as the only difference from uniform MDLM\. The derived schedule improves validity on all three datasets \(QM90\.905→0\.9440\.905\\\!\\to\\\!0\.944, MOSES0\.920→0\.9380\.920\\\!\\to\\\!0\.938, GuacaMol0\.787→0\.8410\.787\\\!\\to\\\!0\.841\) and lowers FCD wherever it is measured \(QM91\.701→1\.6091\.701\\\!\\to\\\!1\.609, MOSES2\.125→1\.8502\.125\\\!\\to\\\!1\.850; FCD and novelty are not computed on GuacaMol\)\. Uniqueness improves on MOSES and stays high on QM9 and GuacaMol, while the small novelty dip on QM9 and MOSES reflects the fidelity–diversity trade\. Figure[2](https://arxiv.org/html/2607.21634#Sx7.F2)shows these gains are stable across five seeds, with tight, largely non\-overlapping intervals\. Since the schedule is derived from measured criticalities rather than chosen arbitrarily \(RQ2\), this consistent improvement across three datasets is in line with the prediction and reflects informed structural allocation\. A matched\-compute staged ablation \(Supplementary Table S2\) separates the mSENT tokenizer gain from the additional gain of the full role\-aware configuration; since that configuration enables the derived schedule and inverse\-exposure weighting jointly, we do not isolate their individual effects\.
00\.250\.50\.751012ValidUniqueNovelFCDrate↑\\uparrowFCD↓\\downarrowQM900\.250\.50\.751012ValidUniqueNovelFCDrate↑\\uparrowFCD↓\\downarrowMOSES00\.250\.50\.751012n/an/aValidUniqueNovelFCDrate↑\\uparrowFCD↓\\downarrowGuacaMolUniform MDLMMotifRole\-Diff\(risk\-optimal\)
Figure 2:Multi\-seed robustnesson QM9, MOSES, and GuacaMol \(55seeds,50005000samples/seed; whiskers are95%95\\%confidence intervals\)\. Validity, uniqueness, and novelty \(*left*axis, higher is better\) and FCD \(*right*axis, lower is better\) for uniform MDLM \(grey\) vs\.MotifRole\-Diff\(amber\)\.MotifRole\-Diffimproves validity on*all three*datasets and FCD on QM9 and MOSES, with tight, largely non\-overlapping intervals; the novelty dip on QM9 and MOSES reflects the fidelity–diversity trade discussed in the text\. Novelty and FCD are not computed on GuacaMol \(marked n/a\), where uniqueness is already saturated\. FCD is computed against the held\-out test reference\. Exact five\-seed means and95%95\\%confidence intervals are reported in Supplementary Table S1\.#### Stability and drug\-likeness\.
Table[5](https://arxiv.org/html/2607.21634#Sx7.T5)completes the metric set:MotifRole\-Diffmatches or improves atom stability \(QM90\.946→0\.9550\.946\\\!\\to\\\!0\.955; MOSES0\.902→0\.9150\.902\\\!\\to\\\!0\.915; GuacaMol0\.796→0\.8390\.796\\\!\\to\\\!0\.839\) and QED \(QM90\.419→0\.4550\.419\\\!\\to\\\!0\.455; MOSES0\.822→0\.8370\.822\\\!\\to\\\!0\.837; GuacaMol0\.438→0\.6340\.438\\\!\\to\\\!0\.634\), consistent with a regularizer that preserves difficult\-role accuracy, with the largest QED gain on the more diverse GuacaMol set\.
Table 5:Atom stability and QED results on QM9, MOSES, and GuacaMol\. We compare Uniform MDLM andMotifRole\-Diffusing50005000generated samples\.DatasetMethodAtom stab\.↑\\uparrowQED↑\\uparrowQM9Uniform MDLM0\.946±\.0820\.946\{\\scriptscriptstyle\\pm\.082\}0\.419±\.0060\.419\{\\scriptscriptstyle\\pm\.006\}MotifRole\-Diff0\.955±\.0040\.955\{\\scriptscriptstyle\\pm\.004\}0\.455±\.0050\.455\{\\scriptscriptstyle\\pm\.005\}MOSESUniform MDLM0\.902±\.0120\.902\{\\scriptscriptstyle\\pm\.012\}0\.822±\.0090\.822\{\\scriptscriptstyle\\pm\.009\}MotifRole\-Diff0\.915±\.0070\.915\{\\scriptscriptstyle\\pm\.007\}0\.837±\.0030\.837\{\\scriptscriptstyle\\pm\.003\}Guacamol \(100k\)Uniform MDLM0\.796±\.0020\.796\{\\scriptscriptstyle\\pm\.002\}0\.438±\.0050\.438\{\\scriptscriptstyle\\pm\.005\}MotifRole\-Diff0\.839±\.0070\.839\{\\scriptscriptstyle\\pm\.007\}0\.634±\.0060\.634\{\\scriptscriptstyle\\pm\.006\}
#### Where the gain comes from\.
The role\-wise diagnostics \(Table[4](https://arxiv.org/html/2607.21634#Sx7.T4)\) showMotifRole\-Diffredistributes reconstruction capacity toward high\-criticality roles rather than improving all roles equally\. On QM9 all three roles improve, most on the protected interface \(NLL0\.497→0\.3540\.497\\\!\\to\\\!0\.354, top\-10\.166→0\.1320\.166\\\!\\to\\\!0\.132\); on MOSES interface NLL is essentially flat \(0\.712→0\.7080\.712\\\!\\to\\\!0\.708\) while its top\-1 improves \(0\.247→0\.2020\.247\\\!\\to\\\!0\.202\) and interior and syntax improve on both metrics\. GuacaMol shows the same pattern: interface stays near baseline \(NLL0\.903→0\.8990\.903\\\!\\to\\\!0\.899\) while the more frequently corrupted syntax role improves sharply \(NLL0\.435→0\.3140\.435\\\!\\to\\\!0\.314\), and interior improves on both metrics\. One plausible mechanism is that protecting the harder interface tokens, which are corrupted less and therefore remain visible longer during denoising, improves the context available for reconstructing the remaining interior and syntax tokens, so the gains are not confined to the protected role\. This redistribution is exactly the reduction in modeled role\-weighted risk the theory targets; the FCD, validity, and uniqueness gains remain empirical downstream effects\.
### Does mSENT Improve Motif Locality? \(RQ4\)
0246Intra\-edge distortion\+7\.5%Boundary alignment\+7\.3%Same\-motif adjacency\+6\.7%Motif span\+4\.2%Compactness\+2\.9%Switch rate\+1\.3%Fragmentation\+1\.0%Motif transitions\+0\.7%Relative improvement of mSENT over SENT \(%\)Figure 3:Motif\-aware serialization \(mSENT\) improves*every*one of the eight motif\-locality metrics over plain SENT on5,0005\{,\}000tokenized QM9 molecules \(bars sorted by relative gain; each metric oriented so that higher is better\)\. The largest gains are in intra\-motif edge distortion \(22\.27→20\.5922\.27\\\!\\to\\\!20\.59\), boundary alignment \(0\.124→0\.1330\.124\\\!\\to\\\!0\.133\), and same\-motif adjacency \(0\.165→0\.1760\.165\\\!\\to\\\!0\.176\), so motif\-coherent substructures stay contiguous in the token stream\. Metric definitions and complete values are in Supplementary Section S4 and Table S4\.On5,0005\{,\}000tokenized QM9 molecules, mSENT improves all eight motif\-locality metrics over plain SENT \(Fig\.[3](https://arxiv.org/html/2607.21634#Sx7.F3); example values in the caption\), keeping motif\-coherent substructures contiguous in the token stream\. Definitions and complete values are in Supplementary Section S4 \(Table S4\)\.
## Related Work
#### Masked and discrete diffusion\.
Discrete diffusion extends denoising diffusion\(Ho, Jain, and Abbeel[2020](https://arxiv.org/html/2607.21634#bib.bib14); Song et al\.[2020](https://arxiv.org/html/2607.21634#bib.bib34)\)to categorical tokens\. D3PM\(Austin et al\.[2021](https://arxiv.org/html/2607.21634#bib.bib1)\)introduced structured transition kernels and found the*absorbing*\(\[MASK\]\) kernel especially effective for text, with multinomial and argmax diffusion\(Hoogeboom et al\.[2021](https://arxiv.org/html/2607.21634#bib.bib15)\)as the uniform\-kernel counterpart; continuous\-time\(Campbell et al\.[2022](https://arxiv.org/html/2607.21634#bib.bib4)\)and score\-based\(Lou, Meng, and Ermon[2023](https://arxiv.org/html/2607.21634#bib.bib25)\)formulations followed\. MDLM\(Sahoo et al\.[2024](https://arxiv.org/html/2607.21634#bib.bib31)\)and simplified objectives\(Shi et al\.[2024](https://arxiv.org/html/2607.21634#bib.bib33)\)show that a clean absorbing\-state ELBO recovers much of the quality of autoregressive models while retaining parallel generation, the diffusion counterpart of masked language modelling\(Devlin et al\.[2019](https://arxiv.org/html/2607.21634#bib.bib8)\)\. For decoding we adopt confidence\-based samplers\(Ghazvininejad et al\.[2019](https://arxiv.org/html/2607.21634#bib.bib11); Chang et al\.[2022](https://arxiv.org/html/2607.21634#bib.bib5)\)\. All of these corrupt every token with a single token\-agnostic rateαt\\alpha\_\{t\}; we reparameterize only that rate\.
#### Graph diffusion and serialization\.
Native graph diffusion perturbs adjacency and feature tensors: EDP\-GNN\(Niu et al\.[2020](https://arxiv.org/html/2607.21634#bib.bib27)\)and GDSS\(Jo, Lee, and Hwang[2022b](https://arxiv.org/html/2607.21634#bib.bib22)\)use score\-based SDEs, DiGress\(Vignac et al\.[2023a](https://arxiv.org/html/2607.21634#bib.bib35)\)runs discrete diffusion over node and edge categories, 3D\-aware variants target geometry\(Hoogeboom et al\.[2022](https://arxiv.org/html/2607.21634#bib.bib16); Vignac et al\.[2023b](https://arxiv.org/html/2607.21634#bib.bib36)\), and autoregressive models generate nodes and edges sequentially\(You et al\.[2018](https://arxiv.org/html/2607.21634#bib.bib38); Shi et al\.[2020](https://arxiv.org/html/2607.21634#bib.bib32)\); these must handle permutation symmetry and quadratic edge tensors\. We instead*serialize*the graph as strings, from SMILES\(Weininger[1988](https://arxiv.org/html/2607.21634#bib.bib37)\)and SELFIES\(Krenn et al\.[2020](https://arxiv.org/html/2607.21634#bib.bib24)\)to sequence\-model generation\(Gómez\-Bombarelli et al\.[2018](https://arxiv.org/html/2607.21634#bib.bib12)\), while fragment\- and motif\-based methods build from substructures\(Jin, Barzilay, and Jaakkola[2018b](https://arxiv.org/html/2607.21634#bib.bib19),[2020](https://arxiv.org/html/2607.21634#bib.bib20)\)\. Our lossless SENT/mSENT serialization additionally emits a per\-token*role*label from the motif partition, so corruption can depend on structural function while the decoder stays unchanged\.
#### Adaptive corruption schedules\.
Non\-uniform noise is established: variational and improved models reshape the*time*schedule\(Kingma et al\.[2021](https://arxiv.org/html/2607.21634#bib.bib23); Nichol and Dhariwal[2021](https://arxiv.org/html/2607.21634#bib.bib26)\), and continuous relaxations tune per\-dimension noise for discrete data\(Dieleman et al\.[2022](https://arxiv.org/html/2607.21634#bib.bib10)\)\. Closest to us, DiffusionBERT\(He et al\.[2023](https://arxiv.org/html/2607.21634#bib.bib13)\)makes the per\-token mask rate non\-uniform via corpus*frequency*; ours is keyed instead to a token’s*structural role*and its measured*difficulty*and*impact*, which frequency cannot see\. It is also not hand\-designed: it solves a fixed\-budget risk\-allocation problem \(Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)\) whose optimum is a water\-filling rule, with uniform MDLM the flat\-criticality special case \(Corollary[1](https://arxiv.org/html/2607.21634#Thmcorollary1)\)\.
## Conclusion
Serialized molecular\-graph diffusion is not role\-agnostic but*role\-sensitive*: the tokens of a lossless serialization play distinct structural roles that differ systematically in reconstruction difficulty and in how much an error costs the decoded graph\. Because a uniformly masked probe shows these roles separating sharply, we derive the corruption schedule as the risk\-optimal allocation of a fixed masking budget \(Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)\), with uniform MDLM its flat\-criticality special case\. At matched architecture, optimizer, and training budget, the derived schedule improves distributional fidelity and validity on QM9 and MOSES, and mSENT improves every motif\-locality metric over SENT while leaving the lossless decoder unchanged\.MotifRole\-Diffreparameterizes only the corruption process, adding neither noise nor compute, so measured structural risk rather than corpus frequency or a hand\-set curve decides where a graph is corrupted; the theorem certifies only the modeled role\-weighted risk, so the generation gains remain empirical\.
## Limitations and Broader Impact
Our single\-token impact estimator may understate multi\-token interface fragility; the inference\-time role priorP\(r∣v\)P\(r\\mid v\)could be sharpened with context\-conditioned estimation; and the full\-scale MOSES run transfers the QM9\-derived schedule rather than re\-estimating it, so a per\-dataset schedule can only help while absolute FCD still trails full\-training graph\-native baselines\. Theorem[1](https://arxiv.org/html/2607.21634#Thmtheorem1)is risk\-optimal only for the modeled objective and does not guarantee that every generation metric improves\. On broader impact,MotifRole\-Diffonly reparameterizes how a fixed serialization is corrupted and adds no capability beyond the uniform\-MDLM baseline; as with any molecular generator, improved sample quality carries a general dual\-use risk\. All datasets are public and contain no sensitive information\.
## References
- Austin et al\.\(2021\)Austin, J\.; Johnson, D\. D\.; Ho, J\.; Tarlow, D\.; and van den Berg, R\.2021\.Structured denoising diffusion models in discrete state\-spaces\.InAdvances in Neural Information Processing Systems, volume 34, 17981–17993\.
- Bickerton et al\.\(2012\)Bickerton, G\. R\.; Paolini, G\. V\.; Besnard, J\.; Muresan, S\.; and Hopkins, A\. L\.2012\.Quantifying the chemical beauty of drugs\.Nature chemistry4\(2\):90–98\.
- Brown et al\.\(2019\)Brown, N\.; Fiscato, M\.; Segler, M\. H\.; and Vaucher, A\. C\.2019\.Guacamol: benchmarking models for de novo molecular design\.Journal of chemical information and modeling59\(3\):1096–1108\.
- Campbell et al\.\(2022\)Campbell, A\.; Benton, J\.; De Bortoli, V\.; Rainforth, T\.; Deligiannidis, G\.; and Doucet, A\.2022\.A continuous time framework for discrete denoising models\.InAdvances in Neural Information Processing Systems, volume 35\.
- Chang et al\.\(2022\)Chang, H\.; Zhang, H\.; Jiang, L\.; Liu, C\.; and Freeman, W\. T\.2022\.Maskgit: Masked generative image transformer\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition, 11315–11325\.
- Chen et al\.\(2023\)Chen, X\.; He, J\.; Han, X\.; and Liu, L\.\-P\.2023\.Efficient and degree\-guided graph generation via discrete diffusion modeling\.InProceedings of the 40th International Conference on Machine Learning, volume 202 ofProceedings of Machine Learning Research\.
- Chen, Krimmel, and Borgwardt \(2026\)Chen, D\.; Krimmel, M\.; and Borgwardt, K\.2026\.Flatten graphs as sequences: Transformers are scalable graph generators\.Advances in Neural Information Processing Systems38:69071–69109\.
- Devlin et al\.\(2019\)Devlin, J\.; Chang, M\.\-W\.; Lee, K\.; and Toutanova, K\.2019\.Bert: Pre\-training of deep bidirectional transformers for language understanding\.InProceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 \(long and short papers\), 4171–4186\.
- Diamant et al\.\(2023\)Diamant, N\.; Tseng, A\. M\.; Hase, P\.; Bansal, M\.; and Biancalani, T\.2023\.Improving graph generation by restricting graph bandwidth\.InProceedings of the 40th International Conference on Machine Learning, volume 202 ofProceedings of Machine Learning Research\.
- Dieleman et al\.\(2022\)Dieleman, S\.; Sartran, L\.; Roshannai, A\.; Savinov, N\.; Ganin, Y\.; Richemond, P\. H\.; Doucet, A\.; Strudel, R\.; Dyer, C\.; Durkan, C\.; et al\.2022\.Continuous diffusion for categorical data\.arXiv preprint arXiv:2211\.15089\.
- Ghazvininejad et al\.\(2019\)Ghazvininejad, M\.; Levy, O\.; Liu, Y\.; and Zettlemoyer, L\.2019\.Mask\-predict: Parallel decoding of conditional masked language models\.InProceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing \(EMNLP\-IJCNLP\), 6112–6121\.
- Gómez\-Bombarelli et al\.\(2018\)Gómez\-Bombarelli, R\.; Wei, J\. N\.; Duvenaud, D\.; Hernández\-Lobato, J\. M\.; Sánchez\-Lengeling, B\.; Sheberla, D\.; Aguilera\-Iparraguirre, J\.; Hirzel, T\. D\.; Adams, R\. P\.; and Aspuru\-Guzik, A\.2018\.Automatic chemical design using a data\-driven continuous representation of molecules\.ACS central science4\(2\):268–276\.
- He et al\.\(2023\)He, Z\.; Sun, T\.; Tang, Q\.; Wang, K\.; Huang, X\.\-J\.; and Qiu, X\.2023\.Diffusionbert: Improving generative masked language models with diffusion models\.InProceedings of the 61st annual meeting of the association for computational linguistics \(volume 1: Long papers\), 4521–4534\.
- Ho, Jain, and Abbeel \(2020\)Ho, J\.; Jain, A\.; and Abbeel, P\.2020\.Denoising diffusion probabilistic models\.Advances in neural information processing systems33:6840–6851\.
- Hoogeboom et al\.\(2021\)Hoogeboom, E\.; Nielsen, D\.; Jaini, P\.; Forré, P\.; and Welling, M\.2021\.Argmax flows and multinomial diffusion: Learning categorical distributions\.Advances in neural information processing systems34:12454–12465\.
- Hoogeboom et al\.\(2022\)Hoogeboom, E\.; Satorras, V\. G\.; Vignac, C\.; and Welling, M\.2022\.Equivariant diffusion for molecule generation in 3d\.InInternational conference on machine learning, 8867–8887\.PMLR\.
- Jang, Lee, and Ahn \(2024\)Jang, Y\.; Lee, S\.; and Ahn, S\.2024\.A simple and scalable representation for graph generation\.InInternational Conference on Learning Representations\.
- Jin, Barzilay, and Jaakkola \(2018a\)Jin, W\.; Barzilay, R\.; and Jaakkola, T\.2018a\.Junction tree variational autoencoder for molecular graph generation\.InProceedings of the 35th International Conference on Machine Learning, volume 80 ofProceedings of Machine Learning Research\.
- Jin, Barzilay, and Jaakkola \(2018b\)Jin, W\.; Barzilay, R\.; and Jaakkola, T\.2018b\.Junction tree variational autoencoder for molecular graph generation\.InInternational conference on machine learning, 2323–2332\.PMLR\.
- Jin, Barzilay, and Jaakkola \(2020\)Jin, W\.; Barzilay, R\.; and Jaakkola, T\.2020\.Hierarchical generation of molecular graphs using structural motifs\.InProceedings of the 37th International Conference on Machine Learning, volume 119 ofProceedings of Machine Learning Research, 4839–4848\.
- Jo, Lee, and Hwang \(2022a\)Jo, J\.; Lee, S\.; and Hwang, S\. J\.2022a\.Score\-based generative modeling of graphs via the system of stochastic differential equations\.InProceedings of the 39th International Conference on Machine Learning, volume 162 ofProceedings of Machine Learning Research\.
- Jo, Lee, and Hwang \(2022b\)Jo, J\.; Lee, S\.; and Hwang, S\. J\.2022b\.Score\-based generative modeling of graphs via the system of stochastic differential equations\.InInternational conference on machine learning, 10362–10383\.PMLR\.
- Kingma et al\.\(2021\)Kingma, D\.; Salimans, T\.; Poole, B\.; and Ho, J\.2021\.Variational diffusion models\.Advances in neural information processing systems34:21696–21707\.
- Krenn et al\.\(2020\)Krenn, M\.; Häse, F\.; Nigam, A\.; Friederich, P\.; and Aspuru\-Guzik, A\.2020\.Self\-referencing embedded strings \(selfies\): A 100% robust molecular string representation\.Machine Learning: Science and Technology1\(4\):045024\.
- Lou, Meng, and Ermon \(2023\)Lou, A\.; Meng, C\.; and Ermon, S\.2023\.Discrete diffusion modeling by estimating the ratios of the data distribution\.arXiv preprint arXiv:2310\.16834\.
- Nichol and Dhariwal \(2021\)Nichol, A\. Q\., and Dhariwal, P\.2021\.Improved denoising diffusion probabilistic models\.InInternational conference on machine learning, 8162–8171\.PMLR\.
- Niu et al\.\(2020\)Niu, C\.; Song, Y\.; Song, J\.; Zhao, S\.; Grover, A\.; and Ermon, S\.2020\.Permutation invariant graph generation via score\-based generative modeling\.InInternational conference on artificial intelligence and statistics, 4474–4484\.PMLR\.
- Polykovskiy et al\.\(2020\)Polykovskiy, D\.; Zhebrak, A\.; Sanchez\-Lengeling, B\.; Golovanov, S\.; Tatanov, O\.; Belyaev, S\.; Kurbanov, R\.; Artamonov, A\.; Aladinskiy, V\.; Veselov, M\.; et al\.2020\.Molecular sets \(moses\): a benchmarking platform for molecular generation models\.Frontiers in pharmacology11:565644\.
- Preuer et al\.\(2018\)Preuer, K\.; Renz, P\.; Unterthiner, T\.; Hochreiter, S\.; and Klambauer, G\.2018\.Fréchet chemnet distance: a metric for generative models for molecules in drug discovery\.Journal of chemical information and modeling58\(9\):1736–1741\.
- Ramakrishnan et al\.\(2014\)Ramakrishnan, R\.; Dral, P\. O\.; Rupp, M\.; and Von Lilienfeld, O\. A\.2014\.Quantum chemistry structures and properties of 134 kilo molecules\.Scientific data1\(1\):1–7\.
- Sahoo et al\.\(2024\)Sahoo, S\. S\.; Arriola, M\.; Schiff, Y\.; Gokaslan, A\.; Marroquin, E\.; Chiu, J\. T\.; Rush, A\. M\.; and Kuleshov, V\.2024\.Simple and effective masked diffusion language models\.InAdvances in Neural Information Processing Systems, volume 37\.
- Shi et al\.\(2020\)Shi, C\.; Xu, M\.; Zhu, Z\.; Zhang, W\.; Zhang, M\.; and Tang, J\.2020\.Graphaf: a flow\-based autoregressive model for molecular graph generation\.arXiv preprint arXiv:2001\.09382\.
- Shi et al\.\(2024\)Shi, J\.; Han, K\.; Wang, Z\.; Doucet, A\.; and Titsias, M\.2024\.Simplified and generalized masked diffusion for discrete data\.Advances in neural information processing systems37:103131–103167\.
- Song et al\.\(2020\)Song, Y\.; Sohl\-Dickstein, J\.; Kingma, D\. P\.; Kumar, A\.; Ermon, S\.; and Poole, B\.2020\.Score\-based generative modeling through stochastic differential equations\.arXiv preprint arXiv:2011\.13456\.
- Vignac et al\.\(2023a\)Vignac, C\.; Krawczuk, I\.; Siraudin, A\.; Wang, B\.; Cevher, V\.; and Frossard, P\.2023a\.DiGress: Discrete denoising diffusion for graph generation\.InInternational Conference on Learning Representations\.
- Vignac et al\.\(2023b\)Vignac, C\.; Osman, N\.; Toni, L\.; and Frossard, P\.2023b\.Midi: Mixed graph and 3d denoising diffusion for molecule generation\.InJoint European Conference on Machine Learning and Knowledge Discovery in Databases, 560–576\.Springer\.
- Weininger \(1988\)Weininger, D\.1988\.Smiles, a chemical language and information system\. 1\. introduction to methodology and encoding rules\.Journal of chemical information and computer sciences28\(1\):31–36\.
- You et al\.\(2018\)You, J\.; Liu, B\.; Ying, Z\.; Pande, V\.; and Leskovec, J\.2018\.Graph convolutional policy network for goal\-directed molecular graph generation\.Advances in neural information processing systems31\.
This document collects the technical appendices for the main paper\. Section, figure, and equation numbers are prefixed with “S”; references to Lemmas, Theorems, Equations, Tables, and Algorithms without an “S” prefix refer to the main paper\.
## Appendix S1Proofs and Additional Theoretical Results
Throughout,πr\>0\\pi\_\{r\}\>0denotes the fraction of role\-rrtokens,cr:=Irϵr≥0c\_\{r\}:=I\_\{r\}\\epsilon\_\{r\}\\geq 0denotes role criticality, andρ¯∈\(ρmin,ρmax\)\\bar\{\\rho\}\\in\(\\rho\_\{\\min\},\\rho\_\{\\max\}\)denotes the uniform\-MDLM exposure budget\.
###### Information preservation\.
We show the identityH\(Zπ\)=H\(G\)H\(Z^\{\\pi\}\)=H\(G\)stated inline in the main text\. BecauseTπT\_\{\\pi\}is deterministic and invertible with inverse decoderDD, it is a bijection between the support ofGGand the support ofZπ=Tπ\(G\)Z^\{\\pi\}=T\_\{\\pi\}\(G\)\. Therefore, for every graphgg,
Pr\(Zπ=Tπ\(g\)\)=Pr\(G=g\)\.\\Pr\(Z^\{\\pi\}=T\_\{\\pi\}\(g\)\)=\\Pr\(G=g\)\.The two random variables have the same multiset of probability masses, only relabeled\. Shannon entropy is invariant under deterministic bijections, soH\(Zπ\)=H\(G\)H\(Z^\{\\pi\}\)=H\(G\)\. ∎
###### Proof of Lemma 1 \(Decoding preserved\)\.
MotifRole\-Diffmodifies only the forward corruption process:
αt\(r\)=1−exp\(−γrΛ\(t\)\)\.\\alpha\_\{t\}^\{\(r\)\}=1\-\\exp\(\-\\gamma\_\{r\}\\Lambda\(t\)\)\.It does not modify the clean sequence distribution, the grammar, or the SENT decoder\. The reverse model still predicts clean tokensz0\(i\)z\_\{0\}^\{\(i\)\}\. Hence, whenever the generated clean sequencez0z\_\{0\}is grammar\-valid, the lossless decoder reconstructs the same graph object,DSENT\(z0\)≅GD\_\{\\text\{SENT\}\}\(z\_\{0\}\)\\cong G\. Settingγr≡1\\gamma\_\{r\}\\equiv 1recovers the uniform MDLM corruption process exactly\. ∎
###### Derivation of Eq\. 5\.
For a token of rolerr, graph distortion occurs only when the token is corrupted and then incorrectly reconstructed\. Averaged over time, the probability of corruption isρr\\rho\_\{r\}\.*Conditional on corruption*, the residual probability that the finite\-capacity denoiser reconstructs the token incorrectly isϵr\\epsilon\_\{r\}\(estimated from the probe difficultyDrD\_\{r\}\), and conditional on such an error the expected graph\-level distortion is at mostIrI\_\{r\}\. Multiplying these factors and weighting by the role frequencyπr\\pi\_\{r\}, the expected contribution of rolerris bounded by
πrIrϵrρr\.\\pi\_\{r\}\\,I\_\{r\}\\,\\epsilon\_\{r\}\\,\\rho\_\{r\}\.Summing over roles gives the role\-decomposed graph risk of the main text,
𝔼\[Δgraph\]≤∑rπrIrϵrρr=:ℛ\(ρ\),\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]\\leq\\sum\_\{r\}\\pi\_\{r\}\\,I\_\{r\}\\,\\epsilon\_\{r\}\\,\\rho\_\{r\}=:\\mathcal\{R\}\(\\rho\),which is Eq\. 5 of the main paper and the exact objective minimized by Theorem 1\. Hereϵr\\epsilon\_\{r\}is the residual error*conditional on corruption*and does not depend onρr\\rho\_\{r\}; the exposure enters only through the linear factorρr\\rho\_\{r\}\.
*Remark \(exposure\-dependent form\)\.*If the residual error itself varies with exposure,ϵr=ϵr\(ρr\)\\epsilon\_\{r\}=\\epsilon\_\{r\}\(\\rho\_\{r\}\), the same argument yields the more general bound𝔼\[Δgraph\]≤∑rπrIrϵr\(ρr\)\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]\\leq\\sum\_\{r\}\\pi\_\{r\}I\_\{r\}\\epsilon\_\{r\}\(\\rho\_\{r\}\)\. The linear objective above is the special case in whichϵr\\epsilon\_\{r\}is exposure\-independent; the generalization result \(Theorem S2\) is stated for this exposure\-dependent form\. ∎
###### Proof of Theorem 1\.
We minimize
ℛ\(ρ\)=∑rπrcrρr\\mathcal\{R\}\(\\rho\)=\\sum\_\{r\}\\pi\_\{r\}c\_\{r\}\\rho\_\{r\}over the compact feasible set
P=\{ρ:∑rπrρr=ρ¯,ρmin≤ρr≤ρmax\}\.P=\\left\\\{\\rho:\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}=\\bar\{\\rho\},\\;\\rho\_\{\\min\}\\leq\\rho\_\{r\}\\leq\\rho\_\{\\max\}\\right\\\}\.SincePPis nonempty and compact andℛ\\mathcal\{R\}is continuous, a minimizer exists\.
Introduce a Lagrange multiplierλ\\lambdafor the budget constraint and nonnegative multipliersμ¯r,μ¯r\\underline\{\\mu\}\_\{r\},\\overline\{\\mu\}\_\{r\}for the lower and upper box constraints\. The Lagrangian is
ℒ\(ρ,λ,μ¯,μ¯\)\\displaystyle\\mathcal\{L\}\(\\rho,\\lambda,\\underline\{\\mu\},\\overline\{\\mu\}\)=∑rπrcrρr−λ\(∑rπrρr−ρ¯\)\\displaystyle=\\sum\_\{r\}\\pi\_\{r\}c\_\{r\}\\rho\_\{r\}\-\\lambda\\left\(\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}\-\\bar\{\\rho\}\\right\)−∑rμ¯r\(ρr−ρmin\)\+∑rμ¯r\(ρr−ρmax\)\.\\displaystyle\\quad\-\\sum\_\{r\}\\underline\{\\mu\}\_\{r\}\(\\rho\_\{r\}\-\\rho\_\{\\min\}\)\+\\sum\_\{r\}\\overline\{\\mu\}\_\{r\}\(\\rho\_\{r\}\-\\rho\_\{\\max\}\)\.\(S1\)Stationarity gives
πr\(cr−λ\)−μ¯r\+μ¯r=0\.\\pi\_\{r\}\(c\_\{r\}\-\\lambda\)\-\\underline\{\\mu\}\_\{r\}\+\\overline\{\\mu\}\_\{r\}=0\.Ifcr\>λc\_\{r\}\>\\lambda, then stationarity requiresμ¯r\>0\\underline\{\\mu\}\_\{r\}\>0, so complementary slackness givesρr⋆=ρmin\\rho\_\{r\}^\{\\star\}=\\rho\_\{\\min\}\. Ifcr<λc\_\{r\}<\\lambda, thenμ¯r\>0\\overline\{\\mu\}\_\{r\}\>0, soρr⋆=ρmax\\rho\_\{r\}^\{\\star\}=\\rho\_\{\\max\}\. Ifcr=λc\_\{r\}=\\lambda, both box multipliers may vanish and any value in the interval is allowed\. This proves the water\-filling form and monotonicity\.
Ifcr≡cc\_\{r\}\\equiv c, then
ℛ\(ρ\)=c∑rπrρr=cρ¯,\\mathcal\{R\}\(\\rho\)=c\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}=c\\bar\{\\rho\},so every feasible allocation, including uniform MDLM, is optimal\.
Conversely, suppose there exist rolesa,ba,bwithca<cbc\_\{a\}<c\_\{b\}\. Becauseρ¯∈\(ρmin,ρmax\)\\bar\{\\rho\}\\in\(\\rho\_\{\\min\},\\rho\_\{\\max\}\), choose
0<δ≤min\{πa\(ρmax−ρ¯\),πb\(ρ¯−ρmin\)\}\.0<\\delta\\leq\\min\\\{\\pi\_\{a\}\(\\rho\_\{\\max\}\-\\bar\{\\rho\}\),\\pi\_\{b\}\(\\bar\{\\rho\}\-\\rho\_\{\\min\}\)\\\}\.Move exposure from the higher\-criticality rolebbto the lower\-criticality roleaa:
ρa=ρ¯\+δ/πa,ρb=ρ¯−δ/πb,\\rho\_\{a\}=\\bar\{\\rho\}\+\\delta/\\pi\_\{a\},\\qquad\\rho\_\{b\}=\\bar\{\\rho\}\-\\delta/\\pi\_\{b\},and keep all other roles atρ¯\\bar\{\\rho\}\. This preserves the budget and remains inside the box\. The change in risk is
ℛ\(ρ\)−ℛ\(ρuni\)=caδ−cbδ=\(ca−cb\)δ<0\.\\mathcal\{R\}\(\\rho\)\-\\mathcal\{R\}\(\\rho^\{\\mathrm\{uni\}\}\)=c\_\{a\}\\delta\-c\_\{b\}\\delta=\(c\_\{a\}\-c\_\{b\}\)\\delta<0\.Therefore uniform MDLM is not optimal whenever criticalities are heterogeneous and a feasible exposure shift exists\.
Finally,
ρ\(γ\)=𝔼t\[1−exp\(−γΛ\(t\)\)\]\\rho\(\\gamma\)=\\mathbb\{E\}\_\{t\}\[1\-\\exp\(\-\\gamma\\Lambda\(t\)\)\]is strictly increasing because
∂∂γ\(1−exp\(−γΛ\(t\)\)\)=Λ\(t\)exp\(−γΛ\(t\)\)\>0\\frac\{\\partial\}\{\\partial\\gamma\}\\left\(1\-\\exp\(\-\\gamma\\Lambda\(t\)\)\\right\)=\\Lambda\(t\)\\exp\(\-\\gamma\\Lambda\(t\)\)\>0wheneverΛ\(t\)\>0\\Lambda\(t\)\>0\. Thusρ−1\\rho^\{\-1\}exists on the feasible interval, and the ordering of optimal exposures transfers to the ordering of optimal exponentsγr⋆=ρ−1\(ρr⋆\)\\gamma\_\{r\}^\{\\star\}=\\rho^\{\-1\}\(\\rho\_\{r\}^\{\\star\}\)\. ∎
###### Proof of Corollary 1\.
The result follows immediately from Theorem 1\. If all criticalities are equal, the objective is constant on the fixed\-budget feasible set, so uniform MDLM is optimal\. If any strict improvement over uniform is possible, the criticalities cannot all be equal\. ∎
###### Proof of Proposition 1\.
Consider the regularized objective
J\(ρ\)=∑rπrcrρr\+1η∑rπrKL\(ρr∥ρ¯\)\.J\(\\rho\)=\\sum\_\{r\}\\pi\_\{r\}c\_\{r\}\\rho\_\{r\}\+\\frac\{1\}\{\\eta\}\\sum\_\{r\}\\pi\_\{r\}\\mathrm\{KL\}\(\\rho\_\{r\}\\\|\\bar\{\\rho\}\)\.For Bernoulli means,
KL\(ρ∥ρ¯\)=ρlogρρ¯\+\(1−ρ\)log1−ρ1−ρ¯\.\\mathrm\{KL\}\(\\rho\\\|\\bar\{\\rho\}\)=\\rho\\log\\frac\{\\rho\}\{\\bar\{\\rho\}\}\+\(1\-\\rho\)\\log\\frac\{1\-\\rho\}\{1\-\\bar\{\\rho\}\}\.Its second derivative is
d2dρ2KL\(ρ∥ρ¯\)=1ρ\(1−ρ\)\>0,\\frac\{d^\{2\}\}\{d\\rho^\{2\}\}\\mathrm\{KL\}\(\\rho\\\|\\bar\{\\rho\}\)=\\frac\{1\}\{\\rho\(1\-\\rho\)\}\>0,so the objective is strictly convex on\(0,1\)\(0,1\)\. Hence the constrained minimizer is unique\.
The derivative is
ddρKL\(ρ∥ρ¯\)=σ−1\(ρ\)−σ−1\(ρ¯\),\\frac\{d\}\{d\\rho\}\\mathrm\{KL\}\(\\rho\\\|\\bar\{\\rho\}\)=\\sigma^\{\-1\}\(\\rho\)\-\\sigma^\{\-1\}\(\\bar\{\\rho\}\),whereσ−1\\sigma^\{\-1\}is the logit\. The Lagrangian stationarity condition is
cr\+1η\(σ−1\(ρr\)−σ−1\(ρ¯\)\)−λ=0\.c\_\{r\}\+\\frac\{1\}\{\\eta\}\\left\(\\sigma^\{\-1\}\(\\rho\_\{r\}\)\-\\sigma^\{\-1\}\(\\bar\{\\rho\}\)\\right\)\-\\lambda=0\.Solving forρr\\rho\_\{r\}gives
σ−1\(ρr⋆\)=σ−1\(ρ¯\)\+η\(λ−cr\),\\sigma^\{\-1\}\(\\rho\_\{r\}^\{\\star\}\)=\\sigma^\{\-1\}\(\\bar\{\\rho\}\)\+\\eta\(\\lambda\-c\_\{r\}\),and therefore
ρr⋆=σ\(σ−1\(ρ¯\)\+η\(λ−cr\)\)\.\\rho\_\{r\}^\{\\star\}=\\sigma\\\!\\left\(\\sigma^\{\-1\}\(\\bar\{\\rho\}\)\+\\eta\(\\lambda\-c\_\{r\}\)\\right\)\.Since the sigmoid is strictly increasing and the argument is strictly decreasing incrc\_\{r\}, the optimal exposure is strictly decreasing in role criticality\. ∎
The following two guarantees are stated informally in the main text \(end of the Theoretical Analysis section\) and are formalized and proved here as supplement\-only results\.
###### Theorem S1\(Stability to noisy criticality estimates\)\.
Letc^r\\widehat\{c\}\_\{r\}be plug\-in criticality estimates withmaxr\|c^r−cr\|≤δ\\max\_\{r\}\|\\widehat\{c\}\_\{r\}\-c\_\{r\}\|\\leq\\delta\. Then the smooth allocation of Eq\. 9 satisfiesmaxr\|ρ^r⋆−ρr⋆\|=O\(ηδ\)\\max\_\{r\}\|\\widehat\{\\rho\}\_\{r\}^\{\\star\}\-\\rho\_\{r\}^\{\\star\}\|=O\(\\eta\\delta\), and ifρ−1\\rho^\{\-1\}is locally Lipschitz on the feasible interval,maxr\|γ^r⋆−γr⋆\|=O\(ηδ\)\\max\_\{r\}\|\\widehat\{\\gamma\}\_\{r\}^\{\\star\}\-\\gamma\_\{r\}^\{\\star\}\|=O\(\\eta\\delta\)\.
###### Proof\.
For the smooth allocation,
ρr⋆\(c\)=σ\(σ−1\(ρ¯\)\+η\(λ−cr\)\)\.\\rho\_\{r\}^\{\\star\}\(c\)=\\sigma\\\!\\left\(\\sigma^\{\-1\}\(\\bar\{\\rho\}\)\+\\eta\(\\lambda\-c\_\{r\}\)\\right\)\.The sigmoid derivative satisfies
0<σ′\(x\)≤14\.0<\\sigma^\{\\prime\}\(x\)\\leq\\frac\{1\}\{4\}\.Therefore, for fixedλ\\lambda,
\|ρr⋆\(c^\)−ρr⋆\(c\)\|≤η4\|c^r−cr\|\.\\left\|\\rho\_\{r\}^\{\\star\}\(\\widehat\{c\}\)\-\\rho\_\{r\}^\{\\star\}\(c\)\\right\|\\leq\\frac\{\\eta\}\{4\}\|\\widehat\{c\}\_\{r\}\-c\_\{r\}\|\.If the multiplier also changes fromλ\\lambdatoλ^\\widehat\{\\lambda\}, then
\|ρ^r⋆−ρr⋆\|≤η4\(\|c^r−cr\|\+\|λ^−λ\|\)\.\\left\|\\widehat\{\\rho\}\_\{r\}^\{\\star\}\-\\rho\_\{r\}^\{\\star\}\\right\|\\leq\\frac\{\\eta\}\{4\}\\left\(\|\\widehat\{c\}\_\{r\}\-c\_\{r\}\|\+\|\\widehat\{\\lambda\}\-\\lambda\|\\right\)\.The budget multiplier is determined implicitly by
∑rπrσ\(σ−1\(ρ¯\)\+η\(λ−cr\)\)=ρ¯\.\\sum\_\{r\}\\pi\_\{r\}\\sigma\\\!\\left\(\\sigma^\{\-1\}\(\\bar\{\\rho\}\)\+\\eta\(\\lambda\-c\_\{r\}\)\\right\)=\\bar\{\\rho\}\.The left\-hand side is smooth and monotone inλ\\lambda\. By the implicit function theorem, perturbing allcrc\_\{r\}by at mostδ\\deltaperturbsλ\\lambdabyO\(δ\)O\(\\delta\)\. Thus
maxr\|ρ^r⋆−ρr⋆\|=O\(ηδ\)\.\\max\_\{r\}\|\\widehat\{\\rho\}\_\{r\}^\{\\star\}\-\\rho\_\{r\}^\{\\star\}\|=O\(\\eta\\delta\)\.Ifρ−1\\rho^\{\-1\}is locally Lipschitz on the feasible interval, then
\|γ^r⋆−γr⋆\|=\|ρ−1\(ρ^r⋆\)−ρ−1\(ρr⋆\)\|=O\(ηδ\)\.\|\\widehat\{\\gamma\}\_\{r\}^\{\\star\}\-\\gamma\_\{r\}^\{\\star\}\|=\|\\rho^\{\-1\}\(\\widehat\{\\rho\}\_\{r\}^\{\\star\}\)\-\\rho^\{\-1\}\(\\rho\_\{r\}^\{\\star\}\)\|=O\(\\eta\\delta\)\.∎
###### Theorem S2\(Schedule generalization\)\.
Assume the per\-example graph\-risk loss is bounded byBB\. With probability at least1−δ1\-\\delta, for allpθ∈ℱp\_\{\\theta\}\\in\\mathcal\{F\},𝔼\[Δgraph\]≤ℛ^\(ρ\)\+2BℜN\(ℱ\)\+Blog\(1/δ\)/\(2N\)\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]\\leq\\widehat\{\\mathcal\{R\}\}\(\\rho\)\+2B\\mathfrak\{R\}\_\{N\}\(\\mathcal\{F\}\)\+B\\sqrt\{\\log\(1/\\delta\)/\(2N\)\}, whereℛ^\(ρ\)=∑rπrI^rϵ^r\(ρr\)\\widehat\{\\mathcal\{R\}\}\(\\rho\)=\\sum\_\{r\}\\pi\_\{r\}\\widehat\{I\}\_\{r\}\\widehat\{\\epsilon\}\_\{r\}\(\\rho\_\{r\}\)\. At matched architecture, data size, and training budget the complexity terms are shared, so schedule comparisons are governed by the empirical role\-weighted riskℛ^\(ρ\)\\widehat\{\\mathcal\{R\}\}\(\\rho\)\.
###### Proof\.
For a fixed exposure allocationρ\\rho, define the population role\-weighted graph risk
ℛ\(ρ\)=∑rπrIrϵr\(ρr\)\.\\mathcal\{R\}\(\\rho\)=\\sum\_\{r\}\\pi\_\{r\}I\_\{r\}\\epsilon\_\{r\}\(\\rho\_\{r\}\)\.The empirical counterpart is
ℛ^\(ρ\)=∑rπrI^rϵ^r\(ρr\)\.\\widehat\{\\mathcal\{R\}\}\(\\rho\)=\\sum\_\{r\}\\pi\_\{r\}\\widehat\{I\}\_\{r\}\\widehat\{\\epsilon\}\_\{r\}\(\\rho\_\{r\}\)\.Assume the per\-example graph\-risk loss is bounded byBB\. Standard uniform convergence using Rademacher complexity gives, with probability at least1−δ1\-\\delta, for allpθ∈ℱp\_\{\\theta\}\\in\\mathcal\{F\},
ℛ\(ρ\)≤ℛ^\(ρ\)\+2BℜN\(ℱ\)\+Blog\(1/δ\)2N\.\\mathcal\{R\}\(\\rho\)\\leq\\widehat\{\\mathcal\{R\}\}\(\\rho\)\+2B\\mathfrak\{R\}\_\{N\}\(\\mathcal\{F\}\)\+B\\sqrt\{\\frac\{\\log\(1/\\delta\)\}\{2N\}\}\.Combining this with the exposure\-dependent form of Eq\. 5 \(see the Remark above\) gives
𝔼\[Δgraph\]≤ℛ^\(ρ\)\+2BℜN\(ℱ\)\+Blog\(1/δ\)2N\.\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]\\leq\\widehat\{\\mathcal\{R\}\}\(\\rho\)\+2B\\mathfrak\{R\}\_\{N\}\(\\mathcal\{F\}\)\+B\\sqrt\{\\frac\{\\log\(1/\\delta\)\}\{2N\}\}\.When comparing schedules at matched architecture, data size, and training budget, the complexity terms are shared\. Therefore differences in the bound are driven by the empirical role\-weighted risk term\. ∎
## Appendix S2Reproducibility and Experimental Details
All results use existing checkpoints with no per\-experiment retraining beyond the trained models described below\. Every generation script is deterministic given its seed: each run is seeded once at start\-up, and the multi\-seed results \(Table[S1](https://arxiv.org/html/2607.21634#A3.T1)\) average over five fixed generation seeds held constant across both methods, soMotifRole\-Diffand uniform MDLM are compared on matched seeds\.
#### Backbone and training\.
A single masked discrete\-diffusion transformer is shared by all conditions: hidden size256256,44layers,∼3\.32\\sim\\\!3\.32M parameters, absorbing \(\[MASK\]\) forward process\. Optimizer Adam, learning rate3×10−43\\times 10^\{\-4\}, batch size6464\. Only the per\-role exponentsγr\\gamma\_\{r\}and the inverse\-exposure weight differ across schedules; the architecture, optimizer, and budget are held fixed \(matched compute\)\.
#### Schedule derivation\.
Role criticalitiesCrC\_\{r\}are measured once on the uniform\-MDLM checkpoint \(difficultyDrD\_\{r\}from probe NLL, impactIrI\_\{r\}from single\-token perturbation\) and the smooth schedule \(Eq\. 9\) is solved withη=2\\eta\{=\}2andγr\\gamma\_\{r\}clipped to\[0\.25,4\.0\]\[0\.25,4\.0\]; special is fixed atγ=1\\gamma\{=\}1\.
#### Metrics\.
Validity is the RDKit\-sanitizable fraction; uniqueness and novelty use canonical SMILES \(novelty vs\. the training set\)\. QED follows\(Bickerton et al\.[2012](https://arxiv.org/html/2607.21634#bib.bib2)\)\. FCD\(Preuer et al\.[2018](https://arxiv.org/html/2607.21634#bib.bib29)\)is computed withfcd\_torchagainst a*held\-out*reference split, never the training set\.
#### Hardware\.
All timing and generation runs use a single NVIDIA TITAN RTX \(2424GB\); no multi\-GPU or mixed\-precision tricks are required\. Code, configs, and the per\-experiment scripts will be released publicly upon publication\.
## Appendix S3Additional Experimental Results and Ablations
#### Multi\-seed robustness \(numeric values for Figure 2 of the main paper\)\.
Table[S1](https://arxiv.org/html/2607.21634#A3.T1)gives the exact per\-metric means and95%95\\%confidence intervals summarized as grouped bars in the main paper’s multi\-seed figure\.
Table S1:Multi\-seed robustness on QM9, MOSES, and GuacaMol\. Values are mean±\\pm95%95\\%confidence interval over55generation seeds using50005000samples/seed; FCD is computed against the held\-out test reference\. Novelty and FCD are not computed on GuacaMol \(marked “–”\)\. Best values per dataset are highlighted\.DatasetMethodValid↑\\uparrowUnique↑\\uparrowNovel↑\\uparrowFCD↓\\downarrowQM9Uniform MDLM0\.905±0\.0060\.905\\pm 0\.0060\.977±0\.0040\.977\\pm 0\.0040\.781±0\.0180\.781\\pm 0\.0181\.701±0\.0801\.701\\pm 0\.080MotifRole\-Diff0\.944±0\.0080\.944\\pm 0\.0080\.954±0\.0030\.954\\pm 0\.0030\.736±0\.0150\.736\\pm 0\.0151\.609±0\.0641\.609\\pm 0\.064MOSESUniform MDLM0\.920±0\.0080\.920\\pm 0\.0080\.945±0\.0030\.945\\pm 0\.0030\.996±0\.0070\.996\\pm 0\.0072\.125±0\.0962\.125\\pm 0\.096MotifRole\-Diff0\.938±0\.0030\.938\\pm 0\.0030\.986±0\.0050\.986\\pm 0\.0050\.916±0\.0150\.916\\pm 0\.0151\.850±0\.0321\.850\\pm 0\.032GuacaMol \(100k\)Uniform MDLM0\.787±0\.0160\.787\\pm 0\.0161\.000±0\.0001\.000\\pm 0\.000––MotifRole\-Diff0\.841±0\.0080\.841\\pm 0\.0080\.984±0\.0010\.984\\pm 0\.001––Table S2:Matched\-compute*staged*component ablation on QM9\. Row 1→\\toRow 2 isolates the mSENT tokenizer contribution \(SENT→\\tomSENT, both uniform, reweighting off\); Row 2→\\toRow 3 gives the additional contribution of the*complete*role\-aware configuration \(derived schedule*and*inverse\-exposure weighting together\)\. Because the final configuration turns on both components jointly, this table does not isolate their individual effects; a schedule\-only \(reweighting\-off\) row would be required for a full decomposition\. Best available value for each metric is shown in bold\.\#TokenizerScheduleReweight\-ingValid\.↑\\uparrowUnique\.↑\\uparrowNovel\.↑\\uparrowAtomStab\.↑\\uparrow1SENTUniform \(γ=1\\gamma=1\)Off0\.8530\.8530\.7500\.7500\.7330\.7330\.8870\.8872mSENTUniform \(γ=1\\gamma=1\)Off0\.9050\.9050\.977\\mathbf\{0\.977\}0\.781\\mathbf\{0\.781\}0\.9460\.9463mSENTMotifRole \(γ∗\\gamma^\{\*\}\)On0\.9440\.9440\.9540\.9540\.7360\.7360\.955\\mathbf\{0\.955\}
#### Computational overhead\.
Table[S3](https://arxiv.org/html/2607.21634#A3.T3)gives the wall\-clock and parameter counts referenced in the main paper’s “Computational overhead” paragraph:MotifRole\-Diffadds no parameters and no extra passes, and matches uniform MDLM in training and sampling time within run\-to\-run noise\.
Table S3:Computational overhead ofMotifRole\-Diffvs\. uniform MDLM \(QM9,44\-layer backbone, one TITAN RTX\)\. The schedule adds no parameters and no passes; wall\-clock matches within noise\. Role\-aware sampling \(Algorithm 1 of the main paper\) adds only twoO\(L\)O\(L\)table operations per reverse step\.MethodParamsTrain \(min/epoch\)Sample10001000\(s\)Uniform MDLM3\.323\.32M∼1\.25\\sim\\\!1\.2591\.891\.8MotifRole\-Diff3\.323\.32M∼1\.25\\sim\\\!1\.2587\.487\.4
## Appendix S4Motif\-Preservation Metric Definitions and Full Results
We compare SENT and mSENT on5,0005\{,\}000tokenized QM9 molecules\. For each molecule we compute atom\-level motif assignments with the CAMT5\-style motif partition, then serialize the same graph with either SENT or mSENT\. Letπ\(v\)\\pi\(v\)be the first sequence position at which atomvvappears andc\(v\)c\(v\)the motif ID ofvv\. Lower span, motif transitions, switch rate, intra\-edge distortion, and fragmentation indicate better motif locality; higher compactness, same\-motif adjacency, and boundary alignment indicate better motif preservation\. The per\-molecule metrics, listed in the same order as the rows of Table[S4](https://arxiv.org/html/2607.21634#A4.T4), are:
Span\(m\)\\displaystyle\\mathrm\{Span\}\(m\)=maxv∈Vmπ\(v\)−minv∈Vmπ\(v\)\+1,\\displaystyle=\\max\_\{v\\in V\_\{m\}\}\\pi\(v\)\-\\min\_\{v\\in V\_\{m\}\}\\pi\(v\)\+1,\(S2\)Compact\(m\)\\displaystyle\\mathrm\{Compact\}\(m\)=\|Vm\|Span\(m\),\\displaystyle=\\frac\{\|V\_\{m\}\|\}\{\\mathrm\{Span\}\(m\)\},\(S3\)SMAR\\displaystyle\\mathrm\{SMAR\}=∑i=1L−1𝟏\[c\(zi\)=c\(zi\+1\)\]L−1,\\displaystyle=\\frac\{\\sum\_\{i=1\}^\{L\-1\}\\mathbf\{1\}\\\!\\left\[c\(z\_\{i\}\)=c\(z\_\{i\+1\}\)\\right\]\}\{L\-1\},\(S4\)Trans\\displaystyle\\mathrm\{Trans\}=∑i=1L−1𝟏\[c\(zi\)≠c\(zi\+1\)\],\\displaystyle=\\sum\_\{i=1\}^\{L\-1\}\\mathbf\{1\}\\\!\\left\[c\(z\_\{i\}\)\\neq c\(z\_\{i\+1\}\)\\right\],\(S5\)Switch\\displaystyle\\mathrm\{Switch\}=TransL−1,\\displaystyle=\\frac\{\\mathrm\{Trans\}\}\{L\-1\},\(S6\)IBA\\displaystyle\\mathrm\{IBA\}=∑i=1L−1𝟏\[c\(zi\)≠c\(zi\+1\)\]1\[\(zi,zi\+1\)∈E\]Trans,\\displaystyle=\\frac\{\\sum\_\{i=1\}^\{L\-1\}\\mathbf\{1\}\[c\(z\_\{i\}\)\\neq c\(z\_\{i\+1\}\)\]\\,\\mathbf\{1\}\[\(z\_\{i\},z\_\{i\+1\}\)\\in E\]\}\{\\mathrm\{Trans\}\},\(S7\)Dintra\\displaystyle D\_\{\\mathrm\{intra\}\}=1\|Eintra\|∑\(u,v\)∈Eintra\|π\(u\)−π\(v\)\|,\\displaystyle=\\frac\{1\}\{\|E\_\{\\mathrm\{intra\}\}\|\}\\sum\_\{\(u,v\)\\in E\_\{\\mathrm\{intra\}\}\}\|\\pi\(u\)\-\\pi\(v\)\|,\(S8\)Frag\(m\)\\displaystyle\\mathrm\{Frag\}\(m\)=\#\{contiguous blocks of motifm\}\.\\displaystyle=\\\#\\\{\\text\{contiguous blocks of motif \}m\\\}\.\(S9\)HereEintra=\{\(u,v\)∈E:c\(u\)=c\(v\)\}E\_\{\\mathrm\{intra\}\}=\\\{\(u,v\)\\in E:c\(u\)=c\(v\)\\\}denotes the set of intra\-motif edges \(both endpoints in the same motif\)\. Equations[S2](https://arxiv.org/html/2607.21634#A4.E2)–[S9](https://arxiv.org/html/2607.21634#A4.E9)correspond one\-to\-one, in order, to the mean motif span, compactness, same\-motif adjacency, motif transitions, switch rate, boundary alignment, intra\-edge distortion, and fragmentation rows of Table[S4](https://arxiv.org/html/2607.21634#A4.T4)\.
Table[S4](https://arxiv.org/html/2607.21634#A4.T4)reports the full per\-metric values summarized as relative gains in Figure 3 of the main paper; arrows indicate the preferred direction, and mSENT improves every metric\.
Table S4:Serialization\-level motif preservation metrics comparing SENT and mSENT on QM9\. Arrows indicate the preferred direction for each metric; mSENT improves all eight\.MetricSENTmSENTMean motif span↓\\downarrow18\.0917\.33Compactness↑\\uparrow0\.6860\.706Same\-motif adjacency↑\\uparrow0\.1650\.176Motif transitions↓\\downarrow22\.0321\.87Switch rate↓\\downarrow0\.8350\.824Boundary alignment↑\\uparrow0\.1240\.133Intra\-edge distortion↓\\downarrow22\.2720\.59Fragmentation↓\\downarrow1\.7961\.778Similar Articles
Adaptive Order Policies for Masked Diffusion
Proposes learning the unmasking order in masked diffusion models using a lightweight policy network, with a weighted loss that outperforms heuristics on combinatorial tasks and protein design.
Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models
This paper introduces Reflective Masking, a lightweight post-training method that enables mask diffusion models to perform multi-turn self-revision through token-level revision policies and history references, improving performance on reasoning tasks like Sudoku, math, code generation, and image editing.
Controllable Molecular Generative Foundation Models
Proposes CoMole, a controllable molecular generative foundation model using motif-aware graph diffusion and reinforcement learning, achieving superior controllability across materials and drug discovery benchmarks.
The Confidence Shortcut: A Reasoning Failure Mode of Masked Diffusion Models
This paper identifies a failure mode in masked diffusion language models where confidence-based decoding leads to high-confidence errors on complex reasoning tasks, and shows that confidence-aligned training exacerbates this issue while random masking preserves reasoning performance.
AnchorDiff: Topology-Aware Masked Diffusion with Confidence-based Rewriting for Radiology Report Generation
AnchorDiff proposes a topology-aware masked diffusion framework for radiology report generation, integrating RadGraph-derived clinical anchors and confidence-based rewriting to achieve state-of-the-art results on MIMIC-CXR and MIMIC-RG4 benchmarks.