Graph DiT for Multi-Conditional Molecular Generation

• The paper demonstrates that Graph DiT integrates discrete score-based diffusion with composable, per-property guidance to enable multi-conditional molecular generation.
• Multi-conditional molecular generation leverages continuous-time Markov chains and discrete concrete scores to ensure valid and diverse molecule synthesis.
• Empirical results show up to a 15.3% improvement in controllability and near-unity validity, underscoring its robustness in addressing complex, multi-objective design challenges.

Multi-conditional molecular generation is the paradigm in which generative models synthesize molecular graphs under the simultaneous control of multiple property constraints. "Multi-conditional" refers to the use of both numerical and categorical properties—such as solubility, bioactivity class, synthetic accessibility, or gas permeability—as conditioning signals, enabling the design of molecules that satisfy complex, multi-objective design criteria. In recent years, score-based diffusion models on molecular graphs, and in particular the Graph DiT and subsequent Composable Score-based Graph Diffusion (CSGD) frameworks, have established a flexible, high-fidelity solution to this problem. These approaches explicitly address the challenge of handling property subsets, enable explicit per-property guidance, and mitigate the distributional challenges arising from training–test mismatches in the discrete molecular domain (Qiao et al., 11 Sep 2025, Liu et al., 2024).

1. Mathematical Formulation of Multi-Conditional Molecular Generation

Given a molecular graph $G = (\mathbf{X}, \mathbf{E})$ with node-type sequence $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$ and adjacency or bond-type matrix $\mathbf{E} \in \mathcal{E}^{n \times n}$, the objective is to sample $G$ from the conditional distribution

$$p(G \mid \mathcal{C}) \propto p_{\mathrm{data}}(G) \prod_{m=1}^M p(c_m \mid G)$$

where $\mathcal{C} = \{c_1, ..., c_M\}$ denotes the full set of property constraints (numerical or categorical). Critically, in practical applications, sampling may target an arbitrary subset $\mathcal{C}' \subseteq \mathcal{C}$, motivating architectures that support compositional, subset-based conditioning.

2. Score-Based Diffusion Models for Discrete Graphs

2.1 Forward–Reverse Processes

In CSGD, the forward (diffusion) process for nodes and edges is governed by continuous-time Markov chains, with independent categorical transitions:

$$\frac{\mathrm{d}}{\mathrm{d} t} p_t^{X} = Q_t^X p_t^{X}, \qquad \frac{\mathrm{d}}{\mathrm{d} t} p_t^E = Q_t^E p_t^E$$

where $Q_t^X = \sigma(t) Q^X$, $Q_t^E = \sigma(t) Q^E$ for fixed transition matrices $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$0 and a scalar noise schedule $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$1. The forward kernels have closed-form, with

$\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$2

and analogously for edges.

The reverse (denoising) process is parameterized not by continuous gradients but by discrete "concrete scores," which are ratios of marginal conditional probabilities (e.g., $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$3 for nodes). These are learned using a diffusion-weighted denoising score entropy loss, aligning with the categorical structure of molecular graphs.

2.2 Model Architecture

The neural denoiser $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$4 is instantiated as a stack of Graph Transformer layers (the GraphDiT backbone). Each layer employs adaptive layer normalization (AdaLN), with scaling and bias determined by the sum of a learned embedding of the joint property vector $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$5 and a time embedding $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$6. Property embedding is performed individually for each property type: categorical via look-up, numerical via an MLP. The final transformer layer branches into node-wise and edge-wise heads, outputting the concrete scores $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$7, $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$8 (Qiao et al., 11 Sep 2025, Liu et al., 2024).

3. Conditional Guidance and Probability Calibration

3.1 Composable Guidance (CoG)

CSGD introduces Composable Guidance as an extension of classifier-free guidance to the multi-property, discrete-graph domain. Rather than requiring joint encoding of all conditions, the denoiser is trained to predict both unconditional and single-property-guided scores. At sampling, for any subset $\mathbf{X} = (x^1, ..., x^n) \in \mathcal{X}^n$9 with per-condition guidance strengths $\mathbf{E} \in \mathcal{E}^{n \times n}$0, the composite score is

$\mathbf{E} \in \mathcal{E}^{n \times n}$1

This enables arbitrary, subset-wise control of molecular properties at test time.

3.2 Probability Calibration (PC)

High guidance scales can push scores out of the support encountered during training, leading to negative or extreme transition probabilities and diminished generation fidelity. PC applies dynamic thresholding: negative probabilities are zeroed, and lower/upper quantiles are clipped and rescaled to $\mathbf{E} \in \mathcal{E}^{n \times n}$2, followed by a temperature transform to restore valid kernels: $\mathbf{E} \in \mathcal{E}^{n \times n}$3 The percentile and temperature parameters $\mathbf{E} \in \mathcal{E}^{n \times n}$4 are set to empirically robust values, e.g., $\mathbf{E} \in \mathcal{E}^{n \times n}$5, $\mathbf{E} \in \mathcal{E}^{n \times n}$6, $\mathbf{E} \in \mathcal{E}^{n \times n}$7 (Qiao et al., 11 Sep 2025).

4. Algorithmic Workflow

The multi-conditional sampling procedure, as realized in CSGD, proceeds by:

1. Initializing from a base noisy distribution.
2. At each reverse step, computing the unconditional and single-property-guided concrete scores for the relevant subset.
3. Composing these via the CoG rule.
4. Building the guided reverse kernel and applying probability calibration.
5. Independently resampling each node and edge token from the calibrated reverse kernel.
6. Iterating until reaching $\mathbf{E} \in \mathcal{E}^{n \times n}$8 and returning the denoised graph.

The formal pseudocode is specified exactly in (Qiao et al., 11 Sep 2025), which demonstrates highly modular, scalable conditional sampling.

5. Empirical Performance and Comparisons

Comprehensive experiments on polymer and drug-likeness datasets demonstrate the quantitative superiority of CSGD and related Graph DiT-class architectures over prior joint-conditioning and continuous-relaxation methods. Typical results for CSGD (Uniform schedule) include:

Dataset	Validity (%)	Diversity	Similarity	FCD	Controllability (MAE/Accuracy)
Polymers	99.38	0.8683	0.9541	6.44	Synth MAE 1.19, O$\mathbf{E} \in \mathcal{E}^{n \times n}$9 0.61
HIV	83.36	0.8945	0.9519	6.05	MAE 0.356, Acc 0.990

Averaged across tasks, CSGD achieves a 15.3% improvement in controllability (property-matching) relative to previous best methods, while validity rates approach unity under the uniform diffusion kernel. Ablations confirm that discrete score-based training provides a 7–13% gain in validity versus mean-prediction, and the addition of PC yields a further 1–3% gain in both validity and controllability. CoG outperforms joint CFG by enabling explicit, per-property subset guidance (Qiao et al., 11 Sep 2025).

6. Comparative Landscape with Other Conditional Generative Models

CSGD and its Graph DiT progenitors contrast with alternate frameworks:

• Sequential Graph Generators: MolRNN and similar architectures perform autoregressive graph construction with condition vectors injected additively into graph-convolutional layers (Li et al., 2018). While achieving high validity and enrichment under multi-conditional tasks, they lack the inherent parallelism, explicit subset guidance, and principled calibration of CSGD.
• Bernoulli Diffusion Models: GraphGUIDE utilizes discrete Bernoulli diffusion on binary graphs. While enabling a high degree of manual interpretability and the possibility of structural intervention, these models are limited to edge-based (structural) properties and do not implement learned embeddings for arbitrary property types (Tseng et al., 2023).
• String-based Generation: STGG+ leverages spanning-tree SMILES representations and leverages random masking, classifier-free guidance, and an auxiliary property-prediction loss with self-criticism (Jolicoeur-Martineau et al., 2024). This approach yields state-of-the-art performance on in/out-of-distribution conditional generation, but the underlying molecular representation and property integration mechanisms differ substantially from CSGD’s combinatorial score-based conditioning.
• Dual-protein and Feature Fusion Models: DiffDTM extends conditional generation to structure-free, dual-protein contexts via fused graph and sequence embeddings, demonstrating high validity and bioactivity for joint protein targets (Huang et al., 2023).
• Loop-guided Score-based Models: Diffusion Twigs decomposes the marginals over structure and properties into "trunk" and "stem" processes with explicit loop guidance, orchestrating property–structure dependencies stepwise during sampling (Mercatali et al., 2024).

CSGD is distinguished as the first model to extend score matching for truly discrete graphs via concrete scores, supporting principled, compositional subset guidance and robust transition calibration. Empirically, this enables state-of-the-art controllability, validity, and distributional fidelity under both in-distribution and out-of-distribution multi-property constraints (Qiao et al., 11 Sep 2025).

7. Significance and Implications

The development of CSGD and Graph DiT establishes a general framework for multi-conditional, property-controllable molecular graph generation. The combination of concrete discrete score estimation, composable guidance, and dynamic probability calibration solves key theoretical and practical limitations of prior approaches—enabling end-users to flexibly specify complex property subsets, control guidance strengths, and ensure robust, chemically valid output distributions. Applications extend across material discovery, drug design, and inverse design scenarios where fidelity to multiple property targets is essential (Qiao et al., 11 Sep 2025, Liu et al., 2024). A plausible implication is that the general principles of composable, property-aware, score-based generative modeling are applicable beyond chemistry, to other discrete, condition-structured graph domains.