MotifRole-Diff: 面向掩码分子图扩散的风险最优角色感知破坏

# 风险最优的角色感知掩码分子图扩散源: https://arxiv.org/html/2607.21634 ###### 摘要 分子图生成的掩码离散扩散通常对所有词符应用统一的破坏调度，这些词符来自无损图到序列表示，隐含地将结构异质的分子组件视为同等难度和同等重要进行重建。然而，不同分子图词符角色在去噪难度及其对解码分子影响方面表现出显著差异，这激发了角色特定的破坏策略。我们引入MotifRole-Diff，一种角色感知的破坏过程，该过程根据经验测量的去噪难度和图级扰动影响分配掩码率，同时保留模型架构、干净序列空间和无损分子图解码器。我们将调度选择形式化为在词符角色间固定掩码预算的风险最优分配。我们的定理刻画了所建模的角色加权残差风险的最优性，而下游生成性能则通过经验评估。在匹配的架构、训练预算和采样计算下，MotifRole-Diff在QM9上将有效性从0.905提升至0.944，同时将FCD从1.701降至1.609；在MOSES上将有效性从0.920提升至0.938，同时将FCD从2.125降至1.850。角色诊断进一步显示，跨分子图词符类别的重建质量提升。这些匹配计算的结果表明，对于序列化分子图扩散，结构信息引导的破坏是比统一调度更有效的掩码策略。 ## 引言 生成有效、新颖且多样化的分子图是计算分子发现中的一个基本问题，而扩散模型已成为分子图生成的最强范式之一(Jo, Lee, and Hwang2022a (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))。基于离散去噪扩散(Austin et al.2021 (https://arxiv.org/html/2607.21634#bib.bib1); Campbell et al.2022 (https://arxiv.org/html/2607.21634#bib.bib4))构建的掩码离散语言建模(MDLM)(Sahoo et al.2024 (https://arxiv.org/html/2607.21634#bib.bib31))提供了一种可扩展的公式，其通过将分子图序列化为词符序列，并行去噪，并通过无损解码器重建原始图，结合了精确重建和序列建模的效率。这种方法根本上依赖于图的序列化。先前工作已表明，序列化顺序影响序列模型的效率(Diamant et al.2023 (https://arxiv.org/html/2607.21634#bib.bib9); Jang, Lee, and Ahn2024 (https://arxiv.org/html/2607.21634#bib.bib17))，而分层分子生成器则证明了保留化学上有意义的基序的重要性(Jin, Barzilay, and Jaakkola2018a (https://arxiv.org/html/2607.21634#bib.bib18),2020 (https://arxiv.org/html/2607.21634#bib.bib20))。然而，现有的掩码离散扩散模型对每个词符位置应用相同的破坏调度(Austin et al.2021 (https://arxiv.org/html/2607.21634#bib.bib1); Sahoo et al.2024 (https://arxiv.org/html/2607.21634#bib.bib31))，隐含地假设所有序列化的词符在重建难度和恢复分子图的重要性上是相同的。这个假设通常是错误的。序列化的分子图包含结构上不同的词符角色：边界标记、序列化语法、基序内部原子和键，以及连接子结构的接口词符。这些角色上的错误后果并不相同，因为有些角色可以可靠地重建，对解码后的图影响很小，而其他角色则更难恢复或导致有效性、连通性严重受损，然而现有的图到序列扩散方法并未考虑破坏是否应取决于这些差异。我们称这种性质为*角色异质性*，并假设掩码分子扩散因此不仅对*顺序敏感*，而且对*角色敏感*：只有当角色具有可比的结构重要性时，统一破坏才是合适的，否则破坏应遵循测量的结构风险。为研究这一假设，我们引入了MotifRole-Diff，一个角色敏感的掩码离散分子扩散框架。我们扩展了无损的SENT表示(Chen, Krimmel, and Borgwardt2026 (https://arxiv.org/html/2607.21634#bib.bib7))，加入了基序感知的序列化*mSENT*，公开了四种结构词符角色，同时保留了解码器的兼容性。我们使用受控的扰动实验量化每个角色的重建难度和图级后果，然后根据测量的结构风险分配固定的掩码预算，推导出角色感知的破坏过程。所得的吸收扩散公式推广了标准MDLM，并在所有词符角色具有相同结构关键性时简化为统一破坏。我们的假设引发了四个研究问题。(RQ1)序列化的词符角色在重建难度和图级结构后果上是否存在显著差异？(RQ2)角色感知的破坏能否被推导为固定掩码预算的约束优化分配？(RQ3)经验信息引导的角色感知破坏是否优于统一掩码，其收益是否源于信息引导的结构分配而非仅仅是非均匀性？(RQ4)基序感知的序列化是否能在保留SENT无损解码保证的同时改善基序局部性？整个过程中，架构、优化、计算和期望掩码率保持恒定，因此差异仅来自序列化和破坏设计。我们的贡献有四个方面： 1. 1.刻画序列化分子扩散中的角色异质性 (RQ1)。我们识别出结构词符角色是图序列化中一个被忽视的性质，并显示它们在重建难度和图级结构后果上系统性地不同。 2. 2.一个原则性的角色感知破坏过程 (RQ2)。我们将破坏调度形式化为固定掩码预算下的约束结构风险最小化，得到一个角色感知的吸收扩散过程，该过程推广了标准MDLM，并将统一破坏作为特例恢复。 3. 3.角色感知破坏的受控评估 (RQ3)。在匹配设置下，MotifRole-Diff在相等计算下比统一MDLM提升了生成质量；由于调度是从测量的关键性而非随意设置推导的，收益反映了信息引导的结构分配而非仅仅是非均匀破坏。 4. 4.一个基序感知且保留角色的序列化方法 (RQ4)。我们引入mSENT，一种SENT的基序感知扩展，在保留原始无损解码器的同时改善了基序局部性，并公开了四种可解释的结构词符角色。 ## 预备知识与问题形式化 #### 符号。一个分子图G=\(V,E,X,A\)G=\(V,E,X,A\)包含原子VV、键EE、特征XX和邻接矩阵AA。无损序列化T\(⋅;π\)T\(\\cdot\\,;\\pi\)在遍历顺序π\\pi下将GG映射到词符序列zπ=T\(G;π\)z^\{\\pi\}=T\(G;\\pi\)，并带有解码器D\(zπ\)≅GD\(z^\{\\pi\}\)\\cong G以及\|zπ\|=O\(\|V\|\+\|E\|\)\|z^\{\\pi\}\|=O\(\|V\|\+\|E\|\)。基序感知的分词器将原子分配到基序块μ:V→\{1,...,K\}\\mu:V\\to\\\{1,\\dots,K\\\}，并为每个词符标注角色ri∈R=\{special,syntax,interior,interface\}r\_\{i\}\\in\\mathcal\{R\}=\\\{\\text\{special\},\\text\{syntax\},\\text\{interior\},\\text\{interface\}\\\}(表1 (https://arxiv.org/html/2607.21634#Sx2.T1))；πr\\pi\_\{r\}表示角色-rr词符的占比。 表 1:基序感知序列化诱导的词符角色。| 角色 | 含义 | 错误后果 |
|------|------|----------|
| Special | 控制 / [MASK] 词符 | 无效边界或填充 |
| Syntax | 语法、重置、分隔符 | 格式错误或无法解码的图 |
| Interior | 一个基序内的原子/键 | 局部基序破坏；偶尔的骨架/原子-键变化 |
| Interface | 跨基序的连接词符 | 连接歧义和可能的连通性错误；单词符扰动影响最大 | #### 掩码离散扩散。一个吸收前向过程将干净序列z0z\_0中的每个词符以概率αt\\alpha\_t在扩散时间tt替换为[MASK]，并且一个去噪器pθp\_\{\\theta\}在吸收状态ELBO下重建被掩码的词符， LMDLM=Ez0,t,Mt\[w\(t\)∑i∈Mt−log⁡pθ\(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\)速率αt\\alpha\_t是词符无关的：每个位置以相同方式被破坏，不论角色如何，并且一个基于置信度的采样器在反向过程中揭示词符。 #### 角色暴露与预算。对于一个角色相关的调度\{αt\(r\)\}\\\{\\alpha\_\{t\}^\{\(r\)\}\\\}，令ρr=Et\[αt\(r\)\]\\rho\_\{r\}=\\mathbb\{E\}\_\{t\}\[\\alpha\_\{t\}^\{\(r\)\}\]为角色rr的平均暴露，ρ ̄=∑rπrρr\\bar\{\\rho\}=\\sum\_\{r\}\\pi\_\{r\}\\rho\_\{r\}为总的掩码预算。统一MDLM是当ρr≡ρ ̄\\rho\_\{r\}\\equiv\\bar\{\\rho\}的情况。固定ρ ̄\\bar\{\\rho\}使期望的掩码词符数量相等，因此共享ρ ̄\\bar\{\\rho\}的调度在预算上匹配，并在共享骨干网络下在计算上也匹配。 #### 问题陈述。给定R\\mathcal\{R\}、频率\{πr\}\\\{\\pi\_\{r\}\\\}和固定预算ρ ̄\\bar\{\\rho\}，我们寻求一个角色相关的调度，在等预算下最小化期望的图级重建误差， minρ\\displaystyle\\min\_\{\\rho\}E\[Δ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\}.这重新分配了一个固定预算跨角色，而不是增加噪声或容量；理论分析使E\[Δgraph\]\\mathbb\{E\}\[\\Delta\_\{\\mathrm\{graph\}\}\]具体化并推导了风险最优分配。 ## 角色敏感性：一种诊断 在一个轻量级的*探针*去噪器pθ0p\_\{\\theta\_\{0\}\}上训练使用统一掩码，我们为每个角色rr测量难度Dr=Ei:ri=r\[−log⁡pθ0\(z0\(i\)∣zt,t\)\]D\_\{r\}=\\mathbb\{E\}\_\{i:r\_\{i\}=r\}\[\-\\log p\_\{\\theta\_\{0\}\}\(z\_\{0\}^\{\(i\)\}\\mid z\_\{t\},t\)\]和top-1误差ErE\_\{r\}，其中t∼U\(0.05,0.95\)t\\sim\\mathcal\{U\}\(0.05,0.95\)。图1 (https://arxiv.org/html/2607.21634#Sx4.F1)\(a\)报告了一个单调的大顺序，interface≫interior\>syntax≫special\\text\{interface\}\\gg\\text\{interior\}\>\\text\{syntax\}\\gg\\text\{special\}(探针NLL Dr=0.11/0.23/0.35D\_\{r\}=0.11/0.23/0.35 nats 和 top-1误差 Er=4.2/8.8/12.9%E\_\{r\}=4.2/8.8/12.9\\% 对于 syntax/interior/interface；special≈0\\approx 0)。完全训练后的统一MDLM在更大的绝对幅度上重现了*相同*的角色顺序(表4 (https://arxiv.org/html/2607.21634#Sx7.T4))，因此我们依赖的是顺序，而不是探针的绝对尺度。这种异质性是MotifRole-Diff利用的空间。 ## MotifRole-Diff 受这种角色异质性的启发，我们使破坏调度依赖于角色。尽管我们在分子图上评估，但该框架适用于任何具有可识别词符角色的无损序列化图。图1 (https://arxiv.org/html/2607.21634#Sx4.F1)在真实QM9值上预览了构造。令Λ\(t\)=∫0tβ\(s\)ds\\Lambda\(t\)=\\int\_\{0\}^\{t\}\\beta\(s\)\\,ds为基础调度，γr\>0\\gamma\_\{r\}\>0为每个角色的指数。吸收过程变为角色相关， αt\(r\)=1−exp⁡\(−γrΛ\(t\)\)，\\alpha\_\{t\}^\{\(r\)\}=1\-\\exp\\\!\\big\(\-\\gamma\_\{r\}\\,\\Lambda\(t\)\\big\),\(3\)因此γr<1\\gamma\_\{r\}<1掩码角色rr更少(可见更长，更早揭示)，而γr\>1\\gamma\_\{r\}\>1掩码更多。统一γr≡1\\gamma\_\{r\}\\equiv 1精确恢复MDLM。干净的序列、语法和解码器不变，因此对于任何语法有效的z0z\_\{0\}，DSENT\(z0\)≅GD\_\{\\text\{SENT\}\}\(z\_\{0\}\)\\cong G（引理1 (https://arxiv.org/html/2607.21634#Thmlemma1)\)。 00.10.20.3051015DrD\_\{r\}\(NLL\)ErE\_\{r\}\(%\)SpecSynIntrIntf\(a\) 测量难度 γr⋆\\gamma\_\{r\}^\{\\star\}\(η=2\\eta\{=\}2\)γ=1\\gamma\{=\}1Intf0.44Intr1.07Syn2.37\(b\) 推导的速率 图 1:在真实QM9值上的测量难度和推导的调度。(a)探针去噪器上的测量角色难度：探针NLL DrD\_\{r\}(青色，左轴)和top-1重建误差ErE\_\{r\}(深蓝，右轴)。两者都将interface排为最难，special排为微不足道。(b)来自方程9 (https://arxiv.org/html/2607.21634#Sx5.E9)的风险最优速率γr⋆\\gamma\_\{r\}^\{\\star\}，在η=2\\eta\{=\}2(表2 (https://arxiv.org/html/2607.21634#Sx5.T2))：最难的角色(interface)被*保护*(γ<1\\gamma<1)，而更容易的角色被破坏更多(γ\>1\\gamma\>1)，同时保持固定掩码预算。无损解码器未受影响。#### 逆暴露损失加权。被保护的低γ\\gamma角色被掩码的频率较低，否则会训练不足。我们通过逆暴露对每个词符的损失进行重新加权， w\(t,ri\)=1Pr⁡\(i∈Mt∣ri\)\+ε，w\(t,r\_\{i\}\)=\\frac\{1\}\{\\Pr\(i\\in M\_\{t\}\\mid r\_\{i\}\)\+\\epsilon\},\(4\)这在γr≡1\\gamma\_\{r\}\\equiv 1时为恒等操作(w≡1w\\equiv 1)。 #### 角色感知采样。在生成时，被掩码词符的基序角色并不直接观察到，但也不必须：去噪器已经在每个位置上诱导了干净词符的分布，而分词器提供了训练词符-角色先验P\(r∣v\)P\(r\\mid v\)。因此，我们在每个被掩码位置ii估计一个*软*角色分布qi\(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\)，并将训练中使用的相同预定角色速率γr\\gamma\_\{r\}带入反向过程，赋予每个位置一个期望指数γi=∑rqi\(r\)γr\\gamma\_\{i\}=\\sum\_\{r\}q\_\{i\}\(r\)\\,\\gamma\_\{r\}和一个角色感知的反向解掩概率piunmask=\(as\(i\)−at\(i\)\)/\(1−at\(i\)\)p^\{\\mathrm\{unmask\}\}\_\{i\}=\(a\_\{s\}^\{\(i\)\}\-a\_\{t\}^\{\(i\)\}\)/\(1\-a\_\{t\}^\{\(i\)\}\)，其中at\(i\)=exp⁡\(−γiΛ\(t\)\)a\_\{t\}^\{\(i\)\}=\\exp\(\-\\gamma\_\{i\}\\Lambda\(t\)\)。然后，位置通过组合的置信度和调度分数log⁡maxv⁡pθ\(xi=v∣zt\)\+log⁡piunmask\\log\\max\_\{v\}p\_\{\\theta\}\(x\_\{i\}\{=\}v\\mid z\_\{t\}\)\+\\log p^\{\\mathrm\{unmask\}\}\_\{i\}被揭示。算法1 (https://arxiv.org/html/2607.21634#alg1)给出了完整过程；γr\\gamma\_\{r\}在采样前固定，因此反向采样端到端地重用推导的调度，而不是将其仅延迟到训练。 算法 1使用预定角色γ的角色感知采样0:去噪器 pθp\_\{\\theta\}；长度 LL；反向步骤 \{t→s\}\\\{t\\\!\\to\\\!s\\\}；角色速率 γspecial,γsyntax,γinterior,γinterface\\gamma\_\{\\text\{special\}\},\\gamma\_\{\\text\{syntax\}\},\\gamma\_\{\\text\{interior\}\},\\gamma\_\{\\text\{interface\}\}；词符–角色先验 P\(r∣v\)P\(r\\mid v\)；噪声调度 Λ\(⋅\)\\Lambda\(\\cdot\)；揭示数量 ktk\_\{t\}；[MASK]，[BOS]，[EOS]；裁剪下限 ε\\epsilon 0:生成的图 G^\\hat\{G\} 1: z←\[\[BOS\],\[MASK\],...,\[MASK\],\[EOS\]\]z\\leftarrow\[\\texttt\{\[BOS\]\},\\texttt\{\[MASK\]\},\\dots,\\texttt\{\[MASK\]\},\\texttt\{\[EOS\]\}\] 2:对于每个反向步骤 t→st\\to s 执行 3: M←\{M\\leftarrow\\\{被掩码位置在 z\}∖\{BOS,EOS\}z\\\}\\setminus\\\{\\text\{BOS\},\\text\{EOS\}\\\} 4:如果 M=