Original: Modular Flows: Differential Molecular Generation
Revised: Modular Flows: Differential Molecular Generation

Original: Abstract
Revised: Abstract

Original: Generating new molecules is fundamental to advancing critical applications suchas drug discovery and material synthesis.
Revised: Generating new molecules is fundamental to advancing critical applications such as drug discovery and material synthesis.

Original: Flows can generate molecules effectivelyby inverting the encoding process, however, existing flow models either requireartifactual dequantization or specific node/edge orderings, lack desiderata such aspermutation invariance, or induce discrepancy between  encoding and decodingsteps that necessitates post hoc validity correction.
Revised: 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 the encoding and the decoding steps that necessitates post hoc validity correction.

Original: We circumvent these issueswith novel continuous normalizing E(3)-equivariant flows, based on a system ofnode ODEs coupled as a graph PDE, that repeatedly reconcile locally towardglobally aligned densities.
Revised: 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.

Original: Our models can be cast as message passing temporalnetworks, and result in superlative performance on the tasks of density estimationand molecular generation.
Revised: Our models can be cast as message passing temporal networks, and result in superlative performance on the tasks of density estimation and molecular generation.

Original: In particular, our generated samples achieve state of theart on both the standard QM9 and ZINC250K benchmarks.
Revised: In particular, our generated samples achieve state of the art on both the standard QM9 and ZINC250K benchmarks.

Original:  Introduction
Revised:  Introduction

Original: Generative models have rapidly become ubiquitous in machine learning with advances from image synthesis(Ramesh et al., 2022) to protein design (Ingraham et al.
Revised: 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).

Original: Molecular generation (Stokes et al., 2020) hasalso received significant attention owing to its promisefor discovering new drugs and materials.
Revised: Molecular generation (Stokes et al., 2020) has also received significant attention owing to its promise for discovering new drugs and materials.

Original: Searching forvalid molecules in prohibitively large discrete spaces is,however, challenging: estimates for drug-like structuresrange between 10 23 and 10 60 but only a tiny fractionon the order of 10 8 - has been synthesized (Polishchuket al., 2013; Merz et al., 2020).
Revised: Searching for valid molecules in prohibitively large discrete spaces is, however, challenging: estimates for drug-like structures range between 10 23 and 10 60 but only a tiny fractionon the order of 10 8 - has been synthesized (Polishchuk et al., 2013; Merz et al., 2020).

Original: Thus, learning representations that exploit appropriate molecular inductivebiases (e.g., spatial correlations) becomes crucial.
Revised: Thus, learning representations that exploit appropriate molecular inductive biases (e.g., spatial correlations) becomes crucial.

Original: Earlier models focused on generating sequences basedon the SMILES notation (Weininger, 1988) used in
Revised: Earlier models focused on generating sequences based on the SMILES notation (Weininger, 1988) used in

Original: Chemistry to describe the molecular structure as strings.
Revised: Chemistry to describe the molecular structures as strings.

Original: However, they were supplanted by generative models that capture valuable spatial information such asbond strengths and dihedral angles, e.g., by embedding molecular graphs via some graph neural network(GNNs)
Revised: However, they were supplanted by generative models that capture valuable spatial information such as bond strengths and dihedral angles, e.g., by embedding molecular graphs via some graph neural network (GNNs)

Original: (Scarselli et al., 2009; Garg et al., 2020).
Revised: (Scarselli et al., 2009; Garg et al., 2020).

Original: Suchmodels primarily include variants of Generative Adversarial Networks (GANs), Variational Autoen- coders (VAEs), and Normalizing Flows (Dinh et al., 2014, 2016).
Revised: Such models primarily include variants of Generative Adversarial Networks (GANs), Variational Autoencoders (VAEs), and Normalizing Flows (Dinh et al., 2014, 2016).

Original: Besides known issues with  training, GANs (Goodfellow et al., 2014; Maziarka et al., 2020) suffer from the well-documented problem ofmode collapse, thereby generating molecules that lack diversity.
Revised: Besides known issues with their training, GANs (Goodfellow et al., 2014; Maziarka et al., 2020) suffer from the well-documented problem of mode collapse, thereby generating molecules that lack diversity.

Original: VAEs (Kingma and Welling, 2013; Lim et al., 2018; Jin et al., 2018), on the other hand, are susceptible to a distributional shift betweenthe training data and the generated samples.
Revised: VAEs (Kingma and Welling, 2013; Lim et al., 2018; Jin et al., 2018), on the other hand, are susceptible to a distributional shift between the training data and the generated samples.

Original: Moreover, optimizing for likelihood via a surrogatelower bound is likely insufficient to capture the complex dependencies inherent in  molecules.
Revised: Moreover, optimizing for likelihood via a surrogate lower bound is likely insufficient to capture the complex dependencies inherent in the molecules.

Original: Flows are especially appealing since, in principle, they enable estimating (and sampling from)complex data distributions using a sequence of invertible transformations on samples from a moretractable continuous distribution.
Revised: Flows are especially appealing since, in principle, they enable estimating (and sampling from) complex data distributions using a sequence of invertible transformations on samples from a more tractable continuous distribution.

Original: Molecules are discrete, so many flow models (Madhawa et al. ,2019; Honda et al., 2019; Shi et al., 2020) add noise during encoding and later apply a dequantizationprocedure.
Revised: Molecules are discrete, so many flow models (Madhawa et al., 2019; Honda et al., 2019; Shi et al., 2020) add noise during encoding and later apply a dequantization procedure.

Original: However, dequantization begets distortion and issues related to convergence (Luo et al.,2021).
Revised: However, dequantization begets distortion and issues related to convergence (Luo et al., 2021).

Original: Moroever, many methods segregate the generation of atoms from bonds, so the assembledstructure is often not a valid molecule and requires post hoc correction to ensure validity (Zangand Wang, 2020), effecting a discrepancy between the encoding and the decoded distributions.
Revised: Moroever, many methods segregate the generation of atoms from bonds, so the decoded structure is often not a valid molecule and requires post hoc correction to ensure validity (Zang and Wang, 2020), effecting a discrepancy between the encoding and the decoded distributions.

Original: Permutation dependence is another undesirable artifact of these methods.
Revised: Permutation dependence is another undesirable artifact of these methods.

Original: Some alternatives havebeen explored to avoid dequantization, e.g., (Lippe and Gavves, 2021) encodes molecules in acontinuous latent space via variational inference and jointly optimizes a flow model for generation.
Revised: Some alternatives have been explored to avoid dequantization, e.g., (Lippe and Gavves, 2021) encodes molecules in a continuous latent space via variational inference and jointly optimizes a flow model for generation.

Original: Discrete graph flows (Luo et al., 2021) also circumvent the many pitfalls of dequantization byresorting to discrete latent variables, and performing validity checks during the generative process.
Revised: Discrete graph flows (Luo et al., 2021) also circumvent the many pitfalls of dequantization by resorting to discrete latent variables, and performing validity checks during the generative process.

Original: However, discrete flows follow an autoregressive procedure that requires a specific ordering of nodesand edges during training.
Revised: However, discrete flows follow an autoregressive procedure that requires a specific ordering of nodes and edges during training.

Original: In general, one shot methods can generate much faster than discrete flows.
Revised: In general, one shot methods can generate much faster than discrete flows.

Original: We offer a different flow-based perspective tailored to molecules.
Revised: We offer a different flow-based perspective tailored to molecules.

Original: Specifically, we suggest coupledcontinuous normalizing E(3)-equivariant flows that bestow generative capabilities from neural partialdifferential equation (PDE) models on graphs.
Revised: Specifically, we suggest coupled continuous normalizing E(3)-equivariant flows that bestow generative capabilities from neural partial differential equation (PDE) models on graphs.

Original: Graph PDEs have been known to enable designingnew embedding methods such as variants of GNNs (Chamberlain et al., 2021), extending GNNsto continuous layers as Neural ODEs (Poli et al., 2019), and accommodating spatial information(Iakovlev et al., 2020).
Revised: Graph PDEs have been known to enable designing new embedding methods such as variants of GNNs (Chamberlain et al., 2021), extending GNNs to continuous layers as Neural ODEs (Poli et al., 2019), and accommodating spatial information (Iakovlev et al., 2020).

Original: We instead seek to bring to the fore their efficacy and elegance as tools tohelp generate complex objects, such as molecules, viewed as outcomes resulting from an interplay ofco-adapting latent trajectories (i.e., underlying dynamics).
Revised: We instead seek to bring to the fore their efficacy and elegance as tools to help generate complex objects, such as molecules, viewed as outcomes resulting from an interplay of co-adapting latent trajectories (i.e., underlying dynamics).

Original: Concretely, a flow is associated with eachnode of the graph, and these flows are conjoined as a joint ODE system conditioned on neighboringnodes.
Revised: Concretely, a flow is associated with each node of the graph, and these flows are conjoined as a joint ODE system conditioned on neighboring nodes.

Original: While these flows originate independently as samples from simple distributions, they adjustprogressively toward more complex joint distributions as they repeatedly interact with the neighboringflows.
Revised: While these flows originate independently as samples from simple distributions, they adjust progressively toward more complex joint distributions as they repeatedly interact with the neighboring flows.

Original: We view molecules as samples generated from the globally aligned distributions obtained aftermany such local feedback iterations.
Revised: We view molecules as samples generated from the globally aligned distributions obtained after many such local feedback iterations.

Original: We call the proposed method Modular Flows ( ModFlow ) tounderscore that each node can be regarded as a module that coordinates with other modules.
Revised: We call the proposed method Modular Flows ( ModFlow s) to underscore that each node can be regarded as a module that coordinates with other modules.

Original: Table 1summarizes the capabilities of ModFlow compared to some previous generative works.
Revised: Tablesummarizes the capabilities of ModFlow compared to some previous generative works.

Original: Contributions.
Revised: Contributions.

Original: We propose to learn continuous-time, flow based generative models, grounded ongraph PDEs, for generating molecules without  any validity correction. In particular, 
Revised: We propose to learn continuous-time, flow based generative models, grounded on graph PDEs, for generating molecules without resorting to any validity correction. In particular, • we propose ModFlow , a novel generative model based on coupled continuous normalizing

Original: [Table]
Revised: E(3)-equivariant flows.

Original: JT
Revised: ModFlow encapsulates essential inductive bias using PDEs, and defines multiple flows that interact locally toward a globally consistent joint density;

Original: [Table]
Revised: [Table]

Original:  ✓ ✓ ✗ ✗
Revised: JT-VAE ✓ ✓ ✗ ✗

Original: [Table] al.[Table] [Table]
Revised: Jin et al.

Original:  ✗ ✗ ✗ ✗
Revised: MRNN ✗ ✗ ✗ ✗

Original: [Table] et al.
Revised: Popova et al.

Original: GraphAF ✗ ✗ ✓ ✗
Revised: GraphAF ✗ ✗ ✓ ✗

Original: Shi et al.
Revised: Shi et al.

Original: GraphDF ✗ ✗ ✓ ✗
Revised: GraphDF ✗ ✗ ✓ ✗

Original: Luo et al.
Revised: Luo et al.

Original: [Table]

Original:  ✓ ✗ ✓ ✗
Revised: MoFlow ✓ ✗ ✓ ✗

Original: Zang and Wang (2020)
Revised: Zang and Wang (2020)

Original: GraphNVP ✓ ✗ ✓ ✗
Revised: GraphNVP ✓ ✗ ✓ ✗

Original: Madhawa et al.
Revised: Madhawa et al.

Original: [Table]

Original:  ✓ ✓ ✓ ✓ this work
Revised: ModFlow ✓ ✓ ✓ ✓ this work

Original: • we propose ModFlow , a novel generative model based on coupled continuous normalizingE(3)-equivariant flows.

Original: ModFlow encapsulates essential inductive bias using PDEs, anddefines multiple flows that interact locally toward a globally consistent joint density;

Original: • we encode permutation, translation, rotation, and reflection equivariance, with E(3) equivariant GNNs adapted to molecular generation, and can leverage 3D geometric information; • ModFlow is end-to-end trainable, non-autoregressive, and obviates the need for any external validity checks or correction; • empirically, ModFlow achieves state-of-the-art performance on both the standard QM9 (Ramakrishnan et al., 2014) and ZINC250K (Irwin et al., 2012) benchmarks.
Revised: • we encode permutation, translation, rotation, and reflection equivariance with E(3) equivariant GNNs adapted to molecular generation, and can leverage 3D geometric information; • ModFlow is end-to-end trainable, non-autoregressive, and obviates the need for any external validity checks or correction; • empirically, ModFlow achieves state-of-the-art performance on both the standard QM9 (Ramakrishnan et al., 2014) and ZINC250K (Irwin et al., 2012) benchmarks.

Original:  Related works
Revised:  Related works

Original: Generative models.
Revised: Generative models.

Original: Earlier attempts for molecule generation (Kusner et al., 2017; Dai et al., 2018)aimed at representing molecules as SMILES strings (Weininger, 1988) and developed sequencegeneration models.
Revised: Earlier attempts for molecule generation (Kusner et al., 2017; Dai et al., 2018) aimed at representing molecules as SMILES strings (Weininger, 1988) and developed sequence generation models.

Original: One of the challenges of these approaches is to learn complicated grammar rulesthat can generate syntactically valid sequences of molecules.
Revised: A challenge for these approaches is to learn complicated grammar rules that can generate syntactically valid sequences of molecules.

Original: Recently, representing molecules asgraphs has inspired new deep generative model for molecular generation (Segler et al., 2018; Samantaet al., 2018; Neil et al., 2018) ranging from VAEs (Jin et al., 2018; Kajino, 2019) to flows (Madhawaet al., 2019; Luo et al., 2021; Shi et al., 2020).
Revised: Recently, representing molecules as graphs has inspired new deep generative models for molecular generation (Segler et al., 2018; Samanta et al., 2018; Neil et al., 2018), ranging from VAEs (Jin et al., 2018; Kajino, 2019) to flows (Madhawa et al., 2019; Luo et al., 2021; Shi et al., 2020).

Original: The core idea is to learn to encode molecular graphsinto a latent space, and subsequently decode samples from the latent space to generate new molecules(Atwood and Towsley, 2016; Xhonneux et al., 2020; You et al., 2018).
Revised: The core idea is to learn to encode molecular graphs into a latent space, and subsequently decode samples from this latent space to generate new molecules (Atwood and Towsley, 2016; Xhonneux et al., 2020; You et al., 2018).

Original: Graph partial differential equations.
Revised: Graph partial differential equations.

Original: Graph PDEs is an emerging area that studies PDEs onstructured data encoded as graphs.
Revised: Graph PDEs is an emerging area that studies PDEs on structured data encoded as graphs.

Original: For instance, one can define a PDE on graphs to track the evolutionof signals defined over the graph nodes under some dynamics.
Revised: For instance, one can define a PDE on graphs to track the evolution of signals defined over the graph nodes under some dynamics.

Original: Graph PDEs have enabled, amongothers, design of new graph neural networks; see, e.g., works such as GRAND (Chamberlain et al.,2021),
Revised: Graph PDEs have enabled, among others, design of new graph neural networks; see, e.g., works such as GNODE (Poli et al., 2019),

Original: GNODE (Poli et al., 2019), NeuralPDE (Iakovlev et al., 2020), Neural operator (Li et al.,2020) and PDE-GCN (Eliasof et al., 2021).
Revised: NeuralPDE (Iakovlev et al., 2020), Neural operator (Li et al., 2020), GRAND (Chamberlain et al., 2021), and PDE-GCN (Eliasof et al., 2021).

Original: Different from all these works, we focus on using PDEsfor generative modeling of molecules (graph-structured objects).
Revised: Different from all these works, we focus on using PDEs for generative modeling of molecules (as graph-structured objects).

Original: Interestingly, ModFlow proposedin this work may be viewed as a new equivariant temporal graph network.
Revised: Interestingly, ModFlow proposed in this work may be viewed as a new equivariant temporal graph network (Rossi et al., 2020; Souza et al., 2022).

Original: Validity oracles.
Revised: Validity oracles.

Original: A key challenge of molecular generative models is to be able to generate validmolecules, according to various criteria for molecular validity or feasibility.
Revised: A key challenge of molecular generative models is to be able to generate valid molecules, according to various criteria for molecular validity or feasibility.

Original: It is a common practiceto use external chemical software as rejection oracles to reduce or exclude invalid molecules, or dovalidity checks as part of autoregressive generation.
Revised: It is a common practice to call on external chemical software as rejection oracles to reduce or exclude invalid molecules, or do validity checks as part of autoregressive generation (Luo et al., 2021;

Original: (Luo et al., 2021; Shi et al., 2020; Popova et al., 
Revised: Shi et al., 2020; Popova et al., 2019).

Original: An important open question has been whether generative models can learn to achieve highgenerative validity intrinsically , i.e., without being aided by oracles or performing additional checks.
Revised: An important open question has been whether generative models can learn to achieve high generative validity intrinsically , i.e., without being aided by oracles or resorting to additional checks.

Original: ModFlow represents a major step forward toward that goal.
Revised: ModFlow takes a major step forward toward that goal.

Original:  Modular Flows
Revised:  Modular Flows

Original: We focus on unsupervised learning of an underlying graph density p ( G ) using a datasetD ={ G n } Nn =1 of observed molecular graphs  G n  .
Revised: We focus on unsupervised learning of an underlying graph density p ( G ) using a dataset of observed molecular graphs D = { G n } Nn =1 .

Original: We learn a generative flow model p θ ( G ) specified byflow parameters θ , and use it to sample novel high-probability molecules.
Revised: We learn a generative flow model p θ ( G ) specified by flow parameters θ , and use it to sample novel high-probability molecules.

Original: 3.1 Molecular Representation
Revised: 3.1 Molecular Representation

Original: Graph representation.
Revised: Graph representation.

Original: We represent the molecular graph G =
Revised: We represent each molecular graph G =

Original: ( V, E ) of size M as a tuple ofvertices V
Revised: ( V, E ) of size M as a tuple of vertices V

Original: =
Revised: =

Original: ( v 1 , . .
Revised: ( v 1 , . .

Original: . , v M ) and edges E ⊂ V × V .
Revised: . , v M ) and edges E ⊂ V × V .

Original: Each vertex takes  value from an alphabet onatoms: v ∈ A
Revised: Each vertex takes a value from an alphabet on atoms: v ∈ A

Original: = { C , H , N , O , P , S , . . .
Revised: = { C , H , N , O , P , S , . . .

Original:  while the edges e ∈ B = { 1 , 2 , 3 } abstract the type of bond(i.e. , single, double, or triple).
Revised: } ; while each edge e ∈ B = { 1 , 2 , 3 } abstracts some bond type (i.e., single, double, or triple).

Original: We assume that, conditioned on the edges, the graph likelihoodfactorizes as a Categorical distribution over vertices given their latent representations:
Revised: We assume that, conditioned on the edges, the graph likelihood factorizes as a product of categorical distributions over vertices given their latent representations:

Original: [Equation]
Revised: [Equation]

Original: where z i = ( z i C , z
Revised: where z i = ( z i C , z

Original: i H , . . . )
Revised: i H , . . . )

Original: ∈ R |A| is a set of atom score parameters, where z ik ∈ R pertains tonode iand type k ∈ A , and σ is the softmax function
Revised: ∈ R |A| is a set of atom scores for node i such that z ik ∈ R pertains to type k ∈ A , and σ is the softmax function

Original: [Equation]
Revised: [Equation]

Original: which turns the real-valued scores
Revised: which turns the real-valued scores

Original: z iof a node i into normalized probabilities.
Revised: z i into normalized probabilities.

Original: ModFlow alsosupports 3D molecular graphs that contain atomic coordinates and angles.
Revised: ModFlow also supports 3D molecular graphs that contain atomic coordinates and angles as additional information.

Original: Tree representations.
Revised: Tree representations.

Original: We can obtain an alternative representation for molecules based on decomposing a molecule into a tree-like structure, by contracting certain vertices into a single node(denoted as a cluster) such that the molecular graphG becomes acyclic.
Revised: We can obtain an alternative representation for molecules: we can decompose each molecule into a tree-like structure, by contracting certain vertices into a single node (denoted as a cluster) such that the molecular graph becomes acyclic.

Original: Followingthe decomposition investigated by Jin et al.
Revised: Following Jin et al.

Original: (2018), we restrict these clusters to ring-substructures presentin the molecular data, in addition to the atom alphabet.
Revised: , we restrict these clusters to ring substructures present in the molecular data, in addition to the atom alphabet.

Original: Thus, we obtain an extended alphabet A tree =
Revised: Thus, we obtain an extended alphabet A tree =

Original:  { C, H , N , . . . , C 1 , C 2 , . . . } , where each cluster label C r corresponds tothe some ringsubstructure in the label vocabulary χ .
Revised: A ∪ { C 1 , C 2 , . . . } , where each cluster label C r corresponds to some ring substructure in the label vocabulary χ .

Original: We  reduce the vocabulary to the 30 most commonly occurringsubstructures of A tree .
Revised: We then reduce the vocabulary to the 30 most commonly occurring substructures of A tree .

Original: For further details, see Appendix A.1.
Revised: For further details, see Appendix A.2.

Original: 3.2 Differential modular flows
Revised: 3.2 Differential modular flows

Original: Normalizing flows (Kobyzev et al., 2021) provide a general recipe for constructing flexible probabilitydistributions, used in density estimation (Cramer et al., 2021;
Revised: Normalizing flows (Kobyzev et al., 2021) provide a general recipe for constructing flexible probability distributions, used in density estimation (Cramer et al., 2021;

Original: Huang et al., 2018) and generativemodeling (Zhen et al., 2020; Zang and Wang, 2020).
Revised: Huang et al., 2018) and generative modeling (Zhen et al., 2020; Zang and Wang, 2020).

Original: We propose to model the atom scores z i ( t ) as a
Revised: We propose to model the atom scores z i ( t ) as a

Original: Continuous-time Normalizing Flow (CNF) (Grathwohl et al., 2018) over time t ∈ R + .
Revised: Continuous-time Normalizing Flow (CNF) (Grathwohl et al., 2018) over time t ∈ R + .

Original: We assumethe initial scores at time t = 0 follow an uninformative Gaussian base distribution z i (0) ∼ N (0 , I )for each node i .
Revised: We assume the initial scores at time t = 0 follow an uninformative Gaussian base distribution z i (0) ∼ N (0 , I ) for each node i .

Original: Node scores evolve in parallel over time by a differential equation,
Revised: Node scores evolve in parallel over time according to the differential equation

Original: [Equation]
Revised: [Equation]

Original: where N i = {z j : ( i, j ) ∈ E } is the set of neighbor scores  at time t , x i , x N i is the positionalinformation (2D/3D)  of node and the neighbours, and θ are the parameters of the flow function f tobe learned.
Revised: where N i = { j : ( i, j ) ∈ E } is the set of neighbors of node i and z N i ( t ) = { z j ( t ) : j ∈ N i } the scores of the neighbors at time t ; x i and x N i denote, respectively, the positional (2D/3D) information of i and its neighbours; and θ denotes the parameters of the flow function f to be learned.

Original: By collecting all node differentials we obtain a modular joint, coupled ODE, which is our key contribution:,
Revised: Stacking together all node differentials, we obtain a modular system of coupled ODEs:

Original: [Equation]
Revised: [Equation]

Original: The coupled system of ODEs is equivalent to a graph PDE (Iakovlev et al., 2020; Chamberlain et al.,2021), where the evolution of each nodeonly depends  on itsimmediate neighbors.
Revised: This coupled system of ODEs may be viewed as a graph PDE (Iakovlev et al., 2020; Chamberlain et al., 2021), where the evolution of each node depends only on its neighbors.

Original: The  flow induces  corresponding density changes by the  divergence of f, (Chen et al., 2018)
Revised: The joint flow induces a corresponding change in the individual densities in terms of divergence of f (Chen et al., 2018),

Original: [Equation]
Revised: [Equation]

Original: starting from the base distribution p 

Revised: starting from the base distribution p

Original: ( z i (0))
Revised: ( z i (0))

Original: =
Revised: =

Original: N ( z i (0) | 0 , I ) .
Revised: N ( z i (0) | 0 , I ) .

Original: The trace picks only the diagonalelements of the Jacobian ∂f ∂ z , which interprets theneighborhood input  z N i term as a ‘control’ foreach node z i at each instant t .
Revised: The trace picks only the diagonal elements of the Jacobian ∂f∂ z , which interprets the input from neighbors, z N i , as a ‘control’ for each node z i at each instant t .

Original: The above system is usually solved via an ODE solver where gradientsare computed via  adjoint sensitivity method (Kolmogorov et al., 1962).
Revised: An ODE solver is used for such systems, and the gradients are computed via the adjoint sensitivity method (Kolmogorov et al., 1962).

Original: This approach incurs alow memory cost, and explicitly controls  numerical error.
Revised: This approach incurs a low memory cost, and explicitly controls the numerical error.

Original: Notably, moving towards modular flowstranslates sparsity also to the adjoints.
Revised: Notably, moving towards modular flows translates sparsity also to the adjoints.

Original: Preposition 1: Modular adjoints are sparser than regular adjoints.
Revised: Preposition 1: Modular adjoints are sparser than regular adjoints.

Original: They can be computed by
Revised: They can be computed as d λ i

Original: [Equation]
Revised: [Equation]

Original: i ∪{ i } λ j ( t ) ⊤ ∂f (cid:0) t, z i ( t ) , z N i ( t ) (cid:1) ∂ z j , (7) where the partial derivatives ∂f∂ z =
Revised: where the partial derivatives ∂f∂ z =

Original: [ ∂f i ∂ z j ] ij are sparse (See Appendix A.5 for  derivation).
Revised: [ ∂f i ∂ z j ] ij are sparse (see Appendix A.1 for the derivation).

Original: 3.3 Equivariant local differential
Revised: 3.3 Equivariant local differential

Original: Our goal is to have a differential function f
Revised: Our goal is to have a differential function f

Original: that is a PDE operator used in Eq. 4, and that satisfies thenatural equivariances and invariances of  molecules.
Revised: that is a PDE operator used in Equation 4, and that satisfies the natural equivariances and invariances of the molecules.

Original: Specifically, the function must be (i) translationequivariant: translating the input results in an equivalent translation of the output; (ii) rotational(and reflection) equivariant: rotating the input results in an equivalent rotation of the output; (iii)permutation equivariant: permuting the input results in the same permutation of the output.
Revised: Specifically, this function must be (i) translation equivariant: translating the input results in an equivalent translation of the output; (ii) rotational (and reflection) equivariant: rotating the input results in an equivalent rotation of the output; and (iii) permutation equivariant: permuting the input results in the same permutation of the output.

Original: Therefore,we chose to use E(3)-Equivariant GNN (EGNN) (Satorras et al., 2021), which is translation, rotationand reflection equivariant (E(n)), and permutation equivariant with respect to an input set of points(See Appendix A.2 for details).
Revised: Therefore, we chose to use E(3)-Equivariant GNN (EGNN) (Satorras et al., 2021), which is translation, rotation and reflection equivariant (E(n)), and permutation equivariant with respect to an input set of points (see Appendix A.3 for details).

Original: The input to the EGNN are the node embeddings and the geometricinformation (polar coordinates (2D) and spherical polar coordinates (3D)).
Revised: EGNN takes as input the node embeddings as well as the geometric information (polar coordinates (2D) and spherical polar coordinates (3D)).

Original: Interestingly, ModFlow canbe viewed as a message passing temporal graph network (Rossi et al., 2020) as shown next.
Revised: Interestingly, ModFlow can be viewed as a message passing temporal graph network (Rossi et al., 2020; Souza et al., 2022) as shown next.

Original: Preposition 2: Modular Flows can be cast as  Temporal Graph Networks (TGN).
Revised: Preposition 2: Modular Flows can be cast as message passing Temporal Graph Networks (TGNs).

Original: The operationsare listed in Table 2, where ModFlow is subjected to a single layer of EGNN.
Revised: The operations are listed in Table 2, where ModFlow is subjected to a single layer of EGNN.

Original: (See Appendix A.3 formore details). 
Revised: (See Appendix A.4 for more details). [Table]

Original: 3.4 Training objective
Revised: ∆ t, e ij ( t )) m ij ( t ) = ϕ e (cid:16) z i ( t ) , z j ( t ) , ∥ r ij ( t ) ∥ 2 ,a ij (cid:17) m ′ i ( t ) = agg (cid:0) { m ′ ij ( t ) | j ∈ N i } (cid:1) m i ( t )

Revised: = (cid:80) j ∈N ( i ) m

Revised: ij ˆ m ij ( t ) =

Revised: r ij ( t ) · ϕ x ( m ij ( t ))

Revised: ˆ m i ( t ) =

Revised: C (cid:80) j ∈N ( i ) ˆ m ij ( t )

Revised: Memory state s i ( t ) = mem( m ′ i ( t ) , s i ( t − ))

Revised: x i ( t + 1) =

Revised: x i ( t )

Revised: +

Revised: ˆ m i ( t )

Revised: Node z ′ i ( t ) =

Revised: (cid:80) j ∈N i h

Revised: ( s i ( t ) , s j ( t ) , e ij ( t ) , v i ( t ) , v j ( t ))

Revised: z i ( t + 1) =

Original:  Normalizing flows are predominantly trained to minimizeKL divergence KL[ p data || p
Revised: ϕ h ( z i ( t ) , m i ( t )) 3.4 Training objective Normalizing flows are predominantly trained to minimize KL[ p data || p

Original: θ ]  between theunknown data distribution p  and the flow-generated distribution p θ .
Revised: θ ] , i.e., the KL-divergence between the unknown data distribution p data and the flow-generated distribution p θ .

Original: This  is equivalent to maximizingtheir cross-entropy E p data [log p θ ]
Revised: This objective is equivalent to maximizing E p data [log p θ ]

Original: (Papamakarios et al., 2021).
Revised: (Papamakarios et al., 2021).

Original: However, this requires the observationsto live in the flow space, which is not true for us: the discrete graphs G and  continuous atom scoresz ( t ) reside in different spaces.
Revised: However, note that the discrete graphs G and the continuous atom scores z ( t ) reside in different spaces.

Original: To applynormalizing flows, a mapping between the spaces is needed.
Revised: Thus, in order to apply flows, a mapping between the observation space and the flow space is needed.

Original: [Figure] Earlier approaches use dequantisation to turn a graph G into a distribution of latent  states 
Revised: Earlier approaches use dequantisation to turn a graph G into a distribution of latent states, and argmax to deterministically map latent states to graphs (Zang and Wang, 2020).

Original: [Figure]
Revised: [Figure]

Original: and argmaxto deterministically map latent states to graphs (Zang and Wang, 2020).

Original: We  reduce the learning problem to maximizingthe score cross-entropy E 
Revised: We instead reduce the learning problem to maximizing E ˆ p data ( z ( T ))

Original: [Figure] z ( T )) [log p θ ( z ( T ))] , where we turn the observed set of graphs { G n } into a set of scores { z n }by usingone-hot encoding
Revised: [log p θ ( z ( T ))] , where we turn the observed set of graphs { G n } into a set of scores { z n } using

Original: [Equation]
Revised: [Equation]

Original: where onehot( G n ) is a matrix of size M ( n ) × |A
Revised: where onehot( G n ) is a matrix of size M ( n ) × |A

Original: ∫ | (i.e., rows equal to  number of nodes of G n andcolumns equal to the number of possible node labels) such that G n ( i, k )
Revised: ∫ | (i.e., rows equal to the number of nodes in G n and columns equal to the number of possible node labels) such that G n ( i, k )

Original: = 1 if v i = a k ∈ A ∫ , thatis if the vertex i is labeled with atom k , and 0 otherwise; 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 from G to z ( T ) skipping the intermediate dependencies,as shown in the  plate diagram of Figure 3.
Revised: = 1 if v i = a k ∈ A ∫ , that is if the vertex i is labeled with atom k , and 0 otherwise; 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 from G to z ( T ) skipping the intermediate dependencies, thereby inducing an unconditional distribution ˆ p data that is slightly different from the true data distribution p data . The plate diagram in Figure 3 summarizes the overall procedure.

Original: Effectively, we exploit the (non-reversible) composition of the argmax and softmax operations totransition from the continuous flow space to the discrete graph space, but skip this compositionaltogether in the reverse direction.
Revised: Effectively, we exploit the (non-reversible) composition of the argmax and softmax operations to transition from the continuous flow space to the discrete graph space, but skip this composition altogether in the reverse direction.

Original: Importantly, this short-circuiting allows ModFlow to keep theforward and backward flows between z (0) and z ( T ) completely aligned  unlike previous approaches.
Revised: Importantly, this short-circuiting allows ModFlow to keep the forward and backward flows between z (0) and z ( T ) completely aligned (i.e., reversible) unlike previous approaches.

Original: Wethus maximize an objective over N training graphs,
Revised: We maximize the following objective over N training graphs:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: which factorises over the size M ( n ) of the n ’th training molecule
Revised: which factorizes over the size M ( n ) of the n ’th training molecule

Original: The encoding probability followsfrom the flow Eq. 6, where z (0) is the reverse solution from z n at time  Tfollowing the flow fbackward in time until t = 0 .
Revised: The encoding probability follows from Equation 6, where z (0) can be traced by traversing the flow f backward in time starting from z n at time t = T until t = 0 .

Original: In practice we solve ODE integrals using  numerical solvers, such as
Revised: In practice we solve ODE integrals using a numerical solver such as

Original: Runge-Kutta.
Revised: Runge-Kutta.

Original: We  delegate this task to a general solver of the form ODESolve ( z , f θ , T ) , where mapf θ is applied for T steps starting on z .
Revised: We thus delegate this task to a general solver of the form ODESolve ( z , f θ , T ) , where map f θ is applied for T steps starting with z .

Original: An optimizer optim is also required for updating θ . 
Revised: An optimizer optim is also required for updating θ . 3.5 Molecular generation

Original: Require: Dataset D , iterations n iter , batch size
Revised: Given a molecular structure, we can generate novel molecules by sampling an initial state z (0) ∼ N (0 , I ) , and running the modular flow forward in time for T steps and obtain z ( T ) . This procedure maps a tractable base distribution p 0 to a more complex distribution p T . We follow argmax to pick the most probable label assignment for each node (Zang and Wang, 2020).

Original:  B , number of batches n B 1: Initialise parameters θ of ModFlow (EGNN) 2: for k = 1 , . .
Revised: We outline the procedures for training and generation in Algorithm 1 and Algorithm 2 respectively. Require: Dataset D , iterations n iter , batch size B , number of batches n B 1: Initialise parameters θ of ModFlow (EGNN) 2: for k = 1 , . .

Original: . , n iter do 3: for b = 1 , . .
Revised: . , n iter do 3: for b = 1 , . .

Original: . , n B do 4:
Revised: . , n B do 4:

Original: Sample D b = ( G 1 , . . .
Revised: Sample D b = { G 1 , . . .

Original: , G B ) ∼ D 5:
Revised: , G B } from D 5:

Original: Init z b ( T ) to z b (  G  b ; ϵ ) 6:z b (0) =
Revised: Define z b ( T ) := { z r ( T ) : G r ∈ D b } 6:

Original: ODESolve ( z b ( T ) , f θ , T ) 7: 
Revised: Set z b ( T ) to z b ( G b ; ϵ ) 7: L b = 1 B (cid:88)

Original: L b= log p θ ( G b  ) 8: end for 9: θ ←− optim ( 1 n B (cid:80) n B b =1 L b ; θ ) 10: end for
Revised: G r ∈ D b log p θ ( z r ( T )) , using z b (0) = ODESolve ( z b ( T ) , f − 1 θ , T ) 8: end for 9: θ ←− optim ( 1 n B (cid:80) n B b =1 L b ; θ ) 10: end for

Original: Algorithm  1:  z (0) ∼ N (0 , I ) 2:z ( T ) =
Revised: Algorithm 2 Generating with ModFlow 1: Sample z (0) ∼ N (0 , I ) 2:

Original: ODESolve ( z (0) , f − 1 θ , T ) 3: Assign labels by argmax( σ ( z ( T))) 3.5 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 basedistribution p 0 to some more complex distribution p T .
Revised: z ( T ) =

Original: We follow argmax to pick the most probablelabel assignment for each node (Zang and Wang, 2020). The pseudo algorithms are shown in Algorithms 1 and 2.
Revised: ODESolve ( z (0) , f θ , T ) 3: Assign labels by argmax( σ ( z ( T )))

Original:  Experiments
Revised:  Experiments

Original: We  evaluated ModFlow models trained, variously, on (i) 2D coordinates (ii) 3D coordinates (iii)2D coordinates + Tree representation,  (iv) 3D coordinates + Tree representation on the tasks ofdensity estimation and molecular generation.
Revised: We first demonstrate the ability of Modular Flows ( ModFlow ) to learn highly discontinuous synthetic patterns on 2D grids. We also evaluated ModFlow models trained, variously, on (i) 2D coordinates, (ii) 3D coordinates, (iii) 2D coordinates + tree representation, and (iv) 3D coordinates + tree representation on the tasks of molecular generation and optimization.

Original: We show that ModFlow compares favorably to otherprominent flow and non-flow based molecular generative models including GraphDF (Luo et al.,2021), GraphNVP (Madhawa et al., 2019), MRNN(Popova et al., 2019), GraphAF
Revised: Our results show that ModFlow compares favorably to other prominent flow and non-flow based molecular generative models, including GraphDF (Luo et al., 2021), GraphNVP (Madhawa et al., 2019), MRNN(Popova et al., 2019), and

Original:  (Shi et al., 2020)etc.
Revised: GraphAF (Shi et al., 2020).

Original: Notably, ModFlow achieves state-of-the-art results without validity checks or post hoc correction.
Revised: Notably, ModFlow achieves state-of-the-art results without validity checks or post hoc correction.

Original: We also provideablation results  to underscore the relevance of geometric features and equivariancetoward this superlative empirical performance.
Revised: We also provide results of our ablation studies to underscore the relevance of geometric features and equivariance toward this superlative empirical performance.

Original: 4.1 Density Estimation
Revised: 4.1 Density Estimation

Original: [Figure] We now demonstrate the high ex- [Figure]
Revised: We generated our synthetic data in the following way.

Original: pressive power of Modular Flows( ModFlow ) on learning highly discontinuous patterns on 2D grid graphs.

Original: We considered two variants of  chessboard pattern as (i) 4{+
+} × 4 grid whereevery node has opposite value to itsneighbors (ii) 16 × 16 grid whereblocks of 4 × 4 nodes have uniformvalues, but opposite across blocks, and(iii) a 20 × 20 grid describing alternating stripes.
Revised: 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.

Original: Figure 4 shows that wecan learna neural differential  f θfunction that reproduces the patterns almost perfectly, indicating sufficient capacity to model complexpatterns.
Revised: Figure 4 shows that ModFlow can learn neural differential functions f θ that reproduce the patterns almost perfectly, indicating sufficient capacity to model complex patterns.

Original: Our model is able to transform the initial Gaussian  into  multi-modal and discontinuousdistributions.
Revised: That is, ModFlow is able to transform the initial Gaussian distribution into different multi-modal and discontinuous distributions.

Original: 4.2 Molecule Generation
Revised: 4.2 Molecule Generation

Original: Data.
Revised: Data.

Original: We train and evaluate the models on ZINC250k (Irwin et al., 2012) and QM9 (Ramakrishnanet al., 2014) datasets.
Revised: We trained and evaluated all the models on ZINC250k (Irwin et al., 2012) and QM9 (Ramakrishnan et al., 2014) datasets.

Original: The ZINC250k set contains 250,000 drug-like molecules up to 38 atoms.
Revised: The ZINC250k set contains 250,000 drug-like molecules, each consisting of up to 38 atoms.

Original: The

Original:  QM9 set contains 134,000 stable small organic molecules consisting of CHONF atoms.
Revised: The QM9 set contains 134,000 stable small organic molecules with atoms from the set { C , H , O , N , F } .

Original: The moleculesare in  kekulized form with hydrogens removed by the RDkit software (Landrum et al., 2013). 
Revised: The molecules are processed to be in the kekulized form with hydrogens removed by the RDkit software (Landrum et al., 2013). [Table]

Original: Setup.
Revised: Setup.

Original: We adopt common quality metrics to evaluate molecular generation.
Revised: We adopt common quality metrics to evaluate molecular generation.

Original: (i)

Original: Validity is thefraction of molecules that satisfy  chemical valency rule.
Revised: Validity is the fraction of molecules that satisfy the respective chemical valency of each atom.

Original: In this work we strictly limit ourselves tocomparing all methods on validity without external correction procedures (ii) Uniqueness  is thefraction of non-duplicate generations, and (iii) Novelty is the fraction  not present in  training data.(iv)
Revised: Uniqueness refers to the fraction of generated molecules that is unique (i.e, not a duplicate of some other generated molecule). Novelty is the fraction of generated molecules that is not present in the training data.

Original: Reconstruction is the fraction of molecules that can be reconstructed from their encoding.
Revised: Reconstruction is the fraction of molecules that can be reconstructed from their encoding.

Original:  We trainthe model over 5 different seed initialization, and generate 50,000 molecular graphs for evaluation.
Revised: Here, we strictly limit ourselves to comparing all methods on their validity scores without resorting to external correction . We trained each model with 5 random weight initializations, and generated 50,000 molecular graphs for evaluation.

Original: We report the mean and the standard deviation of the results from different seeds.
Revised: We report the mean and the standard deviation scores across these multiple runs.

Original: Implementation.
Revised: Implementation.

Original: The models are implemented in PyTorch (Paszke et al., 2019).
Revised: The models were implemented in PyTorch (Paszke et al., 2019).

Original: The EGNNmethod used only a single layer with embedding dimension of 32.
Revised: The EGNN method used only a single layer with embedding dimension of 32.

Original: The input concatenates timeand scalar vocabulary scores, per node. We train with  Adam optimizer (Kingma and Ba, 2014), for50-100 epochs  with batch size 1000 and learning rate 0 . 
Revised: We trained with the Adam optimizer (Kingma and Ba, 2014) for 50-100 epochs (until the training loss became stable), with batch size 1000 and learning rate 0 . 001 . ModFlow is significantly faster compared to autoregressive models such as GraphAF and GraphDF.

Original: For more details, see Appendix A.4.
Revised: For more details, see Appendix A.5.

Original: Results.
Revised: Results.

Original: Tables 3 and 4 show the results on QM9  and ZINC250K. ModFlow achieves state-of-theart results across all metrics.
Revised: Table 3 reports the performance on QM9 (top) and ZINC250K (bottom) respectively. ModFlow achieves state-of-the-art results across all metrics.

Original: Notably, its reconstruction rate is 100% similar to other flow models;in addition, however,  novelty (100%) and uniqueness scores (  99%) are also very high.
Revised: Notably, its reconstruction rate is 100% (similar to other flow models); in addition, however, the novelty (100%) and uniqueness scores ( ≈ 99%) are also very high.

Original: Moreover, ModFlow surpassed all early methods on validity (95%-99%).
Revised: Moreover, ModFlow surpassed the other methods on validity (95%-99%).

Original: These results underscore the promiseof ModFlow as an effective tool forgraph density estimation and molecular generation. Additionalevaluation metrics, property optimization and comparison with properties distribution are shown in
Revised: In Appendix A.6, we document additional evaluations with respect to the MOSES metrics that access the overall quality of generated molecules, as well as the distributions of chemical properties. All these results substantiate the promise of ModFlow as an effective tool for molecular generation. 4.3 Property-targeted Molecular Optimization

Original: [-Appendix A.6 and A.
-]
Revised: [-Appendix A.6 and A.
-]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

Original: [Table]
Revised: [Figure]

Original: [Figure]
Revised: 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.

Original: 4.3 Ablation Studies
Revised: The final embedding thus obtained is decoded as a new molecule via the reverse mapping f − 1 . [Table] ModFlow (2D-EGNN) [Table]

Revised: [Table]

Revised: ModFlow (JT-2D-EGNN)

Revised: [Table]

Original: We performed ablation experiments to gain further insights about ModFlow , as described next.
Revised: ModFlow (JT-3D-EGNN) [Table] [Table] Figure 6 and Figure 7 show examples of the molecules decoded from the learned latent space using this procedure, starting with molecules having a low QED score.

Original: [Table]
Revised: 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.

Revised: 4.4 Ablation Studies

Revised: We also performed ablation experiments to gain further insights about ModFlow , as we describe next.

Original: E(3) Equivariant vs Not Equivariant.
Revised: E(3)-equivariant versus not equivariant.

Original: Molecules exhibit translational and rotational symmetries,so we conducted an ablation study to quantify the effect of incorporating these symmetries in ourmodel.
Revised: Molecules exhibit translational and rotational symmetries, so we conducted an ablation study to quantify the effect of incorporating these symmetries in our model.

Original: We compare the results obtainedvia using  EGNN to non-equivariant graph convolutionalnetworks (GCN).
Revised: We compare the results obtained using an EGNN with a non-equivariant graph convolutional network (GCN).

Original: We used a 3-layer GCN with layer sizes 64-32-32.
Revised: For our purpose, we used a 3-layer GCN with layer sizes 64-32-32.

Original:  Table 5 provides strong evidencein favor of modeling  symmetries explicitly in  Modular Flows
Revised: The validity scores in Table 5 provide strong evidence in favor of modeling the symmetries explicitly in the proposed Modular Flows

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: 2D vs 3D. Finally, we study whether including3D coordinate information  improves the model.
Revised: 2D versus 3D. Finally, we study whether including information about the 3D coordinates improves the model.

Original: The EGNN coupled differential function obtains either the 2D or 3D positions as polar coordinates,whereas 3D has an extra degree of freedom.
Revised: Note that the EGNN-coupled differential function obtains either the 2D or 3D positions as polar coordinates, where the 3D positions have an extra degree of freedom.

Original: Table 6 shows that transitioning  to 3D gives a consistent≈ 2% percentage unit performance improvement with respect to validity.
Revised: Table 6 shows that transitioning from 2D to 3D improves the mean validity score.

Original: [Table]
Revised: [Table]

Original:  Conclusion
Revised:  Conclusion

Original: We proposed ModFlow , a new generative flow model where multiple flows interact locally accordingto a coupled ODE [Table] resulting in accurate modeling of graph densities and high quality moleculargeneration without any validity checks or correction.
Revised: We proposed ModFlow , a new generative flow model where multiple flows interact locally according to a coupled ODE, resulting in accurate modeling of graph densities and high quality molecular generation without any validity checks or correction.

Original: Interesting avenues open up, including the design of (a) more nuanced mappings between discreteand continuous spaces
Revised: Interesting avenues open up, including the design of (a) more nuanced mappings between discrete and continuous spaces, and (b) extensions of modular flows to (semi-)supervised settings.

Original:  [Table]
Revised:  Acknowledgments

Original: and (b) extensions of modular flows to (semi-)supervised settings.
Revised: 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).

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu- tions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu- tions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] c) Did you discuss any potential negative societal impacts of your work?
Revised: [Yes] (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

Original: If you are including theoretical results... a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] b) Did you include complete proofs of all theoretical results?
Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Original: [Yes] 
Revised: [Yes] 3.

Original: If you ran experiments... a) Did you include the code, data, and instructions needed to reproduce the main experimentalresults (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] See the supplementary material.
Revised: [Yes] See the Appendix.

Original: b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: Yes] See the Appendix
Revised: [Yes] See the Appendix.

Original: c ) Did you report error bars (e.g. with respect to the random seed after running experimentsmultiple times
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See the results section.
Revised: [Yes] See the results section.

Original: d)

Original:  Did you include the total amount of compute and the type of resources used ( e .g. , type of GPUs,internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: Yes] See the Appendix4.
Revised: [Yes] See the Appendix 4.

Original: If you are using existing assets (e.g., code, data, models) or curating releasing new assets... a) If 
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: your work

Original: uses

Original: existing

Original: assets, did you cite the creators?

Original: [Yes] b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [N/A]

Original: c) Did you include

Original: any

Original:  new assets either in the supplemental material or 
Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: as

Original: a URL

Original: [Yes] d) Did you discuss whether and how consent was obtained from people whose data you’re using cu rating
Revised: [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/cu- rating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable informa
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content?

Original: tion

Original: or

Original: offensive content

Original: [N/A
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human 
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: subjects...
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [Table] a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Original: [Table] b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB)
Revised: We present a standard adjoint gradient derivation (Bradley, 2019), and show that the adjoint of a graph neighbor- hood differential is sparse.

Original: approvals, if applicable? [N/A]
Revised: For completeness, we define an ODE system

Original: [Table]
Revised: [Equation]

Original: c) Did you include the estimated hourly wage paid to participants and the total amount spent on [Table]
Revised: 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

Revised: [Equation]

Revised: where G is the total loss that consists of instant loss functionals g .

Revised: We desire to compute the gradients of the system ∇ θ G .

Revised: We optimise the constrained problem by solving the Lagrangian

Revised: [Equation]

Revised: The constraints are satisfied by the ODE definition.

Revised: Hence, ∇ θ L = ∇ θ G , and we can set values θ and µ freely.

Revised: We use a shorthand notation ∂a∂b = a b , and omit parameters from the functions for notational simplicity.

Revised: Applying the chain rule, we note that the gradient becomes

Revised: [Equation]

Revised: where the µ term drops out since it does not depend on parameters θ .

Revised: We apply integration by parts to swap the differentials in term

Revised: λ ⊤ ˙

Revised: z θ , resulting in

Revised: [Equation]

Revised: Substituting this into previous equation and rearranging the terms results in

Revised: [Equation]

Revised: The last term is removed since z (0) not depend on θ as a constant, and thus z θ (0) = 0 .

Revised: The difficult term in the equation is z θ .

Revised: We remove it by choosing

Revised: [Equation]

Revised: [Equation]

Revised: Finally

Revised: [Equation]

Revised: we choose λ ( T ) = 0 which also removes the second-to-last term.

Revised: The choices lead to a final term

Revised: [Equation]

Revised: In the derivation the adjoint λ ( t ) =

Revised: ∂

Revised: L ∂ z ( t ) ∈

Revised: R D represents the change of loss with respect to instant states, and is another ODE system that runs backwards from λ ( T )

Original: {+= 
+}
Revised: {+= 
+}

Revised: until λ (0)

Revised: The adjoint depends on the instant loss state derivatives g z .

Revised: These are often only available for observations y j at observed timepoints t j .

Original: compensation?
Revised: This can be represented by having a convenient loss [Equation] [Equation] and now the term g z induces discontinuous jumps at observations.

Revised: This does not pose problems in practice, since we can integrate the ODE in continuous segments between the observation instants.

Revised: The sparsity of the adjoint evolution is evident from Equation 23, where the ˙

Revised: λ i is an inner product between λ and one column of ∂f∂ z , which is invariant to non-neighbors.

Revised: This gives the result

Revised: [Equation]

Revised: A.2 Tree Decomposition

Revised: For tree decomposition of the molecules, we followed closely the procedure described in Jin et al.

Revised: The rings as well as the nodes corresponding to each ring substructure were extracted using RDKit’s functions,

Revised: GetRingInfo and GetSymmSSSR .

Revised: We restricted our vocabulary to the unique ring substructures in the molecules.

Revised: The vocabulary of clusters follows a skewed distribution over the frequency of appearance within the dataset.

Revised: In particular, only a subset ( ∼ 30) of ring substructures (labels) appear with high frequency in molecules within the vocabulary.

Revised: Therefore, we simplify the vocabulary by only representing the 30 commonly occurring substructures of A tree .

Revised: In Figure 8, we show some examples of these ring substructures for the two datasets.

Original: [N/A]
Revised: [Figure]

Revised: A.3 Equivariant Graph Neural Networks

Revised: Equivariant Graph Neural Networks (EGNN) (Satorras et al., 2021) are E(3)-equivariant with respect to an input set of points.

Revised: The E(3) equivariance accounts for translation, rotation, and reflection symmetries, and can be extended to E( n ) group equivariance.

Revised: The inherent dynamics governing the EGNN can be described, for each layer l , as follows.

Revised: 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 .

Revised: [Equation]

Revised: Initially, messages m ij are computed between the neighboring nodes via ϕ e .

Revised: 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 .

Revised: Finally, the node embeddings are updated based on the aggregated messages m i via ϕ h .

Revised: The aggregated message can be computed based on only the neighbors of a node by simply replacing the sum over j ̸

Revised: = i with a sum over j ∈ N i in these equations.

Revised: [Table]

Revised: [Table]

Revised: z i ( t + 1) =

Revised: ϕ h

Revised: ( z i ( t ) , m i ( t ))

Revised: A.4 Connection to Temporal Graph Networks

Revised: Temporal Graph Networks (Rossi et al., 2020; Souza et al., 2022) are state-of-the-art neural models for embedding dynamic graphs.

Revised: 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 .

Revised: Specifically, adopting the notation from Rossi et al.

Revised: (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.

Revised: 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).

Revised: 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 ′

Revised: i ( t ) .

Revised: 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).

Revised: 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 .

Revised: It turns out (see Table 7) that ModFlow can be viewed as an equivariant message passing temporal graph network.

Revised: Interestingly, the coordinate embedding x i plays the role of the memory s i .

Revised: A.5 Implementation Details

Revised: We implemented the proposed models in PyTorch (Paszke et al., 2019).

Revised: 1 We used a single layer for EGNN with embedding dimension 32 and aggregated information for each node from only its immediate neighbors.

Revised: For geometric (spatial) information, we worked with the polar coordinates (2D) or the spherical polar coordinates (3D).

Revised: We solved the ODE system with the Dormand–Prince adaptive step size scheme (i.e., the dopri5 solver).

Revised: The number of function evaluations lay roughly between 70 and 100 .

Revised: The models were trained for 50epochs with the Adam (Kingma and Ba, 2014) optimizer.

Revised: Time comparisons.

Revised: 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).

Revised: However, ModFlow is faster to train than the auto-regressive methods.

Revised: Note that computation is a crucial aspect of generative modeling for application domains with a huge search space, as is true for the molecules.

Revised: We report the computational effort (excluding the time for preprocessing) for generating 10000 molecules averaged across 5 independent runs in Table 8.

Revised: 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.

Revised: 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.

Revised: [Table]

Revised: A.6 Additional Evaluation Metrics

Revised: We invoked additional metrics

Revised: [Table]

Revised: namely the MOSES metrics (Polykovskiy et al., 2020)

Revised: [Table]

Revised: to compare the different models in terms of their ability to generate molecules.

Revised: These metrics, described below, access the overall quality of the generated molecules.

Revised: • FCD : Fréchet Chemnet Distance (FCD)

Revised: (Preuer et al., 2018) is a general purpose metric that measures diversity of the generated molecules, as well as the extent of their chemical and biological property alignment with a reference set of real molecules.

Revised: Specifically, the last layer activations of ChemNet are used for this purpose.

Revised: Lower is better.

Revised: Frag : Fragment similarity (Frag), measures the cosine distance between the fragment frequencies of the generated molecules and a set of reference molecules.

Revised: Higher is better.

Revised: • SNN :

Revised: Nearest Neighbor Similarity (SNN) quantifies how close the generated molecules are to the true molecule manifold.

Revised: Specifically, it computes the average similarity of a generated molecule to the nearest molecule from the reference set.

Revised: Higher is better.

Revised: • IntDiv : As the name suggests, Internal Diversity (IntDiv) accounts for diversity by computing the average pairwise similarity of the generated molecules.

Revised: Higher is better.

Revised: For our purpose, we evaluated these metrics with QM9 and ZINC250K as the reference sets.

Revised: As shown in Table 9 and Table 10, ModFlow achieves better performance results across all metrics.

Revised: Notably, ModFlow registers lower FCD and higher IntDiv scores compared to other methods, suggesting that ModFlow is able to generate diverse set of molecules similar to those present in the real datasets.

Revised: [Table]

Revised: [Table]

Revised: We also evaluated the generated structures via distributions of their important properties.

Revised: Specifically, we obtained kernel density estimates of these distributions to aid in visualization.

Revised: We consider the following standard properties

Revised: • Weight : sum of the individual atomic weights of a molecule.

Revised: The weight provides insight into the bias of the generated molecules toward lighter or heavier molecules.

Revised: • LogP : ratio of concentration in octanol-phase to the aqueous phase, also known as the octanol-water partition coefficient.

Revised: It is computed via the Crippen (Wildman and Crippen, 1999) estimation.

Revised: • Synthetic Accessibility (SA) : an estimate for the synthesizability of a given molecule.

Revised: It is calculated based on contributions of the molecule fragments Ertl and Schuffenhauer (2009).

Revised: • Quantitative Estimation of Drug-likeness (QED

Revised: ) : describes the likeliness of a molecule as a viable candidate for a drug.

Revised: It ranges between [0,1] and captures the abstract notion of aesthetics in medicinal chemistry (Bickerton et al., 2012).

Revised: Figure 9 and Figure 10 show that barring some dispersion in QED and logP (especially on Zinc250K), the property distributions of the molecules generated by ModFlow generally match the corresponding distributions on the reference datasets quite closely.

Revised: These results demonstrate the effectiveness of ModFlow in generating molecules that have properties similar to the molecules in the reference set.

Revised: Additional

Revised: molecules generated by ModFlow .