Modular Flows: Differential Molecular Generation
Abstract
Generating new molecules is fundamental to advancing critical applications such as drug discovery and material synthesis.
Flows can generate molecules effectively by inverting the encoding process, however, existing flow models either require artifactual dequantization or specific node/edge orderings, lack desiderata such as permutation invariance, or induce discrepancy between encoding and decoding steps that necessitates post hoc validity correction.
We circumvent these issues with novel continuous normalizing E(3)-equivariant flows, based on a system of node ODEs coupled as a graph PDE, that repeatedly reconcile locally toward globally aligned densities.
Our models can be cast as message passing temporal networks, and result in superlative performance on the tasks of density estimation and molecular generation.
In particular, our generated samples achieve state of the art on both the standard QM9 and ZINC250K benchmarks.
Introduction
Generative models have rapidly become ubiquitous in machine learning with advances from image synthesis (Ramesh et al., 2022) to protein design (Ingraham et al., 2019).
Molecular generation (Stokes et al., 2020) has also received significant attention in discovering new drugs and materials.
However, searching for valid molecules in large discrete spaces is challenging: drug-like structures range between 10 23 and 10 60 but only a tiny fraction ( ∼ 10 8 ) - has been synthesized (Polishchuk et al., 2013; Merz et al., 2020).
Thus, learning representations that exploit appropriate molecular inductive biases (e.g., spatial correlations) becomes crucial.
Earlier models focused on generating sequences based on the SMILES notation (Weininger, 1988), which were supplanted by generative models that capture valuable spatial information via embedding molecular with some graph neural network (GNNs) (Scarselli et al., 2009; Garg et al., 2020).
This include variants of GANs (Goodfellow et al., 2014; Maziarka et al., 2020) which suffer the problem of mode collapse,
Variational Auto Encoders (VAE) (Kingma and Welling, 2013; Lim et al., 2018; Jin et al., 2018) which are susceptible to a distributional shift between the training data and the generated samples, and Normalizing Flows (Dinh et al., 2014, 2016).
Flows are appealing since they enable estimating (and sampling from) complex data distributions using a sequence of invertible transformations from a tractable continuous distribution.
Molecules are discrete, so many flow models (Madhawa et al., 2019; Honda et al., 2019; Shi et al., 2020) add noise during encoding leading to dequantization procedure which begets distortion and issues related to convergence (Luo et al., 2021).
Moreover, many methods employ post hoc correction to ensure validity (Zang and Wang, 2020), effecting a discrepancy between the encoding and the decoded distributions.
We propose a coupled continuous normalizing E(3)-equivariant flows tailored to molecule generation, that bestow generative capabilities from neural partial differential equation (PDE) models on graphs (Chamberlain et al., 2021; Poli et al., 2019; Iakovlev et al., 2020).
We seek to bring their efficacy and elegance as tools to generate complex objects, such as molecules.
Specifically, a flow is associated with each node, which is conjoined as a joint ODE system conditioned on neighboring nodes.
While these flows originate independently, they adjust progressively toward more complex joint distributions via interacting with the neighboring flows.
We call the proposed method Modular Flows ( ModFlow ) to underscore that each node can be regarded as a module that coordinates with other modules.
Modular Flows
We focus on unsupervised learning of an underlying graph density p ( G ) using a dataset D
= { G n } Nn =1 of observed molecular graphs G n .
We learn a generative flow model p θ ( G ) specified by flow parameters θ and use it to sample novel high-probability molecules.
2.1 Molecular Representation
Graph representation.
We represent the molecular graph G =
( V, E ) as a tuple of vertices
V =
( v 1 , . .
. , v M ) and edges E ⊂ V × V .
Each vertex takes value as: v ∈ A
= { C , H , N , . . .
} ; while the edges e ∈ B = { 1 , 2 , 3 } represent the type of bond.
We assume that conditioned on the edges, the graph likelihood factorizes as a Categorical distribution over vertices given their latent representations:
[Equation]
where z i = ( z i C , z
i H , . . . )
∈ R |A| is a set of atom score parameters of node i, and σ is the softmax function.
We can obtain an alternative tree representation of molecules, similar to Jin et al.
(2018) where we restrict clusters to be ring sub-structures.
We obtain an extended alphabet A tree =
{ C , H , N , . . .
, C 1 , C 2 , . . .
} , where each cluster label C r corresponds to the some ring-substructure in the label vocabulary χ .
For more details see Appendix A.2.
2.2 Differential modular flows
We propose to model the atom scores z i ( t ) as a Continuous-time Normalizing Flow (CNF) (Grathwohl et al., 2018) over time t ∈ R + .
We assume the initial scores at time t = 0 as z i (0) ∼ N (0 , I ) for each node i .
Node scores evolve in parallel over time by a differential equation, where N i is the set of neighbors, x i , x N i is the positional information, and θ are the parameters of the flow function.
By collecting all node differentials we obtain a modular joint, coupled ODE, which is our key contribution:,
[Equation]
[Equation] The above system is usually solved via an ODE solver where gradients are computed via adjoint sensitivity method (Kolmogorov et al., 1962), which incurs a low memory cost, and explicitly controls numerical error.
Notably, moving towards modular flows translates sparsity also to the adjoints.
Proposition 1: Modular adjoints are sparser than regular adjoints.
They can be computed by
[Equation]
where the partial derivatives ∂f∂ z =
[ ∂f i ∂ z j ] ij are sparse (See Appendix A.1 for further details).
The differential function f must satisfy the natural equivariances and invariances of molecules like translation, rotational (and reflection) equivariant and permutation equivariances.
Therefore, we chose to use E(3)-Equivariant GNN (EGNN) (Satorras et al., 2021), which satisfies all the above criteria (See Appendix A.3 for details).
The input to the EGNN are the node embeddings and the geometric information (polar coordinates (2D) and spherical polar coordinates (3D)).
Interestingly, ModFlow can be viewed as a message passing temporal graph network (Rossi et al., 2020) as shown next.
Proposition 2: Modular Flows can be cast as message-passing Temporal Graph Networks (TGNs).
The operations are listed in Table 2, where ModFlow is subjected to a single layer of EGNN. (See Appendix A.4 for more details).
2.3 Training objective
Normalizing flows are trained to minimize KL divergence KL[ p data || p θ ] between the unknown data distribution p data and the flow-generated distribution p θ , which is equivalent to maximizing their cross-entropy E p data [log p θ ] (Papamakarios et al., 2021).
However, this requires a mapping from discrete graphs G to continuous atom scores z ( t ) .
We reduce the learning problem to maximizing the score cross-entropy E ˆ p data ( z ( T ))
[log p θ ( z ( T ))] , where we turn the observed set of graphs { G n } into a set of scores { z n } by using one-hot encoding z (0)
z ( T ) σ ( z ( T ))
G flow f f −softmax argmax
[Equation]
where onehot( G n ) is a matrix of size M ( n ) × |A ∫ | such that G n ( i, k )
= 1 if v i = a k ∈ A ∫ , 1 q is a vector with q entries each set to 1; A ∫ ∈ {A , A tree } ; and ϵ ∈ [0 , 1] is added to model the noise in estimating the posterior p ( z ( T ) | G ) due to short-circuiting the inference process as shown in Fig.2.
We exploit the (non-reversible) composition of the argmax and softmax operations to short-circuit the process, which keeps the forward and backward flows aligned.
We thus maximize an objective over N training graphs,
[Equation]
which factorizes over the size M ( n ) of the n ’th training molecule
In practice, we solve ODE integrals using numerical solvers, such as Runge-Kutta.
We delegate this task to a general solver of the form ODESolve ( z , f θ , T ) , where map f θ is applied for T steps starting on z .
An optimizer optim is also required for updating θ .
2.4 Molecular generation
We generate novel molecules by sampling an initial state z (0) ∼ N (0 , I ) based on structure, and running the modular flow forward in time until z ( T ) .
This procedure maps a tractable base distribution p 0 to some more complex distribution p T .
We follow argmax to pick the most probable label assignment for each node (Zang and Wang, 2020).
Experiments
We evaluated ModFlow models trained, variously, on 2D and 3D coordinates, and their tree representation respectively on the task of a molecular generation.
Notably, ModFlow achieves state-of-the-art results without validity checks or post hoc correction when evaluated against many existing state-ofthe-art methods.
We train and evaluate on ZINC250k (Irwin et al., 2012) and QM9 (Ramakrishnan et al., 2014) datasets.
The molecules are in kekulized form with hydrogens removed by the RDkit software (Landrum et al., 2013).
We used Validity , Uniqueness , Novelty and Reconstruction as metrics.
These metrics measure the molecules obeying valency rules, non-duplicate generations, not present in the dataset, and the fraction of molecules that can be reconstructed from their encoding.
We report the mean and the standard deviation from five different initialization and generate 50,000 molecules for evaluation.
All the implementation is done in PyTorch (Paszke et al., 2019).
The input concatenates time and scalar vocabulary scores, per node.
For more details, see Appendix A.5.
Tables 1 shows the results on QM9 and ZINC250K. ModFlow achieves state-of-the-art results across all metrics.
Notably, its reconstruction rate is 100% similar to other flow models; novelty and uniqueness scores are also very high.
Moreover, ModFlow surpassed all early methods on validity (95%-99%).
Additional results on property optimization and density estimation are shown in
Appendix A.6 and A.7.
[Table]
Conclusion
We proposed ModFlow , a new generative flow model where multiple flows interact locally according to a coupled ODE
[Table]
resulting in accurate modeling of graph densities and high-quality molecular generation without any validity checks or correction.
Interesting avenues open up, including the design of (a) more nuanced mappings between discrete and continuous spaces
[Table]
and (b) extensions of modular flows to (semi-)supervised settings
[Table]
Acknowledgments
The calculations were performed using resources within the Aalto University Science-IT project.
This work has been supported by the Academy of Finland under the HEALED project (grant 13342077)
A Appendix A.1 Derivation of modular adjoint
We present a standard adjoint gradient derivation (Bradley, 2019), and show that the adjoint of a graph neighbor- hood differential is sparse.
For completeness, we define an ODE system
[Equation]
where z ∈ R D is a state vector, ˙ z ∈ R D is the state time differential defined by the vector field function f and parameterised by θ .
The starting state is z 0 , and t, τ ∈ R + are time variables.
Our goal is to solve a constrained problem
[Equation]
where G is the total loss that consists of instant loss functionals g .
We desire to compute the gradients of the system ∇ θ G .
We optimise the constrained problem by solving the Lagrangian
[Equation]
The constraints are satisfied by the ODE definition.
Hence, ∇ θ L = ∇ θ G , and we can set values θ and µ freely.
We use a shorthand notation ∂a∂b = a b , and omit parameters from the functions for notational simplicity.
Applying the chain rule, we note that the gradient becomes
[Equation]
where the µ term drops out since it does not depend on parameters θ .
We apply integration by parts to swap the differentials in term λ ⊤ ˙
z θ , resulting in
[Equation]
Substituting this into previous equation and rearranging the terms results in
[Equation]
[Equation]
The last term is removed since z (0) not depend on θ as a constant, and thus z θ (0) = 0 .
The difficult term in the equation is z θ .
We remove it by choosing
[Equation]
[Equation]
Finally
[Equation]
we choose λ ( T ) = 0 which also removes the second-to-last term.
The choices lead to a final term
[Equation]
[Equation]
In the derivation the adjoint λ ( t ) =
∂
L ∂ z ( t ) ∈
R D represents the change of loss with respect to instant states, and is another ODE system that runs backwards from λ ( T )
=
until λ (0)
The adjoint depends on the instant loss state derivatives g z .
These are often only available for observations y j at observed timepoints t j .
This can be represented by having a convenient loss
[Equation] [Equation]
and now the term g z induces discontinuous jumps at observations.
This does not pose problems in practice, since we can integrate the ODE in continuous segments between the observation instants.
The sparsity of the adjoint evolution is evident from Equation 19, where the ˙
λ i is an inner product between λ and one column of ∂f∂ z , which is invariant to non-neighbors.
This gives the result
[Equation]
For tree decomposition of the molecules, we followed closely the procedure described in Jin et al. The rings as well as the nodes corresponding to each ring substructure were extracted using RDKit’s functions, GetRingInfo and GetSymmSSSR .
We restricted our vocabulary to the unique ring substructures in the molecules.
The vocabulary of clusters follows a skewed distribution over the frequency of appearance within the dataset.
In particular, only a subset ( ∼ 30) of ring substructures (labels) appear with high frequency in molecules within the vocabulary. Therefore, we simplify the vocabulary by only representing the 30 commonly occurring substructures of A tree . In Figure 3, we show some examples of these ring substructures for the two datasets. [Figure]
A.3 Equivariant Graph Neural Networks (EGNN
Equivariant Graph Neural Networks (EGNN) (Satorras et al., 2021) are E(3)-equivariant with respect to an input set of points.
The E(3) equivariance accounts for translation, rotation, and reflection symmetries, and can be extended to E( n ) group equivariance.
The inherent dynamics governing the EGNN can be described, for each layer l , as follows.
Here, h li and x li pertain, respectively, to the embedding for the node i and that for its coordinates; and a ij abstracts the information about the edge between nodes i and j .
[Equation] Initially, messages m ij are computed between the neighboring nodes via ϕ e . Subsequently, the coordinates of each node i are updated via a weighted sum of relative position vectors { ( x i − x j ) : j ̸ = i } with the aid of ϕ x . Finally, the node embeddings are updated based on the aggregated messages m i via ϕ h .
The aggregated message can be computed based on only the neighbors of a node by simply replacing the sum over j ̸
= i with a sum over j ∈ N i in these equations.
A.4 Connection to Temporal Graph Networks
Temporal Graph Networks (Rossi et al., 2020; Souza et al., 2022) are state-of-the-art neural models for embedding dynamic graphs.
A prominent class of these models consists of a combination of (recurrent) memory modules and graph-based operators, and rely on message passing for updating the embeddings based on node-wise or edge-wise events .
Specifically, adopting the notation from Rossi et al.
(2020), an interaction e ij ( t ) between any two nodes i and j at time t triggers an edge-wise event leading to the following steps.
First, a message m ′ ij ( t ) is computed based on the memory s i (cid:0) t − (cid:1) and s j (cid:0) t − (cid:1) of the two nodes just before time t via a learnable function msg (such as multilayer perceptron).
For each node i , the messages thus accrued over a small period due to interactions of with neighbors j are combined (via agg ) into an aggregate message m ′
i ( t ) .
This message, in turn, is used to update the memory of i to s i ( t ) via mem (implemented e.g., as a recurrent neural network).
Finally, the node embedding of i is revised based on its memory s i ( t ) , interaction e ij ( t ) and memory s j ( t ) of each neighbor j ∈ N i , as well as any additional node-wise events v i ( t ) involving i or any node in N i .
the spatial distance x
i − x j , and a ij to denote the attributes of the edge between i and j .
The functions ϕ e , ϕ x , and ϕ h are as defined in Satorras et al.
(2021), and summarized in A.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
It turns out (see Table 2) that ModFlow can be viewed as an equivariant message passing temporal graph network.
Interestingly, the coordinate embedding x i plays the role of the memory s i .
A.5 Implementation Details
We implemented the proposed models in PyTorch (Paszke et al., 2019).
1 We used a single layer for EGNN with embedding dimension 32 and aggregated information for each node from only its immediate neighbors.
For geometric (spatial) information, we worked with the polar coordinates (2D) or the spherical polar coordinates (3D).
We solved the ODE system with the Dormand–Prince adaptive step size scheme (i.e., the dopri5 solver).
The number of function evaluations lay roughly between 70 and 100 .
The models were trained for 50epochs with the Adam (Kingma and Ba, 2014) optimizer.
Time comparisons.
We found the training time of ModFlow to be slightly worse than one-shot discrete flow models that characterize the whole system using a single flow (recall that, in contrast, ModFlow associates an ODE with each node).
However, ModFlow is faster to train than the auto-regressive methods.
Note that computation is a crucial aspect of generative modeling for application domains with a huge search space, as is true for the molecules.
We report the computational effort (excluding the time for preprocessing) for generating 10000 molecules averaged across 5 independent runs in Table 3.
Notably, largely by virtue of being one-shot, ModFlow is able to generate significantly faster than the auto-regressive models such as GraphAF and GraphDF.
ModFlow also owes this speedup, in part, to obviate the need for multiple decoding (unlike, e.g., JT-VAE) as well as any validity checks.
A.6 Density Estimation
[Table] We generated our synthetic data in the fol- lowing way. We considered two variants of a chessboard pattern, namely, (i) 4 ×grid where every node takes a binary value, 0 or 1, and neighboring nodes have different values; and (ii) 16 × 16 grid where nodes in each block of 4 × 4 all take the same value (0 or 1), different from the adjacent blocks. We also experimented with a 20 × 20 grid describing alternating stripes of 0 s and 1 s.
Figure 4 shows that ModFlow can learn neural differential functions f θ that reproduce the patterns almost perfectly, indicat- ing sufficient capacity to model complex patterns.
That is, ModFlow is able to transform the initial Gaussian distribution into different multi-modal and discontinuous distributions
A.7 Property-targeted Molecular Optimization
The task of molecular optimization is to search for molecules that have better chemical properties.
We choose the standard quantitative estimate of drug-likeness (QED) as our target chemical property.
QED measures the potential of a molecule to be characterized as a drug.
We used a pre-trained ModFlow model f to encode molecules M into their embeddings Z = f ( M ) , and applied linear regression to obtain QED scores Y from these embeddings.
We then interpolate in the latent space of each molecule along the direction of increasing QED via several gradient ascent steps, i.e., updates of the form Z ′ =
Z + λ ∗ d Y d Z , where λ denotes the length of the search step.
The final embedding thus obtained is decoded as a new molecule via the reverse mapping f − 1 .
Figure 5 and Figure 6 show examples of the molecules decoded from the learned latent space using this procedure, starting with molecules having a low QED score.
Note that the number of valid molecules decoded back varies on the query molecule.
We report the discovered novel molecules sorted by their QED scores in Table 4.
Clearly, ModFlow is able to find novel molecules with high QED scores.
[Table]