MotifRole-Diff: Risk-Optimal Role-Aware Corruption for Masked Molecular Graph Diffusion

2026-07-20

AI Summary

This work addresses the limitation of conventional molecular graph diffusion models that employ a uniform masking strategy, disregarding the varying denoising difficulty and reconstruction impact across distinct structural units. To overcome this, the authors propose a role-aware masking diffusion approach that, for the first time, categorizes molecular graph tokens according to their chemical roles and dynamically allocates masking rates under a fixed masking budget by modeling role-weighted residual risk. This yields an optimal corruption schedule that minimizes reconstruction risk without altering the model architecture or decoder. Evaluated on QM9 and MOSES benchmarks, the method significantly improves molecular validity—reaching up to 0.944—and reduces Fréchet ChemNet Distance (FCD) to as low as 1.609, while enhancing reconstruction performance across all token roles.

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 introduce MotifRole-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-Diff improves 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.

Problem

• molecular graph generation
• masked diffusion
• corruption schedule
• role-aware masking
• denoising difficulty

Innovation

• role-aware corruption
• masked molecular graph diffusion
• risk-optimal scheduling
• molecular generation
• denoising difficulty