Manipulating 3D Molecules in a Fixed-Dimensional E(3)-Equivariant Latent Space
Authors: Zitao Chen
Yinjun Jia
Zitong Tian
Wei-Ying Ma
Yanyan Lan
Venue: NeurIPS 2025 poster ()
Keywords: Drug Design;VAE;Deep Learninng;Molecule Generation;
Primary area: machine_learning_for_sciences
TL;DR: We present the first variational autoencoder (VAE) model for 3D molecules that operates within a fixed-dimensional, E(3)-equivariant latent space, independent of atom counts.

Abstract: Medicinal chemists often optimize drugs considering their 3D structures and designing structurally distinct molecules that retain key features, such as shapes, pharmacophores, or chemical properties. Previous deep learning approaches address this through supervised tasks like molecule inpainting or property-guided optimization. In this work, we propose a flexible zero-shot molecule manipulation method by navigating in a shared latent space of 3D molecules. We introduce a Variational AutoEncoder (VAE) for 3D molecules, named MolFLAE, which learns a fixed-dimensional, E(3)-equivariant latent space independent of atom counts. MolFLAE encodes 3D molecules using an E(3)-equivariant neural network into fixed number of latent nodes, distinguished by learned embeddings. The latent space is regularized, and molecular structures are reconstructed via a Bayesian Flow Network (BFN) conditioned on the encoder’s latent output. MolFLAE achieves competitive performance on standard unconditional 3D molecule generation benchmarks. Moreover, the latent space of MolFLAE enables zero-shot molecule manipulation, including atom number editing, structure reconstruction, and coordinated latent interpolation for both structure and properties. We further demonstrate our approach on a drug optimization task for the human glucocorticoid receptor, generating molecules with improved hydrophilicity while preserving key interactions, under computational evaluations. These results highlight the flexibility, robustness, and real-world utility of our method, opening new avenues for molecule editing and optimization.

Manipulating 3D Molecules in a Fixed-Dimensional E(3)-Equivariant Latent Space
Zitao Chen 1.2-
Yinjun Jia
Zitong Tian
1,3+ WengYing Ma
Yanyan Lan 44,5f
Institute for Al Industry Research (AIR), Tsinghua University
Department of Computer Science and Technology, Tsinghua University
Qiuzhen College, Tsinghua University
Beijing Frontier Research Center for Biological Structure, Tsinghua University Beijing Academy of Artificial Intelligence {chen-zt23, tttt3tzttamamaaa (iteen-tt2,,ztt2)mmmmi tsinghua. edu. cn
(jiayinjun, maweiying, lanyanyan] Cair tsinghua, edu. on
Abstract
Medicinal chemists often optimize drugs considering their 3D structures and de- signing structurally distinct molecules that retain key features, such shapes, pharmacophores, or chemical properties. Previous deep learning approaches ad- dress this through supervised tasks like molecule inpainting property- -ggide optimization. In this work, we propose a flexible zero-shot molecule manipulation method by navigating in shared latent space of 3D molecules. We introduce a Variational AutoEncoder (VAE) for 3D molecules, named MolFLAE, which learns fixed-dimensional E(3)-equivariant latent space independent of atom counts. MolFLAE encodes 3D molecules using an E(3)-cquivariant neural network into fixed number of latent nodes, distinguished by learned embeddings. The latent space is regularized, and molecular structures are reconstructed via a Bayesian Flow Network (BFN) conditioned on the encoder's latent output. MolFLAE achieves competitive performance standard unconditional 3D molecule gen- eration benchmarks. Moreover, latent space of MolFLAE enables zero-shot molecule manipulation, including atom number editing, structure reconstruction, and coordinated latent interpolation for both structure and properties. We further demonstrate approach on a drug optimization task for the human glucocorticoid receptor, generating molecules with improved hydrophilicity while preserving key interactions, under computational evaluations. These results highlight the flexi- bility, robustness, and real-world utility of our method, opening new avenues for molecule editing and optimization.
Introduction
Structure-guided molecule optimization is crucial task in drug discovery. Medicinal chemists edit molecular structures to improve binding affinity, selectivity, and ADMET (absorption, distribution, metabolism, excretion, and toxicity) properties. These modifications can range from subtle changes, such introducing chlorine and reconstruction of known ligands
methyl 22] more extensive transformations like deconstruction or designing chimera molecules that combine beneficial
"qqual contribution.
Corresponding author.
'The code is available wttttpp:///shhab.bb.oa//ha//a2//
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

These diversified tasks present exciting opportunities for deep learning models to accelerate real-world drug design.
features of different scaffolds 4.55
Previous generative approaches typically decompose 3D molecule editing into a set of narrowly defined subtasks. Notable progress has been made in molecular inpainting 66 7 [. property-gu -uuiddd optimization [0]00, and shape-conditioned regeneration [[]]. While effective, these models often rely on task-specific supervision and architectures, limiting their flexibility and generalizability. Moreover, not all molecule editing tasks align well with the supervised learning paradigm. For example, adding substituents may be too trivial to justify training specialized model, while complex tasks like scaffold hopping by integrating known actives are often data sata-s for supervised approaches. These limitations call for a more flexible framework capable of supporting broad spectrum of molecule editing tasks in unified, tata-efficient manner.
13 14 show that latent space navigation allows for powerful, general-purpose manipulations by perturbing latent vectors. However, 3D molecule generation presents unique challenges not encountered in image domains. Molecules consist of variable numbers of atoms, and they exhibit permutation invariance to the atom order and -euuvvequv to the spatial translation and rotation. These characteristics make latent space modeling significantly more challenging, and most existing 3D generative models operate on the product of latent spaces of each atom or functional group [15] 16 variable dimensional representations. This variability prohibits common operations on vectors, such as interpolation or extrapolation, which are common to image generative models.
Previous successes in image editing and style transfer 12
T8 [9, resulting in
To address these challenges, we propose MolFLAE (Molecule Fixed Length AutoEncoder). a Varia- tional AutoEncoder (VAE) for 3D molecules that learns fixed-dimensional, E(3)-vquivariant latent space, independent of atom counts, Our encoder employs an E(3) E(3)-quiivarian neural network that updates a fixed number of virtual nodes initialized with learnable embeddings, transforming them into fixed-length latent codes for 3D molecules. The latent space is regularized under the standard VAE framework, and a Bayesian Flow Network (BFN) serves as the decoder, reconstructing full molecular structures conditioned on latent codes.
Our autoencoder framework supports wide range of downstream applications. We first demonstrate that our model can unconditionally generate diverse, valid molecules, achieving competitive perfor- mance standard 31) molecular generation benchmarks. More importantly, our fixed-dimensional latent space enables rich, semantically meaningful manipulations. We show that molecular analogs be created with varying atom counts, covering simple substitutions ring contractions. Molecules can also be reconstructed on the shape and orientation of other molecules, yielding chemically plau- sible outputs. Furthermore, interpolating between two latent codes produces chimera molecules that combine substructures and properties from both parents, Finally, we demonstrate a real-world application of our method in a drug optimization task targeting the human glucocorticoid receptor. We design new molecules that preserve the key binding interactions of known actives while achieving better balance of potency and hydrophilicity. These results illustrate the flexibility, robustness, and practical utility of our model, highlighting the promise of latent space manipulation as a powerful tool for molecular editing and optimization. Our main contributions are:
We propose a VAE model that learns a fieedddimensiona,, E(3)-qquivariant latent space for 3D molecular structures;
The learned latent space enables a wide range of molecule manipulation tasks, including analog design, molccule reconstruction, and crructure-properiie co-interpolation:
We introduce quantitative metrics to evaluate the disentanglement of spatial and semantic latent components, as well as the quality of structural and property interpolation.
2
Related Works
Unconditional Generation for 3D Molecules Unconditional 3D molecule generation has achieved rapid advancements, driven by progress in deep generative models. Early works explored auto-
regressive models to construct molecules atom-by-atom 20 21). More recently, inspired by diffusion-
like models 222 23 24 25 26 models like EDM 15 EquiFM [66], and GeoBFN [17], have sig-
nificantly improved generation quality. VAFS 27 offer an alternative by decoding latent embeddings
into molecules, but modeling 3D molecules is challenging due to variable atom counts and the need

---

BFN
Prior
Original Mol
Recon Mol
[(0.0.0)'N2
SE(3)-NN
2
SE(3)-NN
Recon Loss
Learnable Embedding
Regularization Loss
Virtual Nodes
Latent Codes
Figure I: The architecture MoIFLAE, 3D molecules are transformed into latent codes and decoded with BFN. MolFLAE is trained with the recon. loss and the regularization loss of the latent code.
for equivariance of coordinates to rotations and translations. Existing 3D molecular VAEs such as PepGLAD [19] and UniMoMo [28) have made progress by encoding molecular blocks (c.g., amino acid residues) into independent latent nodes. This design ensures -equvvvrvee but has two key limitations: (1) ) the number of latent nodes depends on the molecular composition, complicating cross-sample comparisons and interpolation tasks; and (2) spatial relationships between molecular blocks are tightly preserved, restricting the flexibility of the latent space for generative modeling, especially when latent codes are repurposed for tasks like autoregressive modeling by finetuning LLMs.
A key limitation of most unconditional 3D generative models is that their latent spaces vary in length. This variability prohibits operations such as interpolation. Consequently, zero-shot molecule editing becomes non-trivial even infeasible. To address this, researchers have designed specific models for molecule editing or constructed exed-dimensional latent spaces, which will discussed in the following sections.
3D Molecule Editing 3D molecule editing has traditionally been divided into several specialized subtasks. Linker design focuses on connecting fragments to form valid new compounds [6] 77. while scaffold inpainting involves masking molecule cores and regenerating them (8. These two tasks share similarities and can be unified under mask-prediction framework. However, other editing tasks are harder to generalize. For example, DecompOpt (29] decomposes ligands into substructures and trains a diffusion model conditioned these substructures, enabling deconstruction-.cconttuttoooooo reconstruction of 3D molecules. Another task is property-guided molecule optimization, where molecules are directly edited to improve specific properties using explicit guidance signals. Approaches such as gradient-based optimization [9] and classifier-free guidance [10 have been explored to this end.
While these task- specific methods have demonstrated strong performance, they often lack flexibility and generality. This raises the question if we can achieve more elegant and general 3D molecule manipulation by navigating in the latent space. This motivation has led researchers to explore unified, fixed-length, and semantically meaningful latent spaces for 3D molecules.
Latent Spaces Fixed-length autoencoders for molccules have been explored through voxcl-based models, which discretize 3D space into uniform grids and apply 3D CNNs [30] 31 or neural fields [32]. While straightforward, these methods are E(3)-equivariant and struggle disentangle semantie features from orientations, making molecule manipulation in the latent space not flexible. Local-frame-basee models [3 represent conformations with SE(3) invariant features, including distances, bond angles, and dihedrals, making outputs ireettationaagnsste the inductive bias of geometric equivariance, yet their latent space suffers from similar feature entanglement problem as voxel-based models. A recent work 335 also constructs fixed-length autoencoder, but applies global pooling over atom features, discarding spatial information that critical for reconstruction and interaction modeling.
Generate 3D Molecules from
UAE-3D 34 has proposed 3D fixed-length latent space by discarding

---

In contrast, our method preserves equivariance without requiring data augmentation, contains spatial information, avoids voxelization, and partially disentangles spatial and semantic information which enables unconditional generation and zero-shot molecule editing.
Methodology
We encode a variable-size 3D molecule by concatenating with set learnable virtual nodes and updating them together via an E(3)-equivariant neural network to obtain fixed-length embeddings. The regularization loss Lrgg comes from an Multi- Multi-Layer Perceptron (MLP) predicting the means and variances form variational posterior, from which sample latent codes that condition Bayesian Flow Network (BFN) decoder for reconstruction producing the reconstruction loss Creon The MoIFLAB trained end-to-end by minimizing
Lreeom Lreg-
We recall the classical Variational AutoEncoder (27) (VAE) in Sec. 3.1 that inspires the above loss. Then introduce BFN in Sec. 3.2 And introduce the encoder and decoder of MolFLAE in Sec. 3.3
Sec. 3.4
3.1
Variational AutoEncoder
Let go(2 x) denote the encoder, (x 2) the decoder, and p(z) the prior distribution of the latent codes, which is typically chosen as a standard Gaussian. The VAE loss is the negative Evidence Lower Bound (ELBO):
LVAB (0, x;x) (x)(()))) |- Po z)]
DKL (go x) |1((z))
(2)
reconstruction loss
regularization loss
Our training loss Eq. [ is inspired by the above VAE loss.
3.2 Bayesian Flow Network
The Bayes Flow Network (BFN) 26 incorporates Bayesian inference to modify the parameters of collection of independent distributions and using neural network integrate the contextual information. Unlike standard diffusion models [22 23 that primarily handle continuous data via Gaussian noise, BFN extends this paradigm to support both continuous and discrete data types, including categorical variables. This flexibility makes BFN particularly suitable for modeling 31) molecules [17), where the data naturally consists of mixed modalities of continuous coordinates and discrete atom types.
Unlike traditional generative models that operate directly data, BFN performs inference in the space of distribution parameters. Given a molecule m, the sender distribution (s(y m; a)) transforms into parameterized noisy distribution by adding noise analogous to the forward process in diffusion models. A brief introduction BFN can found in Appendix C.I
3.3 MoIFLAE Encoder
Inspired by the autoencoder in natural languages models [36] who uses of (CLS) tokens for context compression, we append Nz learnable virtual nodes the molecule's point cloud and treat the concatenation of their final embeddings our fixed-length latent codes.
We denote the coordinates and atom type feature of the i-th atom in molecule M by x M E R3 and V
R Pxr respectively. The full molecular input and the learnable virtual nodes are represented as M XM, and Z [xz, vz]. The virtual nodes some artificial atoms with the same size as the real atoms.
We remark that, in order to effectively encode the 3D configuration and chirality of molecule, the number of virtual nodes Nz must be at least so that they can form non-degenerate simplex 3D space. To ensure sufficient capacity for capturing complex spatial structures, set Nz 10.

---

We employ an -queequurar neural network Øg to jointl encode the original molecular point cloud and the appended virtual nodes (M, Z). After rounds of update, we discard the embeddings of M and retain only those of 2 as our fixed-length latent representation. In explicit,
(-, znl) = Φ (IxM, V.l. vzl)
(3)
The latent code is denoted as Z, E RN feature components, D, is the embedding dimension of features. Note that only keep the fixed- length part of the Φy output. So we obtain a fixed-length encoding of the molecules. The network structure and Equivaaive discussion
where 2μ and 2д represent the spatial and
be found in the Appendix A
To regularize the latent space, we adopt a VAE formulation. While the initial output [Zg, zn] is deterministic, this can lead to irregular latent geometry and poor interpolatability. To address this, predict a coordinate-we wise Gaussian distribution for each latent dimension:
[o Mhs T = Linear(z)).
μg = Zz:
The resulting latent posterior is regularized via KL divergence to a fixed spherical Gaussian prior (N 0), [var,, var, I), giving rise to the regularization loss:
Crog = KL (Ni( un), lo ( || N[O 0J, [varz, vara]
where var,, var,
fixed scale parameters. This regularization encourages the latent space be smooth and continuous, facilitating interpolation between molecules and improving the robustness and diversity of samples generated from the prior distribution. In practice, we project (using the linear layer) the feature embedding ZĄ E NR to μ4 E RN DDz For notational simplicity, we continue denote the sampled latent code from M(Hn, I) Zn.- The full expression of the regularization loss is provided in Appendix B
3.4 MoIFLAE Decoder
The encoder defines a Gaussian posterior over the latent space. We sample a latent code (Zas Zn, from this distribution and use it the conditioning input to the BFN decoder. By comparing the coordinates and atom type of reconstructed molecule with the original input, we compute the reconstruction loss:
Lrecom
C
In our molecular BFN setup, we must jointly model both continuous and discrete aspects of atomic data. This requires unified representation that enables neural networks to propagate information across modalities while maintaining compatibility with Bayesian updates. is enough to define suitable sender distribution ps allowing the additivity of precision (26].
We model continuous atomic coordinates using Gaussian distributions. Given ground-truth coordi- nates X and a noise level a pー the sender generates noisy observation by adding isotropic Gaussian noise:
ps(y" ( (xMaa) a) Ny
XM, a II)
For discrete atom types, model each atom using a categorical distribution over K classes. This distribution is parameterized by a continuous matrix 0" e R Na X which is transformed into proba-
bilities via a softmax function. Given the ground-truth atom type matrix e = [v (1) ..... v
(Nw)
NM x where each v M is the column one-hot vector representing one of the K atom categories, the sender perturbs it with artificial Gaussian noise scaled by a', producing:
Ps P Varaa) N (y" a a'(Kevar 1), a' KI)
(
For the initial prior d,. follow 26 and adopt standard Gaussian priors for continuous variables and uniform distributions for categorical ones.
Then we can derive the two loss L LT respectively
The total forward pass can be found in the Algorithm. 2 The inference process is parallel with the general BFN inference but taking two data modalities into consideration. See Appendix C.4
37
5

---

Experiments
We train and evaluate MolFLAE three datasets: QM9 38). Grug-Drugs 39 and ZINC-9M (the in-stock subset of ZINC [40] with 9.3M molecules). QM9 contains 134k small molecules with up to heavy atoms, and GEOM-Drugs is a larger-scale dataset featuring 430k drug-like molecules. On both QM9 and GEOM-Drugs experiment, hydrogens are treated explicitly. We use QM9 and GEOM-Drugs evaluate MolFLAE in unconditional 3D molecule generation task, and demonstrate other applications on the more comprehensive scale dataset ZINC-9M, where hydrogens are treated implicitly.
Unconditional Molecule Generation To assess the capability of MolFLAE generate stable, diverse molecules, first focus on 3D molecule generation task following the setting of prior works We conduct 10,000 random samplings in the latent space, then decode them into molecules using MoIFLAE decoder, subsequently evaluating qualities of these molecules. We sample the atom number from the prior of the training set as previous works like [3]]. Table iillstraate the benchmark results of unconditional generation with MolFLAE. We also provide results on drug- likeness metrics on GEOM- --Dgg in Appendix [E] confirming MolFLAE's outstanding performance in generating structurally reasonable and drug-like like molecules compared to previous methods.
15
16 17
Table 1: Performance comparison of different methods on the QM9 and GEOM- --ugg dataset.
QM9
GEOM-Drugs
# Metrics
Atom Sta (%)
Mol Sta (%)
Valid (%)
VxU (%)
Novelty (%)
Atom Sta (%)
Valid (%)
Data
99.0
95.2
97.7
97.7
86.5
99.9
ENF
41
85.0
4.9
40.2
39.4
G-Schnel
423
95.7
68.1
85.5
80.3
GDM-AUG
15
97.6
71.6
90.4
89.5
74.6
77.7
91.8
1
98.7
90.7
81.3
92.6
B0M-Bridge
43
98.8
84.6
90.7
82.4
92.8
GEOLDM 18
98.9
89.4
93.8
92.7
57.0
84.4
99.3
GROBEN 17
98.3
85.1
92.3
90.7
72.9
75.1
91.7
GROBPN 100
98.6
87.2
93.0
91.5
70.3
78.9
93.1
UniGEM
99.0
85.1
MolFLAE 60
99.3
90.4
92.1
77.1
86.9
99.2
MolFI.AF
99.4
92.0
96.8
88.9
74.5
86.7
99.7
Compared with several baseline models, MolFLAE achieves competitive performance across atom stability, molecular stability, and validity metrics both QM9 and GEOM-Drup 1roM-- dataset, while requiring fewer sampling steps. These results suggest that our latent space is well-structured, supporting efficient and reliable molecular generation.
Generating Analogs with Different Atom Numbers First, we probe the smoothness of MolFLAE decoder atom numbers by forcing the generation with increased or decreased atom numbers based on original latent code. We examine the similarities between generated molecules and original molecules with MCS-IoU (Maximum Common Substructure Intersection-over-unne molecules share similar orientations, shapes and 2D structures with the original input, validating the desired smoothness. Detailed results presented in Table 2 and three examples provided in Fig. 2 for better illustration.
Generated
Table 2: Evaluating 2D similarities between generated analogs and original molecules.
Atom Number
-2
-1
MCS-IoU similarity
69.79
76.69
84.08
76.05
69.95
Valid(%)
100.0
99.89
99.76
99.89
99.68
Atom Sta(%)
84.58
83.28
82.48
82.38
82.53
Exploring the disentanglement of the latent space via molecule reconstruction The latent space of MoIFLAE consists of two parts, the E(3))equivariant component and the E(3)-invariant component 2₁, Ideally, spatial and semantic features of molecules disentangle spontaneously, with zy encoding the shape and orientation and n encodes substructures of molecules.

---

N₂. -2
N
N1 +1
N₂. +2
Figure 2: Examples for analog generation with variable atom numbers.
(2), 2.)
(z/, z
(z) z!)
Figure 3: An example for molecule reconstruction with new shape and orientation.
In this section, we explore this disentanglement hypothesis by swapping and 2u between molecules and observe decoded molecules. Formally, with two molecules Mo and M,, and their latent (z x2) and (z/, ) we decode molecules with (z) k z!!) (named preserving 2;)) and (z) 2 (named preserving 2.)). respectively.
In experiments, we have observed a partial disentanglement the MolFLAE latent space. As shown in Fig. following the disentanglement hypothesis, substructures information of M1 should be able to be extracted by isolating 2/,- As z) is not sufficient to reconstruct M1, we view it as deconstructed molecule. Then, reconstruct these substructures into the shape and orientation of Mo. by decoding the hybrid latent code (z), 2 The resulted molecule shares a similar shape and orientation with My. indicated by the dash lines in Fig Moreover, it also shares similar substructures as M, like amide, chlorobenzyl, sulfamide, indicated by rectangles corresponding colors.
Quantitatively, we compute the MACCS fingerprint 45) similarity (considering substructure overlap- ping) and situ shape similarity (considering shape, orientation and relative position) of 1000 hybrid molecules (z),, zb) and (2% 2-) with the original molecule (z), 22) (Table 3). Under the setting of preserving Z,, the shape similarity is significantly higher than the preserving 2n (0.394 vs 0. 0.14:
indicating z indeed encodes shape and orientation information. Similarly, MACCS similarity is

---

higher under the preserving 2₁, setting than the preserving 2₂ setting (0.580 vs 0.421). These results support that the latent space of MolFLAE is partially disentangled, with 2 representing substructure composition and 2g representing shape and orientation.
Table 3: Measuring molecule reconstruction similarities under different settings.
Preserving &.
Preserving N.
Metries
MACCS Sin
Shape
Vasdrao
Atom S
MACCSS Sim
Valdisen
Atom Stark
MoIFLAE
0421
85.20
0.580
0.174
100.0
84.62
Latent Interpolation The diedddimensional latent space allows flexible manipulation of molec- ular representations via vector convex combinations. In MolFLAA,, the regularization loss further encourages smooth transitions between latent codes, facilitating continuous transformations be- tween molecules. Fig. presents the interpolation of three pairs of molecules, indicating smooth transformations of shape and orientation.
Source
Target
Figure 4: Examples for the latent interpolation between molecules.
To further quantify the quality of interpolation, We evaluate the trend of properties of the intermediate molecules during the transformation, considering both structural and physical properties (detailed descriptions provided in Appendix (D). We compute the Pearson correlation coefficient between each property value and its corresponding interpolation index (adjusted with the sign of property difference between the source and target molecule). We also report the associated p-values from the Pearson significance test to assess the statistical evidence for linear trends. A property considered to exhibit significant linear variation along the interpolation trajectory if the null hypothesis of zero correlation is rejected at the 5% significance level, i.c., if log results documented in Table and Table 4 We also report the step-wise molecular validity and atom stability during interpolation Table which indicates that most intermediate molecules are valid and stable.
log (0.00) 1.3. Detailed
Table 4: Monitoring the trend of structural properties along with molecule interpolations.
Similarity Preference
sp3frac
BertzCT
QED
Interpo Num
Pearson's
-logp
Pearson's 's
-logp
Pearson's s
-logr
Pearson's
-logp
8
0.9261
3.3346
0.4314
0.8518
0.6537
1.7888
0.5460
1.2128
10
0.9191
4.1340
0,3982
0.9639
0.6344
2.1138
0.5194
1.4228
0.9122
4.8476
0.3809
1.0889
0.6281
2.4728
0.5059
1.6216

---

Table 5: Monitoring the trend of physical properties along with molecule interpolations.
Labute ASA
TPSA
LogP
MR
Interpo Num
Pearson's
-logp
P
-logp
Pearson's
-logp
Pearson's 's
-logp
8
0.9067
4.6366
0.5711
1,3783
0.5400
1.2363
0.8467
3,5188
10
0.9041
5.7687
0,5620
1.6533
0.5216
1,4533
0.8388
4.3350
12
0.8939
6.8567
0.5425
1.8572
0.4925
1.6259
0.8255
5.1216
Table 6: Step-wise Validity and Atom Stability during interpolation.
Step
Interpo Num
Metrics
10
12
Valid(3))
53.91
99.501
99.50
100.0
99.90
100.0
100.0
100.0
Atom Sta(%) %)
82.23
82.95
84.17
84.57
84.21
83.74
Valid((3))
100.0
99.90
99.90
99.80
99.80
100.0
100.0
100.0
99.80
99.90
Atom Sta(%)
82.23
82.75
83.53
84.07
84.36
84.35
84.16
84.17
82.78
82.67
Valid(3))
99.90
99,90
100.0
Atom Sta( (6)
82.39
82.50
83,76
83.67
83,92
84.65
84,20
84,50
84,27
83.00
82. 84
82.57
Applying MoIFLAE to optimize molecules targeting the human glucocorticoid receptor (hGR) To assess the real-world utility of our method, we applied our method to optimize drug candidates for the hGR, which requires balancing The hGR is a key target for anti- anti-inllammatory. and our optimization starts from two known actives. AZD2906 potent hGR modulator but is poorly soluble [46), while BI-653048 is more soluble but less potent 447.. To showcase the performance of our model, We computationally evaluate the potency and hydrophilicity with Glide docking and QikProp CLogPo/w from the Schrodinger Suite, where lower docking scores indicate better potency, and lower CLogPo/w values indicate better hydrophilicity. These computational metrics align with real-world properties of AZD2906 and BI-653048. AZD2906 has docking score of - 13. 16 and a high CLogPo/w of 5.61, whereas BI-653048 shows a better CLogPo/w of 3.90 but a weaker docking score of 10.62 (Fig. 55 and C). These results facilitate the computational evaluation of MolFLAE generated molecules.
cetroppoobiciiteracttinn
binding with aqueous solubility.
To explore trade-offs, we blended these two molecules latent space using 90% AZD2906 and 10% BI-653048, generating 100 candidates. The top 10 molecules outperformed B1-653048 in docking score, and also improved hydrophilicity (CLogPo/w) compared AZD2906. These candidates preserved AZD2906''s s binding shape while introducing polar groups for better solubility. Sample 34, for instance, retained key pharmacophores of AZID2006 and BI-653048 for interacting with receptor (indicated by colored rectangles for each pharmacophore), achieving balanced property with its docking score 11.15 and CLogPo/w of 3.75 (Fig. FF. Moreover, its docking pose closely matched both its generated conformation (RMSD= 1.35 Å) and AZD2906's 's crystal structure (Fig. and E), representing the advantage of explicitly modeling of 3D coordinates by MolFLAE. These results highlight our method's potential for meaningful molecular optimization and drug design.
Conclusion and Future Works
In this work, we present MolFLAE, flexible VAE framework for manipulating 3D molecules within a fixed-dimensional, E(3)-equivariant latent space. Our method demonstrates strong performance across multiple tasks, including unconditional generation, analog design, substructure reconstruction, and latent interpolation. We further validate the real-world utility of MolFLAE through a case study on generating diug-like molecules targeting the hGR, balancing potency and solubility.
Beyond the reported experiments, MolFLAE naturally extends to a wider range of tasks. For example, molecule inpainting can be achieved by encoding discontinuous fragments and decoding to larger atom sets. Structural superposition can be achieved efficiently via the weighted Kabsch algorithm on latent nodes, avoiding the high complexity of atom-wise bipartite matching. These applications exemplified in Fig 6 and a deeper exploration left for future work due to space constraints.
While MoIFLAE demonstrates strong performance across multiple tasks, there remains room for improvement in the disentanglement and interpretability its latent space. We hypothesize that better disentanglement can be achieved by enforcing the invariance of semantic latent z to molecular

---

conformational changes or other non non-rigid perturbations. Future works may explore incorporating self-contrastive objectives to better capture chemical semantics with latent representations.
In summary, our results highlight the versatility of MolFLAE and its promise as a general-puroose framework for 3D molecule generation and editing. This work opens directions for exploring the broader applications of exed-dimensional, E(3)-equivariant latent spaces in molecular modeling.
A
B
AZD2906
BI-653048 Glide-score =-00.62
Glide-score = 113.16
PDB: 6EL6 (AZD2906)
QPLogPow 5.61
aPLogPolw = 3.90
Sample 34 vs AZD2906
e'
Sample 34 Glide-score = -11..15
Sample 34 (RMS=1.35A)
Sample 34 (Docking)
Docking Generated
OPLogPolw = 3.75
Figure 5: Applying MoIFLEA optimizing AZD2096 targeting the hGR. A, the cystal structure of AZD2096 in complex with hGR. B, C, and F, 2D structures of AZD2096, BI BI-653048, and sample 34, with their docking score and CLogPo/w. D, the docking pose of sample 34. E, comparing the docking pose of sample 34 with AZD2096 and its generated pose before docking.
Molecule inpainting (linker design)
Molecule alignment (superposition)
Fragments
Inpainted
Inpainted
Reference ligand
Before alignment
After alignment
Figure 6: Exemplifying the application of MolFLAE to molecule inpainting and superposition.
10

---

Acknowledgements
This work is supported by Beijing Academy of Artificial Intelligence and Beijing Frontier Research Center for Biological Structure Fundings.
References
[1] Debora Chiodi and Yoshihiro Ishihara. "magic chloro": Profound effects of the chlorine atom in drug discovery. Journal of Medicinal Chemistry, (0(8):5305-533) 2023, PMID: 37014977.
[2] Heike Schönherr and Tim Cernak. Profound methyl effects in drug discovery and a call for new c-h methylation reactions. Angewandte Chemie International Edition, 52(47):12256-1226) 2013.
[3] J. Henry Blackwell, Iacovos N. Michaelides, and Floriane Gibault. A perspective on the strategic application ddcconsruction-reonstuuiiioooooi drug discovery. Journal of Medicinal Chemistry, 0(0):null, O. PMID: 40324045.
[4] Tingting Chen, Jiafu Leng. Jun Tan, Yongjun Zhao, Shanshan Xie, Shifang Zhao, Xiangyu Yan, Liqiao Zhu, Jun Luo, Lingyi Kong, and Yong Yin. Discovery of novel potent covalent glutathione peroxidase 4 inhibitors highly selective ferroptosis inducers for the treatment of triple-negative breast cancer. Journal of Medicinal Chemistry, (6(14): 1003-10059,, 2023. PMID: 37452764.
[5] Kunyu Shi, Jifa Zhang, Enda Zhou, Jiaxing Wang, and Yuxi Wang. Small-molecule receptor- interacting protein I (ripl) inhibitors as therapeutic agents for multifaceted diseases: Current medicinal chemistry insights and emerging opportunities. Journal of Medicinal Chemistry, 65(22):14711-1499) 2022. PMID: 36346971.
[6] Ili Igashov, Hannes Stärk, Clément Vignac, Victor Gareia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design, 2022.
[7] Jiaqi Guan, Xingang Peng, PeiQi Jiang, Yunan Luo, Jian Peng, and Jianzhu Ma. Linkernet: Fragment poses and linker co-design with 3d equivariant diffusion. In Thirty-sevntth Conference on Neural Information Processing Systems, 2023.
[8] Shicheng Chen, Odin Zhang, Chenran Jiang, Huifeng Zhao, Xujun Zhang, Mengting Chen, Yun Liu, Qun Su, Zhenxing Wu, Xinyue Wang. Wanglin Qu Yuanyi Ye, Xin Chai, Ning Wang, Tianyue Wang. Yuan An, Guanlin Wu, Qianqian Yang, Jiean Chen, Wei Xie, Haitao Lin, Dan Li, Chang- Yu Hsich, Yong Huang, Yu Kang, Tingjun Hou, and Peichen Pan. Deep lead optimization enveloped in protein pocket and its application in designing potent and selective ligands targeting LTK protein. Nat. Mac. Intell., 7(3):448-455 2025.
[9] Keyue Qiu, Yuxuan Song, Jie Yu, Hongbo Ma, Ziyao Cao. Zhilong Zhang, Yushuai Wu, Mingyue Zheng, Hao Zhou, and Wei- Wei-Ying Ma. Empower structure-basee molecule optimization with gradient guidance, 2025.
[10] Alex Morehead and Jianlin Cheng, Geometry-complete diffusion for 3d molecule generation and optimization. ArXiv, 2023.
[11] Keir Adams, Kento Abeywardane, Jenna Fromer, and Connor W. Coley. ShEPhERD: Diffusing shape, electrostatics, and pharmacophores for bioisosteric drug design. In The Thirteenth International Conference on Learning Representations, 2025.
[12] Peiye Zhuang, Oluwasanmi O Koyejo, and Alex Schwing. Enjoy your editing: Controllable {gan)s for image editing via latent space navigation. In International Conference on Learning Representations, 20221.
[13] Qiucheng Wu, Yujian Liu, Handong Zhao, Ajinkya Kale, Trung Bui, Tong Yu, Zhe Lin, Yang Zhang, and Shiyu Chang. Uncovering the disentanglement capability in text-to-image diffusion models. In 2023 IEEF/CVE Conference on Computer Vision and Pottern Recognition (CVPR), pages 190001910, 2023.
11

---

[14] Yong-Hyun Park, Mingi Kwon, Jaewoong Choi, Junghyo Jo, and Youngjung Uh. Understanding the latent space of diffusion models through the lens of riemannian geometry. In Thirry-seventh Conference Neural Information Processing Systems, 2023.
[15] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 8867-8887. PMLR, 2022.
[16] Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei- Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[17] Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei-Ying Ma. Unified generative modeling of 3d molecules with bayesian flow networks. In The Twelfth International Conference on Learning Representations, 2024.
[18] Minkai Xu, Alexander S. Powers, Ron O. Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In Andreas Krause. Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 38592-38610. PMLR, 2023.
[19] Xiangzhe Kong, Yinjun Jía, Wenbing Huang, and Yang Liu. Full-atom peptide design with geometric latent diffusion. In The ehirt Annual Conference on Neural Information Processing Systems, 2024.
[20] Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International Conference on Learning Representations, 2022.
[21] Niklas W. A. Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapteer generation of
3d point sets for the targeted discovery of molecules. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc Emily B. Fox, and Roman Gamett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 7564-7576, 2019.
[22] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information P Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[23] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[24] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
[25] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
[26] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks, 2023.
[27] Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
[28] Xiangzhe Kong, Zishen Zhang, Ziting Zhang, Rui Jiao, Jianzhu Ma, Wenbing Huang, Kai Liu, and Yang Liu. Unimomo: Unified generative modeling of 3d molecules for de novo binder design. In Forty- International Conference on Machine Learning, 2025.
12

---

[29] Xiangxin Zhou, Xiwei Cheng, Yuwei Yang. Yu Bao. Liang Wang, and Quanquan Gu. Decom- popt: Controllable and decomposed diffusion models for structure-based molecular optimization. In The Twelfth International Conference on Learning Representations, 2024.
[30] Tomohide Masuda, Matthew Ragoza, and David Ryan Koes. Generating 3d molecular structures conditional a receptor binding site with deep generative models. CoRR, abs/2010. 2020.
[31] Pedro Pinheiro, Joshua Rackers, joseph Kleinhenz, Michael Maser, Omar Mahmood, An- drew Martin Watkins, Stephen Ra, Vishnu Sresht, and Saeed Saremi. 3d molecule generation by denoising voxel grids. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
[32] Matthieu Kirchmeyer, Pedro O. Pinheiro, and Saeed Saremi. Score-based 3d molecule gen- eration with neural fields. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[33] Robin Winter, Frank No'e, and Djork-Arné Clevert. Auto-encoding molecular conformations. ArXiv, abs/2101.016 2021.
[34] Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Kenji Kawaguchi, Tat-Seng Chua, and Xiang Wang. Towards unified latent space for 3d molecular latent diffusion modeling, 2025.
[35] Tianxiao Li, Martin Rengiang Min, Hongyu Guo, and Mark Gerstein. 3d autoencoding diffusion model for molecule interpolation and manipulation, 2024.
[36] Tao Ge, Hu Jing, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model. In The Twelfth International Conference on Learning Representations, 2024.
[37] Yanru Qu, Keyue Qiu, Yuxuan Song, Jingjing Gong, Jiawei Han, Mingyue Zheng, Hao Zhou, and Wei-Ying Ma. MoICRAFT: Structure-baset drug design continuous parameter space. In Forty-first International Conference on Machine Learning, 2024.
[38] Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp. and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 08 2014.
[39] Simon Axelrod and Rafael Gómez-Bombarelli GEOM: energy-annotated molecular conforma- tions for property prediction and molecular generation. CoRR, abs/2006.05531, 2020.
[40] John J. Irwin, Khanh G. Tang, Jennifer Young, Chinzorig Dandarchuluun, Benjamin R. Wong, Munkhzul Khurelbaatar, Yurii S. Moroz, John Mayfield, and Roger A. Sayle. Zinc20-a free ultralarge-scale chemical database for ligand discovery. Journal of Chemical Information and Modeling, 60(12)x6065-6773, 2020. PMID: 33118813.
[41] Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In NeurlPS, pages 4181-4192, 2021.
[42] Niklas W. A. Gebauer, Michael Gastegger, and Kristof T. Schütt. addmetry-adapted generation of 3d point sets for the targeted discovery of molecules. CoRR, abs/1906.00957, 2019.
[43] Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. In NeurIPS, 2022.
[44] Shikun Feng, Yuyan Ni, Lu yan, Zhi-Ming Ma, Wei- Ma, and Yanyan Lan. UniGEM A unified approach to generation and property prediction for molecules. In The Thirteenth International Conference on Learning Representations, 2025.
[45] Joseph L. Durant, Burton A. Leland, Douglas R. Henry, and James G. Nourse. Reoptimization of mdl keys for use drug discovery. Journal of Chemical Information and Computer Sciences, 42(6): (0:6):1273-11 1280, 2002. PMID: 12444722.
13

---

[46] Martin Hemmerling, Karl Edman, Matti Lepistö, Anders Eriksson, Svetlana Ivanova, Jan Dah- mén, Hartmut Rehwinkel, Markus Berger, Ramon Hendrickx, Matthew Dearman, Tina Jelles- mark Jensen, Lisa Wissler, and Thomas Hansson. Discovery of indazole ethers novel, potent, non-steroidal glucocorticoid receptor modulators. Bioorganic & Medicinal Chemistry Letters, 26(23):5741-5748, 2016.
[47] Christian Harcken, Doris Riether, Pingrong Liu, Hossein Razavi, Usha Patel, Thomas Lee, Todd Bosanac, Yancey Ward, Mark Ralph, Zhidong Chen, Donald Souza, Richard M. Nel- son, Alison Kukulka, Tazmeen N. Fadra-Khan, Ljiljana Zuvela- Zuvela-Jelaska, Mita Patel, David S. Thomson, and Gerald H. Nabozny. Optimization of drug-like properties of nonsteroidal gluco- corticoid mimetics and identification of clinical candidate. ACS Medicinal Chemistry Letters, 5(12):1318-1322 2014.
[48] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[49] Kaiwen Xue, Yuhao Zhou, Shen Nie, Xu Min, Xiaolu Zhang. JUN ZHOU, and Chongxuan Li. Unifying bayesian flow networks and diffusion models through stochastic differential equations. In Forty-first International Conference Machine Learning. 2024.
[50] Steven H. Bertz. The first general index of molecular complexity. Journal of the American Chemical Sociery, 103((2):3599-36)) -5090,,,,,, 1981
[51] G. Richard Bickerton, Gaia V. Paolini, Jeremy Besnard, Sorel Muresan, and Andrew L. Hopkins. Quantifying the chemical beauty of drugs. Nature Chemistry, 4(2):90-98, February 2012.
[52] Paul Labute. widely applicable set of descriptors. Journal of Molecular Graphics and Modelling, 18(4):464 477, 2000.
[53] Peter Ertl, Bernhard Rohde, and Paul Selzer. Fast calculation of molecular polar surface area as a sum of fragment-based contributions and its application to the prediction of drug transport properties. Journal of Medicinal Chemistry, 43(20):3714-3711 2000. PMID: 11020286.
[54] Scott A. Wildman and Gordon M. Crippen. Prediction of physicochemical parameters by atomic contributions. Journal of Chemical Information and Computer Sciences, (9(55:868 873, 1999.
14

---

NeurIPS Paper Checklist
The checklist is designed to 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 follow 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 short (1-2 sentence) justifiicton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are 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 reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.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"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or 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.
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]
We do claim our contributions and scope in the abstract and the last part of introduction.
Justification:
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 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 can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as 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? Answer: [Yes]
15

---

Justification:
We discuss the limitations of our work in the Discussion and Conclusion part.
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 in 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 approxinations 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 a 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 low images are taken in low lighting. Or speech-to-text system might 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, a 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 proof'
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:
[NA]
Justification: Our work includes no theoretical analysis.
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 if they appear in the supplemental material, the authors are 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 be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We give the expiriment details in Appendix. A and hyperparameters in Appendix. G
Guidelines:
16

---

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, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 to 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) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If 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 or way to reproduce the model (e.g., with an open- source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
We will release our codes reproduce the main experiment soon.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https c/nips.cc/publi///guu
CodeeSumiiii wwttubbmmssiiioi ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
17

---

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
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?
Question:
Answer: [Yes]
Justification: We give the expiriment details in Appendix. C and hyperparameters in Appendix. G
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: [No]
Justification: Our experiments follow standard deep learning practice and do not report statistical error bars significance tests.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, 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.
It is OK to report 1-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 to show in tables or figures symmetric error bars that would yield results that are out of 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide our computer resources in Appendix. G
Guidelines:
The answer NA means that the paper does not include experiments.
18

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well 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 make 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
Answer: [Yes]
Justification: The paper does conform, in every respect, with the NeurIPS Code of Ethics.
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:
We have discussed the potential impact of our work in Appendix
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 not 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, it 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 is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
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 (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: We have discussed the Safeguards of our work in Appendix H
19

---

Guidelines:
The answer NA means that the paper poses no 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 to 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 challenging, and many papers do not require this, but 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]
Justification: All assets are properly cited in the main text. Their licenses and terms of use are respected in accordance with the original distribution terms.
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 is used and, if 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 (c.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pacerswithcode thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license 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 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: [Yes]
Justification: We will provide new validation set splitted from ZINC-9M, along with our code of latent molecule manipulation experiments.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/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 was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You either create an 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 details about compensation (if any)?
20

---

Answer:
[NA]
Justification:
We do not do such research.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in 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 in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
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?
Question:
Answer: [NA]
Justification:
We do not do such experiments.
Guidelines:
The answer NA means that the paper does 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 we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification:
We do not use LLM on such important parts.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non ntn-standard components.
Please refer to LLM policy (https://neurip: htt.r/ppps..cooooeeeeeeeeeee//// for what should or should not be described.
21

---

A Neural Network Details
We adopt the network structure of MoICRAFT 37 as our E(3)-Equivariant graph neural network (GNN) backbone. Originally designed to model interactions between ligand and protein pocket atoms, the network distinguishes between update nodes (whose coordinates are updated) and condition nodes (which provide contextual information and the coordinates not updated).
A.1 Neural Network Architecture
To align with our autoencoder framework, we adapt this formulation by assigning roles to nodes based on the encoding or decoding stage. In the encoder, the update nodes correspond to the virtual nodes, while the condition nodes are the atoms of the input ground truth molecule. Conversely, in the decoder, the update nodes are the molecular atoms being generated, and the condition nodes are latent codes produced by the encoder.
To be convenient, we concatenate the spatial part and feature part of the update nodes (upddate Vupdare] and condition nodes [xcondition Vcondition with writing x = X Xcondition] and h = updatee Vc where the superscript represents the 8-th layer of Φ, 0 S f L. The Initial hidden embedding h" is obtained by an MLP embedding layer that encodes the atom feature [h]. No embedding layer for the atom spatial coordinates. The construction of Φe is alternately updating the atom feature embeddings h and coordinates x as
X
h h
+
Σ JEV())
In (dly, h/ h/. exji On),
(9)
x
[+
x
Σ jEV(i)
(xf
xf) for fo bftl, hf" eiyi 0z) upddate
(10)
j
Here, V(i) is the neighbors of i, who could have information communication to i. We choose the nearest nodes from i. dē = ||x - x |l2 denotes the Euclidean distance between atoms and j at layer f, and oiy encodes whether the pair (i, belongs to the update nodes, condition nodes, or the connection between them. The indicator 1 Lupdate ensures that coordinate updates only applied to update nodes, keeping the positions condition nodes fixed.
A.2 E(3)-Equivariance Discussion
In molecular modeling, is essential that the learned distribution over update nodes be invariant to translations, reflections, and rotations of the condition nodes. This E(3)-invariance reflects a fundamental inductive bias in molecular systems [5 48 Since SE(3) is a subgroup of E(3), any -quivvarvaa model is also SE(3) -equivrrina While some prior works (e.g., (37 adopt the SE(3), others (e.g.. [5) use E(3), which more accurately describes the symmetry of their networks. We follow the latter avoid ambiguity.
The full Euclidean group in R3 denoted E(3), consists of rigid-body transformations of the form T(x) = Rr + 4, where E 3x3 is a orthogonal matrix and E IR3 is a translation vector.
If pre-align the condition nodes by centering them at their center of mass (i.c., eliminating the translational degree of freedom), then the resulting likelihood becomes E(3)-invariant under the following condition:
Proposition A.1 (Proposition 4.1 in (37). Let T' E E(3) denote rigid transformation. If the condition nodes are centered at zero and the parameterization (0, c, 1) is E(3)-equivariant, then the likelihood is invariant under T:
Ps (T(n) T(c)) = (o
This property ensures that the decoder's predictions respect the underlying geometric symmetries of molecular structures, which crucial for both sample quality and spatial information learning of latent codes.
22

---

A.3 Encoder Network
In the encoder network, the update nodes are the virtual nodes while the condition nodes are the input ground truth molecule.
Before passing into the network, we apply a linear layer to embed the one-hot atom features VM E R-NMx DM into continuous feature space IR xD,, where D, denotes the embedding dimension. The virtual node features are initialized learnable parameters in the same embedded space IR xxxxxD and are only defined at the embedding level. Their initial spatial positions are set to zero.
After L layers of message passing, the output of the Network ΦB is given by
Φe (xxM, v.l. Xz, vzl) = (RMs V.l [5z, Vz]) (_, Zzz zn.))
(11)
where only the virtual node outputs (Ze> zn| are retained as the final latent code. Since the coordinate updates are designed to be -(q))eiuvva at each layer, the entire encoder Φo is equivariant by
construction.
Importantly, we do not apply softmax projection to convert feature embeddings into one-hot vectors. Instead, we preserve the continuous representations retain richer information for downstream generation.
Then apply VAE layer to furtherly encode the latent code to Gaussian distribution. See Appendix
A.4 Decoder Network
The decoder follows a mirrored architecture, where the update nodes correspond the generated molecule, and the conditioning nodes are the latent virtual codes. Similar to the encoder, discard the outputs corresponding to the latent nodes and retain only the decoded molecule representations for final output. The atom number NM is to be known beforehand. We can sample it from the training set prior when doing unconditional generation [5] or edit in generating analogs with different atom numbers.
B
VAE Details
We adopt the regularization loss to regularize the latent space. Given the deterministic initial output [Z,, 2n), we predict coordinate-wic Gaussian distribution for each latent dimension:
Zez
[o 1hh o Linear(Z...
(12)
where we assume isotropic variance for 3D coordinates (i.e., each atom shares a scalar variance across x, V. 2), while the feature dimensions N assigned independent variances per entry. This design ensures that the latent distribution preserves equivariance in the spatial domain while maintaining expressiveness in the feature domain.
In practice, we project (using the linear layer) the feature embedding 2n N2x, to fs E R 22 For notational simplicity, continue to denote the sampled latent code from ( I.
as Z
The resulting latent posterior is regularized via KL divergence to a fixed spherical Gaussian prior N(O0, 0), [var,, var, I), giving rise the regularization loss:
Lreg = KL (W(( Mn), jo o || |VO,, 0J, [var,, var,
(13)
where var,; vara two fixed scale parameters. Since the KL-divergence between two Gaussian distribution P; ~ (l,, of), i == 2
022 of + (11
12
DKL(Pi = log
(14)
σ1
20
2
23

---

Hence the regularization loss is Lieg = [(h)
((x) where
+
Nex
[1 of
KL
Σ
log 07 log -1).
(15)
2
var
i=1
ck
N2
x3
[1
of
log o7 -1)
(16)
var
This regularization encourages the latent space to be smooth and continuous, facilitating interpolation between molecules and improving the robustness and diversity samples generated from the prior distribution.
C
BFN details
C.1 A Brief Introduction to Bayes Flow Network
BFN can be interpreted as a communication protocol between a sender and a receiver. The sender observes the ground-truth molecule and deliberately adds noise to obtain corrupted version y, which is transmitted to the receiver. Given the known precision level (e.g., a from a predefined schedule (()). the receiver performs Bayesian inference and leverages neural network to incorporate contextual information, producing an estimate of m. The communication cost at each timestep is defined the KL divergence between the sender's noising distribution and the corrupted output of the receiver's 's updated belief. By minimizing this divergence across timesteps, the receiver learns approximate the posterior and generate realistic samples from prior noise using the same learned inference process.
C.1.1 Training of BFN
Concretely, at each communication step t, the sender perturbs m using the Sender distribution (adding noise distribution) ps(y; m: a;) to produce noisy latent analogous to the forward process in diffusion models. The receiver then reconstructs (using Bayes updates and Neural Networks) m via the Output distribution:
Pm y:0) = ((0,-1,,,tt).
(17)
where + is a neural network which is expected to reconstruct the sample m given the Bayes-updated parameters 0,-1, conditioning code z and time t.. Then re-applies the same noising process obtain the receiver distribution ( 0;-1, z; t.) = Enapo Ps(y; m, az). The training objective is to minimize the KL. divergence (T.(pslppa each step, encouraging consistency between sender and receiver.
In practice, use Gaussian for Continuous data and categorical distribution for discrete data. Therefore Bayesian updates has closed form. The details can be found in the Appendix C.3 update distribution pv stems from the Bayesian update function h,
Bayesian
Pr 0 0 m, 2; a;) = Eysops (6, h(0; 110 Yis as)))
(18)
where 5(.) is Dirac delta distribution. This expectation eliminates the randomness of the sent sample from the sender.
According to the nice additive property of accuracy (26). the best prediction of m up to time 4 is the Bayesian flow distribution pp which could obtained by adding all precision parameters:
pr m. z; (4) = Pu (0; Co, m, z; ((t)), where B(t.) = Σ aj.
(19)
JSt
Therefore, the training objective for n steps is to minimize:
L (m, z) (E,U((1))) Ey 6: 1 DKL PR).
(20)
24

---

C.1.2 Inference of BFN
During inference, the sender is no longer available to provide noisy samples to help the receiver improve its belief. However, as training minimizes DKL (ps pn). Thus, can reuse the same communication mechanism by iteratively applying the learned receiver distribution pn to generate samples.
Given prior parameters Oo. accuracies o1,,..., and corresponding times t, = i/n, the n-step sampling procedure recursively generates 0 ......
......, On by sampling x' from po(- 0;-1, 2, 4- y from ~ PR( 0 z, li-1- 0s),, then setting 0, - 1(0,,1, 2, y), and pass the result to the neural network. The final sample is drawn from po 0.. z, 1).
This recursive procedure enables BFN to generate molecule from simple prior, guided solely by the learned receiver network and the latent codes. Algorithm. !
However, explicitly sampling x' from po(- 0,-1, 2, 14-1), particularly for discrete data could introduce unnecessary noise and impair the stability of the generation process. Instead, following we directly operate in the parameter space, avoiding noise injection from sampling and enabling more deterministic and efficient inference. See Algorithm,
37 49
Algorithm 1: Inference of General Bayes Flow Networks
Input: Initial prior parameter Oo. noise schedule {ax)?=n timestep grid {4, = i/n) conditioning latent code z
Output: Final molecular sample Xfinal
for i to do
Sample x' ~ - 0;-1, 2, 14-1)
Sample y ~ PR( 01-1, f 11-1: a,)
Update latent state: 0, + 10, - 2.
Apply the network: 0g $ (0,, 2, ti)
Sample final output: Xfinal ~ p On, z, 1)
C.2 MoIFLAE Reconstruction Loss
BFN can be trained by minimizing the KL-divergence between noisy sample distributions. BFN allows training discrete time and continuous time, and for efficiency we adopt the n-step discrete loss.
Given the ground truth molecule m [x, v] and its latent code z, we can have the reconstruction loss
Lrecon LT + C..
(21)
The above two summands coordinates loss and atom type loss:
Since the atom coordinates and the noise are Gaussian, the loss can be written analytically as follows:
C DrI. (N(x, (M(x.00)))|
a: |x x(,
(22)
The atom type loss can also be derived by taking KL-divergence between Gaussians [26]. assuming DM is the number of atom types, NM is the number of the atom:
C. In (y"
aassDDee - -1),a;DM1)
Nu
Dur Pook 0:1)N (oo (ac,ee-11),,,,aa 1),aDa)))
Σn d=
In
(23)
k=1
25

---

Together with two training losses, we can summarize the forward pass. In practice, we use reconstruction and regularization loss weights to get the final loss, see Table. T0
Algorithm 2: Forward Pass MolFLAE
Input:
Molecule M (XM. Var, Number of Virtual nodes Nz Output: Reconstruction loss Lreon, Regularization Lieg
Introduce Nz virtual nodes:
2 x Vz)
equivvarra encoding:
(-. 2 2,, < (M, 2))
VAE parameterisation:
M, Ze. _ Mh: o.l Linear(Z;)
Get regularization loss Lreg
(Bq. 5
Sample latent code:
f Ma pb. O Un.ohhe c~N10
BFN decoding (Alg. 3):
i f- BFNDecode
Get reconstruction loss Lrecon
(Eq 6
return Lrecon, Lreg
C.3 Bayes Updating Function for Molecular Data
Continuous coordinates The receiver observes the noisy input yd and the corresponding noise level
a. Starting from prior parameterized by 0 {14-1: (p-1)- According to Bayes' rule, it updates its belief by the bayes updating function h 1 ac;) = (is.
P: Pi-1 +
(24)
p. 1i-1 a.y"
(25)
Pt
Discrete atom types Upon receiving the noisy signal y" and the noise factor af, the receiver updates its belief by applying the Bayes formula with the previous parameters 0 0 The Bayes updating function is:
h: 1 MO0-1.y",oo a())
(26)
k= evi 1)s
where c denotes element-wiss wise multiplication.
C.4 MolFLAE Decoding Process
In the inference phase of BFN, is in principle possible to draw samples from the receiver distribution to perform Bayesian updates. However, explicitly sampling such intermediate pariables-partiullle for discrete data-car can introduce unnccessary stochasticity and impair the stability of the generation process. Instead, following [26] 44), we directly operate in the parameter space, avoiding noise injection from sampling and enabling a more deterministic and efficient inference.
To implement this approach, we define (() := T-ク(t) Let m [&, v) denote the neural network's output at given step, where v represents the continuous (pre-sor roff-oof logits for atom types. Instead of sampling noisy observations explicitly, we directly use mn to update the parameters for the next iteration, thereby bypassing the stochastic sampling step in the standard Bayesian update 0, h00 ax).
3(t)
Under this formulation, the Bayes Flow parameter updates simplify to:
PF (μ X 2)() A' (m 7()) &, (() - -(()) I)
(27)
(6" V. 22) (No
(0 softmax(y"))
(28)
26

---

Here, the continuous coordinates are updated using a closed-form Gaussian posterior, while the discrete atom types are updated via an expected categorical distribution induced by a softmax over noisy logits. In practice, this expectation is approximated with single Monte Carlo sample.
Throughout the generation process, updates performed entirely in the parameter space, avoiding noisy sampling exeps- for the final decoding stage, where an actual molecular structure is drawn from the output distribution.
Algorithm 3: Decoder: sampling Molecules conditioned on latent code
Input: Network 4, latent code (R-Nx(3-Dz), total steps N, number of atoms NM, number of types DM. noise levels a1 81 Output: Sampled molecule [&, V]
/* Define update function
Function update (& E TRNN V NarxDN B(t), (t), EER R+): (
A(t) I-0(c)
7
~ N(x, (1 1)
y" Nyy ((t)(Deee ())
0" (oftmax((y"))())))
return 0"
/* Initialize parameters
0, 0"
Dis Na DMM
for i = 1 do
t-
Sample &, v ~ po 2², z.
Update latent: μu, 0" + update($, v1,1,
/* Final sampling
Sample &, V- po (H, 0", 1)
return & |]
D Molecules Property Metrics
Following the setup of EDM[13) in unconditional generation experiments, employed the following metrics to evaluate the quality generated molecules:
Atom Stability: Proportion of atoms with valid bond counts.
Molecular Stability: Proportion of molecules where all atoms stable.
Validity: Proportion of molecules with RDKit-parsable SMILES.
Novelty: Proportion of molecules whose SMILES are not in the training set.
Uniqueness: Proportion of unique molecules in the generated set.
In our interpolation experiments, to fully assess the chemical and physical properties of intermediary molecules generated during interpolation, we employed the following metrics:
Similarity Preference: Defined as follows:
Similarity Preference S S
S + S
where S₁ and S, denote the Tanimoto similarity (calculated using MACCS fingerprints 45 to the target and source molecules.
sp3frac: Represents the proportion of 1p3-hybridized carbon atoms in a molecule relative to the total number of carbon atoms.
topological index based on molecule's structure, considering factors like atomic connectivity, ring size, and number.
BertzCT [50]:
27

---

QED [5]]: Evaluates the similarity of molecule to known drug molecules by considering multiple physicochemical properties.
Labute ASA : Measures the surface area each atom a molecule contact with the solvent, reflecting the molecule's 's solvent interaction ability.
TPSA [53] value calculated based on molecule's 's topological structure and atomic polarity, used to predict solubility and biological membrane penetration.
logP [54 : The logarithm of the partition coefficient of an organic compound between octanol and water, indicating the hydrophobicity of a molecule.
MR [54 Measures molecule's ' ability to refract light, related to factors like polarizability, molecular weight, and density.
E
Results on drug-likeness metrics on GEOM-Drugs
To provide more comprehensive evaluation, we conduct experiments on the GEOM-Drugs dataset with additional drug likeness metrics, where MolFLAE showed significant improvements over existing baselines. The results are summarized in Table It should be noted that we were unable to locate open resources for GeoBEN's checkpoints on GFOM-Drugs, so experiments on GeoBEN were not conducted.
In details, sample 10,000 molecules from each model with their default settings, obtaining atom positions and types, and then inferred bond types using OpenBabel. We then fix the bond order using Schordinger due some bugs in OpenBabel. The final molecules are then be evaluated using RDKit for the following metrics: QED, SA, Lipinski, and Strain Energy.
Table 7: Comparison of models QED, SA, Lipinski and Strain Energy on GEOM-Drugs dataset.
Metrics
QED (t)
SA (1)
Lipinski (f)
Strain Energy (L)
Data
0.64
0.84
4.80
80.19
EDM
0.36
0.59
4.26
705.2
GeoLDM
0.40
0.63
4.31
446.1
UniGEM
0.36
0.63
4.24
490.4
MoIFLAE
0.60
0.75
4.75
84.66
F
Ablation Study
We conduct ablation experiments on three key design choices of MolFLAE: the presence of the regularization loss, the feature embedding dimension Dz, and the number of virtual nodes Nz. All models are evaluated using 100-step decoding. Unless otherwise noted, all hyperparameters follow the configuration in Table 10
Table summarizes the results on QM9, Each row varies only one component while keeping all other settings fixed.
We do ablation Regularization loss, embedding dimension Dz and length of latent codes Nz- Our decoder uses 100 steps for sampling. And the MolFLAE means the same hyperparameters as in Table IO. The following three experiments keeps the same settings except for the marked one. Removing the regularization loss leads to drop in molecular stability and novelty, suggesting that
Table 8: Ablation study, performance comparison of different model config on QM9.
Model Contig (steps=100)
Atom Sta (%)
Mol Sta (%)
Valid (%)
VxU (%)
Novelty (%)
MoIFLAE
99.39
92.01
96.81
88.94
74.49
w/o Regularization Loss
98.73
85.01
97.82
86.43
66.04
Dz 16
98.90
87.05
93.09
94.19
70.32
N
99.21
89.24
96.60
73.21
61.54
latent smoothness is crucial for robust generation. Reducing the latent dimensionality (Dz 16)
28

---

or the number of latent nodes (N2 5) also impacts overall performance, particularly in terms of novelty and reconstruction fidelity.
Moreover, as our iterative decoding process is a little complex, we tried to simplify it to a one-step decoding variant. However, this approach failed to generate any valid molecules.
G Hyper-parameter Settings
Hyperparameters for training on QM9, GEOM-DRUG and ZINC-9M are listed in Table I0 We followed prior works in the choice of network structure, and carried out ablation study determine the number of latent codes introduced.
Table 9; Training costs.
Dataset
GPUs
Time
Max Epoch
GEOM-DURG
Nvidia A A100s(80G)
days
15
QM9
Nvidia A (100s(80G)
16h
250
ZINC-9M
Nvidia A800s(80G)
3 days
25
Table 10: Hyperparameteers for training.
Parameter
Value or description
69214221/99//rmmiinnng data for GEOM-DRUG
Train/Val/Test Splittigg
93322660/932rreminnng data for ZINC-9M
10000017748/eemining data for QM9
Batch size
100 for GEOM-DRUG,200 for ZINC-9M,400 for QM9
Optimizer
Adam
31
0.95
B2
0.99
Lr
0.005
Weight decay
0
Learning rate decay policy Learning rate factor
Reducel_ROnPlann
0.4 for GEOM-DRUG, 0.6 for QM9 and ZINC-9M 3 for GEOM-DRUG, 10 for QM9 and ZINC-9M
Patience
Min learning rate
1.00E-06
Embedding dimension D,
128
Head number
16
Layer number
9
k (knn) Activation function
32
ReLU
Nz
10
Dz
32
var, varı
100
Reconstruction loss weight
1
Regularization loss weight
0.1
H Broader Impacts and Safety Discussion
Our work develops an autoencoder model for molecular design, which has potential positive societal impact in areas such drug discovery, materials science, and green chemistry by enabling the
29

---

efficient generation of candidate molecules with desired properties. However, we acknowledge that may also be misused, for instance generate harmful or toxic compounds. While our model is trained and evaluated on general-purpose datasets without any bias toward hazardous compounds, we emphasize that any downstream deployment should include domain-specific safeguards, such as toxicity filters and expert oversight.
30


Reviews:
[{"review_id": "259aY9tAAw", "reviewer": "Reviewer_dc5U", "summary": "This paper introduces a VAE for 3D molecule generation with a fixed-size equivariant latent space. It allows flexible manipulation of molecules in a zero-shot way. The model shows good results on generation, editing, and drug optimization tasks. The fixed-length latent space enables interpolation and disentanglement.", "questions": "-- How sensitive is the method to the number of virtual nodes?\n-- Could the learned latent space be used for other molecular tasks? Which ones?\n-- Are the interpolation results robust across different dimensions of variation (molecular types, etc)?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n-- Novel latent space design for SE(3)-equivariant molecular representations.\n-- Strong experimental results on several datasets.\n-- Demonstrates useful molecule editing tasks (e.g., atom count change, interpolation).\n-- Includes a real drug design use case with docking evaluation.\n\nWeaknesses\n-- Latent disentanglement is only partial and not deeply explored. Doing more here would be very intereting.\n-- Method is complex; requires many architectural components. Can it be simplified and better ablated?\n-- Limited comparison to recent unified editing models. Would really like to see more here.", "was_revised": true, "final_justification": "I thank the authors for a good response.\nThe response further supports my evaluation.", "rebuttal": "We thank the reviewer for the thoughtful and encouraging feedback.  Below, we address your concerns and questions.\n\n>W1:  Latent disentanglement is only partial and not deeply explored. Doing more here would be very interesting.\n\nWe appreciate the reviewer’s insightful observation regarding the extent of latent disentanglement. Our current model does not incorporate any explicit supervision signal or inductive bias specifically enforcing disentanglement between invariant and equivariant features. While partial disentanglement emerges naturally from our architecture and training objective, we agree that further exploration would be valuable. One promising direction is to employ a contrastive learning objective that encourages consistency of the invariant representation across different conformations of the same molecule. This could explicitly encourage the model to separate invariant and equivariant features more effectively. We thank the reviewer for highlighting this important area and will pursue this line of investigation in future work.\n\n\n---\n>W2: Method is complex; requires many architectural components. Can it be simplified and better ablated?\n\nWe thank the reviewer for pointing out the complexity of the method. Our approach is fundamentally built upon a standard VAE framework, consisting of an encoder, a latent space with regularization, and a decoder. The main architectural difference lies in the decoder, where we introduce an iterative decoding mechanism (Bayesian flow network, BFN) similar to diffusion models. This modification is critical for generating valid molecular structures. We experimented with a simplified, one-step decoding variant (i.e., without iterative decoding), but it failed to produce any valid molecules. Due to the poor quality and instability of this baseline, we excluded it from the main ablation results to maintain focus on more informative comparisons. Nonetheless, we acknowledge that including such a baseline, even as a failure case, could be instructive, and we will consider adding it in the appendix for completeness. In summary, the current architecture reflects the minimal design necessary to achieve effective and valid molecular generation, and we are actively exploring ways to further streamline the components without compromising performance. \n\n\n---\n>W3: Limited comparison to recent unified editing models. Would really like to see more here.\n\nThank you for requiring more discussion on other unified editing models. To the best of our knowledge, MolFLAE is the first fixed-length equivariant molecular autoencoder capable of unified latent space manipulation for 3D molecules. This includes zero-shot operations like atom editing, structure inpainting, and property-guided interpolation, all within the same latent space.\n\nA close prior work is the SMILES-based VAE [1], which enables continuous latent space editing but lacks any modeling of 3D structure or equivariant properties. Another recent approach [2] introduces a fixed-length autoencoder but relies on global pooling over atom features, discarding spatial and geometric details critical for reconstruction. UAE-3D [3] has proposed a 3D fixed-length latent space by discarding the inductive bias of geometric equivariance. Furthermore, both [2] and [3] adopt evaluation settings different from ours and do not provide open-source code, making direct comparison difficult.\n\nWhile most prior editing methods target specific tasks (e.g., linker design [4], scaffold inpainting [5], and Gradient-guided or property-based optimization [6,7]), they are not designed as general-purpose frameworks. Notably, MolFLAE can handle such tasks under a unified interface (e.g., see linker design and superposition in Figure 6), and we will explore broader evaluations in our future works.\n\nIn summary, we struggled to find editing models with the same settings as ours for a fair comparison. MolFLAE is a powerful Auto-Encoder, which means a broader application to other editing tasks.\n\n[1]Gómez-Bombarelli et al, “Automatic Chemical Design Using a Data-Driven Continuous Representation of Molecules.” ACS Central Science 4 (2): 268–76, 2018. \n\n[2]Tianxiao Li, Martin Renqiang Min, Hongyu Guo, and Mark Gerstein. 3d autoencoding diffusionmodel for molecule interpolation and manipulation, 2024.\n\n[3]Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Kenji Kawaguchi, Tat-Seng Chua, and XiangWang. Towards unified latent space for 3d molecular latent diffusion modeling, 2025.\n\n[4]Ilia Igashov et al, Equivariant 3d-conditional diffusion models for molecular linker design, 2022.\n\n[5]Shicheng Chen et al, Deep lead optimization enveloped in protein pocket and its application in designing potent and selective ligands targeting LTK protein. Nat. Mac. Intell., 7(3):448–458, 2025.\n\n[6]Keyue Qiu et al, Empower structure-based molecule optimization with gradient guidance, 2025.\n\n[7]Alex Morehead and Jianlin Cheng. Geometry-complete diffusion for 3d molecule generation and optimization. ArXiv, 2023.\n\n\n\n---\n>Q1:  How sensitive is the method to the number of virtual nodes? \n\nWe appreciate the reviewer’s question regarding the sensitivity of our method to the number of virtual nodes. We investigated this by reducing the number of virtual nodes ($N_Z$) from 10 to 5 on the QM9 dataset. As shown in Table 6 (Appendix E) in our paper, this setting yields a molecular stability of 89.24% and validity of 96.60%. While slightly lower than the performance with 10 nodes, these results are still strong and indicate that the method is relatively robust to moderate changes in the number of virtual nodes. Lower numbers of virtual nodes correspond to a smaller latent space capacity, which may restrict the diversity of generated molecules. As a result, the model tends to produce outputs more similar to the training data, potentially leading to lower novelty scores. We show Table 6 here for convenience:\n| Model Config (steps=100)|Atom Sta (%)|Mol Sta (%)|Valid (%)|V×U (%)|Novelty (%)|\n|-|-|-|-|-|-|\n|**MolFLAE**|**99.39**|**92.01**|96.81|88.94|**74.49**|\n|w/o Regularization Loss| 98.73| 85.01| **97.82** | 86.43| 66.04|\n|$D_Z=16$|98.90|87.05|93.09|**94.19**|70.32|\n|$N_Z=5$|99.21|89.24|96.60|73.21|61.54|\n\nWe acknowledge that more virtual nodes may be required for modeling larger and more complex molecular structures, such as peptides or proteins, due to the inherent compression capacity of limited virtual nodes. However, recent work in image generation, such as TiTok [1], demonstrates that even high-resolution images can be compressed into extremely compact latent representations (e.g., 32 tokens) while preserving generation quality. These results suggest promising scalability potential for our approach as well, motivating future extensions toward more efficient tokenization strategies for large biomolecules.\n\n[1] Qihang Yu et al, An image is worth 32 tokens for reconstruction and generation. NeurIPS, 2024.\n\n---\n>Q2: Could the learned latent space be used for other molecular tasks? Which ones?\n\nThank you for this thoughtful question. Yes, the learned latent space in our model offers promising utility beyond unconditional generation and can potentially support a range of downstream molecular tasks:\n1. Conditional Generation: By applying diffusion or flow-matching techniques in the latent space, we can perform conditional latent code generation with given pocket structures. This could enable controllable design with high validity and structural fidelity, because our decoder is trained on much more diversified molecules than the current conditional molecular generation datasets.\n2. Property Prediction: The compact invariant latent representations can be used as input features for supervised property prediction tasks, acting as learned molecular embeddings that capture both global and local structural information.\n3. Language-Guided Molecular Generation: Perhaps the most exciting direction is leveraging the 1D latent codes as a “molecular language” compatible with large language models (LLMs). This opens the door to natural language-driven molecular generation, akin to the recent BAGEL model [1] or BLIP3-o [2], which integrates diverse visual modalities into pretrained LLMs like Qwens. This approach could unify molecular generation with broader language-based interfaces.\n\nWe see significant promise in these directions and plan to explore them further in future work.\n\n[1] Deng, Chaorui, et al. \"Emerging properties in unified multimodal pretraining.\" arXiv:2505.14683 (2025).\n\n[2] Chen, Jiuhai, et al. \"Blip3-o: A family of fully open unified multimodal models-architecture, training and dataset.\" arXiv:2505.09568 (2025).\n\n\n\n---\n>Q3: Are the interpolation results robust across different dimensions of variation (molecular types, etc)?\n\nThank you for the insightful question regarding interpolation robustness. We would appreciate further clarification on what is meant by \"molecular types\".  Does it refer to categories like peptides, antibodies, or other biomolecules? They are not in our current training set.\n\nIn our study, we perform interpolation experiments across diverse molecules with different sizes and structures (see Figure 4, Table 4, and Table 5). These experiments demonstrate that our fixed-length latent space enables relatively smooth transitions in both structure and properties, indicating the robustness and regularity of the latent representation.\n\nIf you could provide further clarification on \"molecular types,\" we would be happy to elaborate on our results in more detail.\n\n---"}, {"review_id": "QTsANuM4w1", "reviewer": "Reviewer_EbXx", "summary": "This paper introduces a novel VAE architecture for learning 3D molecules. It uses learnable virtual nodes to achieve a fixed number of latent nodes. Thus, MolFLAE has an SE(3)-equivariant latent space with fixed dimension independent from the number of atoms, allowing for easier interpolation or extrapolation. The decoder of architecture employs a Bayesian Flow Network. The author evaluated MolFLAE on various downstream tasks, demonstrating model efficacy. The MolFLAE also exhibits the ability to learn disentangled spatial and semantic latent components.", "questions": "1. In formula (3), the model discards the embedding of molecule M and only retains the embedding of virtual nodes Z. Since the number of virtual nodes is fixed-length, will the performance of the model vary for molecules with different numbers of atoms, particularly for very large molecules? \n2. In the GEOBFN paper, the best performance was achieved when sampling the molecules with 2000 steps. Did the author try to compare the performance of MolFLAE when sampling with 2000 steps? \n3. In drug discovery, it is crucial to generate drug-like molecules. Why are metrics like drug-likeness (QED) not used to evaluate the model with the GEOM-drug dataset?", "limitations": "Yes.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. It is a common problem in molecular VAE where the number of atoms varies across molecules. The author provides a novel approach by introducing virtual nodes to effectively achieve a fixed-dimensional latent space. \n2. The authors demonstrate that the MolFLAE can generate similar molecules with the same latent code, despite varying numbers of atoms, thereby illustrating the smoothness of model generation. The disentanglement of spatial and semantic features also shows the possibility of generating molecules with desired properties while maintaining certain constraints. \n3. The paper demonstrated the model's performance in optimizing molecules targeting the human glucocorticoid receptor (hGR). This real-world task shows the effectiveness of the model in drug design applications.\n\n**Weaknesses**\n1. Table 1 shows that MolFLAE outperforms baselines across different metrics. However, it seems the improvement is a bit marginal. The author should have more comparisons and discussions between MolFLAE and baselines other than unconditional molecule generation.", "was_revised": true, "final_justification": "The author has addressed all my questions. The quality of the work meets the standard for a NeurIPS acceptance. However, I still believe the performance boost is a little bit marginal, and I am not completely satisfied with the significance of the work. Thus, I will retain my score as it is.", "rebuttal": "Many thanks to your valuable comments and questions, which help us a lot to improve our work. We address your questions as follows.\n>W1.1:  Table 1 shows that MolFLAE outperforms baselines across different metrics. However, it seems the improvement is a bit marginal. \n\nThank you very much for your insightful comments. We truly appreciate your thoughtful feedback and suggestions.\n\nWe acknowledge that while MolFLAE demonstrates improvements over baseline models in Table 1, these enhancements may appear relatively modest. This is likely because the QM9 tasks are relatively simple. To provide a more comprehensive evaluation, we conducted experiments on the GEOM-Drugs dataset with additional drug likeness metrics, where MolFLAE showed significant improvements over existing baselines. We sincerely appreciate your Question 3 about drug likeness, which inspires us to do these experiments.\n|  Model & sampling steps | QED(↑)   | SA(↑)    | Linpinski(↑) | Median Strain Energy(↓) |\n| ------------- | ------ | ------ | ---------- | --------------------- |\n| data| 0.64 | 0.84 | 4.80     | 80.19|\n| EDM 1000| 0.36 | 0.59 | 4.26     | 705.2|\n| GeoLDM 1000| 0.40 | 0.63 | 4.31     | 446.1|\n| UniGEM 1000 | 0.36 | 0.63 | 4.24     | 490.4|\n| MolFLAE 100 | **0.60** | **0.75** | **4.75**     | **84.66**|\n\nIt should be noted that we were unable to locate open resources for GeoBFN's checkpoints on GEOM-Drugs, so experiments on GeoBFN were not conducted. We will provide the detailed results of these experiments in the camera-ready version.\n\nAlso, it is important to clarify that our primary objective is not to build a state-of-the-art unconditional generative model. Instead, we aim to demonstrate that a fixed-dimensional, equivariant latent space can effectively capture molecular semantics. This capability is crucial for a wide range of downstream tasks, such as molecular modification and optimization.\n\n>W1.2 The author should have more comparisons and discussions between MolFLAE and baselines other than unconditional molecule generation.\n\nTo further contextualize our results, we provide an additional comparison against several recent structure-based drug design (SBDD) models on standard drug-like property metrics including QED, SA and Lipinski score, as shown below. These results indicate that our model could serve as a strong decoder for latent diffusion models in the field of SBDD, providing better drug likeness and synthesizability.\n| Model                        | QED↑     | SA↑      | Linpinski↑|\n|-----------------------------|----------|----------|------------|\n| TargetDiff（from UniMoMo ） | 0.49     | 0.60     | 4.57       |\n| MolCRAFT（from UniMoMo ）   | 0.48     | 0.66     | 4.39       |\n| UniMoMo                     | 0.55     | 0.70     | 4.68      | \n| MolFLAE（100 steps）        | **0.60** | **0.75** | **4.75** |\n\nHowever, we respectfully note that direct comparisons to conditional generation models may not be entirely appropriate, as they often operate on different datasets.\n\nThank you again for your valuable feedback. We are committed to improving this work and greatly appreciate your guidance.\n\n\n---\n>Q1:  In formula (3), the model discards the embedding of molecule M and only retains the embedding of virtual nodes Z. Since the number of virtual nodes is fixed-length, will the performance of the model vary for molecules with different numbers of atoms, particularly for very large molecules?\n\nThank you for your question regarding the potential performance variation of our model for molecules with different numbers of atoms, especially large molecules. We conduct additional experiments to evaluate reconstruction similarity and validity across different heavy atom counts. The results are summarized in the table below:\n\n| Heavy Atom Count         | [0, 20) | [20, 30) | [30, 40) | [40, 50) | [50, +∞) |\n|--------------------------|---------|----------|----------|----------|----------|\n| Averaged Similarity(%)   | 94.16  | 88.62   | 73.43   | 65.99   | 49.30   |\n| Valid(%)                  | 99.19  | 99.80   | 100.0   | 100.0   | 100.0   |\n| Proportion in Train Set(%)| 12.17  | 58.43   | 26.39   | 2.89   | 0.12   |\n\nWe observe that larger molecules tend to be reconstructed with lower similarity. This is expected for two reasons:  \n1. The training set contains very few large molecules, making it difficult for the model to learn their distribution.  \n2. As shown in our ablation study (Table 6 in Appendix E), the number of virtual nodes would slightly affect modeling performance. Larger molecules require more virtual nodes for accurate reconstruction.\n\nIn Computer Vision (CV), extensive compression is often feasible [1,2,3], but it typically relies on a strong decoder. In contrast, our decoder is relatively lightweight, a 9-layer Transformer with only 2.83M parameters. Our model is a prototypical design rather than a fully optimized solution. While modeling large molecules remains a critical and challenging task, we believe our framework provides a promising foundation, and we are keen to explore its extension in this direction in future work.\n\n[1] Qihang Yu et al, An image is worth 32 tokens for reconstruction and generation. NeurIPS, 2024.\n\n[2] Chen, Jiuhai, et al. \"Blip3-o: A family of fully open unified multimodal models-architecture, training and dataset.\" arXiv:2505.09568 (2025).\n\n[3] Sun, Quan, et al. Emu: Generative pretraining in multimodality. ICLR, 2024.\n\n---\n>Q2: In the GEOBFN paper, the best performance was achieved when sampling the molecules with 2000 steps. Did the author try to compare the performance of MolFLAE when sampling with 2000 steps?\n\nThank you very much for your question regarding the sampling steps in our work and the comparison with GEOBFN. \n\nWe have evaluated the performance of MolFLAE across various sampling steps, including 500, 1000, and 2000 steps, and compared it with GEOBFN’s performance at 2000 steps.\n\nIn our experiments, we found that when MolFLAE was sampled at just 100 steps, its performance was already comparable to that of GEOBFN at 2000 steps and significantly outperformed GEOBFN's results at 100 steps. However, while increasing the sampling steps did lead to improvements in Atom Stability, Mol Stability, and Validity, these improvements were rather marginal, and we noticed a decline in Uniqueness and Novelty. \n| Model & sampling steps | Atom Sta(%)      | Mol Sta(%)| Valid(%)| U×V(%)| Novelty(%)   |\n|------------------------|------------------|----------------|----------------|----------------|-----------|\n| Geobfn 100             | 98.64       | 87.21    | 93.03    | 91.53    | 70.3      |\n| Geobfn 2000            | 99.31            | 93.32          | 96.88          | **92.41**      | 65.3      |\n| MolFLAE 100            | 99.39            | 92.01          | 96.81          | 88.94          | **74.49** |\n| MolFLAE 500            | 99.58            | 94.52          | 98.18          | 80.17          | 68.86     |\n| MolFLAE 1000           | **99.61**        | 94.81          | 98.38          | 75.57          | 71.05     |\n| MolFLAE 2000           | **99.61**        | **94.99**      | **98.49**      | 72.47          | 65.37     |\n\nThese observations and the balance between computation efficiency and performance, combined with the fact that our original goal was not to develop a powerful unconditional generation model but rather to create a model capable of flexible molecular modification, led us to present the 100-step results in our paper. The 100-step results effectively demonstrated the model's capabilities without necessitating a large number of sampling steps.\n\n\n\n---\n>Q3: In drug discovery, it is crucial to generate drug-like molecules. Why are metrics like drug-likeness (QED) not used to evaluate the model with the GEOM-drug dataset?\n\nThank you for bringing up the important point regarding the evaluation metrics used in our model with the GEOM-Drugs dataset.\n\nIn our experiments with the GEOM-Drugs dataset, we initially followed the setup of EDM, GeoLDM and UniGEM, which primarily focused on testing Validity and Atom Stability. In response to your inquiry, we have further evaluated our model as well as these three baselines with additional metrics, including drug-likeness (QED), synthetic accessibility and Lipinski's rule of five, where we significantly outperform baseline methods. Detailed results are in answer to W1.1. \n\nWe sincerely hope this information addresses your concerns and provides a clearer picture of our model's performance. Thank you once again for your valuable feedback.\n\n---"}, {"review_id": "6XTKReiAna", "reviewer": "Reviewer_fHNV", "summary": "This paper presents MolFLAE, an SE(3)-equivariant variational autoencoder (VAE) designed for 3D molecular structures. To ensure a fixed-dimensional latent representation, it introduces a fixed number of virtual nodes. A Bayesian Flow Network is then used to reconstruct molecules from these latent codes.\n\nMolFLAE is evaluated across a variety of tasks, demonstrating its versatility in applications such as:\n- Unconditional molecular generation\n- Analog generation with varying atom counts\n- Mixing invariant and equivariant codes of distinct molecules\n- Interpolating between source and target molecules\n- Optimizing drug candidates for hGR\nThe model also shows promise for broader molecular applications, including molecule inpainting and structural alignment.", "questions": "1. Have the authors considered extending the model with a diffusion process in the latent space (i.e., using a latent diffusion model)? I’d be interested to hear the authors’ thoughts on how this might improve performance.\n2. It seems that molecular validity is reported only for the unconditional generation task. I may have missed it, but were validity (and/or stability) checks also performed for the downstream applications?\n3. Beyond the noted issues with disentanglement and interpretability, are there other limitations the authors would like to highlight?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "# Strengths\n- The paper makes a good technical contribution by introducing a more effective method to capture rich molecular features while maintaining a fixed-dimensional latent representation. Instead of relying on simple pooling, it introduces multiple virtual nodes to achieve this.\n- It demonstrates the versatility of MolFLAE through a wide range of molecular tasks, addressing a key limitation of prior models that were typically designed for narrow or specific subtasks. \n\n# Weaknesses\n- According to Appendix A.1, the authors appear to use E(3)-equivariance rather than SE(3)-equivariance, which may limit the model's ability to capture chirality. Please correct me if I’ve misunderstood this point.\n- While the paper offers novel contributions, certain components are similar to prior work, such as UniMoMo [1] and PepGLAD [2], which also model the molecules with E(3)-equivariant VAE. Although MolFLAE may be the first to apply these ideas to 3D small molecules, acknowledging related efforts in other molecular domains would help better contextualize the work and strengthen its positioning.\n- (Minor) While I'm familiar with the molecular domain and the techniques used, the evaluation metrics for 3D small molecules are less familiar. A brief explanation -- e.g., in Appendix D -- of how metrics like atom stability, molecular stability, validity, and novelty are defined would be helpful.\n\n## Typos\n- Line 107: inculding -> including\n- Line 191: Alogorithm -> Algorithm \n- Line 192: modelities -> modalities \n- Line 227-228: Shouldn’t $z_h^1$ correspond to the substructure information?\n- Line 574: Incomplete sentence: “And pass the Neural Network”\n- Equation (12) in Appendix B: The equation $[\\sigma_x^2, \\mu_h, \\sigma_h^2] = \\text{Linear}([z_x, z_h])$ appears misleading, as the actual implementation is $[\\sigma_x^2, \\mu_h, \\sigma_h^2] = \\text{Linear}(z_h)$, without dependence on $z_x$, in order to preserve SE(3)-equivariance.\n\n## References\n[1] Kong, Xiangzhe, et al. \"UniMoMo: Unified Generative Modeling of 3D Molecules for De Novo Binder Design.\" arXiv preprint arXiv:2503.19300 (2025).\n[2] Kong, Xiangzhe, et al. \"Full-atom peptide design with geometric latent diffusion.\" arXiv preprint arXiv:2402.13555 (2024).", "was_revised": true, "final_justification": "The authors have thoroughly addressed my concerns, so I am keeping my acceptance score. However, I did not raise it to 6 due to current limitations, such as the need for atom count during decoding and limited scalability to larger molecules.", "rebuttal": "Dear reviewer fHNV, \n\nThank you for pointing out the typos to help us improve the quality of our work. We sincerely apologize for these oversights. We will fix them in the camera-ready version. We respond to your concerns one by one.\n>W1:  According to Appendix A.1, the authors appear to use E(3)-equivariance rather than SE(3)-equivariance, which may limit the model's ability to capture chirality. Please correct me if I’ve misunderstood this point.\n\nThank you for pointing out this important detail. Our model is E(3)-equivariant, and a careful examination of the proof confirms that it does not require the orthogonal transformation to have determinant 1. We acknowledge that referring to it as SE(3)-equivariant throughout the paper may be misleading. We will revise this terminology in the camera-ready version. \n\nFrom a mathematical perspective, because SE(3) is a subgroup of E(3), any E(3)-equivariant model is also SE(3)-equivariant. We followed prior works like MolCRAFT [1] and TargetDiff [2] in using SE(3), since we adopt the same network architectures.\n\nImportantly, this distinction does not affect the correctness of our results. E(3)-equivariance is more suitable for generative tasks like molecule editing, as it preserves the chirality of latent codes and molecules. SE(3)-only equivariance (i.e., excluding reflections) would be more appropriate for chirality-sensitive property prediction, which is beyond our current scope.\n\nMoreover, our main baselines including GeoBFN[3], EDM[4], and UniGEM[5] are also based on E(3)-equivariant networks. Thus, our comparisons are on fair ground.\n\n[1] Yanru Qu et al, MolCRAFT: Structure-based drug design in continuous parameter space. ICML, 2024.\n\n[2] Jiaqi Guan et al, 3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction. ICLR, 2023.\n\n[3] Yuxuan Song et al, Unified generative modeling of 3d molecules with bayesian flow networks. ICLR, 2024.\n\n[4] Emiel Hoogeboom et al,  Equivariant diffusion for molecule generation in 3d. ICML, 2022.\n\n[5] Shikun Feng et al,  UniGEM: A unified approach to generation and property prediction for molecules. ICLR, 2025.\n\n\n\n---\n>W2: While the paper offers novel contributions, certain components are similar to prior work, such as UniMoMo [1] and PepGLAD [2], which also model the molecules with E(3)-equivariant VAE. Although MolFLAE may be the first to apply these ideas to 3D small molecules, acknowledging related efforts in other molecular domains would help better contextualize the work and strengthen its positioning.\n\nWe thank the reviewer for raising this important point. While MolFLAE shares high-level similarities with UniMoMo [1] and PepGLAD [2] in leveraging E(3)-equivariant VAEs, our approach differs significantly in motivation, architecture, and application scope.\n\nWe have already discussed PepGLAD in the related work section under Unconditional Generation for 3D Molecules, and we will include UniMoMo in the camera-ready version as they are quite similar. Both UniMoMo and PepGLAD compress each molecular block (e.g., an amino acid residue) into a single latent node. This design has two key limitations. First, the number of latent nodes depends on the molecular composition, making cross-sample comparisons and operations like interpolation difficult. Second, the spatial relationships between molecular blocks are tightly preserved, restricting the flexibility of the latent space for generative modeling, e.g. when latent codes are used for autoregressive modeling by finetuning LLMs.\n\nIn contrast, MolFLAE encodes entire molecules into a fixed-length 1D latent sequence, where virtual nodes are associated with a learnable embedding to assign an intrinsic order. This greatly facilitates applications such as molecular interpolation and molecule-to-molecule translation (optimization).\n\nIn comparison with current vision-language models, MolFLAE adopts a CLIP-style encoder with a diffusion-like decoder, akin to recent advances in vision-language models (e.g., Emu2 [3] and BLIP3-o[4]), while UniMoMo or PepGLAD use pixel-level reconstructions. Our latent codes are not atom-aligned but rather semantically rich and highly compressed, similar to how CLIP-style models encode images. This makes our approach more suitable for downstream semantic tasks and aligns with the emerging paradigm in multimodal generation. This paradigm has been proven to be powerful in transforming non-1D modalities like images into 1D tokens.\n\nWe will ensure these distinctions are more clearly highlighted in the related work part of the camera-ready version.\n\n[1] Kong, Xiangzhe, et al. \"UniMoMo: Unified Generative Modeling of 3D Molecules for De Novo Binder Design.\"  arXiv:2503.19300 (2025). \n\n[2] Kong, Xiangzhe, et al. \"Full-atom peptide design with geometric latent diffusion.\" arXiv:2402.13555 (2024).\n\n[3] Sun, Quan, et al. \"Generative multimodal models are in-context learners.\" CVPR, 2024.\n\n[4] Chen, Jiuhai, et al. \"Blip3-o: A family of fully open unified multimodal models-architecture, training and dataset.\" arXiv:2505.09568 (2025).\n\n\n---\n>W3: While I'm familiar with the molecular domain and the techniques used, the evaluation metrics for 3D small molecules are less familiar. A brief explanation -- e.g., in Appendix D -- of how metrics like atom stability, molecular stability, validity, and novelty are defined would be helpful.\n\nThank you very much for your insightful comments. To be self-contained, we will extend Appendix D in the camera-ready version. \nHere is a precise metric definition (all in [0,1]):\n- Atom Stability: Proportion of atoms with valid bond counts.\n- Molecular Stability: Proportion of molecules where all atoms are stable.\n- Validity: Proportion of molecules with RDKit-parsable SMILES.\n- Novelty: Proportion of molecules whose SMILES are not in the training set.\n\n---\n>Q1:  Have the authors considered extending the model with a diffusion process in the latent space (i.e., using a latent diffusion model)? I’d be interested to hear the authors’ thoughts on how this might improve performance.\n\nThank you for the insightful question. We agree that extending MolFLAE with a latent diffusion model is a promising direction. Since our model provides a high-quality molecule decoder, the generation task can be naturally framed as generating latent codes. This offers an opportunity for controllable molecule design with improved validity and structural fidelity, especially because the decoder has been trained on more diversified molecules than typical conditional molecular generation datasets.\n\nFurthermore, your suggestion opens up exciting avenues for future work. For example, integrating latent diffusion or latent flow matching could enable more precise and diverse molecular edits. These techniques may also facilitate conditional generation with strong generalization.\n\nWe appreciate your valuable suggestion and will actively consider this line of research in future iterations of our work.\n\n---\n>Q2: It seems that molecular validity is reported only for the unconditional generation task. I may have missed it, but were validity (and/or stability) checks also performed for the downstream applications?\n\nThank you very much for your careful observation regarding the validity checks. We sincerely apologize for not including them in our initial submission. To address this, we have conducted additional experiments to evaluate the validity and stability of all molecules involved in the three downstream tasks. Here are the detailed results:\n### Generating Analogs with Different Atom Numbers (Figure 2)\n|Atom Number|-2|-1|0|1|2|\n|-|-|-|-|-|-|\n|Valid(%)|100.0|99.89|99.76|99.89|99.68|\n|Atom Sta(%)|84.58|83.28|82.48|82.38|82.53|\n\n### Exploring the disentanglement of the latent space via molecule reconstruction (Figure 3)\n||$(z_{h_{0}},z_{x_{0}})$|$(z_{h_{1}},z_{x_{1}})$|$Preserving\\ z_{x}$|$Preserving\\ z_{h}$|\n|-|-|-|-|-|\n|Valid(%)|100.0|100.0|100.0|100.0|\n|Atom Sta(%)|80.90|80.81|85.20|84.62|\n\n### Latent Interpolation (Table 4 & 5)\n||Start|2|3|4|5|6|7|End|\n|-|-|-|-|-|-|-|-|-|\n|Valid(%)|99.90|99.90|99.90|100.0|99.90|100.0|100.0|100.0|\n|Atom Sta(%)|82.23|82.95|84.17|84.57|84.21|83.74|83.10|82.63|\n\n||Start|2|3|4|5|6|7|8|9|End|\n|-|-|-|-|-|-|-|-|-|-|-|\n|Valid(%)|100.0|99.90|99.90|99.80|99.80|100.0|100.0|100.0|99.80|99.90|\n|Atom Sta(%)|82.23|82.75|83.53|84.07|84.36|84.35|84.16|84.17|82.78|82.67|\n\n||Start|2|3|4|5|6|7|8|9|10|11|End|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|Valid(%)|99.90|100.0|99.90|100.0|100.0|100.0|100.0|100.0|99.90|99.90|100.0|99.80|\n|Atom Sta(%)|82.39|82.50|83.76|83.67|83.92|84.65|84.20|84.50|84.27|83.00|82.84|82.57|\n\nThe validity is high among all experiments. Thank you once again for your valuable feedback and for bringing this to our attention.\n\n---\n>Q3: Beyond the noted issues with disentanglement and interpretability, are there other limitations the authors would like to highlight?\n\nWe appreciate the reviewer’s suggestion to reflect on further limitations. We highlight the following points:\n1. Atom Count Specification During Decoding: Our current decoder requires specifying the number of atoms. Although this is common in diffusion-based molecule generation, it could be eliminated in our setting by a small atom count predictor conditioned on the latent code.\n2. Broader Validation of Editing Capabilities Needed: Figure 6 shows preliminary editing results, but we have not yet systematically evaluated MolFLAE across diverse editing tasks. We plan to explore this in future work.\n3. Scalability to Large Molecules: Our current experiments focus on small to medium-sized molecules. Extending the framework to larger structures such as proteins is a promising but non-trivial extension.\n\n\n---"}]

Author rebuttal:

Meta-review: All reviewers assigned positive scores for this paper, although two reviews are still in the "borderline" regime. 
In all reviews, it was mentioned that this work is of significant novelty, that the experimental setup is convincing, and that this work is relevant for  designing equivariant molecular representations. In particular, one reviewer mentioned that  the experiment al result nicely demonstrate the effectiveness of the proposed model in drug design applications. 

There were also some concerns, mainly about the potential of latent disentanglement,  the model's ability to capture chirality, and some similarity to prior work. In the rebuttal, however, the authors could address most of these concerns in a way that seemed to be plausible for the reviewers. I also think that most of the critical remarks have been addressed in a convincing way by the authors, and therefore I recommend acceptance of this paper.
Decision: Accept (poster)