Traversing Chemical Space with Latent Potential Flows

Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-02 19:31:15

Abstract: We consider the latent traversal problem in studying and exploring the chemical space with the learned latent space of a generative model. We propose a new framework, ChemFlow, that unified previous molecule manipulation and optimization method with a dynamical system perspective. Specifically, we formulate the problem as learning a vector field that transports the mass of the molecular distribution to the region with desired molecular properties or structure diversity. We also propose several alternative dynamics which exhibit various advantages over previous methods. We validate the efficacy of our proposed methods on both supervised and unsupervised molecule manipulation and optimization scenarios.

Main paper: 
 introduction: 
 Unveiling the structure of the chemical space and designing efficient search algorithms to explore it is a long-standing challenge in chemistry (Bohacek et al., 1996; Lipinski et al., 2012) . Recently, with the promising results by deep generative models to generate valid molecules, more attention has been attracted to study the structure of the learned latent space of those models and optimization methods over the latent space to find better molecule candidates (Du et al., 2022a) .Initially, researchers explicitly introduces additional constraints to encourage a disentangled latent space such that any dimension learns a meaningful factor that may correspond to a molecular property (Du et al., 2022b; c) . However, it turns out the assumption is too strong thus cannot learn a meaningful latent space. Inspired by the study of interpretability of generative models on images (Shen et al., 2020) , Du et al. (2023) leverage the observed structure similarity and the "linear separability" assumption (Gómez-Bombarelli et al., 2018) to leverage a linear model on the pre-trained model to find meaningful paths in the latent space corresponding to molecular properties. Another branch of work explicitly focus on leveraging the smooth and low-dimensional latent space for molecule optimization with a commonly used approach is gradient ascent over the latent vectors (Liu et al., 2018; Griffiths & Hernández-Lobato, 2020) .In this paper, we propose a new framework, ChemFlow, based on potential flows to efficiently explore the latent structure of molecule generative models. Specifically, we unify previous approaches (gradient-based optimization, linear latent traversal, and disentangled traversal) under the realm of flow that transforms data density along time via a vector field. In contrast to previous linear models, our framework is flexible to learn nonlinear transformations inspired by popular partial differential equations (PDEs) such as heat and wave equations. We also analyze the special properties of each distinct dynamics. For example, under the mild assumption, Langevin dynamics by Fokker Planck equation exhibit convergence to the global minimum. Our framework can also generalized to both supervised and unsupervised settings. Particularly in the unsupervised setting, we introduce a structure diversity potential to find directions that maximize the structure change of the molecules. We conduct extensive experiments with physicochemical properties and drug-related properties on both molecule manipulation and optimization experiments. The experiment results demonstrate the proposed alternative methods using different dynamical priors achieve better or comparable results with existing approaches. 
background: traversing latent space of molecules: 
 The latent space Z of molecule generative models is often learned through an encoder function f θ (•) and a decoder function g ψ (•) such that the encoder maps the input molecular structures x ∈ X into an (often) low-dimensional and continuous space (i.e. latent space) while the decoder maps the latent vectors z ∈ Z back to molecular structures x ′ . Note that this encoder-decoder architecture is general and can be realized by popular generative models such as VAEs, flow-based models, GANs, and diffusion models (Jin et al., 2018; Madhawa et al., 2019; Cao & Kipf, 2018; Vignac et al., 2023) .For simplicity, we focus on VAE-based methods in this paper. To traverse the learned latent space of molecule generative models, two approaches have been proposed: gradient-based optimization and latent traversal.The gradient-based optimization methods first learn a proxy function h(•) parameterized by a neural network that provides the direction to traverse (Zang & Wang, 2020) . This can be formulated as a gradient flow following the direction of steepest descent of the potential energy function h(•) and discretized, as follows:EQUATIONThe latent traversal approaches leverage the observation of linear separability in the learned latent space of molecule generative models (Gómez-Bombarelli et al., 2018) . Since the direction is assumed to be linear, it can be found easily. ChemSpace (Du et al., 2023) learns a linear classifier that defines the separation boundary of the molecular properties. Then the normal direction of the boundary provides a linear direction n ∈ Z for traversing the latent space:z t = z 0 + nt (2)We notice that the above gradient flow and linear traversal can be analyzed and designed in a dynamical system perspective, e.g. linear traversal can be considered as a special case of wave functions. This connection inspires us to consider designing more dynamical traversal approches.We leave a more substantial background for Wasserstein gradient flow and related work in Appendix Sec. A and Sec. E, respectively. 
methodology: 
 We present ChemFlow as a unified framework for latent traversals in chemical latent space as potential flows (detailed in Appendix A). Motivated by the optimal transport theory of Wasserstein gradient flows, we parameterize a set of scalar potential energies ϕ k = MLP θ k (z, t) ∈ R using Physics-informed Neural Networks (PINNs) and take potential flow ∇ z ϕ to traverse the latent sam-ples:EQUATIONLearning latent potential flows. Given a pre-trained molecule generative model g ψ : Z → X with prior distribution p(z), we would like to model K different latent trajectories that correspond to different molecule properties. The optimal transport is given by Hamilton-Jacobi Equation (HJE):∂ ∂t ϕ(z, t) + 1 2 ||∇ z ϕ(z, t)|| 2 = 0 (4)where the velocity field is defined as the potential flow ∇ϕ. The HJE is usually interpreted as fluid-dynamic optimal transport, i.e., under the velocity field ∇ϕ, the fluid will evolve to the target distribution with the optimal transport cost. Despite the optimality in moving distributions, HJE may not create smooth trajectories due to the advection field ∇ϕ, which may not lead to desired smooth property variations. Alternatively, we can sacrifice the optimal transport property and restrict the advection term to enforce other types of dynamics for smooth spatiotemporal dynamics. For example, we can specify the flow to follow the wave-like dynamics with the coefficient c:EQUATIONThe above constraint empirically produces highly diverse and realistic trajectories. We use a PINN (Raissi et al., 2019) to enforce the PDE constraint. Compared with traditional PDE solvers, PINNs can be orders of magnitude faster. Our PINN objective is to minimize:L r = 1 T T -1 t=0 ||r k (z t , t)|| 2 2 , L ϕ = ||∇ z ϕ k (z 0 , 0)|| 2 2 (6)where T represents the total number of traversal steps, L r restricts the energy to obey our physical constraints, and L ϕ restricts ϕ(z t , t) to match the initial condition.Supervised semantic potential guidance. When an explicit semantic potential or labeled data for the semantic of interest is available, we can use the provided semantic potential to guide the learning of the flow. Firstly, we train a surrogate model h η : X → R (parameterized by a deep neural network) to predict the corresponding molecular property. Then we use the trained surrogate model as guidance to learn flows that drive the increase of the property:d = ⟨-∇ z h η (g ψ (z t )), ∇ z ϕ k (z t , t)⟩, L P = -sign(d)∥d∥ 2 2 (7)The intuition behind this objective is to learn the vector field z t such that it aligns with the direction of the steepest descent (negative gradient) of the objective function. Note that the sign of the dot product matters as it determines minimizing or maximizing the property.Unsupervised structure diversity guidance. When no explicit potential function is provided to learn the flow, we need to define a potential that captures the change of the molecular properties. As molecular properties are determined by the structures, we devise a potential energy that maximizes the continuous structure change of the generated molecules. Inspired by Song et al. (2023b) , we couple the traversal direction with the Jacobian of the generator to maximize the traversal variations in the molecular space. We therefore introduce the guidance L J as follows:g(z t + ϵ∇ z ϕ k (z t , t)) ≈ g(z t ) + ϵ ∂g(z t ) ∂z t ∇ z ϕ k (z t , t), L J = - ∂g(z t ) ∂z t ∇ z ϕ k (z t , t) 2 2 (8)where we perform the first-order Taylor approximation on the left equation. In the unsupervised setting, for sufficiently small ϵ, if the Jacobian-vector product (the underlined term in Eq. ( 8)) can cause large variations in the generated sample, the direction is likely to correspond to certain properties of molecules. Compared to the supervised setting which maximizes the change of the molecular properties, it aims to find the direction that causes the maximal change of the structures. This can in turn effectively push the initial data distribution to the target one concentrated on the maximum property value. The Jacobian guidance will compete with the dynamical regularization (e.g. wave-like form) on the flow to yield smooth and meaningful traversal paths. In high-dimensional space, a trivial solution exists such that all the flows learn the exact same direction. To encourage the model to learn more disentangled paths, we design a disentanglement regularization detailed in Appendix B.In Appendix C, we show one particular property of a specific flow governed by Fokker Planck equation such that it is equivalent to simulating Langevin dynamics in the latent space which has weak convergence to the global minimizer of the molecular property. 
experiments: 
 We conduct experiments by pre-training a VAE model on a commonly used molecular dataset (ZINC250k (Irwin & Shoichet, 2005) ). We compare our proposed methods against two baseline methods with linear traversal (Du et al., 2023) and gradient flow-based traversal (Eckmann et al., 2022) . For all other details, we defer to Appendix F.Due to the space limit, we defer all experiments related to molecule manipulation in Appendix F.5.In the main paper, we only report results for molecule optimization in the latent space. (Brown et al., 2019) . Following the procedures described in JT-VAE (Jin et al., 2018) and LIMO (Eckmann et al., 2022) , we select the 800 molecules with the lowest QED scores in the ZINC250k dataset and perform 1,000 steps of optimization until all methods are converged. Table 1 reports the statistics of QED improvements under different similarity constraints. Among all the approaches, Langevin dynamics achieves the best overall performance. As we analyze above, it is more suitable for the optimization task than other dynamical priors. In addition, we further show it enjoys a faster convergence rate empirically than other dynamical priors as well in Figure 5 . Surprisingly, we observe that the random direction performs also quite well on molecule optimization tasks. This observation motivates us to study the structure of the latent space. We show that the molecular structure distribution on the latent space follows a high-dimensional Gaussian distribution and the random direction increases the norm of the latent vectors that have strong correlations with molecular properties. We analyze this systematically in Appendix G. It is also notable that although random directions could be effective in optimizing molecules, the distribution of the entire molecule sets being optimized does not change accordingly as shown in Figure 2 . We also report additional experiment on unconstrained and multi-objective molecule optimization results in Appendix F.7 and F.8, respectively. 
conclusion, limitation and future work: 
 In this paper, we propose a unifying framework that learns a flow transformation through a vector field for traversing the latent space of molecule generative models. Under this framework, we propose a variety of dynamical regularizations which exhibit different properties.  
 
 Appendix: Gradient flows define the curve x(t) ∈ R n that evolves in the direction of steepest descent of a function F : R n → R. The time evolution of the gradient flow is given by the ODE x ′ (t) = -∇F(x(t)).The choice of the functional F determines the metric space and the associated geodesics. Wasserstein gradient flows describe a special type of gradient flow where F is set to be the Wasserstein distance. For example, as introduced in Benamou & Brenier (2000) , the commonly used L 2 Wasserstein distance has the following dynamic formulation:EQUATIONwhere ρ 0 and ρ 1 are two probability measures at the source and target distributions, respectively. Interestingly, if we take the gradient of a potential energy ∇ϕ as the velocity field applied to a distribution, the time evolution of ∇ϕ can be seen to minimize the Wasserstein distance and thus follow optimal transport.In Appendix A, we give detailed derivations of how the vector fields minimize the L 2 Wasserstein distance and discuss some PDEs of the density evolution that can be interpreted under different learned potentials.As shown in the main paper, based on the dynamic formulation of optimal transport (Benamou & Brenier, 2000) , the L 2 Wasserstein distance can be re-written as:W 2 (ρ 0 , ρ 1 ) = min ρ,v ρ t (z)|v t (z)| 2 dzdt (10)where v t (z) is the velocity of the particle at position z and time t, and ρ t (z) is the density dµ(z) = ρ t (z)dz. The distance can be optimized by the gradient flow of a certain function on space and time. Consider the functional F : R n → R that takes the following form:F(µ) = U (ρ t (z)) dz (11)The curve is considered as a gradient flow if it satisfies ∇F = -d dt ρ t (z) (Ambrosio et al., 2005) . Moving the particles leads to:EQUATIONThe velocity vector satisfies the continuity equation:EQUATIONwhere -∇ • v t (z)ρ t (z) is the tangent vector at point ρ t (z). Eq. ( 12) can be simplified to:EQUATIONOn the other hand, the calculus of differential geometry givesd dt F(µ) = DiffF| ρt (-∇ • v t (z)ρ t (z) ) = ⟨∇F, -∇ • v t (z)ρ t (z) ⟩ f (15)where ⟨, ⟩ f is a Riemannian distance function which is defined as:EQUATIONThis scalar product coincides with the W 2 distance according to Benamou & Brenier (2000) . Then Eq. ( 14) can be similarly re-written as:d dt F(µ) = ⟨-∇ • ∇U ′ (ρ t (z))ρ t (z) , -∇ • v t (z)ρ t (z) ⟩ f (17)So the relation arises as:EQUATIONSince we have ∇F = -d dt ρ t (z), the above equation can be re-written asEQUATIONThe above derivations can be alternatively made by JKO schemes (Jordan et al., 1996) . This explicitly defines the relation between evolution PDEs of ρ t (z) and the internal energy U . For our method, we use the gradient of our scalar energy field ∇u(z, t) to learn the velocity field which is given by U ′ (ρ t (z)). Interestingly, driven by certain specific velocity fields ∇u(z, t), the evolution of ρ(z, t) would become some special PDEs. Here we discuss some possibilities:Heat Equations. If we consider the energy function U as the weighted entropy:EQUATIONWe would have exactly the heat equation:d dt ρ t (z) - d dz 2 ρ t (z) = 0 (21)Injecting the above equation back into the continuity equation leads to the velocity field v t (z) asEQUATIONWhen our ∇u(z, t) learns the velocity field -∇ log(ρ t (z)), the evolution of ρ(z, t) would become heat equations.Fokker Planck Equations. For the energy function defined as:EQUATIONwe would have the Fokker-Planck equation asEQUATIONThe velocity field can be similarly derived asEQUATIONFor the velocity field ∇A -∇ log(ρ t (z)), the movement of ρ(z, t) is the Fokker Planck equation.Porous Medium Equations. If we define the energy function asEQUATIONThen we would have the porous medium equation where m > 1 and the velocity field:EQUATIONWhen the ∇u(z, t) learns the velocity -mρ m-2 ∇ρ, the trajectory of ρ(z, t) becomes the porous medium equations.Disentanglement Regularization. While the above formulation can encourage smooth dynamics and meaningful output variations, the potential flows are likely to mine identical directions which all correspond to the maximum Jacobian change. To avoid such a trivial solution, we adopt an auxiliary classifier l γ to predict the potential index and use the cross-entropy loss to optimize it:EQUATIONWhere x t = g(z t ) is the generated sample from timestep t. We see the extra classifier guidance would encourage each potential flow to be independent and find distinct properties.In scenarios where our flow adheres to the dynamics of the Fokker-Planck equation, our approach may also be interpreted as employing a learned potential energy function to simulate Langevin Dynamics for global optimization (Gardiner et al., 1985) . Notably, the convergence of Langevin dynamics, particularly at low temperatures, tends to occur around the global minimum of the potential energy function (Chiang et al., 1987) . The continuous and discretized Langevin dynamics are as follows: (1991) ). Given a langevin dynamics in the form of z t = z t-1 -a t (∇ z h η (z t-1 ) + u t ) + b t w t (30) where w t is a d-dimensional Brownian motion, a t and b t are a set of positive numbers with a T , b T → 0, and u t is a set of random variables in R n denoting noisy measurements of the energy function h η (•). Under mild assumptions, z t converges to the set of global minima of h η (•) in probability.EQUATIONFollowing Proposition C.1, we show that the learned latent potential can be used to search for molecules with optimal properties and it converges to the global minima of the learned latent potential.Due to space limit, we demonstrate the training and inference procedure for ChemFlow in Alg. 2 and Alg. 2, respectively.Algorithm 1 ChemFlow Inference / Traversal Require: Pre-trained encoder f θ , pre-trained potential function ϕ, (optional) pre-trained proxy function h, timestamps T , step size α, LD strenth β 1: Sampling:z 0 = f θ (x 0 ) 2: for t = 1, . . . , T do 3:if Langevin Dynamics then 4:z t = z t-1 -α∇ z h η (z t-1 ) + β √ 2αN (0, I) 5:else if Gradient Flow then 6:z t = z t-1 -α∇ z h η (z t-1 ) 7:else 8:z t = z t-1 + α∇ z ϕ(z t-1 , t -1) 9:end if 10: end for Algorithm 2 ChemFlow Training Require: Pre-trained encoder f θ , decoder g ψ , (optional) classifier l γ , timestamps T , # of potential functions K 1: Initialize ϕ j (•) ← MLP for j = 1, . . . , K 2: repeat 3:Sampling: z 0 = f θ (x 0 ), t ∼ Categorical(T ), k ∼ Categorical(K) 4:for i = 1, . . . , t do 5:z i+1 = z i + ∇ z ϕ k (z i , i) 6:end for 7:Decode: (Du et al., 2022a) . One idea is to leverage the advancement of deep generative models which approximate the data distribution from a provided dataset of molecules and then sample new molecules from the learned density. This idea inspires a line of work developing deep generative models from variational auto-encoders (VAE) (Gómez-Bombarelli et al., 2018; Jin et al., 2018) , generative adversarial networks (GAN) (Guimaraes et al., 2017; Cao & Kipf, 2018) , normalizing flows (NF) (Madhawa et al., 2019; Zang & Wang, 2020) and more recently diffusion models (Hoogeboom et al., 2022; Vignac et al., 2023; Jo et al., 2022) . However, to respect the combinatorial nature of molecules, another line of work leverage combinatorial optimization to search new molecules including genetic algorithm (GA) (Jensen, 2019) , Monte Carlo tree search (MCTS) (Yang et al., 2017) , reinforcement learning (RL) (You et al., 2018) , but often with sophisticated optimization objectives beyond simple valid molecules.x t = g ψ (z t ), x t+1 = g ψ (z t+1 ) 8: if unsupervised then 9: Classification: k = l γ (x t ; x t+1 ) 10: Loss: L = L r + L ϕ + L J + L k 11:In addition to simply generating valid molecules, a more realistic application is to generate molecules with desired properties (Du et al., 2022a) . For deep generative model-based methods, it is naturally combined with on-the-fly optimization methods such as gradient-based or Bayesian optimization (in low data regime) as it often maps data to a low-dimensional and smooth latent space thus more friendly for these optimization methods (Griffiths & Hernández-Lobato, 2020) . For methods that do not explicitly reduce the dimensionality of data such as diffusion models, Schneuing et al. ( 2022) propose an evolutionary process to iteratively optimize the generated molecules. As it is observed that the learned latent space exhibits explicit structure (Gómez-Bombarelli et al., 2018), Du et al. (2023) leverage such property to learn a linear classifier to find the latent direction to optimize the property of given molecules. In opposition to deep generative models, combinatorial optimization methods are often inherently associated with optimization, e.g. reward function in RL, selection criteria in GA, etc (Fu et al., 2022; Loeffler et al., 2023) .Beyond molecule generation, there is a vast literature on the study of the latent space of generative models on images for image editing and manipulation (Goetschalckx et al., 2019; Jahanian et al., 2020; Voynov & Babenko, 2020; Härkönen et al., 2020; Zhu et al., 2020; Peebles et al., 2020; Shen & Zhou, 2021; Song et al., 2022; 2023b; a; c) . Here we highlight some representative supervised and unsupervised approaches. Supervised methods usually require pixel-wise annotations. We compare with the following baselines:• Random: we take a linear direction that is sampled from Multi-variant Gaussian distribution in the high dimensional latent space and normalized to unit length for all molecules across all time steps.• Random 1D: we take a unit vector where only 1 randomly selected dimension is either 1 or -1 as the linear direction.• ChemSpace (Du et al., 2023) : a separation boundary of the training dataset in latent space w.r.t. the desired property is classified by an Support vector machine (SVM). Then we take the normal vector corresponding to the positive separation as the manipulation direction of control.• LIMO (Eckmann et al., 2022) / Gradient Flow: a VAE-based generative model that encodes the input molecules into SELFIES (Krenn et al., 2019) and auto-regressive on the tokenized molecule. LIMO uses Adam optimizer to reverse optimize on the input latent vector z whereas Gradient Flow is equivalent to using an SGD optimizer for the same purpose.We report the following metrics for our experiments:• Penalized logP: Estimated octanol-water partition coefficient penalized by synthetic accessibility (SA) score and the number of atoms in the longest ring.We noticed that plogP is a commonly reported metric in recent molecule discovery literature but does not share the same normalization scheme. Following Gómez-Bombarelli et al. (2018, Eq. 1), the SA scores and a ring penalty term were introduced into the calculation of penalized logP as the following J logP (m) = logP(m) -SA(m)ring-penalty(m)Each term of logP(m), SA(m), and ring-penalty(m) are normalized to have zero mean and unit standard derivation across the training data. However, no sufficient details were included in their paper or their released source code on how the ring-penalty(m) is computed. Specifically, 3 implementations are widely used in various works.Penalized by the length of the maximum cycle without normalization where ring-penalty(m) is computed as the number of atoms on the longest ring -6 in their implementation. Neither logP(m), SA(m), or ring-penalty(m) is normalized.Penalized by the length of the maximum cycle with normalization where ring-penalty(m) is computed same as without normalization. We report plogP using this metric.Penalized by number of cycles As described by Jin et al. (2018) , ring-penalty(m) is computed as the number of rings in the molecule that has more than 6 atoms. LIMO reports plogP using this metric.Datasets. We consider a commonly used molecular dataset to study drug-related properties, ZINC250k which is a subset of the ZINC database (Irwin & Shoichet, 2005) containing ∼250,000 commercially available compounds for virtual screening. The molecules in the dataset is firstly encoded as SELFIES strings for molecule representations and then tokenized into input IDs following the language modeling commons with [PAD] and other special tokens. All input IDs are padded to the maximum length in the overall dataset.Implementations. We establish our framework by pre-training a VAE model that learns a latent space of molecules and is capable of generating new molecules by decoding latent vectors from the latent space. We adapt the framework in Eckmann et al. ( 2022) which is a basic VAE architecture with molecular SELFIES string representations plusing an MLP for surrogate property predictor. See Appendix F.3 for all implementation and hyper-parameter details.Model variants. As discussed in Section 3, our proposed framework is general to incorporate different dynamical priors to learn the flow. For the experiments, we consider four types of dynamics including gradient flow (GF), Wave flow (Wave, Eq. ( 5)), Hamilton Jacobi flow (HJ, Eq. ( 4)) and Langevin Dynamics or equivalently Fokker Planck flow (LD, Eq. ( 29)).Pre-trained VAE We adjust the VAE architecture from LIMO that consisting a 64-dimension embedding, 1024 latent space size, 3-hidden-layer encoder, and 3-hidden-layer decoder both with 1D batch normalization and ReLU activation functions. The hidden layer sizes are {2000, 1000, 1000} for the encoder and reversely for the decoder. We empirically find that replacing the ReLU activation function with its newer variant Mish activation function (Misra, 2020) results in faster convergence and better validation loss. All the experiments reported in this paper use this Mish-activated variant of VAE. The VAE is trained using an AdamW (Loshchilov & Hutter, 2017) optimizer with default settings of PyTorch implementation, 0.001 initial learning rate, and 1,024 training batch. To better prevent the model from being stacked at a sub-optimal local minimum, a cosine learning rate scheduler with periodic restart is applied. The VAE is trained for 150 epochs with 4 restarts on 90% of the ZINC250k dataset and validated with the rest 10% data. The checkpoint corresponding to the best validation loss epoch is selected.Surrogate Predictor Similarly, we find that the Mish activation function improves the performance of the original MLP predictor described by LIMO. The predictor consists of 3 linear hidden layers of size 1024 with Mish activation function. Similar to the LIMO setups, we find that the choice of optimizer and training hyperparameters like learning rate or learning rate scheduler is crucial for successful training. Like LIMO, we use PyTorch Lightning to choose the optimal learning rate for the Adam optimization. The predictor is trained for 20 epochs on 100,000 randomly generated samples and validated with 10,000 unseen data. The epoch with the best validation loss is selected.We use an MLP structure to parameterize the potential energy function. The spatial time input t is embedded with a sinuous positional embedding followed by a linear layer. The spacial input x is encoded with a linear layer and ReLU activation function. We empirically find that the Tanh and GELU activation function of MLP does not help PINN learn the corresponding Jacobian structure. Therefore, a ReLU activation function is used instead. The training of PINN follows uses 90,000 random data and 10,000 unseen data for validation. For unsupervised settings, 10 disentangled potential energy functions are trained for 70 epochs with a batch size 100. The epoch with the best validation loss is selected.Reproducibility All the experiments including baselines are conducted on one RTX 3090 GPU and one Nvidia A100 GPU. The code implementation will be released upon the acceptance of the paper.We conduct more experiments to analyze the performance of the proposed methods systematically. They are referred and discussed in the main paper. Table 3 : Pearson Correlation. The average Pearson correlation between the sequence of real properties and sequence of time steps along the manipulation trajectory following a learned potential function ϕ k (z, t) using Hamilton Jacobi equations. .28 ± 6.40 (71.6) 4.90 ± 6.56 (22.1) 7.77 ± 6.75 (78.1) 7.12 ± 6.18 (72.0) 7.03 ± 6.82 (71.5) 6.15 ± 7.39 (46.9) 6.41 ± 7.72 (71.1) 4.60 ± 6.22 (30.9) 7.24 ± 6.07 (72.0) 0.4 4. 66 ± 5.56 (42.5) 3.67 ± 5.49 (16.8) 1.81 ± 3.71 (9.0) 3.54 ± 4.46 (20.9) 2.53 ± 3.52 (11.9) 1.81 ± 3.26 (18.9) 1.00 ± 1.16 (10.1) 1.14 ± 1.83 (24.8) 0.56 ± 0.36 (3.1) 2.79 ± 3.80 (12.9)Table 4 : Similarity-constrained plogp maximization. For each method with minimum similarity constraint δ, the results in reported in format mean ± standard derivation (success rate %) of absolute improvement, where the mean and standard derivation are calculated among molecules that satisfy the similarity constraint. The value of QED is scaled by 100 for better presentation. Molecule manipulation is a relatively new task proposed in Du et al. (2023) to study the performance latent traversal methods. Specifically, the main idea of molecule manipulation is to find smooth local changes of molecular structures that simultaneously improve molecular properties which is essential to help chemists systematically understand the chemical space. As our methods are general in both learning the traversal direction with surrogate model (supervised) or with only structure change (unsupervised), we evaluate below the performance of both use cases.Table 5 shows the success rate results of manipulating 1,000 randomly sampled molecules to optimize each desired property. Following Du et al. (2023) , we traverse the latent space for 10 steps in the traversal direction of each method as reported the strict and relaxed success rate. Details of the definition of these metrics can be found in Appendix F.2. Among all the approaches, our method with gradient flow achieves the highest success rates on multiple properties such that it takes the steepest descent of the surrogate model. When the step size is small enough, it is reasonable to learn a smooth path. However, the results vary across properties.One particular example is DRD2 such that supervised approaches all fail to achieve good success rates.As the correspondence between specific molecular properties and learned latent potential flows is not explicitly given in the unsupervised scenario, we use an artificial process to mimic the use case in reality. Specifically, we learned 10 different potential energy functions representing 10 disentanglement potential flows following Algorithm 2 using Wave equation and Hamilton-Jacobi equation and validated them on 1,000 unseen molecules. For each potential flow, we measure the real properties of molecules generated along the manipulation trajectory of 10 steps. We report the average of all Pearson correlations between the sequence of properties and time step sequence from 1 to 10 in Table 2 for Wave eqn. and Table 3 for HJ eqn.The best correction potential flow is selected for each property representing the learned Jacobian structural change that would most effectively optimize the corresponding property.In Table 5 , we can observe that even though it is without supervised training of traversal directions, the flow still learns meaningful directions from molecular structure to property changes. Surprisingly, the results of manipulating molecules for DRD2, JNK3, and GSK3B in unsupervised settings are better than in supervised settings. We hypothesize that this is partially because of the training and generalization errors of the surrogate model, as shown in Appendix Figure 3 . On the contrary, the structure change measurement does not provide supervision but is reliable. We would like to point out that this is an open question in chemistry, often referred as to the structure-activity relationship (Dudek et al., 2006) , such that it is important to know the correspondence between structure and activity. We believe this is a promising result to demonstrate that generative models "realize" molecular property by learning from structures.Among the quantitative results, it is notable that the random direction achieves surprisingly high success rates, we argue that's because of the specific property of the learned latent space. The underneath generative model regulated the latent space to be smooth such that similar molecular structures are often mapped to close areas in the latent space. We also analyze in Appendix G such that we find some molecular properties are highly correlated with their latent vector norms in which a random direction always increases the norm and therefore successfully manipulates a portion of molecules by chance. Unconstrained Molecule Optimization. For unconstrained optimization problems, we randomly sample 100K molecules from the latent space and report the top 3 scores after 100 steps of manipulation of each method in Table 6 . All methods use 0.1 relative step sizes for fair comparison. Among them, we observe that gradient flow and Langevin dynamics are the two best approaches that find better candidates than other methods. This is reasonable since they are taking the steepest descent of the energy function in Euclidean and probability space, respectively. However, we notice that in Figure 2 , Langevin dynamics pushes the entire distribution much further than gradient flow despite the optimal solutions are similar. In addition, we find that despite the random direction seeming to be effective in optimizing molecules, it does not push the distribution much.F.8 MULTI-OBJECTIVE OPTIMIZATION RESULTS Table 7 : Similarity-constrained Multi-objective (QED+SA) maximization. The value of QED and SA is scaled to both have a range from 0 to 100 for an equal-weighted sum. The method with the highest equal-weighted sum score of QED+SA of each structure similarity level is bolded.Multi-objective Molecule Optimization. As we are learning distinct vector fields and potential energy functions for each property, they can be readily added together for multi-objective optimization (Eckmann et al., 2022; Du et al., 2023) . To generate molecules that are optimized on multiple properties, we use a similar setting as similarity-constrained molecule optimization to select 800 molecules from the ZINC250k dataset with the lowest QED and aim to generate molecules with high QED as well as SA simultaneously. At each time step, the latent vector is optimized following the averaged direction of two supposed potential flow directions. This scheme could be seamlessly generalized to m-objectives optimization and is commonly used in related works. Table 7 shows that Langevin dynamics and ChemSpace achieve the best or competitive performance at all similarity cutoff levels. As observed in experiments such that random directions perform surprisingly well on molecule manipulation and optimization tasks, we look into the learned latent space to understand its structure.As the prior of a VAE is an isotropic Gaussian distribution, we first verify if the learned variational poster also follows a Gaussian distribution and we find that it does learn so from the evidence shown in Figure 6 , where the norm of the molecule projected to the latent space concentrate around 32 which is around √ d such that the latent dimension d is 1024. We also visualize in Figure 7 that how the properties of the molecules in the training dataset are related to their latent vector norms. Surprisingly, we find a strong correlation between almost all molecular properties and their latent norms. Combining this two evidences, it is not surprising that a random latent vector taking a random direction will change the molecular property smoothly and monotonically. In addition, we further plot when we traverse along a random direction in the latent space, how the change of the norm may correspond to the change of a certain property. Among them, we find that SA is particularly in strong positive correlation with the traversal in Figure 8 . Though the emergence of the structure in the latent space is interesting and suggests that better algorithms can be developed to exploit the structure, we leave this to future work. In addition to quantitative evaluations, we demonstrate some qualitative evaluations in this section. We showcase one selected manipulation path Figure 9 in by gradient flow (GF) on plogP. We can see the learned path follows the trend to replace the O-containing non-stable macrocycle step by step with six-or three-member rings while optimizing the desired property (i.e, plogP). For molecule optimization, we also showcase one selected path Figure 13 Fokker Planck flow on plogP. We also observe a similar trend such that it replaces the N-containing fused ring step by step with threemember rings while optimizing the desired property (i.e., plogP). Additional molecule manipulation paths are shown in Figure 10 , Figure 11 , Figure 12 , and optimization paths are shown in Figure 14 , Figure 15 , Figure 16 .The original method was trained on a much larger training set (2M vs250K), we report the reproduced results for fair comparison. We also use SGD instead of Adam to align with our scheme.2 Note that we notice there is a misalignment of normalization schemes for the plogP property in the previous literature, so we only rerun and compare with related methods that align with our normalization scheme. Details can be found in Appendix F.2. 


Decision: Thanks to the reviewer for the detailed review and feedback. 
The paper proposes a generative modeling framework (VAE based) , Chemflow, for chemical molecule structure analysis and manipulation, Authors are expected to go through all comments and address them clearly in their final revision. Under this condition, the work can be accepted as poster.

Reviews: review: title: title: The authors consider the problem of studying and exploring the dynamics of the (chemical) latent space from a generative model.
main_review: main_review: The authors consider the problem of studying and exploring the (chemical) latent space of a generative model. They aim to make connections of how one can explore the latent space of a (generative model) and dynamical systems. They use different approaches e.g. computing the gradient (gradient ascent or descent), learn a wave type parametrization of the potential (with the help of PINNS) a HJE and also use Langevin dynamics.

In my comments below I would mostly focus on the dynamic’s aspect of the work and less on the chemical/molecular aspect that is not my field.

I believe this is a sound paper and that the authors have done a lot of work.

The four-page limitation makes the paper difficult to follow therefore going back and forth between the main text and the appendix becomes necessary to get understand the overall frameworks. There are a few parts of the paper that I believe need improvement, I provide more details below. 

Major Comments:

(I)	In many expressions in the main text there is a superscript k e.g. (equations 3,5,6) I do not see where this is defined. Can you clarify/change this?

(II)	Can the authors mention why learning a vector field is better and not evaluating a network in a dense grid and choosing the minimum or maximum on the grid? What cases one can expect the one to be better than the other?

(III)	What happens if the maximum, minimum is close to the boundary? Are the trajectories going to “leave” the latent space and explode? Are there ways they can encode this information to their model?

(IV)	Why the Langevin dynamics might be better? Is it better regardless of the size of the noise? 

(V)	Even though I understand the limitations because of the four-page limit I think is important to improve the section “conclusions, limitations”. Now is more a reiteration of the overall paper and it does not mention what the conclusions and the limitations are.


Minor:

Abstract: (a) Check a little bit the English (e.g. that unified -> unifies?)

Introduction: (a) (Figure 1) The NN in the schematic of the decoder has as input two neurons and one (yellow) output, I think this should have been reversed right? (from less inputs to more outputs in the decoder). 

(b) Perhaps also a change is needed for the green box in the schematic for the predictor (go from the yellow to a function)?

(c ) The yellow color is not so visible.

Background: (Traversing Latent Space of Molecules) (a) The gradient of the potential energy function h is estimated in equation 1. Can you say how this is computed (automatic differentiation, symbolic differentiation, finite differences)? The same holds also for the case of the potential φ later (e.g. equation 3).

(b) The first time PINNs are mentioned is in page 2 perhaps move the citation for the paper there?

(c ) “Compared with traditional PDE solvers PINNs can be orders of magnitude faster”; can you clarify if this includes both the training time or only the inference time?

(d) Please clarify for each case you used PINNs if you solve the forward or inverse problem (fit the constant c)? 

(e ) Do you assume in this case you have some data for the potential or not? If not, how can you guarantee the uniqueness of the solution for the constant?

(f) Equation (6) does not mention a loss function for the Boundary Conditions, is this necessary or not in your case if not can you please explain why?


Supervised semantic potential guide (a) The input for the surrogate model $h_η$ lives in the high dimensional space is this correct? I thought we are working with the latent space of the generative model. 


(b) It is not clear to me, how equation 7 gives you the vector field z_t, can you add a few words in the SI and explain in more detail if you think is important for your overall scheme? The output seems to be just a scalar to me.

(c ) In high-dimensional space, a trivial solution exists such that all the flows learn the exact direction -> Can you justify that why this is the case, intuitively I would expect the opposite?

(d) Table (1) What is delta? Please explain>


Appendix : 
(a) Disentanglement Regularization: Since here there is no limit in the SI can you make things more clear? Why having the maximum Jacobian change is a problem? Why an auxiliary classifier is helpful? Perhaps cite or say a few words here of what is auxiliary classifiers.

(b) Check the text after equation 28, I think it might be wrong.

(c ) Algorithms: For the Langevin dynamics how β was selected. How sensitive the method might be to that value?  Can you mention include the selected parameters in your test cases?

(d) For algorithms 2, I don’t know if I missed it in the text or in the Appendix what sampling the timestamps and the potential functions means and why there are two random variables for that? 

(e ) Check Caption for Figure 4 in the SI.

(f) Figure 9, those are representative molecules, right since you have a generative model. Perhaps check the captions.