Navigating Chemical Space with Latent Flows
Authors: Guanghao Wei
Yining Huang
Chenru Duan
Yue Song
Yuanqi Du
Venue: NeurIPS 2024 poster (2024)
Keywords: Dynamical System
Optimal Transport
Molecular Discovery
Deep Generative Models
Primary area: machine_learning_for_physical_sciences

Abstract: Recent progress of deep generative models in the vision and language domain has stimulated significant interest in more structured data generation such as molecules. However, beyond generating new random molecules, efficient exploration and a comprehensive understanding of the vast chemical space are of great importance to molecular science and applications in drug design and materials discovery.
In this paper, we propose a new framework, ChemFlow, to traverse chemical space through navigating the latent space learned by molecule generative models through flows. We introduce a dynamical system perspective that formulates 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. 
Under this framework, we unify previous approaches on molecule latent space traversal and optimization and propose alternative competing methods incorporating different physical priors. 
We validate the efficacy of ChemFlow on molecule manipulation and single- and multi-objective molecule optimization tasks under both supervised and unsupervised molecular discovery settings.
Codes and demos are publicly available on GitHub at 
[https://github.com/garywei944/ChemFlow](https://github.com/garywei944/ChemFlow).

Navigating Chemical Space with Latent Flows
Guanghao Wei"! Cornell University gw3382cornel1.co
Yining Huang"
Harvard University
yininghuanggénms ha harvard.edu
Chenru Duan Deep Principle, Inc. duanchenzu@gmail
Yue Song!
Yuanqi Du'
California Institute of Technology
Cornell University
yuesong@caltech.cem
yd3920cornel1.
Abstract
Recent progress of deep generative models in the vision and language domain has stimulated significant interest in more structured data generation such as molecules. However, beyond generating new random molecules, efficient exploration and comprehensive understanding of the vast chemical space are of great importance molecular science and applications in drug design and materials discovery. In this paper, we propose a new framework, ChemFlow, to traverse chemical space through navigating the latent space learned by molecule generative models through flows. We introduce dynamical system perspective that formulates the problem as learning vector field that transports the mass of the molecular distribution to the region with desired molecular properties or structure diversity. Under this framework, we unify previous approaches on molecule latent space traversal and optimization and propose alternative competing methods incorporating different physical priors. We validate the efficacy of ChemFlow on molecule manipulation and single- and multi-objective molecule optimization tasks under both supervised and unsupervised molecular discovery settings. Codes and demos are publicly available on GitHub at
1 Introduction
Designing new functional molecules has been long-standn aingsttnding challenge in molecular discovery which concerns wide range of applications in drug design and materials discovery [48, 57]. With the increasing interest in applying deep learning models in scientific problems [60, 70], molecular design has attracted considerable attention given its massively available data and accessible evaluations. Among the developed methods, two paradigms emerge: one paradigm searches for new molecules based combinatorial optimization approaches respecting the discrete nature of molecules: the other paradigm builds upon the success of deep generative models in approximating the molecular distribution with a given dataset and then generating new molecules from the learned models [8]. Both of the approaches have demonstrated promising results in small molecule, protein, and materials design [61, 26, 41, 69]. Despite the promise, the chemical space is tremendously large with the number of diug-like small molecule compounds estimated to be from 1023 to 10 [3, 39]. This necessitates either more efficient searching methods or better understanding about the structure of the chemical space. Following the progress made in studying the latent structure of deep generative models (e.g. generative adversarial networks (GANs) [19], variational autoencoders (VAEs) [35].
Equual Contribution.
I 'Equal Supervision.
*This work was completed while the author was at Cornell University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

and denoising diffusion models [23]) in computer vision [28, 5, 22, 38), decent efforts have recently been made in understanding the learned latent space of molecule generative models.
Initially, disentangled representation learning becomes a popular paradigm to enforce a structured and interpretable representation [9]. Specifically, each latent dimension is expected learn disentangled factor of variation, and tweak the latent vector along the dimension could lead to generating new samples with changes only in one molecular property. However, even if imposing such constraints in the training of molecule generative models, the models still struggle to learn meaningful disentangled factors in the early attempts [10]. In addition to constraining the model training procedure, exploring the structure of pre- pre-terined molecule generative models is more efficient. The main approach developed is to utilize optimization approaches to discover the region in the latent space with the desired molecular property. It often trains a proxy function to map from the latent vector to the property, providing access to gradients for gradient-based optimization [40, 20, 13). The third line of work also builds upon pre-trained models as well. It leverages one interesting finding such that the learned latent space of molecule generative models is linearly separable [18], which is also widely studied and used as priori in computer vision [51, 50]. ChemSpace [1 I] develops highly efficient approach to use linear classifiers to identify the separation boundary and considers the normal direction to the boundary the direction of control. Nevertheless, the linear separability assumption may be too strong. It is worth noting that the first line of work does not require labels and can be trained in an unsupervised manner (referred to as unsupervised discovery), while both the second and third lines of work require access to labels to train/identify a guidance model/direction (referred to as supervised discovery).
In this paper, we propose a new framework, ChemFlow, based on flows in a dynamical system to efficiently explore the latent structure of molecule generative models. Specifically, we unify previous approaches (gradient-basee optimization, linear latent traversal, and disentangled traversal) under the realm of flows that transforms data density along time via a vector field. In contrast to previous linear models, our framework flexible to learn nonlinear transformations inspired by partial differential equations (PDEs) governing real-world physical systems such as heat and wave equations. We then analyze how different dynamics may bring special properties to solve different tasks. Our framework can also generalize both supervised and unsupervised settings under the same umbrella. Particularly in the under-studied unsupervised setting, we demonstrate a structure diversity potential can be incorporated find trajectories that maximize the structure change of the molecules (which in turn leads to property change). We conduct extensive experiments with physicochemical, drug-related properties, and protein-ligand binding affinities on both molecule manipulation and (single- and multi-objective) molecule optimization experiments. The experiment results demonstrate the generality of the proposed framework and the effectiveness of alternative methods under this framework to achieve better or comparable results with existing approaches.
2 Background
2.1
Navigating Latent Space of Molecules
The latent space Z of molecule generative models is often learned through an encoder function fo(-) and a decoder function (-) such that the encoder maps the input molecular structures æ E X into an (often) -iw-dinensional and continuous space (i.e. latent space) while the decoder maps the latent vectors z E back molecular structures z'. Note that this encoder-decoder architecture is general and can be realized by popular generative models such as VAFs, flow-based models, GANs, and diffusion models [31, 43, 6, 58,88 66). For simplicity, we focus on VAF-bbased 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 fs(-) parameterized by a neural network that provides the direction to traverse [67]. This can be formulated as a gradient flow following the direction of steepest descent of the potential energy function h(-) and discretized, as follows:
dz 2 dt
2, = 1 - V (z
where we take a dynamic system perspective on the evolution of latent samples. The latent traversal approaches leverage the observation of linear separability in the learned latent space of molecule gener-
2

---

Paramotizze potental function of (b
Predictor Rift*)
Superwised
I
plogP=-842
Encoder felr)
Decoder sy(s)
Property
Unsupervised
Inpul Molecute
Latent Spane
Reconstructed Molecule a'
Jacobian Structural Change
Decoder soloz)
plopP
QED - 0.508
Optiized Property
Drug
ChemFlow framework: (1) a pre-trat pre-trained encoder fo(-) and decoder (+) that maps between molecules æ and latent vectors 2, (2) we use a property predictor (.) (green box) or "Jacobian control" (yellow box) as the guidance to learn a vector field (t z,) that maximizes the change in certain molecular properties (e.g. plogP, QED) molecular structures, (3) during the training process, we add additional dynamical regularization on the flow. The learned flows move the latent samples change the structures and properties of the molecules smoothly. (Better seen in color). The flow chart illustrates case where a molecule is manipulated into a drug like caffeine.
Figure 1:
ative models [18]. Since the direction assumed to be linear, can be found easily. ChemSpace [11] learns linear classifier that defines the separation boundary of the molecular properties. Then the normal direction of the boundary provides linear direction n E 2 for traversing the latent space:
+ nz
(2)
We notice that the above gradient flow and linear traversal can be analyzed and designed from dynamical system perspective: linear traversal can be considered as a special case of wave functions, i.e., we have 0'2,/02,.- 0²2./02t 0 satisfied by wave functions. This connection inspires to consider designing more dynamical traversal approaches in the latent space.
2.2 Wasserstein Gradient Flows
Gradient flows define the curve att) E IR" that evolves in the negative gradient direction of function F: + R. The time evolution of the gradient flow is given by the ODE x' et(t F(x(((tt)) Wasserstein gradient flows describe special type of gradient flow where F is set to be the Wasserstein distance. For example, as introduced Benamou and Brenier [2], the commonly used 7,2 Wasserstein metric induces dynamic formulation of optimal transport:
{
1 2 ptt, x)t((t x² dt dæ d,plt,x) - (v(t,z))(t,x)) (3)
Wa(u, min V.P
where jr, are two probability measures at the source and target distributions, respectively. Interest- ingly, we take the gradient of potential energy V as the velocity field applied distribution, the time evolution of V can be seen minimize the Wasserstein distance and thus follow optimal transport. In Appendix A, we give detailed derivations Wasserstein distance and discuss alternative PDEs of the density evolution recovered by Wasserstein gradient flow (e.g. Wasserstein gradient flow over the entropy functional recovers heat equation) following the seminal JKO scheme [34].
how vector fields minimize the I
3 Methodology
present ChemFlow as a unified framework for latent traversals chemical latent space as latent flows. We parameterize set of scalar-valued valued energy functions o* MLPPss (t, x) IR and use the

---

learned flow V == to traverse the latent samples. The traversal process can be described by the following equation in Lagrangian way (particle trajectory):
2 +
- 1, 2
(4)
Alternatively, as an Eulerian approach, we can write the time evolution of the density through pushforward map:
P Pt-1
(5)
where 8y defines the time-dependent flow that transforms the densities latent samples through a probability path. The pushforward measure [vx]. induces a change of variable formula for densi- ties [47]:
Oui (=)
- (v- (z)] det
(6)
Oz
In the following, we will introduce how v=0f% matched to some pre-defined velocities for generating different flows.
3.1 Learning Different Latent Flows
Given pre-trained molecule generative model Is 2 + X' with prior distribution p(z), we would like to model K different semantically disentangled latent trajectories that correspond to different properties of the molccules, numbered by superscript k.
HamiltonJJacobi Flows. One desired property for the latent traversal comes from optimal transport theory such that the transport cost is minimized (i.e. shortest path). This property can be enforced by solving Eq. (3) by Karush-Kuhn Kaai--KKnnTTTckkee Tucker (KKT) conditions, which will give the optimal solution
the Hamilton-Jacobi Equation (HJE):
Ot o (t, 2)
1 V 2)11 =0
where the velocity field is defined as the flow VØt. The HJE can also be interpreted as mass transportation in fluid dynamics, i.e., under the velocity field Vot, the fluid will evolve to the target distribution with an optimal transportation cost.
We achieve the HJE constraint by matching our flow fields and define the boundary condition as:
C. = Σ ot, 1 (t,
Lo Lo
Σ
II x (0, 20)1
(8)
where T represents the total number of traversal steps, C, restricts the energy to obey physical constraints, and Cg restricts ott x,) to match the initial condition. Our latent traversal can be thus regarded as dynamic optimal transport between distributions of molecules with different properties.
Wave Flows. Alternatively, we can pivot the optimal transport property to enforce additional physical and dynamical priors. For example, if we specify the flow to follow wave-like dynamics, we can use the second-order wave equation:
02
012 v (1, 2) 2V2 =
The above constraint empirically produces highly diverse and realistic trajectories. Our velocity matching objective and boundary condition then become:
T
,, T 1=0
2
oi (t,z)
L. - - 20)
(10)
k=0
where L, and C. restrict the physical constraints and the initial condition, respectively. Note that ok = 0 is a trivial optimal solution for the above two objectives regarding that both C, and L.s are non-negative. To prevent the parameterized &e from converging to such a trivial solution, we introduce more guidance the loss function in Section 3.2 separately.
Alternative Flows. Besides HJE and wave equations, our framework is also general to include other commonly used PDEs that allow for different dynamics along the flow, such as Fokker Planck equation and heat equation. In the experimental section, we will explore the effectiveness of each latent flow in different supervision settings.

---

3.2 Supervised & Unsupervised Guidance
Supervised Semantic Guidance. When an explicit semantic potential energy function 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 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 for the trajectory of the generated molecules.
d= (z.)), V= (t,)),, Lp = - sign (d) ||d]
(11)
The intuition behind this objective is to learn the vector field 2, 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 it determines minimizing maximizing the property.
The proposed objective function in the supervised scenario is
L = L, + Ls + Lp
Unsupervised Diversity Guidance. When explicit potential energy function is provided to learn the flow, we need to define potential energy function that captures the change of molecular properties. As molecular properties are determined by the structures, we devise potential energy that maximizes the continuous structure change of the generated molecules. Inspired by Song et al. [54], we couple the traversal direction with the Jacobian of the generator to maximize the traversal variations in the molecular space. The perturbation on latent samples can be approximated by the first-order Taylor approximation:
g + eV vV (t, x:)) = g(x,) + 0g(zi) V (1) + RL(g(z)))
(12)
0z.
where denotes perturbation strength, and (.) is the high-order terms. In the unsupervised setting, for sufficiently small c, if the Jacobian-vects vector product can cause large variations the generated sample, the direction is likely to correspond to certain properties of molecules. We therefore introduce such a Jacobian- product guidance:
-1=11001 [gg(z)) Ly =
V=0 (t, 21)
(13)
dz
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 pushes the initial data distribution the target one concentrated 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.
Disentanglement Regularization. While the above formulation can encourage smooth dynamics and meaningful output variations, the flows are likely mine identical directions which all correspond to the maximum Jacobian change. To avoid such a trivial solution, we adopt an auxiliary classifier L following Song al. [54] predict the flow index and use the cross-entropy loss optimize it:
L+ LCE(L, (ge (z)); (z.+1)), k))
(14)
Where (zx) is the generated sample from timestep t. We see the extra classifier guidance would encourage each flow to independent and find distinct properties. For each target property, we compute the Pearson correlation coefficient using a randomly generated test set. This coefficient measures the correlation between the property and natural sequence (from to time step !) along the optimization trajectory. We then select the energy network that achieves the highest correlation score for optimizing molecules with that specific property.
The proposed objective function in the unsupervised scenario is
L - L, Ls Lg La
3.3 Connection with Langevin Dynamics for Global Optimization
In scenarios where our flow adheres to the dynamies of the Bokker-Planch equation, our approach may also be interpreted as employing learned potential energy function to simulate Langevin Dynamics
5

---

for global optimization [15]. Notably, the convergence of Langevin dynamics, particularly at low temperatures, tends oceur around the global minima of the potential energy function [7]. The continuous and discretized Langevin dynamics are as follows:
dz: (, + 2dw,
(15)
21 21-1 V V-h V 2dt
Proposition 3.1. (Global Convergence of Langevin Dynamics, adapted from Gelfand and Mitter [16]). Given a Langevin dynamics in the form of
= == a. av (V + u.) + b.w
where is a d-dimensional Brownian motion, 04 and by are set of positive numbers with ar, br - and ur is a set of random variables in R denoting noisy measurements of the energy function (.).)) Under mild assumptions, converges to the set of global minima of (.) in probability.
Following Proposition 3.1, the learned latent flow be used search for molecules with optimal properties and it converges to the global minimizers of the learned latent potential energy function.
4 Experiments
4.1 Experiment Sct-up
Datasets & Molecular properties. We extract 4,253,577 molecules from the three commonly used datasets for drug discovery including MOSES [46], ZINC250K [27], and ChEMBL [68]. Molecules are represented by SELFIES strings [37]. All input molecules are padded to the maximum length in the dataset before fitting into the generative model. We consider a total of molecular properties which include 3 general drug-related properties - Quantitative Estimate of lrug-likeness (QED), Synthesis Accessibility (SA), and penalized Octanol-water Partition Coefficient (plogP) and machine learning- based target activities - DRD2, JNK3 and GSK3B [25] 2 simulation-basee target activities docking scores for two human proteins ESRI and ACAAI. See Appendix D.3 for details.
Implementations. We establish our framework by pre-training a VAE model that learns latent space of molecules that can generate new molecules by decoding latent vectors from the latent space. We adapt the framework Eckmann et al. [13] which is a basic VAE architecture with molccular SELFIES string representations and an additional MLP model the surrogate property predictor. See Appendix D.5 for all implementation and hyper-parameter details.
Model variants. As discussed in Section 3.1, our proposed framework general to incorporate different dynamical priors learn the flow. For the experiments, including gradient flow (GF), Wave flow (Wave, eq. (9)), Hamilton Jacobi flow (HJ. eq. (7)) and Langevin Dynamics equivalently Fokker Planck flow (LD, eq. (15)).
consider four types of dynamics
For the specific molecular properties and evaluations, the readers are kindly referred Appendix D
for details. We also move qualitative evaluations to Appendix F due to space limit.
4.2 Molecule Optimization
Molecule optimization is key in drug design and materials discovery, aiming to identify molecules with optimal properties [4]. Various machine learning methods have accelerated this process [8]. Our discussion focuses on optimization within the latent space of generative models, primarily using gradient-based optimization as outlined in Section 2.1. We categorize molecule optimization into three scenarios: (1) unconstrained optimization to identify molecules with the best properties, (2) constrained optimization to find molecules with the best-expected property and similar to specific structures- a common step the lead optimization process, and (3) multi-ohjective optimization to simultancously enhance multiple properties of molecule.

---

Table 1: Unconstrained plogP, QED maximization, and docking score minimization. (SPV denotes supervised scenarios, UNSUP denotes unsupervised scenarios). Boldface highlights the highest- perforrip generation for each property within each rank.
METHOD
PLOGP T
QED 1
ESRI DOCKING 1
ACAAI DOCKING
1ST
2ND
1ST
3kD
1ST
2ND
1ST
3RD
RANDOM
3.52
3.43
3.37
0.940
0.933
0.932
-10.32
-10.18
-10.03
-9.86
-9.50
-9.34
CHEMSPACE
3,74
3.69
3.64
0.941
0.936
0.933
-11.66
-10.52
-10.43
-9.81
-9.72
-9.63
GRADIENT
-10.46
-9.90
WAVE (SPV)
4.76
3.78
3.71
0.947
0.934
0.932
-11.05
-10.71
-10.68
-10.48
-10.04
-9.88
WAYE (UNSUP)
5.30
5.22
5.14
0.905
0.902
0.978
-10.22
-10.06
-9.97
-9.69
-9.64
-9.57
(SPV)
HJ (UNSUP)
4.26
4.10
4.07
0.930
0.928
0.927
-10.24
-9.96
-9.92
-9.73
-9.11
-9.24
LD
4.74
3.61
3.55
0.947
0.947
0.942
-00.68
-10.29
-10.28
-10.34
-9.74
-9.64
Baselines. For molecule optimization, we follow the same experiment procedure in Eckmann al. [13]'. To ensure fair comparison, we use the same pre-trained VAE model for all the methods. The details about the baselines are deferred to Appendix D.I. We also propose evolutionary algorithm (EA)-based latent optimization approaches our comparison (Appendix D.7).
Unconstrained Molecule Optimization. In this study, randomly sample 100,000 molecules from the latent space and assess the top three scores after 10 steps of optimization for each method (de- tails in Table 1). For two specific docking scores tasks, however, only 10,000 molecules are sampled due to computational resource constraints. All methods employ a step size of 0.1 to ensure fair comparison. Our findings reveal that the efficacy of optimization methods varies with target prop- erties, highlighting the necessity of employing diverse set of approaches within the optimization framework, rather than depending on a single dominant method. We also visualize some generated ligands docked into protein pockets in Figure 2.
Random
ChemSSpace -11.66 kcalimal
Gradient Flow
Vilvi egn, (SpV) -11.05 scalimol
Wave egn, (unsup) -10.2 kcallmol
Hi egn. (apv) -10,50 kcalmmll
HJ eon, (unsup) -10 kcalmmol
Langavin Dynamics -10.80 kcallmal
-10,32 kcallmmll
-11.00 kcalimal
ESR1
Random -3.86 kcalinol
ChemSpace -9.81 koaimmal
Gradient F low 90 kcalimal
Weve 100.4 kcaliml
Wave (unsup)
HJ
Dynamics
-3.89 acalimol
-9.80 koalmnol
.9 73 Kealimal
-10, 34 kealimal
ACAA1
Figure Visualization of generated ligands docked against target ESRI and ACAAI.
Furthermore, when extending the optimization procedure to 1,000 steps, illustrated in Figure 3, Langevin dynamics significantly pushes the entire distribution to molecules with better properties, surpassing other methods in performance. Although the random direction method is effective in optimizing molecules, it does not consistently produce significant shifts in the distribution. Moreover, the outputs from ChemSpace often converge to just few molecules, which is indicative of the challenge posed by Out-of-Diitrrbution (OoD) generation.
Similarity-coostrranned Molecule Optimization. Adopting methodologies from JT-VAE [31] and LIMO [13], we select 800 molecules with the lowest partition coefficient (plogP) scores from the ZINC250k dataset. These molecules undergo 1,000 steps of optimization until convergence is achieved for all methods. The optimal results for each molecule, adhering predefined similarity constraint (8), is reported as in Table 2. The structural similarity is measured using Tanimoto similarity between Morgan fingerprints with radius of 2.
Without any similarity constraints, the unsupervised approaches significantly improve molecular properties with a high success rate. However, as the similarity constraints increase, both the magnitude
Note that we notice there was 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 be found in Appendix D.3.
7

---

Handom
Grudient Flow
ShamSpace
Ward
Langests Dynanice
6.005
0,0025
0.0025
54005
0,003
6-3000
0.0020
200
100
0.0615
0.024
0.00:0
500
0.0025
0.022
0.0666
0.000
660
0.003
6:3391
01000
0.0010
00000
500
550
Figure 3: Molecular property plogP distribution shifts following the latent flow path.
Table 2:
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.
Method
6=0 0
5=
8= 0.4
6= 0.6
Random
11.76 (99.0)
7.64 6.38 (80.0)
5.03 5.73
2.37 + 3.71 (21.1)
ChemSpace
12.13 (99.8)
9.07 6.80 (90.2)
7.52 6.29 (59.4)
5.80 (20.2) 5.50 (180))
Gradient Flow
7.88 7.28 (60.4)
7.20 + 6.98 (56.5)
5.45 6.45 (41.9)
Wave (spv)
6.83 (59.6)
5.62 3662 (54.9)
4.31 5.55 (41.9)
2.47 4.22 (20.6)
Wave (unsup)
19.76 + 13.62 (99.6)
7.47 + 9.62 (50.2)
2.06 + 4.37 (27.3)
0.77 + 2.21 (16.8)
HJ (spv)
8.58 # 8.08 (68 ()
6.62 + 7.4 (60.0)
4,27 + 5 (40.6)
2.39 + 4,10 (18.5)
(unsup)
20.64 12.93 (98.0)
8.57 9.69 (50.1)
2.12 (19.5)
0.67
(8.6)
L.D
12.98 6.23 (99.6)
9.70 + 6.21 (94.4)
6.14 + 5.99 (70.9)
2.94 + 4.34 (35.4)
and success rate of unsupervised methods decrease notably. The most considerable improvements of ChemSpace are observed when similarity constraints are set at 0.4 and 0.6. Despite this, as shown in Figure and previously discussed, the generation within ChemSpace encounters significant OoD issues after extensive iterations. Among all techniques evaluated, Langevin dynamics stands out for its overall high improvements and high success rate. As further illustrated in Figure 5, Langevin dynamics also demonstrates a notably faster empirical convergence rate. We also report the performance in optimizing the QED property in Appendix D.8.
Surprisingly, we observe that the random direction performs 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 E. 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 3.
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 optimiza- tion [13, 11]. To generate molecules that are optimized on multiple properties, we use a similar setting as mimilaarty-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 low SA simultaneously. At each time step. the latent vector is optimized following the averaged direction of two corresponding flow directions. This scheme could be seamlessly generalized to k-objectives optimization. Table shows that Langevin dynamics and ChemSpace achieve the best or competitive performance at all similarity cutoff levels.
4.3 Molecule Manipulation
Molecule manipulation is a relatively new task proposed in Du et al. [11] to study the performance of 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 essential to help chemists systematically understand the chemical space.
Supervised Molecule Manipulation Table shows the success rate results of manipulating 1,000 randomly sampled molecules to optimize each desired property. Following Du et al. [11], we traverse
8

---

Table 3:
Multi-objective (QED-SA) maximization. The value of QED and SA scaled to both have range from 0 to 100 for an equal-weighted sum. The method with the highest weightee sum improvement of QED and SA of each structure similarity level is bolded.
6 D
0.2
6 0.4
d 0.6
QED
Random
45.5 13.3 (99.5)
20.7 (81.)
12.8+ 10.7 (57.0)
8.0 (29.5)
ChemsSpac
47.0 (996)
25.9
(88.2)
15.5 + 13.2 (634)
9.7 10.1 (31.6)
Gradient Flow
31.9 (8999
23.0 16.4 (80.0)
14.4 (5..))
9.4 (32.2)
Wave (SPV)
14.6
(24.2)
1..
68 (12.4)
Wave (UNSUP)
39.0 18.5 (96.1)
18.3 14.0 (69.4)
10.0 (43.2)
7.0 (24.9)
HJ (SPV)
45.2+ 13.7 (98.9)
22.7
15.6 (84.9)
13.9
(57.5)
9.3 (30.2)
(UNSUP)
40.6
(96.8)
(46.2)
(26.9)
L.D
47.9 13.1 (99.6)
27.6
16.3 (92.1)
15.4
12.5 (71.4)
9.5 9:4 (40.1)
Random
8.04
6.10 (27.0)
5.73
2.61 1.21 (10.9)
ChemSpace
7.92 (42.8)
6.63 6.71 (36.1)
4
2.89 3.22 (15.4)
Gradicnt Flow
9.56 (44.0)
6.66
5.27
5.30 (26.6)
3.04 3.62 (17.1)
Wave (SPV)
6.37
5.77 (12.4)
4.54 (10.5)
3.59
Wave (UNSUP)
15.21 4:
6.51 (70.8)
3.92 (15.9)
2.22 (25.9)
HJ (SPV)
8.93 + 8.39 (52.1)
6.69 6.61 (38.8)
5.21 5.72 (29.2)
3.23 + 3.40 (18.6)
HJ (UNSUP)
16.03 + 10.31 (91.0)
7.35 6.34 (68.5)
4.22 4.09 (46.8)
2.73 + 2.85 (27.8)
(69.2)
7.51 (45.4)
4.95 (33.6)
2.75 3.09 (20.1)
Table 4: Success Rate of traversing latent molecule space to manipulate over a variety of molecular properties. Numbers reported are strict success ratc/relaxed success rate in %. (SPV
denotes supervised scenarios, UNSUP denotes unsupervised scenarios). The ranking is the average
between the ranking of average strict success rate and ranking of the average relaxed success rate.
RANKING
AVERAGE
PLOGP (1)
QED
SA (1)
DRD2 (1)
JNK3
GSK3B (1)
RANDOM-ID
6.00/31.60 1.42
0.00/0.00
0.00/
2.50/9.55 9.55
0.00/ 0.00
RANDOM
0.57 1 42.3
0.00/ 32.60
0.10/3.20
0.40 8.60
87.0/
1.50 / 81.40
0.90/ 40.90
CHIMSPACH
6.17 22.83
5.20 25.00
6.00 / 18.10
6,80 26.50
3.20/ 18.00
8.60 25.80
7,20 23. 23.60
WAVE (UNSUP)
1.18/45.28
0.60
0.60/
1.90/
0.40
L.50/ /..0/78.20
1.80/4..10
WAVE (SPV)
1.85/8.08
0.00
3.40 12.10
4.00 18.60
3.50/
0.00/
0.20/ 0.30
HJ (UNSUP)
3
2.3/25.28
3.00 15.60
0.701 3.20
1.70 13.00
0. 0.20/M7.20
4,80 18.07
3.40 14.
HJ (SPV)
1.9
3.00 / 13.20
3.00/ 7 2.20
3,70 15.00
1.90
0.20/ 0. 50
0.00 /..00/0.00
GP (SPV)
7.62 28.
6.90 / 28
6.60 16.20
6.30//25.10
7.10/36.00
11.70/
7.10/
(SPV)
6.23
26.68
5.90 26.00
15.50
5,20 5.29/22.90
6.00 6.00/33.30 33.30
8 33.80
5.30/288.60
the latent space for 10 steps in the traversal direction of each method and report the strict and relaxed success rate. Details of the definition of these metrics can be found Appendix D.6. Among all approaches, the gradient flow achieves the highest success rates in multiple properties such that it takes the steepest descent of the surrogate model. When the step size is small enough, is reasonable to Iearn smooth path. However, the results still vary across properties.
Unsupervised Molecule Manipulation As the correspondence between specific molecular properties and learned latent flows is not explicitly given in the unsupervised scenario, we use an artificial process to mimic the use case in reality. Specifically, we learn 10 different potential energy functions representing 10 disentangled flows following Algorithm 1 using Wave equation and Hamilton-Jacon equation and validated them on 1,000 unseen molecules. For each flow, we evaluate the properties of molecules generated along the 10-step manipulation trajectory. The learned potential energy function with the highest correlation score is selected for each property representing the learned Jacobian structural change that would most effectively optimize the corresponding property.
In Table 4. 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 relaxed success rate of manipulating molecules for 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. On the contrary, the structure change measurement does not provide supervision but is correlated with key molecular properties. We would like to point out that this an open question in chemistry, often referred as to the structure-activity relationship [12], such that it is important to know the correspondence between structure and activity. We believe this a promising result to demonstrate that generative models "realize" molecular property by learning from structures.

---

Among the quantitative results, it is interesting that the random direction achieves good relaxed success rate for some properties, we argue this is because of the specific property of the learned latent space. The latent space learned by generative models tend to be smooth such that similar molecular structures are often mapped to close areas in the latent space. In Appendix E, we find that some molecular properties are highly correlated with their latent vector norms, in which random direction always increases the norm and thus successfully manipulates a portion of molecules by chance.
5 Conclusion, Limitation and Future Work
In this paper, we propose unified framework for navigating chemical space through the learned latent space of molecule generative models. Specifically, we formulate the traversal process flow that defines a vector to transport the mass of molecular distribution through time to desired concentrations (e.g. high properties). Two forces (supervised potential guidance and unsupervised structure diversity guidance) are derived to drive the dynamics. We also propose a variety of new physical PDEs the dynamics which exhibit different properties. We hope this general framework can open up a new research avenue to study the structure and dynamics the latent space of molecule generative models.
Limitation and future work. This work is a preliminary study on small molecules and it may be interesting to see it transfer to larger molecular systems or more specialized systems and properties. Beyond molecules, this approach has the potential to be extended to languages and other data modalities.
Acknowledgement
Y.D. would like to thank Ziming Liu and Kirill Neklyudov for helpful discussions.
References
[1] L. Ambrosio, N. Gigli, and Savaré. Gradient flows: in metric spaces and in the space of probability measures, Springer Science & Business Media, 2005.
[2] J.-D. Benamou and Y. Brenier. A computational fluid mechanics solution to the Lonne-kantorovidd mass transfer problem. Numerische Mathematik, 84(3):375-333, 2000.
[3] R. Bohacck, C. McMartin, and W. C. Guida. The and practice of aruettre-based drug design: molecular modeling perspective. Medicinal research reviews, 16(1):3-50, 1996.
[4] N. Brown, M. Fiscato, M. H. Segler, and A. C. Vaucher. Guacamol: benchmarking models for de novo molecular design. Journal of chemical information and modeling, 59(3): 1096-118,, 2019.
[5] C. P. Burgess, Higgins. A. Pal, L. Matthey, N. Watters, G. Desjardins, and A. Lerchner. Understanding disentangling fa- ArXiv, ftss/1804.03599 2018.
[6] N. D. Cao and T. Kipf. MoIGAN: An implicit generative model for small molecular graphs, 2018.
[7] T.SS. Chiang, C.-R. Hwang, and S. J. Sheu. Diffusion for global optimization in r'n. SIAM Journal on Control and Optimization, 25(3):737-753, 1987.
[8] Y. Du, Fu, J. Sun, and S. Liu. Molgensurvey: A systematic survey in machine learning models for molecule design. arXiv preprint 24.. 2022.
[9] Y. Du, X. Guo, A. Shehu, and I.. Zhao. Interpretable molecular graph generation via monotonic constraints. In Proceedings of the 2022 SIAM International Conference on Data Mining (SDM), pages 73-81. SIAM, 2022.
[10] Y. Du, X. Guo, Wang, A. Shehu, and L. Zhao. Small molecule generation via disentangled representation learning. Bioinformatics, 38( 38(12):3200 -208,, 2022.
[11] Y. Du, X. Liu, N. M. Shah, S. Liu, J. Zhang, and B. Zhou. Chemspace: Interpretable and interactive chemical space exploration. Trans. Mach. Leam. Res., 2023, 2023.
10

---

[12] A. Z. Dudek, T. Arodz, and J. Gálvez. Computational methods in developing quantitative structure-activity relationships (qsar): review. Combinatorial chemistry & high throughput screening, 933::233222,, 2006.
[13] P. Eckmann, K. Sun, B. Zhao, M. Feng. M. Gilson, and R. Yu. Limo: Latent inceptionism for targeted molecule generation. International Conference on Machine Learning, pages 3777-5792, PMLR, 2022.
[14] T. Fu, W. Gao, Coley, and J. Sun. Reinforced genetic algorithm for structure-based drug design. Advances in Neural Information Processing Systems, 35:12325- 12338, 2022.
[15] C. W. Gardiner et al. Handbook of stochastic methods, volume 3. springer Berlin, 1985.
[16] S. B. Gelfand and S. K. Mitter. Recursive stochastic algorithms for global optimization r'd. SIAM Journal on Control and Optimization, 29(5):999-100) 1011, 1991.
[17] L. Goetschalckx, A. Andonian, A. Oliva, and P. Isola. Ganalyze: Toward visual definitions of cognitive image properties. In ICCV, 2019.
[18] R. Gómez-Bombarelli, J. N. Wei, D. Duvenaud, J. M. Hermández-Lobato B. Sinchez-LLngelligg D. She- berla, J. gquilera-Iparraguin T. D. Hirzel, R. P. Adams, and A. Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science, 4(2):268-276 2018.
[19] I. Goodfellow, Y. Bengio, A. Courville, and Y. Bengio. Deep learning, volume 1. MIT Press, 2016.
[20] R.-R. Griffiths and J. M. Hemándezz-Lobato. Constrained bayesian optimization for automatic chemical design using variational autoencoders. Chemical science, 11(2):777-566 2020.
[21] G. L.. Guimaraes, B. Sanchez-Lengeling, C. Outeiral, P. L. C. Farias, and A. Aspuru-Guzik, Objective- reinforced generative adversarial networks (organ) for sequence generation models. arXiv preprint arXiv:1705. 10843, 2017.
[22] E. Härkönen, A. Hertzmann, J. Lehtinen, and S. Paris. Ganspace: Discovering interpretable gan controls. Advances in neural information processing systems, 33:9841-9850, 2020.
[23] Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840 6851, 2020.
[24] E. Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling. Equivariant diffusion for molccule generation in 3d. In International conference on machine learning, pages 8867-8887. PMLR, 2022.
[25] K. Huang, T. Fu, W. Gao, Y. Zhao, Y. Roohani, J. Leskovec, C. W. Coley, C. Xiao, J. Sun, and M. Zitnik. Artificial intelligence foundation for therapeutic science. Nature chemical biology, 18 10300)10033-1066 2022.
[26] J. B. Ingraham, M. Baranov, Z. Costello, K. W. Barber, W. Wang, A. Ismail, V. Frappier, D. M. Lord, C. Ng- Thow-Ifing, R. Van Vlack, et al. Illuminating protein space with a programmable generative model. Nature, pages 1-9, 2023.
[27] J. J. Irwin and B. K. Shoichet. Zinc- a free database of commercially available compounds for virtual screening. Journal of chemical information and modeling. 45(1):177-182, 2005.
[28] A. Jahanian, L. Chai, and P. Isola. On the" stecrability" of generative adversarial networks. In International Conference Leaminng Representations, 2019.
[29] A. Jahanian, L. Chai, and P. Isola. On the" steerability" of generative adversarial networks. ICLRR 2020.
[30] J. H. Jensen. A graph-based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space. Chemical science, 10(12)3367-33722 2019.
[31] W. Jin, R. Barzilay, and T. Jaakkola. Junction tree variational autoencoder for molecular graph generation. ICML, 2018.
[32] Jin, R. Barzilay, and T. Jaakkola. Hierarchical generation molecular graphs using structural motifs. In International Conference on Machine Learning. 2020.
[33] J. Jo, S. Lee, and S. J. Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International Conference on Machine Learning, pages 10362-1038) PMLR, 2022.

---

[34] R. Jordan, D. Kinderlehrer, and F. Otto. The variational formulation of the fokker-planck equation. Siam Journal Applied Mathematics, 1996.
[35] D. P. Kingma and M. Welling. Auto-encnding variational bayes. arXiv preprint arXiv: 2013.
[36] M. Krenn, F. Hase, A. Nigam, P. Friederich, and A. Aspuru-Guzik. Self-referencing embedded strings (selfies): A 100% robust molecular string representation. Machine Learning: Scienee and Technology. 1, 2019.
[37] M. Krenn, F. Häse. A. Nigam, Friederich, and A. Aspuru-Guzik, Self-referencing embedded strings (selfics): A 100% robust molecular string representation. Machine Learning: Seience and Technology, I (4):045024, 2020.
[38] Kwon, Jeong, and Y. Uh. Diffusion models already have semantic latent space. In The Eleventh International Conference on Learning Representations, 2022.
[39] C. Lipinski, Lombardo, B. Dominy, and P. Feeney. Experimental and computational approaches to estimate solubility and permeability in drug discovery and development settings. Advanced drug delivery reviews, 64:4-17, 2012.
[40] Q. Liu, M. Allamanis, M. Brockschmidt, and A. Gaunt. Constrained graph variational autoencoders for molecule design. Advances in neural information processing systems, 31, 2018.
[41] H. Loeffler, J. He, A. Tibo, J. P. Janet, A. Voronov, L. Mervin, and O. Engkvist. Reinvent4: Modern ai di-driven generative molecule design. 2023.
[42] T. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2017.
[43] K. Madhawa, K. Ishiguro. K. Nakago, and M. Abe. Graphnvp: An invertible flow model for generating molecular graphs. arXiv preprint arXiv: 1600, 2019.
[44] D. Misra. Mish: A self regularized non-monotonic activation function. In British Machine Vision Conference, 2020.
[45] W. Peebles, J. Peebles, J.-Y. Zhu, A. Efros, and A. Torralba. The hessian penalty: A weak prior for unsupervised disentanglement. In ECCV, 2020.
[46] D. Polykovskiy, A. Zhebrak, B. Sánchez-Lengeling, S. Golovanov, O. Tatanov, S. Belyaev, R. Kurbanov, A. A. Artamonov, V. Aladinskiy. M. Veselov, A. Kadurin, I. Nikolenko, A. Aspuru-Guzik, and A. Zha- voronkov. Molecular sets (moses): A benchmarking platform for molecular generation models. Frontiers in Pharmacology. 11, 2018
[47] D. Rezende and Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
[48] B. Sanchez-Lengeling and A. Aspuru-Guzik, Inverse molecular design using machine learning: Generative models for matter engineering. Science, 3640)))3)33 2018.
[49] A. Schneuing, Y. Du, C. Harris, A. Jamasb, Igashov, W. Du, T. Blundell, P. Lió, C. Gomes, M. Welling, et al. Structure- based drug design with cquivariant diffusion models. arXiv preprint Xxii:2210011699 2022.
[50] Y. Shen and B. Zhou. Closed-form factorization of latent semantics in gans. In Proceedings of the JEFF/CV conference on computer vision and pattern recognition, pages 1532-1540, 2021.
[51] Y. Shen, C. Yang, X. Tang, and B. Zhou, Interfacegan: Interpreting the disentangled face representation learned by gans. IEEE transactions pattern analysis and machine intelligence, 44(4):2004 2018,, 2020.
[52] Y. Song, N. Sebe, and W. Wang. Orthogonal svd covariance conditioning and latent disentanglement. IEEE T-PAMI, 2022.
[53] Y. Song, A. Keller, N. Sebe, and M. Welling. Flow factorzied representation learning. In NeurIPS, 2023.
[54] Y. Song, A. Keller, N. Sebe, and M. Welling. Latent traversals in generative models as potential flows. arXiv preprint arXiv:2304. 12944, 2023.
[55] Y. Song, J. Zhang, N. Sebe, and W. Wang. Householder projector for unsupervised latent semantics discovery. In ICCV, 2023.
12

---

[56] Y. Song. T. A. Keller, Y. Yue, P. Perona, and M. Welling, Unsupervised representation learning from sparse transformation analysis. arXiv preprint arXiv:2410.0556. 2024.
[57] J. Vamathevan, D. Clark, P. Czodrowski, I. Dunham, E. Ferran, G. Lee, B. Li, A. Madabhushi, P. Shah, M. Spitzer, et al. Applications of machine learning in drug discovery and development. Nature reviews Drug discovery, 18(6)463-477, 2019.
[58] C. Vignac, I. Krawezuk, A. Siraudin, Wang, V. Cevher, and P. Frossard. Digress: Discrete denoising diffusion for graph generation. In The Eleventh International Conference on Learning Representations, 2023.
[59] A. Voynov and A. Babenko. Unsupervised discovery of interpretable directions in gan latent space, In ICML, 2020.
[60] H. Wang, T. Fu. Y. Du, W. Gao, Huang, Liu, Chandak, Liu, Van Katwyk, A. Deac, et al. Scientific discovery in the age of artificial intelligence. Nature, 620(7972):47-60) 60, 2023.
[61] L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, E. Eisenach, W. Ahem, Borst, R. J. Ragotte, L. F. Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620 1089 1100, 2023.
[62] Y. Xie, C. Shi, H. Zhou, Y. Yang. W. Zhang, Y. Yu, and L. Li. Mars: Markov molecular sampling for multi-objective drug discovery. ArXiv, abs/2103. 10432, 2021.
[63] X. Yang, J. Zhang, K. Yoshizoe, K. Terayama, and K. Tsuda. ChemTS: an efficient python library for de novo molecular generation. Science and technology of advanced materials, 18(1):972-97 976, 2017.
[64] J. You, B. Liu, R. Ying, V. Pande, and J. Leskovec. Graph convolutional policy network for goal-directed molecular graph generation. In NIPS, 2018.
[65] J. You, B. Liu, R. Ying, V. S. Pande, and J. Leskavec. Graph convolutional policy network for goal-directed molecular graph generation. In Neural Information Processing Systems, 2018.
[66] Y. You, R. Zhou, J. Park, H. Xu, C. Tian, Z. Wang, and Y. Shen, Latent 3d graph diffusion. In The Twelfth International Conference on Learning Representations, 2024.
[67] C. Zang and F. Wang. Moflow; an invertible flow model for generating molecular graphs. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pages 617-626. 2020.
[68] B. Zdrazil, E. Felix, F. Hunter, E. J. Manners, J. Blackshaw, S. Corbett, M. de Veij, H. Ioannidis, D. M. Lopez. F. Mosquera, M. P. Magariñiss. N. Bosc, R. Arcila, T. Kizilören, A. Gaulton, A. P. Bento, M. F. Adasme, P. Monecke, G. A. Landrum, and R. Leach. The chembl database 2023: a drug discovery platform spanning multiple bioactivity data types and time periods. Nucleic Acids Research, 52:D1180- 11192, 2023.
[69] C. Zeni, R. Pinsler, D. Zügnner, A. Fowler, M. Horton, X. Fu, S. Shysheya, J. Crabbé, L. Sun, Smith, et al. Mattergen: a generative model for inorganic materials design. arXiv preprint arXiv:2312.03687, 2023.
[70] X. Zhang, L. Wang. J. Helwig, Y. Luo, C. Fu, Y. Xie, M. Liu, Y. Lin, Z. Xu, K. Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint NVV:2337.00422, 2023.
[71] Z. Zhou, S. M. Kearnes, L. R. learning. Scientific Reports, 9, 2018.
R. N. Zare, and P. Riley. Optimization of molecules via deep reinforcement
[72] X. Zhu, C. Xu, and D. Tao, Leamiigg disentangled representations with latent variation predictability. In ECCV, 2020.
13

---

Appendix for ChemFlow
A Wasserstein Gradient Flow
15
PDE-regularized Latent Space Learning
16
Extended Related Work
17
C.I Machine Learning for Molecule Generation
17
C.2 Goal-oriented Molecule Generation
17
C,3 Image Editing in the Latent Space
17
Experiments Details
18
D.I Baselines
18
D.2 Training Dataset
18
D.3 Molecule Properties
18
D.4 Training and inference
19
D.5 Experiments Setup
19
D.6 Evaluation Metrics
20
D.7 Additional Baselines
D.8 More Experiment Results
22
Latent Space Visualization and Analysis
24
Qualitative Evaluations
25
14

---

A Wasserstein Gradient Flow
As shown in the main paper, based on the dynamic formulation of optimal transport [2]. the L2 Wasserstein distance can be re-written as:
J (=(z))ux(2)P dzdl
W2(uo, pi) = min
(16)
where (x))) is the velocity of the particle at position & and time t. and (x is the density dj((z) = Pe(z)dz. The distance can be optimized by the gradient flow of certain function space and time.
Consider the functional F IR" + R that takes the following form:
F U((((2))
F(u) =
(17)
The curve is considered as a gradient flow if satisfies VF= (z)) [1]. Moving the particles
leads to:
F(1) / (=))))))))))))))))) dd((z) dz
d
(18)
dt.
dt
The velocity vector satisfies the continuity equation:
dp(cz) - (v(()))))((())
(19)
dt
where V (m()))((z) the tangent vector at point pe(z). Eq. (18) can be simplified to:
dt F(1) / - ('(((=))VV (<(()))(()) dz
d
/ (((""((())))((())ec
(20)
On the other hand, the calculus of differential geometry gives
d F(n) DiffF (-v (1)(()())()))))))))))))))))))))) - (VF,-V (VE,)) (VML---)))))))))))))
(21)
where (,) is a Riemannian distance function which defined as:
(-V (wi(z)p((z)) -V- (()2(()))((()))) (x)) dz
(22)
This scalar product coincides with the W2 distance according to Benamou and Brenier [2]. Then eq. (20) can be similarly re- re-written as:
d _ dt
F()) = (-v (VU(()(3))((E))
(veee)))-v-v(iee (v((())a(())),,
-v.
(23)
So the relation arises as:
VF VF= -F=-V. (VV((((()))))
(24)
Since have V = -d (z),, the above equation can be re-written
d (=) -v v. (V))()(())))(())
(25)
The above derivations can be alternatively made by the poineering JKO scheme [34]. This explicitly defines the relation between evolution PDEs of pt(z) and the internal energy For our method, we use the gradient of our scalar energy field Vn(z, L) to learn the velocity field which is given by U' ('(p))))). Interestingly, driven by certain specific velocity fields Vu(z, (), the evolution of p(z, I) would become some special PDEs. Here we discuss some possibilities:
15

---

Heat Equations. If we consider the energy function U as the weighted entropy:
(()( - (og(p:(z))
(26)
We would have exactly the heat equation:
d
d
dt Pt(2)
dz (z2P((2)
(27)
Injecting the above equation back into the continuity equation leads the velocity field uz(2)
d d -- (1a(2))((z)) d
dz2P((z)
dt
(28)
ve(z) Vpi(z) log(pt(z))
Pe(z)
When our Vu(z,t) learns the velocity field -V Vlog(p(())), the evolution of p(z,t) would become heat equations. Fokker Planck Equations. For the energy function defined as:
((p((z) - pr(z) + (=z))oo((((())
(29)
we would have the Fokker-Planck equation as
d
di [VAp(())]
d
d [p<[2)]] =0, dz?
t px(())
(30)
dz
The velocity field can be similarly derived
(x(2) = VA- - V Vlo((+(z))
(31)
For the velocity field VA V (p the movement of p(z, t) is the Fokker Planck equation.
Porous Medium Equations. If we define the energy function
U(p((2)) = m - - (z)
(32)
Then we would have the porous medium equation where >> 1 and the velocity field:
Pc(z) dz Pi ( (z) = O, ve(2) mm mp
(33)
When the Vu(z,t) learns the velocity - -mp - the trajectory of p(z,t) t) becomes the porous medium equations.
B PDE-regularized Latent Space Learning
Our framework can be extended to incorporate the PDE dynamics as part of the training procedure to encourage a more structured representation. To validate the effectiveness, we incorporate a PDE loss such that expect any path in the latent space follow specific dynamics (wave equation in our experiment). In addition to the initial setup outlined in Appendix D.5, we further fine-tune the VAE model by applying PDE regularization loss term, defined as
L = LV + L, + Csb
that includes the latent space learning can also be adapted to other generative models by replacing Lv We further optimize the model and energy network for 10 epochs across the full training dataset using an AdamW optimizer with a le-4 learning rate and cosine learning rate scheduler, without warm-up period. All other training parameters remain consistent with those initially described in Appendix D.5.
objective C, and boundary condition Lg. This PDE-regularized
During fine-tuning, the VAE loss drops from 0.2187 to 0.06288. The results of QED optimization surpass those of two other unsupervised methods, indicating potential areas for future research and refinement.
16

---

Table 5: Single-objective Maximization with PDE-regularized Latent Space Learning The results are the same as of Table I, only unsupervised results are presented for fair comparison. (UNSUP denotes unsupervised scenarios). K represents the number of energy networks trained for unsupervised training with disentanglement regularization.
METHOD
PLOOP T
QED 1
1ST
2ND
3RD
1ST
2ND
3RD
RANDOM
N/A
3.52
3.43
3.37
0.940
0.933
0.932
WAVE (UNSUP)
5.30
5.22
5.14
0.905
0.902
0.978
HJ (UNSUP)
4.26
4.10
4.07
0.930
0.928
0.927
WAVE (UNSUP FT)
3,71
3.58
3.46
0.936
0.933
0.933
C Extended Related Work
C.1 Machine Learning for Molecule Generation
Molecules are highly discrete objects and two branches of methods are thus developed to design search new molecules [8]. 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 line of work developing deep generative models from variational auto-encoders ( (VAE) [18, 31J, generative adversarial networks (GAN) [21, 6], normalizing flows (NF) [43, 67] and more recently diffusion models [24, 58, 33]. However, respect the combinatorial nature of molecules, another line of work leverage combinatorial optimization to search new molecules including genetic algorithm (GA) [30], Monte Carlo tree search (MCTS) [63], reinforcement learning (RL) [64], but often with sophisticated optimization objectives beyond simple valid molecules.
C.2 Goal-oriented Molecule Generation
In addition to simply generating valid molecules, a more realistic application is to generate molecules with desired properties [8]. 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 now-dimensional and smooth latent space thus more friendly for these optimization methods [20]. For methods that do not explicitly reduce the dimensionality of data such as diffusion models, Schneuing et al. [49] propose evolutionary process to iteratively optimize the generated molecules. As is observed that the learned latent space exhibits explicit structure [18], Du et al. [11] leverage such property to learn linear classifier to find the latent direction optimize the property of given molecules. In opposition to deep generative models, combinatorial optimization methods often inherently associated with optimization, c.g. reward function in RL, selection criteria in GA, etc [14,41].
C.3 Image Editing in the Latent Space
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 [17, 29, 59, 22, 72, 45, 50, 52, 54, 53, 55, 56). Here we highlight some representative supervised and unsupervised approaches. Supervised methods usually require pixel-wise annotations. InterfaceGAN [51] leverages face image pairs of different attributes to interpret disentangled latent representations of GANs. Jahanian et al. [29] explores linear and non-linear walks in the latent space under the guidance of user-specified transformation. Compared to supervised methods, unsupervised ones mainly focus on discovering meaningful interpretable directions in the latent space through extra regularization. Voynov and Babenko [59] proposes to jointly learn a set of orthogonal directions and classifier to learn the distinct interpretable directions. SeFa [50] and HouseholderGAN [55] propose to use the eigenvectors of the (orthogonal) projection matrices as interpretable directions to traverse the latent space. More relevantly. Song et al. [54] proposes to use wave-like potential flows to model the spatiotemporal dynamies in the latent spaces of different generative models.
17

---

D Experiments Details
D.1 Baselines
We compare with the following baselines:
we take 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:
Random 1D: we take unit vector where only I randomly selected dimension either or -1 as the linear direction.
ChemSpace [11]: separation boundary of the training dataset in latent space w.r.t. the desired property classified by an Support vector machine (SVM). Then we take the normal vector corresponding to the positive separation as the manipulation direction of control.
Gradient Flow (LIMO) [13]: a VAE-based generative model that encodes the input molecules into SELFIES [36] and auto-regressive on the tokenized molecule. LIMO uses Adam optimizer to reverse optimize on the input latent vector whereas Gradient Flow is equivalent using an SGD optimizer for the same purpose.
Evolutionary Algorithm (EA): EA is general framework that leverages random mutation and crossover to select better samples iteratively to search good solutions. In our scenario, we realize an EA-based baseline by perturbing the vector by the directions given by random, ChemSpace and gradient. The pscudocode can be found in Algorithm 3. The results are presented Appendix D.7.
D.2 Training Dataset
ChEMBL [68] a datahase of 2.4M bioactive molecules with drug-like properties, includ- ing features like Natural Product likeness score and annotations for chemical probes and bioactivity measurements.
MOSES [46] is benchmarking dataset derived from the ZINC [27] Clean Leads collection, containing 2M filtered molecules with specific physicochemical properties, organized into training, test, and unique scaffold test sets facilitate the evaluation of model performance on novel molecular scaffolds.
ZINC250K subset of the ZINC database containing ~550,000 commercially available compounds for virtual screening.
D.3 Molecule Properties
We report the following metrics for our experiments:
Penalized logP/plogP: Estimated octanol-water partition coefficient penalized by synthetic accessibility (SA) score and the number of atoms the longest ring.
Quantitative Estimate of Drug-likeness, metric that evaluates the likelihood of a molecule being a successful drug based on its pharmacophores and physicochemical properties.
QED:
SA: Synthetic Accessibility, a score that predicts the of synthesis of molecule, with lower values indicating casier synthesis.
DRD2 activity: Predicted activity against the D2 dopamine receptor, using machine learning models trained known bioactivity data.
JNK3 activity: Predicted activity against the c-Jun N-terminal kinase 3, important for developing treatments for neurodegenerative diseases.
GSK3B activity: Predicted activity against Glycogen Synthase Kinase beta, which plays crucial role in various cellular processes including metabolism and neuronal cell development.
ESR1 docking score: Simulation- based score representing the binding affinity of molecule to Estrogen Receptor 1, relevant in the context of breast cancer therapies.
18

---

ACAA1 docking score: Simulation- based score representing the binding affinity of molecule to Acetyl-CoA Acyltransferase important for metabolic processes in cells.
D.3.1 Misalignment of normalization schemes for penalized logP
We notice that plogP is commonly reported metric in recent molecule discovery literature but does share the same normalization scheme. Following scores and a ring penalty term were introduced into the calculation of penalized logP as the following
et al. [18, Eq. 1]. the SA
flogr (m) logP(m) SA(m) loig-penalty(m)
Each term of lm)) SA(m), and fing-penalty(m) are normalized have zero mean and unit stan- dard derivation across the training data. However, no sufficient details were included in their paper or their released source code on how the are widely used in various works.
(m) is computed. Specifically, implementations
Penalized by the length of the maximum cycle without normalization hing-penalty(m) is computed as the number of atoms on the longest ring minus 6 in their implementation. Neither (m))(m), SA(m), or (ing-penalty(m) is normalized. MolDQN [71] reported their results in this scheme.
Penalized by the length of the maximum cycle with normalization ring- (ing-penaty((m) is computed same as without normalization. MARS [62]. HierVAE [32], GCPN [65], and ours report plog? using this scheme.
Penalized by number of cycles As described by Jin et al. [31, page footnote 3], ring-penalty (emal is computed as the number of rings in the molecule that has more than atoms. LIMO reports plogP using this metric.
D.4 Training and inference
We detail our training and inference workflows in Alg. and Alg, D.4, respectively.
Algorithm 1 ChemFlow Training
Require: Pre-trained encoder fe. decoder Ig. (optional) classifier (,, timestamps T, # of potential
functions K
1: Initialize (²(-) + MLP for j I,...,K
2: repeat
3: Sampling: 2u = fo (zo); t Categorical()). ~ Categorical(K))
4: for ...... . do
5: 2,41 = 2, V = xi)
6: end for
7: Decode: 0, go (2t), #t+1 = go (2++1)
8: if unsupervised then
9: Classification: i: L, (Ii; x(+1)
10:
Loss: L L, Ls + L + Ls
else
12:
Loss: L L. Lp
end if
14:
Back-propagation through the Loss L until Convergence
D.5 Experiments Setup
Pre-trained VAE We follow the VAE architecture from LIMO consisting of 128 dimension em- bedding layer, 1024 latent space size, 3-hidden-layer encoder, and 3-hidden-layer decoder both with 1D batch normalization and non-linear activation functions. The hidden layer sizes are (4096, 2048, 1024) for the encoder and reversely for the decoder. We empirically find that re- placing the ReLU activation function with its newer variant Mish activation function [44] results
19

---

Algorithm 2 ChemFlow Inference Traversal
Require: Pre-trained encoder fe, pre-trained potential function o. (optional) pre-trained proxy function h, timestamps T, step size a, LD strenth 3
1: Sampling: 20 fo(o))
2: for = 1,..., do
3: if Langevin Dynamics then
4: 2, - 1 a (=,n(===1))) + +3 (0,11)
5: else if Gradient Flow then
6:
2: 2t-1 Vi (2
7:
else
8:
2: 21-1 a aV -1-1)
9: end if
end for
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 Adam [42] optimizer, 0.001 initial learning rate, and 1,024 training batch size. To better prevent the model from being stacked at sub-optimal local minimum, a cosine learning rate scheduler with Ic-6 minimum learning rate with periodic restart is applied. The VAE trained for 150 epochs with restarts on 90% of the training data and validated with the rest 10% data. The checkpoint corresponding to the epoch with the lowest validation loss is selected. Training 150 epochs takes ~ hours on a single RTX 3090 desktop.
Surrogate Predictor The performance of the surrogate predictor is erucial to the proposed latent traversal framework. To handle chemical properties different magnitudes, we normalize all chemical properties in the training data have mean and unit variance. Then activation-norm MLP with residual connections as the surrogate predictor. The predictor contains residual blocks of size 1024 and the output dimension is 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. The predictor trained for 20 epochs on 100,000 randomly generated samples and validated with 10,000 unseen data with SGD optimizer, 0.001 learning rate, and batch size 1000. The epoch with the best validation loss selected. Training each predictor takes less than minute.
a pre-
Energy Network We use an MLP structure to parameterize the energy function (the spatial derivative gives the velocity). The time input is embedded with sinusoidal positional embedding followed by linear layer. The special input x is encoded with a linear layer and ReL.U activation function. The training of the network uses 9,000 random data and 1,000 unseen data for validation. For unsupervised settings, 10 disentangled potential energy functions are trained for 310 epochs with batch size of 100. The epoch with the best validation loss is selected. Training an energy network with 10 disentangled potential energy functions takes 400 minutes.
Reproducibility All the experiments including baselines conducted one RTX 3090 GPU and one Nvidia A100 GPU, All docking scores are computed on one RTX 2080Ti GPU, The code implementation is available at http://ghhhhccom/arveee44///ccn
D.6 Evaluation Metrics
Success Rate The success rate is used as the evaluation metric for the molecule manipulation task. It first randomly generates molecules and traverses each of them in the latent space for l steps. The success rate calculated as the percentage k-step trajectories that are successful. In our case, we generate 1000 molecules and traverse for 10 steps. The manipulation is successful if the local change molecular structure is smooth and molecular property increased. Specifically, we showed two success rates: the strict success rate and the relaxed success rate.
20

---

For the strict success rate, manipulation is a success if the molecular property is monotonically increasing, molecular similarity with respect the previous step monotonically decreasing, and molecules are diversity the manipulation trajectory. These constraints are formulated as follows:
Csp(a, k. 1 st..P.)))(z,) P( S 0], Css(x, K) 1 E [k], s.t.S (Zi11( S(z,, - 0),
CsD(æ, k) 1[|24 Vi [a]] 2],
(34)
SSR
Σ ØEX
[Csr(z, k P) (ss(x,K k 8) Csp(z,))
X
where {ri) one k-step manipulation trajectory, X contains n manipulation trajectories, C' is the constraint, P is the property evaluation function, S is the structure similarity function (Taminoto similarity over Morgan fingerprints). The CSp constraints that the property of molecules must monotonically increase. The C'ss constraints that the structure similar regard the starting molecule must monotonically decrease. The Csp constraint that the molecules must at least change twice during the manipulation. SSR calculates the percentage of trajectories that satisfy all success constraints.
The relaxed success rate relaxes some constraints by adding tolerance interval. It formulated as follows:
Csp k, 1 [k], s.t.P((())) P(x++1) $ c),
Css(x, S) = 1[vi E [k], s.t.S( -((r, x1)= 1 71.
(2D(æ, k) 1[|2 Vi E [+]| 2].
(35)
Σ SSR = X ØEX
[Csr(æ, k. P) A Css(r, k 5) A CsD(2,
The relaxed success rate does not require monotonic increase of molecular property but sets tolerance threshold e. This tolerance threshold e defined as 5% the range property in the training dataset. It also does not require a monotonically decrease of structure similarity with tolerance threshold 7 of 0.I.
D.7 Additional Baselines
Evolutionary Algorithm. We present the Evolutionary Algorithm-based (EA) approach optimize molecules in the latent space as an additional baseline for the experiment. The pseudocode is provided as Algorithm 3. Table shows the performance of the evolutionary algorithm- based approach on unconstrained optimization. For fair comparison, all methods in Table have the same number of Oracle calls, The result shows that our methods outperform all EA approaches.
Algorithm 3 Evolutionary Algorithm Augmented Optimization
1: Input: samples select i per iteration, pre-trained ChemSpace/Graddien direction step size o, pre-trained surrogate model h. decoder P
2: Randomly sample n latent vectors { 1n in latent space
3: for each iteration
to T do
4: Evaluate sampled latent vector scores: s² h(z)) for i 1,...,
5: Select top-k scored latent vectors: [=
6: for each selected vector = to k do
7: Update z++1 Zip..j
+ where (oO, ), / the evolution direction guided by
8: end for
9: Randomly sample samples around each = to generate new latent vectors {2 Ti=1
10: end for
Decode latent vectors: x! = (x.")) for I....,n
21

---

Table 6: Unconstrained plogP, QED maximization of Evolutionary Algorithm. (SPV denotes supervised scenarios, UNSUP denotes unsupervised scenarios). Random/ChemSpacefGradeeee are the evolution directions of EA. Since EA (Gradient Flow) has converged to single molecule when optimizing pLogP, only one value is reported.
METHOD
PLOGP 1
QED
1ST
2ND
3RD
1ST
2ND
3RD
EA (RANDOM)
2.29
1.64
1.52
0.836
0.801
0.794
EA (CHEMSPACE)
3.79
3.79
3.79
0.933
0.931
0.931
EA (GRADIENT FLOW)
3.53
0.930
0.929
0.929
WAVE (SPV)
4.76
3.78
3.71
0.947
0.934
0.932
WAVE (UNSUP)
5.30
5.22
5.14
0.905
0.902
0.978
HJ (SPV)
4.39
3.70
3.48
0.946
0.941
0.940
HJ (UNSUP)
4.26
4.10
4.07
0.930
0.928
0.927
LD
4.74
3.61
3.55
0.947
0.947
0.942
Unconstrained Molecular Optimization Additional Results In addition to reporting the top 3 scores as presented in Table I, we computed the mean and standard deviation for the top 100 molecules after unconstrained optimization in Table 7. The table shows that methods have overall the best optimization performance. In addition, HJ exhibits better performance on mean and standard deviation than on top showing that minimizing the kinetic energy efficient in pushing the distribution to desired properties.
Table 7: MSE of Unconstrained plogP, QED maximization, and docking score minimization. (SPV denotes supervised scenarios, UNSUP denotes unsupervised scenarios). Each entry in the table follows the format mean # std (median). Boldface highlights the highest-performing generation for each property within each rank.
METHOD
PLOGP1 1
QED f
ESRI DOCKING
ACAAI DOCKING
RANDOM
2.345 0.386 (2.259)
0.903 + 0.014 (0.902)
-9.127 + 360 (-9.015)
-8.454 + 0.316 (-8.390)
GRADIENT FLOw
2.664 = 0.382 (2,537)
0.910 # 0.012 (0. (0.908)
-9.452 = 0.338 (-9,365)
-8.735 # 0.337 (-8.650)
CHEMSPACE
2.580 0.406 (2.446)
0.907 1 0.014 (0.906)
-9.523 0.409 (-9.395)
-8.749 0.356 (-8.640)
WAVE (SPV)
2.536 0.439 (2.388)
0.015 (0.898)
-9.630
0.399 (-9.525)
8.764 0.344 (-8.650)
WAVE (UNSUP)
1.736 0.401 (1.610)
0.845 1 0.014 (0.840)
-9.074 : 0.329 (-9.000)
-8,813 0.265 (-8.745)
HJ (SPV)
2,482 # 0.397 (2,382)
0.899 # 0.017 (0.894)
-9.544 0.322 (-9. 460)
-8.792 # 0.332 (-8.675)
(UNSUP)
0.254 (3.377)
(0.909)
9.112 0.321 (-9.090)
0.243 (-8.630)
LD
2.463 0,388 (2.399)
0.905 + (0.903)
-9.400 + 0.360 (-9.300)
-8.709 0.372 (-8.585)
D.8
More Experiment Results
We conduct more experiments to analyze the performance of the proposed methods systematically. They are referred and discussed the main paper.
Pearson correlation score for unsupervised diversity guidance with disentanglement regular- ization. Tables 8 and 9 present the Pearson correlation scores for the trained energy networks of wave equations and Hamilton-Jacee properties other than synthetic accessibility (SA), we select the network with the highest correlation score to maximize these properties. Conversely, for SA, the network with the lowest correlation score (most negative score) is chosen.
equations with disentanglement regularization, respectively, For
Distribution shift and convergence for plogP optimization. Figure 4 illustrates the distribution shift in plogP optimization, complementing the analysis in Figure 3. Similar to the findings discussed in Section 4.2, both unsupervised methods encounter out-of-distribution (OOD) issues after 400 steps, consistent with those observed with ChemSpace. The Random ID method does not achieve the expected distribution shift, as manipulates only one dimension of the latent vector. Figure depicts the convergence trends of each method's improvements. Consistent with the predictions in Proposition 3. Langevin Dynamics demonstrates the fastest and most effective convergence among all methods.
22

---

Table 8: Pearson Correlation of trained Wave PDE Energy Network. The average Pearson correlation between the sequence of real properties and sequence of time steps along the manipulation trajectory following learned potential function e (t, 2) using wave equations.
PLOGP
SA
QED
DRD2
JNK3
JSK3B
0
0.019
0.003
0.016
0.029
0.015
0.051
1
0.160
-0.451
0.275
-0.074
-0.153
-0.272
0,035
-0.003
0.011
0.002
-0.006
0.017
0.072
-0.096
0.065
0.011
-0.017
-0,028
0.042
0.003
0.039
-0.010
0.025
0.018
5
-0.036
-0.022
0.150
0.009
-0.017
0.008
6
0.032
-0.045
0.006
0.002
-0.011
0.002
0.023
-0.023
0.054
-0.002
-0.013
-0.017
0.075
-0.085
0.064
-0.040
-0.007
-0.054
9
0.013
0.020
-0.011
0.031
0.005
0.014
INDEX
0
Table 9: Pearson Correlation of trained Hamilton-Jacobia PDE Energy Network. The average Pearson correlation between the sequence of real properties and sequence of time steps along the manipulation trajectory following learned potential function @ (t, using Hamilton-Jacobi equations.
PLOGP
SA
QED
DRD2
JNK3
GSK3B
0.345
-0.453
0.141
-0.210
-0.127
-0.350
0.257
-0.289
0.057
-0.154
-0.121
-0.276
2
0.203
-0.284
0.050
0.020
-0.044
-0.164
-0.029
0.041
0.034
-0.001
0.001
0.002
0.304
-0.343
0.066
-0.225
-0.179
-0.336
0.008
-0.036
0.029
-0.009
-0.021
0.015
6
0.316
-0.370
0.173
-0.208
-0.163
-0.309
7
0.305
-0.429
0.046
-0.291
-0.191
-0.352
8
0.003
-0.011
0.016
-0.021
-0.035
0.026
0.311
-0.386
0.209
-0.222
-0.209
-0.342
INDEX
0
QED optimization and distribution shift. To further explore optimization tasks additional properties, also attempted to enhance the QED of molecules. Table 10 presents the results of efforts to maximize the QED of 800 molecules from the ZIC250K dataset that initially had the lowest QED scores. Despite encountering OoD issues with ChemSpace and two unsupervised methods, as illustrated in Figure 6, the performance of Langevin Dynamics surpasses other methods across various similarity levels, which is consistent with the result from Table 2.
Table 10: Similarity-constrainee constrained QED maximization. For each method with minimum similarity constraint 5, the results in reported format mean standard derivation (success rate %) of absolute improvement, where the mean and standard derivation are calculated among molecules that satisfy the similarity constraint.
Method
6=0 0
5 = 0.2
5= 0.4
8 0.6
Random
0.16 (98.0) 0.13 (40.1)
0.19 (78.4)
0.11 (54.4)
0.08 + 0.08 (29.9)
Random ID
0.1200.1((88.).
0.09 (29.2)
0.07 # 0.06 (18.0)
Gradicnt Flow
0.48 +.13 (99.))
0..16 (84.9)
0.13 (53.8)
0.10 + 0.10 (24.9)
ChemSpace
0.47 + 0.13 (99.8)
0.188 (90.1)
0.18 # (62.9)
0.12 0.12 (33.5)
Wave (spv)
0.38 = 0.16 (97.))
0.23 (84.9)
0.13 (62.6)
0.08 # 0.08 (35.0)
Wave (unsup)
0.54 = 0.18 (99.1)
0..10
0.05 0.06 (27.4)
0.03 0.04 (14.0)
HJ (spv)
0.16 (74.2)
0.17 (69.4)
0.12 (57.0)
0.08
(35.0)
HI (unsup)
0.52 + 0.19 (98.5)
0.11 (60.4)
0.05 0.05 (28.6)
0.03 + 0.03 (13.4)
LD
0.53 + 0.12 (99.8)
0,31 0.17 (96.2)
0.16 (77.8)
0.10 0.09
23

---

Random
Razdom 1D
Gredieat Flow
0.0014
0.0014
0. 0.0020
0.0012
0.0012
0.0015
0.0010
0.0010
Deastty
0.0008
0.0008
0.0010
0.0006
0.0006
0.0004
0.0004
0.0005
0.0002
0.0002
0.0000
0.0000
0.0000
ChemSpace
Wave (spv)
Wave (unsup)
0.0014
0.0200
0.005
0.0012
0.0175
0.0010
0.0150
100
0.0125
200
0.0008
300
0.0100
0.0006
500
0.0075
600
0.0004
0.0050
0.001
800
0.0002
0.0025
900
999
0.000
0.0000
0.0000
HJ ogn. Espr)
HJ ogm. (unsup)
Langevin Dynamtes
0.0012
0.025
0.0014
0.0010
0.0012
0.020
0.0010
0.0008
0.015
0.0008
0.00 00086
0.0006
0.010
0.0004
0.0004
0.0002
0.005
0.0000
0.000
0.0000
-10
-5
0
-10
-5
-10
-5
5
plogp
plogp
plogp
Figure 4: Distribution shift for plogP optimization
E
Latent Space Visualization and Analysis
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 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 7, where the norm of the molecule projected to the latent space concentrate around 32 which is around Vd such that the latent dimension d is 1024. We also visualize in Figure 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 these two evidences, it is not surprising that a random latent vector taking random direction will change the molccular property smoothly and monotonically. In addition, we further plot when we traverse along 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 9. 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.
Additionally, we visualize the traversal trajectory for each different property using both supervised and unsupervised wave flow in Figure 10. The plot shows that almost all trajectories grow towards unique direction in the t-SNE plot, which implies the disentanglement of learned directions and, thus, molecular properties. In addition, the figures display sinusoidal wave-shape trajectories, indicating the
24

---

Imgrovmmmnt Imprrveeenet plagg
Random
Random 1D
Gradlent Flow
agn.
HJ eqn. (pc)
Langesth Dynomics
0.4
0.2
0.0
200
400
600
800
1000
200
400
000
800
1000
Stegs
Figure 5: Optimization Convergence Langevin Dynamics shows faster convergence and achieves
greater improvement in plogP.
flow follows wave-like dynamics. the unsupervised t-SNE plot, the trajectories of some properties overlap, such as plogP and SA. It is because some properties correlate with the same disentangled direction, so their traversal follows the same direction, thus the same trajectories.
Qualitative Evaluations
In addition quantitative evaluations, we demonstrate some qualitative evaluations in this section. We showcase three manipulation trajectories in Figures 11 to 13. Each of these paths is 6-step manipulation for different molecular properties using different flows. Specifically, the supervised wave flow and gradient flow improves the molecular property by the conversion of large heterocyclic rings for better synthesizability. We also show three optimization trajectories in Figures 14 to 16. From left to right, each molecule is snapshot selected from trajectory of 1000-step optimization. It is notable that the supervised Hamilton-Jacobi flow optimizes the property by reducing the number of nitrogen atoms. This leads to more chemically stable molecule, whereas the original molecule, with a large number of nitrogen atoms, is unstable and potentially explosive. The supervised wave flow optimizes the molecular property by simplifying the poly-cyclic molecule which enhances its synthesizability and overall stability.
that
25

---

Random
Random
Gradient
0.07
0.05
0.06
0.06
0.05
0.05
Dessiry
0.04
0.04
0.03
0.02
0.02
0.01
ChemSpace
Wave (spr)
Wore egn. (unsup)
0.30
0.06
0.5
0.25
100
0.4
0.20
300
0.3
0.15
0.10
0.2
600
800
0.05
900
0.00
0.00
0.0
HJ egn. (spv)
HJ ego. (unsup)
Langevin Dynamics
0.06
0.20
0.06
0.15
0.04
0.10
0.05
1.00
0.75 QED
1.00
QED
QED
Figure 6: Distribution shift for QED optimization
15000
Count
1 10000
5000
30
32
34
Latent Space Norm
Figure 7: Latent Vector Norm. Distribution of the norm of the latent vectors projected from the training dataset onto the learned latent space.
26

---

DRD2
05:38
0030
-32
3.03
3.00
25
43
10
Letert Embwed ng Norm
Figure 8: Embedding Norm against Property Value of each path. Norm and property value of molecules along the direction of latent traversal with random direction. The middle curve shows the mean property value and latent embedding norm for all paths. The shaded area is the standard deviation of property value.
0.10
20
33
Figure 9: Embedding Norm against Property Value of each Molecule. Scatter plot of norm and property value of individual molecules in the training set encoded in the latent space.
Wave Supervised Trajectory for Each Properties
Wave Unsupervised Trajectory for Each Properties
Proparties plogo
Proparties
plggo. cep.ooed.oo
drd2
TS-M Dimenztoe
drd2
grldb
TE Dimandiur
25
jk3
Pk3
0-
25-
-50
-100
-75
-30
-25
75
-50
-25
100
TSNE Dimonsiion
Figure 10: t-SNE Visualization of Optimization Trajectories. Optimization Trajectories following supervised and unsupervised wave flow visualized using t-SNE.
plogp -3pp-33.550 1.000
plogp 36.654
plagp -16.554 0.583
plage -16.505
plogp-1.00 -1006 0.042
plegp-1.015
plogp -1.016
Figure 11: Molecule Manipulation Trajectory Molecule manipulation by chemspace plogP.
plagp .29 985 sim 1.000
plagp 112.162 sim 0.460
plage -12.075 sim 0.351
plagp -11.878 sim 0.283
plagp -2.803 sim 0.089
plago -2.003 sim 0.089
Figure 12: Molecule Manipulation Trajectory Molecule manipulation by gradient flow plogP.
g5k3b 0.050
gbk3b 1.100 sim
gik3bb sim 0.287
gik3ab gok3b sim
gb
1.000
sim 0.210
Figure 13: Molecule Manipulation Trajectory Molecule manipulation by supervised wave flow GSK3B.
27

---

plogp -12.77 736 sim 1.000
plogp -22 126 sim 0.646
plagp 4.333 sim 0.246
plagp -4.310 sim 0.174
plagp -3.904
plogp 33904 sim 0.174
sim 0.174
Figure 14: Molecule Optimization Trajectory Molecule optimization by random direction on plogP.
qed 0.215
qed 0.678
qed 0.678
oed 0.699
qed 0.726
qed 0.726 sim 0.198
sim 1.000
sim 0.268
sim 0.268
sim 0.226
sim 0.198
Figure 15: Molecule Optimization Trajectory Molecule optimization by supervised wave flow on QED.
qed
Figure 16: Molecule Optimization Trajectory Molecule optimization by supervised Hamilton- Jacobi flow QED.
28

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide a short (1-2 sentence) justificatin right after your answer (even for NA).
integral part of your paper submission. They visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The checklist answers
The reviewers of your paper will be asked to use the checklist one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] " "rrovided a proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" " or "we were unable to find the license for the dataset we used"). "[No] " or "[NA] is not grounds for rejection. While the questions are phrased in binary way, acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification elaborate. All supporting evidence can appear either in the main paper the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
general, answering
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The claims in the abstract and introduction reflect the scope and contributions of the paper, including the development of the ChemFlow framework for navigating chemical latent spaces and its validation across multiple molecule manipulation and optimization tasks.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
29

---

Answer: [Yes]
Justification: The paper discusses limitations such as scaling with billions of available molecules in large databases and generalizing beyond small molecules.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
The paper does not invent any new theory but replies on existing theoretical results. The details including assumptions and derivations (referenced in Section and Appendix A), discussed.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
30

---

All experimental setups are clearly described, including data sources, model parameters, and evaluation criteria. The source code is released at results presented.
Justification:
This should allow for the reproducibility of the
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the data used in our experiments are open-sourced and we include our codes and instructions to run with README.nd in
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niil co/
pub lublic//e quides/A (CoeeSuum ss asionoolicy for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips cc/ wttpc//clcccgcccceceeeeeeeopppce www.a/eieeeeeeeeeeeeeeeeoeaeeeee Code wwwssssooooiiiooosoooooooiisoioi icy for more details.
31

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper specifies all necessary experimental details, such as data splits, hyperparameters, and types of optimizers, which essential for understanding and repro- ducing experimental results, referred in Appendix D.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Statistical significance is discussed with appropriate measures (std) and success rate provided for key results, ensuring the reliability of the findings.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
32

---

Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: The paper details the computational resources used, including the type of GPUs and the estimated compute time for experiments in Appendix D.5, which aids in the replication and understanding of the resource requirements.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: //euriuupppp cc cccppccccicccccccccccocccccccccc wwtpwccpcpppcccpccccccteccccccct wtt.eieeeeeiiiiieeeeeeeee.eeeeee
Answer: [Yes]
Justification: The research adheres to the NeurlPS Code of Ethics, and ethical considerations, especially concerning dual-use risks and responsible Al practices, are discussed Section 5.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e-g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Both positive and negative societal impacts are considered, with discussions how the ChemFlow framework could impact drug discovery and potential misuse scenarios.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
33

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks. mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
The paper does not introduce new models or datasets with high risks misuse at this moment; thus, specific safeguards are not necessary. However, general practices for responsible Al are discussed.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The paper correctly credits all used assets, including datasets and codes, with appropriate references and discusses the terms of use and licenses where applicable.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datases has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not introduce new datasets or tools that require additional documentation licensing information.
34

---

Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the datasct/codc/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The research does not involve human subjects or crowdsourcing, making this question not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subject research is conducted; hence IRB approval is not required.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
35


Reviews:
[{"review_id": "CliNyLn2WF", "reviewer": "Reviewer_kzEx", "summary": "This paper presents a novel gradient flow-based method to traverse the latent space of molecular generation models, known as ChemFlow. The authors instantiate their framework with a number of different flows inspired by dynamical systems. They also investigate the use of supervised and unsupervised guidance for the flow methods with the goal of optimising molecular properties of interest. The authors perform many experiments focused on molecular optimisation in order to evaluate their method. Specifically, they investigate unconstrained optimisation, similarity-constrained optimisation and multi-property optimisation.", "questions": "- For the unsupervised guidance, when you match flows with properties, what the correlations are computed between? Does this require you to have an existing dataset of molecule-property pairs?\n\n- For the unsupervised cases did you experiment with different values of k? Which values of k were used? It seems to me that k might need to be very large in general in order to find a flow which matches with an arbitrary property.\n\nOther suggestions: I assume equation 8 should have $\\phi^k$ instead of $\\phi$? Additionally, the ordering of t and z in $\\phi(.,.)$ or $\\phi^k(.,.)$ is inconsistent - compare figure 1, the caption for figure 1 and equation 11 with all the other places were $\\phi$ is used.", "limitations": "- While the authors evaluate their approach on optimising two properties simultaneously it's unclear how well this would work with a larger number, since their approach relies on simply adding the guidance terms together.\n\n- In this study the authors used a VAE with a fixed input size. It's unclear how easy it would be to apply a similar approach to a model with an arbitrary number of elements in an input or latent sequence, such as chemical language models or many diffusion-based models, which are much more commonly used in practice for molecular generation.\n\n- As far as I can tell the method as it stands requires a latent space optimisation to be done for every sample. This could lead to much longer sampling times than other methods such as RL fine-tuning which allows samples to be generated as normal but from an optimised model.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The authors present a novel formulation for exploring chemical latent spaces and optimising molecular properties which is an important and relevant task within the pharmaceutical and molecular design domains. They present a few different versions of their method, including an implementation which can use a surrogate property prediction model to guide the flow to optimised chemical space, as well as an unsupervised implementation which aims to maximise structural changes to the molecule.\n\n- The authors perform an extensive evaluation on molecule optimisation related tasks, comparing a number of different instantiations of their framework. They also benchmark against a previously introduced model for latent space traversal and a random traversal strategy.\n\n- ChemFlow methods show very promising results in comparison to baselines, especially when only looking at top performing molecules or applying similarity-constrained optimisation.\n\n- The authors also present a useful analysis of the latent changes under the random traversal strategy and an explanation for why random traversal can work reasonably well.\n\n- The paper is mostly very well written and, the evaluations in particular, are very clearly presented and easy to follow.", "weaknesses": "- I find the methodology section on its own quite unclear since it's not clear how to actually use the objective functions that are outlined. Appendix sections D4 and D5 are helpful but ideally it would be possible to follow the main text on its own. Particularly, I think the text would benefit from showing the full loss function in the methodology and including a short outline of the training and sampling procedures and referring to the appendix.\n\n- The baselines for some tasks are a bit weak, particularly for the unconstrained molecular optimisation. For this task techniques like evolutionary algorithms and reinforcement learning fine-tuning have been proposed and widely used before. It would be very useful to see a comparison of ChemFlow with methods like these, as well as an evaluation of the training and sampling time for each.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: The baselines for some tasks are a bit weak, particularly for the unconstrained molecular optimisation. It would be very useful to see a comparison of ChemFlow with methods like EA and RL fine-tuning.**\n\nA: Thanks for the suggestion. As our method focuses on the latent space of deep generative models, we compared mostly with methods for molecular optimization and traversal in a similar setup. However, we agree with the reviewer that the paper would benefit from additional baselines. We add an evolutionary algorithm-based (EA) approach to optimize molecules in the latent space. The pseudocode is provided as Algorithm 1 in the uploaded PDF in the general response. For a fair comparison, all methods in Table 1 in the uploaded PDF have the same number of oracle calls. The results in Table 1 show that our methods outperform all EA approaches. \n\nIt is possible to use reinforcement learning to guide the search in the latent space of molecular generative models but the main reason to use the latent space of generative model is to avoid the discrete nature of molecular structures and instead conduct optimization over a continuous space. We believe it will be nontrivial to propose a new reinforcement learning algorithm and thus leave it as a future study.\n\n**C2: For the unsupervised guidance, when you match flows with properties, what the correlations are computed between? Does this require you to have an existing dataset of molecule-property pairs?**\n\nA: We compute “the correlation between the property and a natural sequence (from 1 to time step t) along the optimization trajectory” (line 177). Specifically, we compute the following measurement:\n\n$$Spearman([P(m_1), P(m_2), …, P(m_t)], [1,2,...,t])$$\n\nwhere $P(\\cdot)$ is the function to measure the real chemical property of a given molecule $m_t$ at time $t$.\n\nEven though we do not require access to a dataset of molecule-property pairs, we do rely on minimal supervision to match the direction to the list of properties it may control. In reality, we do this by scoring the molecules using oracle functions.\n\n**C3: For the unsupervised cases did you experiment with different values of k? Which values of k were used?**\n\nWe observe as long as the value of k is larger than the actual number of properties, the results will not be impacted much. We agree having an approximation of the number of properties is a prior knowledge of determining the hyperparameters, but a good practice is always to start with a relatively large k and select the active ones after training.\n\n**C4: While the authors evaluate their approach on optimising two properties simultaneously it's unclear how well this would work with a larger number, since their approach relies on simply adding the guidance terms together.**\n\nA: Thanks for the constructive advice. We get motivated by the disentanglement literature so we mainly focus on optimizing individual properties. Assuming each property is an energy $\\phi^k$ so the stationary distribution is a Boltzmann distribution $p^k(x) \\sim exp(\\phi^k)$, then the summation of the guidance terms (i.e. energies) corresponds to sampling from the product distribution $\\pi \\sim exp(\\sum \\phi^k) = \\prod exp(\\phi^k)$. Even though this implicitly assumes independence among different properties, it is commonly used (also known as product of expert [1]) in machine learning, e.g. energy-based models [2]. We will leave how to leverage the correlation between different objectives as future work. \n\n> [1] Hinton, G.E., 2002. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8), pp.1771-1800.\n>\n> [2] Du, Y. and Mordatch, I., 2019. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32.\n\n**C5: In this study the authors used a VAE with a fixed input size. It's unclear how easy it would be to apply a similar approach to a model with an arbitrary number of elements in an input or latent sequence, such as chemical language models or many diffusion-based models.**\n\nA: Thanks for the question. The fixed-length VAE model supports any input with a length less than its maximum limit. Indeed it can be less efficient for varied-length input such as text, however, it has still been widely used in languages and other sequence data. Moreover, despite that we validate the proposed method on a specific problem, molecular design, and select a fixed-length VAE model as the generative model, the method is not limited to fixed-length input.\n\nAs long as the model architecture has a well-defined latent space, e.g. [3] uses the bottleneck of U-Net as the latent space for diffusion models and [4] similarly in the attention head of large language models, it is possible to adopt our method in other networks to discover meaningful properties.\n\n> [3] Kwon, M., Jeong, J. and Uh, Y., Diffusion Models Already Have A Semantic Latent Space. In The Eleventh International Conference on Learning Representations. \n>\n> [4] Li, K., Patel, O., Viégas, F., Pfister, H. and Wattenberg, M., 2024. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36.\n\n\n**C6: The method as it stands requires a latent space optimisation to be done for every sample. This could lead to much longer sampling times than other methods such as RL fine-tuning which allows samples to be generated as normal but from an optimised model.**\n\nWe focus on pre-trained generative models, where multiple objectives can be composed or removed at any time, while we admit this necessitates the extra sampling time as the model itself does not directly sample from the desired distribution, RL-based fine-tuning would limit the flexibility of the pre-trained model once it has been tuned for a specific task. The extra inference time is also minimal as it only requires a surrogate model evaluation (which is a relatively simple MLP model).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6icCLuKMsm", "reviewer": "Reviewer_piJE", "summary": "The authors propose a new method called ChemFlow, which navigates molecular distributions in chemical space through flow.", "questions": "1. Why were only two specific targets evaluated in the docking score experiment?\n2. I would like to get some intuitive understanding : What problems arise if training a molecular property predictor in latent space and updating it directly based on its gradient? How is this issue typically addressed in the proposed method?", "limitations": "This work is a preliminary study.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The method demonstrates high generality, applicable to various molecular optimization tasks.\n2. Based on the experimental results presented, the method shows significant optimization of molecular properties.", "weaknesses": "1. The description of the experimental section lacks detail, such as which software was used to measure the docking scores?\n2. Table 1 displays properties of several indicators, but showcasing only the top 3 among numerous sampled molecules may lack sufficient \n persuasiveness. It would be better to include a broader distribution, such as the mean and so on.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: The description of the experimental section lacks detail, such as which software was used to measure the docking scores?**\n\nA: Thanks for pointing it out. We will thoroughly revise the experimental section in the revised manuscript to make sure all details are fully explained. Specifically, we used AutoDock [1] to calculate the docking score. We also used RDKit [2] and TDC [3] to calculate other molecular properties and structure similarity. \n\n> [1] Morris, G.M. et al. (2009) ‘AutoDock4 and AutoDockTools4: Automated docking with selective receptor flexibility’, Journal of computational chemistry, 30(16), pp. 2785–2791.\n> \n> [2] RDKit: Open-source cheminformatics. RDKit. Available at: https://www.rdkit.org.\n> \n> [3] Huang, K. et al. (2021) ‘Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development’, arXiv [cs.LG]. Available at: http://arxiv.org/abs/2102.09548.\n\n\n**C2: Why were only two specific targets evaluated in the docking score experiment?**\n\nA: Thanks for the comment. To respect previous literature, we follow the set-up as in [4] that target the binding sites of two human proteins ESR1 and ACAA1. Human estrogen receptor (ESR1) is chosen because it is a well-characterized protein with known disease relevance. Human peroxisomal acetyl-CoA acyl transferase 1 (ACAA1) is chosen to demonstrate the model’s de novo drug design ability as ACAA1 has no known binders. \n\nIn addition, computing docking scores requires significant computational resources, as calculating the docking score for 10,000 molecules generated by a single method takes approximately 20 hours on one GPU.\n\n> [4] Eckmann, P. et al. (2022) ‘LIMO: Latent Inceptionism for Targeted Molecule Generation’, Proceedings of machine learning research, 162, pp. 5777–5792.\n\n**C3: I would like to get some intuitive understanding : What problems arise if training a molecular property predictor in latent space and updating it directly based on its gradient? How is this issue typically addressed in the proposed method?**\n\nA: Thanks for the comments. Training a property predictor and updating the latent vector directly based on its gradient is the exact method used in the previous work LIMO [4] as gradient-based optimization can be viewed as discretization of a gradient flow. In our work, we generalize this method as different types of flows. Gradient-based optimization may suffer from several challenges, such as stuck in local minima and poor convergence, especially in a high-dimensional space with noisy gradient guidance. \n\nIn our proposed method, these issues can be improved using techniques like Langevin dynamics, which introduces diffusion noise into the gradient updates. Intuitively, this approach helps the model escape local minima by injecting stochasticity into the optimization process, thus promoting better exploration of the latent space.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "to6NWlbn34", "reviewer": "Reviewer_BiGQ", "summary": "Designing new functional molecules within the vast chemical space is challenging, which necessitates efficient exploration and understanding of this space. The paper introduces a new framework called ChemFlow, which leverages latent space learned by molecule generative models and navigates it using flows. ChemFlow formulates the problem as a vector field that guides the molecular distribution to regions with desired properties or structure diversity. The paper conducts extensive empirical studies and justifies the effectiveness of the proposed method.", "questions": "- Is it possible to visualize the vector field using a tool like t-SNE or UMAP to provide further insight into the entanglement of the molecular properties?\n- Could you discuss the connection between flow learning and downstream tasks?\n- Could you discuss the extension of ChemFlow to macromolecular tasks like protein?", "limitations": "The paper could provide further insights into the connection between flow learning and downstream tasks.\nAdditionally, the author can consider discussing the extension of ChemFlow to macromolecular tasks like protein, which could make the proposed framework have a broad impact.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The proposed ChemFlow unifies the previous approaches via the vector field, which is novel and effective for learning a latent space with rich nonlinearity information. This can benefit various downstream tasks, including drug-related properties and protein-ligand binding.\n- Extensive experiments have been conducted to provide a good insight into the components of the proposed method. ChemFlow achieves faster empirical convergence and higher success rates, especially using Langevin dynamics.\n- The paper is generally well-written, with clear illustrations and tables.", "weaknesses": "- Despite the novelty of the proposed method, the ChemFlow mainly focuses on small molecules. This might hinder the border impact of the learned latent space for macromolecular tasks like protein.\n- ChemFlow employs multiple approaches to learning different latent flows. However, the experiment's results show that different methods have different specialties, and the paper does not discuss the connection between flow learning and downstream tasks.\n- The paper mentioned the out-of-distribution generation problem in Appendix D.7 and Sec. 4.2. ChemFlow has encountered such a problem in an unsupervised manner. This could hinder the utility of the learned latent space for scenarios with distribution shifts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: ChemFlow employs multiple approaches to learning different latent flows. However, the experiment's results show that different methods have different specialties, and the paper does not discuss the connection between flow learning and downstream tasks.**\n\nA: Thanks for the question. This is indeed a good point. For molecular optimization task, we might have better properties of one flow vs another (e.g. Langevin dynamics vs gradient flow for optimization). However, for the latent traversal or unsupervised setting, we often do not know which flow is better. Nevertheless, we can use this as a prior to perform traversal while simultaneously structuring the latent space. As indicated in Table 5 of the Appendix, we demonstrate that the different flows work as various inductive biases to improve the performance.  \n\n\n**C2: The paper mentioned the out-of-distribution generation problem in Appendix D.7 and Sec. 4.2. ChemFlow has encountered such a problem in an unsupervised manner. This could hinder the utility of the learned latent space for scenarios with distribution shifts.**\n\nA: We appreciate reviewer points out that out-of-distribution generalization hinder the utility of the learned latent space. The out-of-distribution generalization is not a particular problem for our proposed method, but rather a limitation of the underlying generative models and the surrogate model used. We thus leave this for future study.\n\n\n**C3: Is it possible to visualize the vector field using a tool like t-SNE or UMAP to provide further insight into the entanglement of the molecular properties?**\n\nA: Thanks for suggesting visualizing the vector field. We provided the t-SNE visualization of the traversal trajectory for each different property using both supervised and unsupervised wave flow in Figure 1 of the uploaded PDF. The plot shows that almost all trajectories grow towards a unique direction in the t-SNE plot. This implies the disentanglement of learned directions and, thus, molecular properties. In addition, the figures display sinusoidal wave-shape trajectories, indicating the flow is following the wave-like dynamics. \n\nIn the unsupervised t-SNE plot, the trajectories of some properties overlap, such as plogP and sa. It is because some properties correlate with the same disentangled direction, so their traversal follows the same direction, thus the same trajectories.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zsP1gUPhrm", "reviewer": "Reviewer_D6r5", "summary": "The authors built a general latent flow-based framework unifies traversal and optimization in the molecular latent space. The flow is trained by utilizing energy functions so that the vector field aligns with the gradients, with regularization imposed by an auxiliary classifier that tries to differentiate each distinct flow. Under multiple evaluation settings, ChemFlow outperforms or is generally on par with previous SOTA.", "questions": "- In Table 2 and 3, why did ChemFlow perform best under mild similarity constraint, while suboptimal when the similarity threshold gets higher than 0.4? I think the similarity constraint is important, since usually in practical drug design people would expect to develop new therapeutics based on some drugs whose effects are already known, and the newly designed drug molecules are preferred to be similar so as to keep the effect.\n- For Figure 4, why is the predictor deviating so much from ground truth? LogP, QED and SA don't seem very hard to learn as far as I know.\n- Please consider elaborating on why \"the learned variational poster[ior] also follows a Gaussian distribution\" and \"a strong correlation between almost all molecular properties and their latent norms\" would contribute to the observed result that \"a random latent vector taking a random direction will change the molecular property smoothly and monotonically\".", "limitations": "This paper does not include a discussion on efficiency as compared with ChemSpace.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- This paper is generally well-written and easy to follow.\n- It is a novel contribution to formulate the manipulation and optimization of molecules in latent space as learning the vector fields toward optimal distribution.", "weaknesses": "- Typo: L747 \"we first verify if the learned variational poster also follows a Gaussian distribution and we find that it does learn so\", poster -> posterior\n- For Table 1, it risks not fully revealing the optimization ability if only TOP3 results are reported. Please consider adding mean and median values, too.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: Typo: L747 \"we first verify if the learned variational poster also follows a Gaussian distribution and we find that it does learn so\", poster -> posterior.**\n\nA: Thanks for pointing it out. It was a typo and we will fix it in the revised manuscript.\n\n\n**C2: In Table 2 and 3, why did ChemFlow perform best under mild similarity constraint, while suboptimal when the similarity threshold gets higher than 0.4? I think the similarity constraint is important, since usually in practical drug design people would expect to develop new therapeutics based on some drugs whose effects are already known, and the newly designed drug molecules are preferred to be similar so as to keep the effect.**\n\nA: Thank you for the question. We did not explicitly trained our methods for similarity constraint optimization. In the future work, we will explicitly encode more constraints to encourage our method to maintain similarity while optimizing the molecular properties. In addition, even though the absolute improvement of our method is not optimal for similarity threshold gets higher than 0.4, it has a optimal success rate compared to baseline methods. \n\n**C3: For Figure 4, why is the predictor deviating so much from ground truth? LogP, QED and SA don't seem very hard to learn as far as I know.**\n\nA: Thanks for the question. In Figure 4 (original paper), we set up an out-of-distribution scenario where the ground truth/prediction was the ZINC250k dataset (our VAE model was trained on a mixed of MOSES, ChEMBL and ZINC250k dataset where ZINC250k dataset was only a small fraction). However, the training set for the surrogate model was 10k randomly sampled molecular structures from the VAE model. We also only sampled 10k data points from ZINC250k which might not be representative of ZINC250k either. We have uploaded a new Figure 4 (Figure 2 in the uploaded PDF in the general response) showing the comparison on the full ZINC250k dataset (still out-of-distribution but slightly better). \n\nWe have further reported the training/test performance for the surrogate model, see Table 1 and 2 (training MAE and test MAE). To further study the out-of-distribution scenario we set up, we include Figure 3 in the uploaded PDF to show our training/test set for the surrogate model and ZINC250k dataset.\n\n**Table 1** Training Error (In-Distribution)\n\n|       | plogp | qed  | sa    | drd2   | jnk3  | gsk3b |\n|-------|-------|------|-------|--------|-------|-------|\n| **MAE**   | 9.840  | 0.189| 0.956 | 0.006| 0.016| 0.041|\n| **RMSE**  | 12.948| 0.231| 1.205 | 0.011 | 0.021| 0.052|\n\n**Table 2** Test Error (Out-of-Distribution)\n\n|       | plogp | qed  | sa    | drd2   | jnk3  | gsk3b |\n|-------|-------|------|-------|--------|-------|-------|\n| **MAE**   | 9.976 | 0.341| 2.030 | 0.010 | 0.017| 0.039|\n| **RMSE**  | 11.076| 0.366| 2.203 | 0.038 | 0.025| 0.049|\n\n\n\n**C4: Please consider elaborating on why \"the learned variational poster[ior] also follows a Gaussian distribution\" and \"a strong correlation between almost all molecular properties and their latent norms\" would contribute to the observed result that \"a random latent vector taking a random direction will change the molecular property smoothly and monotonically\".**\n\nA: Thanks for the question. As VAE enforces the variational posterior to be a Gaussian distribution and we know that the geometry of high-dimensional Gaussian distribution is spheraical and the mass concentrates on the shell (if it is zero-centered, the norm of the sampled data is around $\\sqrt{d}$ where d is the dimension of the data). In $R^d$, any random direction would eventually take it to the outer shell of the Gaussian ball with larger norms and further increase or decreases the property value. \n\nThus the observation of the correlation between property values and latent norms is an emergent geometry of the learned latent space. Under this geometry, it is reasonable that any random direction could lead to monotonic and smooth change of the property.\n\n\n**C5: This paper does not include a discussion on efficiency as compared with ChemSpace. Similar to the table 2 in chemspace**\n\nA: Thanks for suggesting discussing the efficiency. Below is a table that summarize the efficiency of all baselines and our methods. It is benchmarked by training the model described in the unconstrained optimization task. The inference time is the time for optimizing 100,000 molecules for 1 step using a batch size of 10,000.\n\nEven though our method has a slower training time compared to ChemSpace because learning the flows requires learning a neural network instead of a linear model (e.g. linearSVM), they have similar inference time. This fast inference time ensures that our method is also capable of conducting high-throughput molecule optimization and screening for drug discovery. \n\n\n| Method                   | Training | Inference/Iter (without oracle) |\n|--------------------------|---------------|---------------------------------|\n| **ChemSpace**            | <1 min        | 0.01 s                          |\n| **Gradient Based**       | 7 min         | 0.03 s                                |\n| **Supervised Guidance**  | 20 min        | 0.03 s                                |\n| **Unsupervised Guidance**| 32 min        | 0.03 s                                  |\n| **Langevin Dynamics**    | 7 min         | 0.03 s                                  |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their valuable feedback that helps us improve the manuscript. We appreciate reviewer’s common sentiment that our work is novel, general, applicable, and well written. We are also glad that reviewers compliment our work has extensive experiments showing the effectiveness and significance of proposed methods.


We will first address the points raised by more than one reviewer in this general response and then provide individual responses to each reviewer. 


**1. Broader applications of the proposed approach to other tasks.**

This paper mainly focus on the problem of molecular design and optimization. However, we do not foresee any issues attached to our proposed method to be applied to other tasks. As long as the generative model architecture has a well-defined latent space, e.g. similarly in the attention head of large language models [1], it is possible to adopt our method in other networks to discover meaningful properties.

For example, our framework can be applied to protein design tasks. Previous work uses gradient-based method to optimize protein in the latent space [2]. This is generalized as traversing with gradient flow in our framework. Diffusion has deomonstrated its powerful ability in generating de novo proteins with desired properties [3]. By define a latent space for diffusion models, such as using the bottleneck of U-Net [4], it is also possible to extend our method to diffusion models for de novo protein generation.

> [1] Li, K., Patel, O., Viégas, F., Pfister, H. and Wattenberg, M., 2024. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36.
>
> [2] Castro, E. et al. (2022) ‘Transformer-based protein generation with regularized latent space optimization’, Nature Machine Intelligence, 4(10), pp. 840–851.
>
> [3] Watson, J.L. et al. (2023) ‘De novo design of protein structure and function with RFdiffusion’, Nature, 620(7976), pp. 1089–1100.
>
> [4] Kwon, M., Jeong, J. and Uh, Y., Diffusion Models Already Have A Semantic Latent Space. In The Eleventh International Conference on Learning Representations.


**2. Mean/STD/median of the scores in Table 1.**

In addition to reporting the top 3 scores, we computed the mean and standard deviation for the top 100 molecules after unconstrained optimization. Each entry in the table below follows the format `mean ± std (median)`. The table shows that our methods have overall the best optimization performance. In addition,  HJ exhibits better performance on mean and standard deviation than on top 3, showing that minimizing the kinetic energy is efficient in pushing the distribution to desired properties. 

| Model          | plogP $\uparrow$    | QED $\uparrow$     | ESR1 Docking $\downarrow$ | ACAA1 Docking $\downarrow$ |
|-------------|--------------|------------|-------------|-------------|
| **Random**      | 2.345 ± 0.386 (2.259) | 0.903 ± 0.014 (0.902) | -9.127 ± 0.360 (-9.015) | -8.454 ± 0.316 (-8.390) |
| **Chemspace**   | 2.580 ± 0.406 (2.446) | 0.907 ± 0.014 (0.906) | -9.523 ± 0.409 (-9.395) | -8.749 ± 0.356 (-8.640) |
| **Gradient Flow**        | 2.664 ± 0.382 (2.537) | 0.910 ± 0.012 (0.908) | -9.452 ± 0.338 (-9.365) | -8.735 ± 0.337 (-8.650) |
| **Wave (spv)**    | 2.536 ± 0.439 (2.388) | 0.903 ± 0.015 (0.898) | **-9.630 ± 0.399** **(-9.525)** | -8.764 ± 0.344 (-8.650) |
| **Wave (unsup)**  | 1.736 ± 0.401 (1.610) | 0.845 ± 0.014 (0.840) | -9.074 ± 0.329 (-9.000) | **-8.813 ± 0.265** **(-8.745)** |
|  **HJ (spv)**       | 2.482 ± 0.397 (2.382) | 0.899 ± 0.017 (0.894) | -9.544 ± 0.322 (-9.460) | -8.792 ± 0.332 (-8.675) |
| **HJ (unsup)**     | **3.405 ± 0.254** **(3.377)** | **0.911 ± 0.009** **(0.909)** | -9.132 ± 0.321 (-9.090) | -8.668 ± 0.243 (-8.630) |
| **LD**          | 2.463 ± 0.388 (2.399) | 0.905 ± 0.014 (0.903) | -9.400 ± 0.360 (-9.300) | -8.709 ± 0.372 (-8.585) |

Meta-review: All four reviewers agree to accept the paper (although three of them only weakly), with some reviewers raising their score after the author response. For the camera-ready version, the authors should be certain to include the new experiments and results described in the rebuttal, as well as the additional baselines promised for reviewer kzEx, and a clear presentation of the overall objective pointing to the appropriate appendix material as necessary for the new algorithm block.
Decision: Accept (poster)