D²G-TO: Task-aware and OOD-guided Discrete Graph Diffusion for Robust CNS Drug Discovery

1 Introduction

The therapeutic landscape of central nervous system (CNS) disorders remains one of the most intractable frontiers in biomedical research. CNS drug discovery success rates are markedly lower than in most other disease areas, while clinical demand vastly exceeds the available treatment options. This persistent gap reflects three tightly coupled challenges:

• The chemical space is vast while truly viable structures are extremely sparse [Paul et al., 2010; Scannell et al., 2012; Stumpfe and Bajorath, 2012];
• CNS drugs must finely balance blood-brain barrier permeability and overall efficacy, where even slight deviations can readily erode clinical benefit [Pardridge, 2005; Cummings et al., 2014; DiChiara et al., 2024];
• Complex molecular scaffolds markedly amplify risks in synthesis and development, forcing synthetic accessibility to become a hard constraint already in the early design stage [Dichiara et al., 2009];

Jointly optimizing chemical space exploration, CNScritical pharmacological properties, and synthetic accessibility at the hit and lead discovery stages has thus become a central challenge that ultimately shapes clinical translation success.

In recent years, computer-aided drug design (CADD) and deep learning have advanced rapidly, with graph neural networks (GNNs) and related representation learning methods achieving strong performance in property prediction and virtual screening, thereby helping to shorten development cycles and improve success rates [Cheng et al., 2021]. Building on this, different generative models have been introduced for molecular design, aiming to directly propose molecules with desired properties in vast chemical spaces [Yang et al., 2023; Du et al., 2024]. However, in practical drug discovery, property control and molecular generation remain weakly integrated, and the combined challenge of multi-objective design and distributional robustness has yet to be systematically addressed [Cheng et al., 2021]. Most existing frameworks still implicitly assume that training and deployment share similar data distributions, making it difficult to explicitly handle distribution shift and out-of-distribution (OOD) structures in realistic CNS scenarios [Liu et al., 2023; Wu et al., 2022].

To address these challenges, we propose D2G-TO (Taskaware and OOD-guided Discrete Graph Diffusion), a taskaware and distribution-consistency discrete diffusion framework for CNS-oriented molecular graph generation. D2G-TO injects multiple CNS-relevant pharmacological objectivessuch as target inhibition, brain penetration, and synthetic accessibility-directly into the diffusion process. While the structural guidance pulls samples toward in-distribution (ID) graphs and pushes them away from potential OOD modes. D2G-TO achieves strong performance in molecular validity, diversity, and distributional robustness across CNS benchmarks. In an Alzheimer's disease (AD) case study, D2GTO prioritizes early-stage candidate compounds that jointly satisfy efficacy and developability requirements, followed by ADMET profiling and structure-based docking to assess binding-mode plausibility. Thus, D2G-TO is practical in silico design to accelerate CNS drug discovery.

2 RelatedWork

2.1 Multiple Challenges in CNS Drug Discovery

Drug discovery remains one of the most formidable challenges in CNS disorders. Despite decades of research and investment, progress in clinical translation has been slowsuch as AD, where repeated Phase III failures and success rate reduced to below 0.4% [Cummings et al., 2014]. This stagnation arises from three intertwined challenges: chemical space, pharmacological property, and synthetic accessibility.

The vastness of chemical space provides a fundamental challenge for CNS drug discovery. With an estimated > 1060 drug-like molecules, exploration through traditional screening is infeasible [Paul et al., 2010; Scannell et al., 2012; Stumpfe and Bajorath, 2012; Vamathevan et al., 2019; Cheng et al., 2021; Mullard, 2017; Berdigaliyev and Aljofan, 2020; Zhang et al., 2025]. And pharmacological properties impose intrinsic constraints on molecular optimization. CNS drug design must simultaneously balance blood-brain barrier (BBB) permeability and precise target selectivity, both serving as key determinants in diseases such as AD [Pardridge, 2005; Cummings et al., 2014]. However, these requirements often conflict, narrowing the drug design space and hindering balanced optimization. In addition, synthetic accessibility further increases these challenges, as complex or unstable scaffolds demand costly synthesis cost and validation, prolonging development and limiting scalability [DiChiara et al., 2024; Dichiara et al., 2009].

These challenges need for computational frameworks capable of navigating complex drug design trade-offs while jointly optimizing pharmacological, synthetic, and structural objectives. CADD has shown substantial potential in supporting multiple stages of drug development, significantly enhancing the efficiency of drug discovery [Button et al., 2019]. With the rapid advancement of artificial intelligence, deep learning models have emergedparticularly GNNs-demonstrating strong capabilities in encoding molecular topology and predicting key pharmacological properties[Cheng et al., 2021; Choung et al., 2023; Huang et al., 2021; Jing et al., 2025]. However, most related studies on CNS remain focused on property prediction and optimization. These models effectively evaluate existing compounds, but they limit capability in generating novel molecules that simultaneously satisfy pharmacological, synthetic, and structural constraints. To address these limitations, generative modeling offers new possibilities for advancing CNS drug design.

2.2 Generative Diffusion and OODGeneration

Generative modeling has emerged as a transformative paradigm for drug discovery, providing an approach to explore the vast and uncharted chemical space and to design novel, diverse, pharmacological compounds beyond existing libraries [Cheng et al., 2021; Zhavoronkov et al., 2019; Du et al., 2024; Ren et al., 2023].

Early based sequence and graph generative frameworks, such as SMILES-based RNNs and LSTMs [Moret et al., 2023], as well as variational autoencoders including JTVAE [Jin et al., 2018] and its Bayesian variant JTVAE-BO [G´omez-Bombarelli et al., 2018], demonstrated the feasibility of learning molecular distributions for de novo design. However, these models often depend on heuristic latent sampling or post hoc property optimization. Therefore, they limit the generation capacity about multi-objective constraints, particularly those related to pharmacological property, synthetic accessibility, and structural rationality. Subsequent graph-based genetic algorithms GB-GA expanded chemical diversity through stochastic graph mutations, but lacked mechanisms to integrate domain priors for guided optimization [Jensen, 2019]. Additionally, fragment-based methods like MARS improved sampling efficiency though adaptive graph edits and MCMC exploration, yet remained limited in jointly handling complex, multi-objective constraints[Xie et al., 2021].

Diffusion models, based in probabilistic formulation, have recently emerged as a dominant approach for generation. Graph diffusion models further extend the capability to capture complex structural distributions [Liu et al., 2023; Yang et al., 2023]. Continuous diffusion models, such as EDPGNN and GDSS, generate molecular graphs through iterative denoising with score-based Langevin dynamics, while discrete approaches like DiGress introduces a discrete denoising diffusion framework that reformulates graph generation as node and edge classification tasks [Lee et al., 2023; Niu et al., 2020; Vignac et al., 2023; Jo et al., 2022]. Furthermore, GraphDiT extends discrete diffusion framework with a denoiser based on transformer and a graph-dependent noise model, and DiffMC-Gen introduces a dual denoising architecture that combines discrete graph diffusion with SE(3)- equivariant denoising in continuous 3D space [Liu et al., 2024; Yang et al., 2025].

Despite these advances, most diffusion models assume a consistent distribution between training and testing graphs. This condition is rarely satisfied in the real world for CNS drug design. Newly generated molecules often deviate from training distributions, resulting in OOD samples that affect chemical validity and pharmacological reliability. Addressing this distributional robustness challenge has become crucial for trustworthy generation models [Wu et al., 2022; Liu et al., 2024]. Recent studies such as MOOD and PGR-MOOD introduce OOD generation in continuous diffusion process. But these methods remain external to the generative process and absent from discrete diffusion models [Lee et al., 2023; Shen et al., 2024]. Therefore, developing distribution-aware diffusion frameworks that integrate property control, structural rationality, and distributional robustness constitutes an essential step toward reliable CNS drug design.

3 Method

We propose Task-aware and OOD-guided Discrete Graph Diffusion (D2G-TO) method, a three-stage discrete diffusion framework for molecular graph generation. As illustrated in Figure 1, the framework unifies task-aware modeling and distribution-guided diffusion within a single generative pipeline. In Stage 1, we pre-train a discrete diffusion model with multi-pharmacological conditioning to learn propertyaware molecular graph representations. In Stage 2, we generate candidate ID/OOD graphs and introduce a Structural Similarity Guidance (SSG) module to select guidance structures that remain close to ID regions while repelling OOD regions. In Stage 3, we train denoising network based on Graph Transformer for noise molecular graphs to generate new molecules. The reverse diffusion process is governed by dual constraints that couple property-aware task optimization with distribution-consistent structural guidance, enabling task- and structure-guided joint molecular generation for precise and robust CNS drug discovery.

3.1 Task-Aware Pre-training

In Stage 1, we introduce multi-pharmacological properties (e.g., permeability, synthetic accessibility) into a discrete diffusion framework to learn molecular graph representations conditioned on these properties, providing a stable generative prior for subsequent stages (Figure 1, Stage 1).

Given a molecular graph G = (X,E), with X the node/atom feature and E the edge/bond feature tensorthe, the forward diffusion process incrementally perturbs node and edge states:

where xti and etij denote node and edge states at timestep t. The reverse process reconstructs clean molecules under multi-pharmacological conditions {ym}Mm=1 :

where fθ(·)is a denoising network conditioned on multipharmacological properties. These property-aware representations serve as the basis for distribution-guided diffusion in later stages.

3.2 Candidate Structural Generation and Guidance

In Stage 2, we simulate pseudo-OOD structures using training set and construct guidance graphs in the discrete diffusion space. The stage learns explicit structural distribution boundaries, thereby avoids anomalies and enhances structural robustness during the generation.

In Figure 1 (Stage 2), we use the pretrained model (Stage 1) to generate candidate guidance structures and inject targeted perturbations at node and edge levels to construct a set of candidate OOD graphs GOOD. And they are used only as repulsive structural references rather than desired generation targets. These graphs differ markedly from the training distribution (ID) in topological patterns, thereby approximating potential OOD regions.

To distinguish ID and OOD samples, we introduce the Structural Similarity Guidance (SSG) module, which evaluates, for each candidate graph, its structural deviation from both ID and OOD reference sets. SSG is defined as a joint topology-feature distance. This design preserves global structural information, making it amenable to end-to-end optimization within diffusion training [Shen et al., 2024; Vayer et al., 2019]:

where A1,A2 are graph adjacency matrices, i = 1, . . . ,N1, j = 1, . . . , N2, and xk(i) denotes node features in graph Gk. The first term captures global topological divergence, while the second term measures cross-graph feature deviation.

Based on SSG distances to ID and OOD sets, we define a bilateral selection criterion:

where γ > 0 is a weighting hyperparameter. The selected graphs GID and GOOD are used as structural guidance in the subsequent generation stage.

3.3 Dual-Constrained Joint Generation

As illustrated in Figure 1 (Stage 3), the reverse diffusion process is governed by dual constraints that couple task-guided property optimization and distribution-guided structural constraint.

Forward diffusion

The forward process follows the same Markovian noising scheme in discrete graph space:

which gradually corrupts a clean molecular graph G0 into a noisy graph GT that serves as the start for denoising. Before reverse diffusion, we encode M pharmacological attributes (y1, . . . , yM) into a property-aware vector:

where cluster encoder maps numerical attributes (e.g., synthetic accessibility) and one-hot encoder maps categorical attributes (e.g., permeability). The resulting vector y is used as a property-aware control signal for the denoising network.

Reverse denoising with joint constraints

The reverse diffusion starts from GT and iteratively reconstructs molecular structures under the joint control of multi-pharmacological properties y and structural guidance (GID, GOOD):

where fθ(·) is a denoising network based on Graph Transformer with joint constraints.

Task-guided property optimization. At each step t to t - 1, the denoiser employs a Condition Norm module to incorporate the multi-property vector y into feature updates:

where hl denotes hidden features at layer l. This process steers the denoising towards molecules to satisfy target design, such as high permeability and low synthetic accessibility. The task loss is defined as:

where Lm denotes the supervised loss associated with each pharmacological property.

Distribution-consistency constraint. We train distribution consistency using the SSG module dSSGintroduced in Eq. 3. And we define the structural distribution consistency loss as:

where dSSG ID (·)and dSSG OOD (·) denote SSG diversity to ID and OOD guidance sets, and β1, β2 > 0 control the balance between attraction to ID and repulsion from OOD.

Dual-Constraint Training Objective. The overall loss in Stage 3 is:

where Ldiff denotes the diffusion reconstruction loss and λ1, λ2 > 0 are the weights.

From Noise to Molecule Generation. After training, molecule generation corresponds to the Sampling stage in Figure 1. Starting from a randomly initialized noisy molecular graph GT , we use the denoising network fθ(·), and simulate the reverse Markov chain pθ(Gt-1 | Gt ,C),C = (y,GID,GOOD). From t = T down to t = 0 to obtain the final sample G0 .

Denoising Network Architecture. The denoising network is implemented as a Graph Transformer that updates node and edge features. In Figure 2, it combines Multi-head Selfattention Block (MSA) [Peebles and Xie, 2023; Dwivedi and Bresson, 2021] with Gated condition Norm Block to support task-aware, distribution-guided generation. The MSA is:

Multi-pharmacological properties are incorporated via the Condition Norm module in Eq. 8, keeping latent representations aligned with the CNS drug design objectives throughout denoising. To fuse property-aware signals with structuredistribution information, we employ a gated fusion block:

where σ(MLP(h)) acts as an element-wise gate applied to the attention output MSA(h), enabling an adaptive balance between property condition and structure features along the denoising trajectory.

Consequently, we enforce a dual constraint regime of property-aware task optimization and distribution-consistent structural guidance, enabling the model to generate molecules that satisfy multi-pharmacological property requirements while remaining structurally plausible distributions.

4 Experiments

4.1 Experiment Setup

Datasets. We evaluate on twoCNS-related benchmarks and one large-scale small molecule dataset. For CNS drug design, we adopt the BBBP and BACE subsets of MoleculeNet [Wu et al., 2018], which provide categorical property labels for BBB permeability and β-site amyloid precursor proteincleaving enzyme 1 (BACE-1) inhibition, and we further assign each molecule a synthetic accessibility (SAS) score as a numerical property [Dichiara et al., 2009]. To assess generation over a broader region of chemical space, we additionally include the QM9 dataset comprising approximately 130k small organic molecules [Ramakrishnan et al., 2014].

Evaluationmetrics. For CNS benchmarks (BBBP, BACE), we evaluate the following metrics: Validity (fraction of chemically valid molecules), Diversity (internal structural diversity among generated molecules), Similarity (fragmentbased similarity to the reference set), Distance (Fr´echet ChemNet Distance to the reference distribution), Synthe.MAE (mean absolute error between generated and target synthetic accessibility scores), and Property.Acc (classification accuracy on task-related categorical conditions such as BBBP/BACE labels, computed by an oracle). For QM9, we focus on Validity, Diversity, and Distance to assess basic chemical correctness, structural variability, and distributional differences in a large-scale small molecule setting. The evaluation oracle is implemented as a random forest trained on task-related molecules; therefore, Property.Acc should be interpreted as a surrogate indicator of conditional consistency rather than direct experimental evidence of pharmacological activity. Lower Distance and Synthe.MAE, and higher Validity, Diversity, Similarity, and Property.Acc indicate better performance.

Baselines. We compare with a broad set of molecular generative models, including traditional methods GraphGA [Jensen, 2019], MARS [Xie et al., 2021], JTVAE with Bayesian optimization (JTVAE-BO) [Jin et al., 2018], and LSTM-HC on SMILES [Brown et al., 2019]; continuous diffusion models GDSS [Jo et al., 2022] and MOOD [Lee et al., 2023]; and discrete diffusion models DiGress [Vignac et al., 2023] and GraphDiT [Liu et al., 2024]. On QM9 experiments, we further include DiffMC-Gen [Yang et al., 2025]. All baselines are implemented following their original settings and evaluated under the same protocol as ours.

4.2 Comparison with Baseline Methods

Comparisons between D2G-TO and baseline models on the CNS benchmarks and the QM9 dataset are reported in Tables 1-3 and Figure 3-4. On BBBP and BACE, D2G-TO delivers strong performance and improves Property Acc. by 3% over the best competing method. On the large-scale QM9 dataset, D2G-TO also achieves the best overall rank among all baselines. The bar plots, radar and petal charts in Figure 34 further show that D2G-TO attains a more favorable balance among Validity, Diversity, and distributional Distance metrics than traditional generators, continuous diffusion models, and discrete diffusion models.

4.3 Model Analysis

Ablation Studies

To verify the effectiveness of the components in our D2GTO model, we design ablation studies on BBBP and BACE benchmarks. In Table 4, introducing SSG guidance consistently improves Property.ACC and reduces both Distance and Synthe MAE. These changes indicate that the generated molecules more closely follow the target chemical-space distribution and exhibit enhanced synthetic accessibility. While Validity, Coverage and Diversity are at least maintained compared with the w/o SSG guidence. These results show that the proposed guidance improves distributional quality without harming the basic generative capability.

Analysis of Structural Distribution Consistency

To further analyze structural distribution consistency, we project each generated molecule into the Structural Distribution State Space (SDSS) as a 2D descriptor [dID, dOOD]. This SDSS effectively filters out semantic redundancy and high dimensional interference, and we then compute pairwise Euclidean distances in this space. On BBBP dataset, we visualize (i) similarity maps between generated molecules and guidance graphs, and (ii) sample×sample structural similarity matrices based on ∆ = dOOD - dID for Stage 1 (without SSG guidance) and Stage 3 (D2G-TO with SSG guidance). These plots characterize the distribution of generated molecules under an ID/OOD structural reference.

In Figures. 5, the Stage 1 (a) similarity map is dominated by warm colors and indicates that many molecules violate the "close to ID, far fromOOD" requirement. But the Stage 3 (b) map is cooler color and shows that most samples now satisfy this constraint under SSG guidance. For the sample×sample matrices, molecules are first sorted in descending order of ∆. The Stage 1 (c) matrix appears nearly homogeneous blue with weak clustering, making it difficult to distinguish structurally good and bad samples. In contrast, the Stage 3 (d) matrix exhibits clear block structures that distinctly separate higher quality samples like guidance from lower quality samples like OOD. These observations demonstrate that SSG guidance substantially improves distributional consistency.

4.4 Case Study

We conduct a case study on AD by targeting BACE-1. D2GTO generates 10,000 molecules per design space from two complementary spaces (BACE-focused and BBBP-focused). After threshold filtering (1,689), ADMET screening (118), and multi-objective Pareto selection [Swanson et al., 2024; Deb et al., 2002], we obtain seven candidates. To assess structural plausibility, we dock these molecules into human BACE-1 receptor (PDB: 7MYI). And we additionally docked the clinical-stage BACE inhibitor LY3202626 as a reference [McKinzie et al., 2021].

Docking results suggest that the seven generated candidates are predicted to adopt plausible binding poses in the BACE-1 pocket, but due to limited manuscript space, we show only four representative binding poses in Figure 6. LY3202626 shows a docking score of -8.7kcal/mol, featuring hydrogen bonds with THR231/ASP228/ASP32 and TYR71-centered aromatic interactions, with additional contacts from VAL69/ILE118 and SER35. Molecule (a) adopts a compact pose stabilized by key interactions with TYR71 with additional hydrophobic stabilization from LEU30. Molecule (b) retains the key ASP32-TYR71-ILE118 interaction pattern and appears more favorable from a developability perspective. Molecule (c) shows the best docking score (-10.5kcal/mol), featuring an extensive polar contact network with ASP32/ASN37/ASP228 and aromatic anchoring near TYR71. Molecule (d) provides an alternative scaffold, characterized by TYR71-centered aromatic interactions and supporting contacts from TRP76/PHE108. Overall, these candidates span diverse chemotypes with competitive docking scores and better developability, supporting D2G-TO as a computational prioritization pipeline for early-stage virtual screening.

5 Conclusion

In this paper, we propose D2G-TO, a task-aware and OOD-guided discrete diffusion framework for CNS-oriented molecular graph generation. We introduce Structural Similarity Guidance (SSG) module that explicitly steers the generation process toward ID regions and away from OOD structures. By jointly coupling multi-pharmacological properties with distribution-aware structural guidance, D2G-TO enforces a dual task-structure guidance scheme within a unified generative model. The AD case study further shows that D2G-TO can prioritize candidates predicted to satisfy key CNS constraints, supporting its utility for early-stage in silico drug discovery.

Acknowledgments

This work was supported by National Natural Science Foundation of China (22273010).