Towards Unified and Lossless Latent Space for 3D Molecular Latent Diffusion Modeling
Authors: Yanchen Luo
Zhiyuan Liu
Yi Zhao
Sihang Li
Hengxing Cai
Kenji Kawaguchi
Tat-Seng Chua
Yang Zhang
Xiang Wang
Venue: NeurIPS 2025 poster ()
Keywords: 3D molecule generation
3D variational auto-encoder
latent diffusion model
Primary area: machine_learning_for_sciences

Abstract: 3D molecule generation is crucial for drug discovery and material science, requiring models to process complex multi-modalities, including atom types, chemical bonds, and 3D coordinates. A key challenge is integrating these modalities of different shapes while maintaining SE(3) equivariance for 3D coordinates. To achieve this, existing approaches typically maintain separate latent spaces for invariant and equivariant modalities, reducing efficiency in both training and sampling. 
In this work, we propose **U**nified Variational **A**uto-**E**ncoder for **3D** Molecular Latent Diffusion Modeling (**UAE-3D**), a multi-modal VAE that compresses 3D molecules into latent sequences from a unified latent space, while maintaining near-zero reconstruction error. This unified latent space eliminates the complexities of handling multi-modality and equivariance when performing latent diffusion modeling. We demonstrate this by employing the Diffusion Transformer--a general-purpose diffusion model without any molecular inductive bias--for latent generation. Extensive experiments on GEOM-Drugs and QM9 datasets demonstrate that our method significantly establishes new benchmarks in both *de novo* and conditional 3D molecule generation, achieving leading efficiency and quality. On GEOM-Drugs, it reduces FCD by 72.6% over the previous best result, while achieving over 70% relative average improvements in geometric fidelity. Our code is released at [https://github.com/lyc0930/UAE-3D/](https://github.com/lyc0930/UAE-3D/).

Towards Unified and Lossless Latent Space for 3D Molecular Latent Diffusion Modeling
Yanchen Luo', Zhiyuan Liu Yi Zhao', Sihang Li', Hengxing Cai², Kenji Kawaguchi", Tat-Seng Chua², Yang Zhang?, Xiang Wang"
'University of Science and Technology of China
National University of Singapore
DD Technology
luoyanchen@mail..c ustc.edu. edu. cn
acharkq@gmail,co com
ciangwang1220@mail.com
Abstract
3D molecule generation is crucial for drug discovery and material science, requir- ing models process complex multi-modalities, including atom types, chemical bonds, and 3D coordinates. A key challenge is integrating these modalities different shapes while maintaining SE(3) equivariance for 3D coordinates. To achieve this, existing approaches typically maintain separate latent spaces for in- variant and equivariant modalities, reducing efficiency in both training and sam- pling. In this work, we propose Unified Variational Auto-Encoded for 3D Molec- ular Latent Diffusion Modeling (UAE-3D), multi-modal VAE that compresses 3D molecules into latent sequences from a unified latent space, while maintaining near-zero reconstruction error. This unified latent space eliminates the complex- ities of handling multi-modalits and equivariance when performing latent diffu- sion modeling. We demonstrate this by employing the Diffusion Transformer-a general-purpose diffusion model without any molecular inductive bias-for latent generation. Extensive experiments on GEOM-Drugh and QM9 datasets demon- strate that our method significantly establishes new benchmarks in both de novo and conditional 3D molecule generation, achieving leading efficiency and quality. On GEOM-Drugs, it reduces FCD by 72.6% over the previous best result, while achieving over 70% relative average improvements in geometric fidelity. Our code is released at https://ght wttp:///ggttt.....//eee/......bt httpttbbbccm//yy099//0//33//
Introduction
The discovery of novel molecules is fundamental to various scientific fields, particularly in drug and material development. Given significant progress has been made in designing 2D molecular graphs [[]-А], recent research has increasingly focused on the generation of 3D molecules [5, 6). Unlike 2D molecular generation, which focuses on forming valid molecular structures based on chemical bonds, 3D generation must also predict 3D atomic coordinates that align with the 2D structure. Accurate 3D molecule generation is essential to power many important applications, such as structure-based drug design [M] and inverse molecule design targeting quantum properties [X]]
3D molecule generation is challenging due to its multi-modal nature. As shown in Figure Ilaa), a 3D molecule consists of features of three distinct modalities: atom types, atomic coordinates, and chem- ical bonds. This requires the generation model to handle both discrete (e.g., atom types) and con- tinuous features (e.g., coordinates), while also addressing differences in feature shapes (atom-wise
CCoresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Malecule
Invariant
Bond Recuncy 25.2%
Aceuneyy
RMKDO
equivariont
UAC30 (9unk)
RINSO 0L08
Gloms
Board Accuracy
100
56
26
1:0
011
Molecular Space Encoder Lotent Space Decodem
Alom/oond type acturney (96 (J) Coordi-ata RMSD (A).
coordinatis
xna
fidd
Ged.oV
X
X7.
bonds
J0D0
x2.7
lossless
UPW-5D (3h000
(a)
255
400
0.00
0,25
0.50
0.75
(6) Trairing Time (hrs).
(4) Bampping Time (stmolocul)).
Figure 1: (a) A 3D molecule has multi-modal features. (b)
Figure 2:
Comparing UAE-3D and UDM-3D with other methods on the QM9 dataset. (a:b) Reconstruction er- rors on the test set. (c:d) Comparing training and inference time.
Prior methods use separate latent spaces for equivariant (3D) and invariant (2D) modalities, inducing unnecessary complexity for model architecture. (c) UAE-3D reduces
this complexity by establishing unified and near-lossless latent space that integrates all molecular modalities.
v.S. edge-wise features). Worse still, the modality of 3D coordinate requires special care to respect rotational and translational equivariance, whereas other modalities do not, further complicating the task. Mitigating this latent spaces: some maintain separate diffusion processes for each modality [R] g Ii], while oth- ers interleave the prediction of different modalities across autoregressive generation steps [12,, [3] However, we argue that handling each modality separately complicates the model design, reducing both the training and sampling efficiency. Moreover, processing each modality separately risks com- promising the consistency between them. These issues raise a critical question: Can we design a unified generative model that seamlessly integrates all three modalities of 3D molecule generation?
challenge, prior works mostly process each modality in separate
To answer this research question, we propose build multi-modal latent diffusion model (LDM) [14]]] for the unified generative modeling of 3D molecules. LDM extends the diffusion paradigm by operating in compressed latent space learned through variational autoencoders (VAEs) [ offering improved computational efficiency and generation quality. For 3D molecules, we can build a multi-modal VAE that compresses all three modalities in single unified latent space. Scrutinizing the previous works, we identify the following challenges to achieve this purpose:
The Challenge of Compressing Multi-modality in One Latent Space. While previous studies have explored LDMs for 3D molecule and protein generation [IZ I8], they fail to build a unified latent space that integrates all modalities. Their difficulty arises from the reliance on neural net- works with baked-in 3D equivariance [19 16], which mostly maintain separate latent spaces for 3D equivariant and invariant modalities (cf. Figure [(bb)..
The Challenge of Ensuring Near-Lossless Compression. Ensuring low reconstruction error for the molecular VAE is critical, as even small eтrors can result in invalid or unstable 3D struc- tures. Furthermore, imprecise latents that cannot reproduce the original molecule can propagate their errors to the LDM, easily disrupting the generative modeling. Despite its critical importance, reconstruction error has been largely overlooked in previous works [117] [X] (cf. Figure (aa:))).
In this work, we introduce UAE-3D, Unified Variational Auto-Encoder for 3D Molecular Latent Dif- fusion Modeling, a VAE that can compress the multi-modal features of 3D molecules into a unified latent space while maintaining near zero (100% atom/bond accuracy and 2E-4 coordinate RMSD) reconstruction error. To obtain a latent space integrating both 3D equivariant and invariant modal- ities, we draw inspiration from the "bitter lesson" [20] of 3D molecules: rather than using models with intricate, baked-in 3D equivariance, UAE-3D trains a neural network to "Iearn" t 3D equivari- ance through our tailor-made SE(3) augmentations, encouraging the transformation on the input coordinates to be reflected equivariantly on the output coordinates. Moreover, UAE-3D employs the Relational Transformer [[[]] as its encoder, leveraging its scalability and flexibility to incorporate both atom-wise and edge-wise features. A Transformer-basee [22]] decoder is jointly trained with the encoder to reconstruct both 3D invariant and equivariant molecular features.
Despite its simplicity, UAE-3D can compress 3D molecules into token sequences in a unified latent space, eliminating the complexities of handling multi-modalities and 3D equivariance in latent diffusion modeling. Figure (aab))) show that UAE-3D reaches near-lossless reconstruction of atom and bond types, with a near-zero coordinate RMSD. To further demonstrate its effectiveness,

---

employ the Diffusion Transformer (DiT) [X] general-purpose diffusion backbone without any molecular inductive bias, to model UAE-3D's latents. We show that DiT can successfully generate stable and valid 3D molecules, with significantly improved training (by 2.7 times) and sampling (by 7.3 times) efficiency (ef. Figure Dic:))). We refer to this LDM pipeline as UDM-3D: Unified Latent Diffusion Modeling for 3D Molecule Generation.
Extensive experiments demonstrate that UDM-3D, powered by UAE-3D, achieves state-of-the results in both de novo and conditional 3D molecule generation the QM9 [ and GEOM- Drugs [25] datasets. On GEOM-Drugs our model achieves a 72.6% reduction in Fréchet ChemNet Distance [3], while significantly reducing the MMD) of bond length, bond angle, and dihedral angle by 88.4%, 55.6%, and 74.0%, respectively. In conditional generation on QM9, it achieves the lowest MAE in five out of six properties. These results confirm the effectiveness of our unified modeling approach. Ablation studies show the effectiveness our key components. Further analysis reveals that UAB-3D's latents present structured variations across geometric movement.
2 Background: 3D Molecular Latent Diffusion Models
A 3D molecular LDM involves 3D molecular VAE to compress 3D molecules into the latent space, where a diffusion model performs generative modeling. Below, we introduce both components.
Notations. A 3D molecule is represented by G = (F, E, X), where F E RRI is the atom feature matrix (e.g,, atomic numbers), E E / is the bond feature matrix (e.g., bond connectivity and type), and X E R|V]*3 is the 3D atom coordinate matrix. d1 and d2 are dimensions of atom and bond features. V is G's set of atoms and | denotes the number of atoms.
3D Molecular Variational Auto-Encoding. A molecular VAE consists of an encoder & and a decoder D. Given a 3ID molecule G, the encoder & maps it into sequence of latent tokens E(G) Z [z, E i E V). Each latent z; is sampled from Gaussian distribution W(z,; M4, o;) using the reparameterization trick [IG]] where μ4 and o; are learned parameters for atom i. The decoder D then reconstructs G from these latents: G = D(Z). The VAE is trained with a reconstruction loss and regularization term:
LVAR = EGrppaan
[IG - -GG DKI. (kI
where DKL. denotes the Kullback-Leiin Leibler divergence; (Z|G) Ilevv N(2;; M4; σx) is the approx- imated posterior distribution; and ||G -Gl measures the difference between the original and the
predicted graph, whose definition varies across different works. In [[Z]], || - Gll combines an
MSE loss for 3D coordinates and a cross-entropy loss for atom types.
Diffusion Model. Building on the molecular VAE's latent Z, a diffusion model [1E]] Ad] performs generative modeling. In forward diffusion process, gradually add noise the original latent g = Z following ( (z z)) (1 a where E (0, 1] the diffusion timestep, and ((*) is a hyperparameter controlling the signal/noise ratio. In practice, Z is sampled
as (t)) = V (z)) + v1 - a(t)e where e ~ NN, Given the noised latents (()) a diffusion model co (t), el1 Once trained. new latents (0) are sampled with the model Eœ via iterative denoising [15], and
()),) t) is trained to predict the added noise e by minimizing the MSE loss (Z t)
the decoder D reconstructs the corresponding 3D molecule via G = (0))
Separated Latent Spaces. The VAE latents in previous works [I8. include two parts: Z [Z: 2). where Z for equivariant features and Z for 3D invariant features. They also rely on diffusion models with baked-in 3D equivariance. In contrast, as Section B shows, our method uses unified latent space, allowing a purpose diffusion model- any geometric or molecu- lar inductive bias-to achieve strong performance.
3 Methodology
In this section, we propose the UAE-3D, a multi-modal variational auto-encoder designed to ef- fectively compress the diverse modalities of 3D molecules into a unified latent space. Based on UAE-3D, we introduce UDM-3D, an LDM for 3D molecule generation.
3

---

MLP
R- Trans
DiT
MLPadeLNN
GBF
MLP
xt
Attn
adaLN
Molecular Space
Unified Lossless
Candition
Latent Space
rrecmn
MLP
MLP
Transformer
DiT
Derarorg
biffusion Process
(-1)
2
2
Figure 3: Overview of the UDM-3D and UAE-3D models. The UAE-3D encodes 3D molecules from molecular space into unified latent space, integrating multi-modal features such as atom types, chemical bonds, and 3D coordinates. Utilizing this latent space, UDM-3D employs DiT to perform generative modeling. Then, the denoised latents are decoded back into 3D molecules.
3.1 Unified Variational Auto-Encoder for 3D Molecular Latent Diffusion Modeling
UAE-3D is designed to address the complexities of modal and equivariance of molecular data by compressing the atomic features, bond features, and atomic coordinates into unified latent space. This is achieved through three key components: (1) a Relatinoal-Transfornee [2]]] that effectively integrates the multi-modal features into token sequences, (2) Transformer-base [22] decoder and
the reconstruction loss for 3D molecule reconstruction; and (3) our tailor-made SE(3)-equivarian data augmentations train the model to learn 3D equivariance.
Compressing 3D Molecules with Relational Transformer (R-Trans) Encoder. Given a 3D molec- ular G (F, E, X), we compute its initial atom-wise and edge-wise embeddings as:
H MLP (XX: (X;
F)) E Mixd
(2)
H MLP((E;DD DJ) E R
(3)
where D X -
R4; GBB(-)) implements the Gaussian basis functions to expand inter-atomic distances as feature vectors [2Z]]: and is the embedding dimension. The initial node embeddings H" combines both atomic features and its 3D coordinates, while the edge embeddings H" incorporate bond features and inter-atomic distances. H and H together effectively represent all 3D molecular information for the subsequent encoding. Given these embeddings, we process them with the molecular encoder {.) of layers of R-Trans [2]]]:
H R-Trans HH" HH) E RRlixd
where H denotes the updated atom embedding for the next layer, computed follows:
Q H H.
W
Rd
(5) Kiy; Vyl = [H): H whk E =dd
Q Vd
aiy softmax,
E E.
(7)
H MLP Σ oig Vo)
where W and Wey are linear projectors. Unlike the original R- R-Trans, our implementation does not update the edge H every layer, which improves efficiency without compromising performance.
Compared to Transformer [22]] R- R-Trans is well-suited for 3D molecule encoding, that it can effec- tively integrate edge embeddings H E IR [ and atom embeddings H E RDV despite their different shapes. This integration occurs during the computation of queries, keys and values (Equation (5) and (G)), ensuring the output atom embeddings fully incorporate edge information.
Decoder. Given encoder's latents Z, employ Transformer [ decoder D(.) to obtain the atom-wise output P Transformer(()) (Z) E IR 1iixd, Unlike the encoder, our decoder includes no molccule-speciallzze design, because it processes sequences of latents. Further, we employ three MLP predictors for atom, bond types, and atom coordinates to reconstruct complete 3D molecules:
X, ML.P (P.) E R F, ML.P2(PP) E No; É.g MLPP + +P,) E N.N. (9)
where N. and N₂, are the number of atom and bond types.

---

Reconstruction Loss. We define the distance between our reconstructed graph and the original G||. It includes multiple components to ensure the fidelity of generated 3D molecules:
graph ||G
futtom
CE(F,, F.):
Leuordinate
MSE(X; X;),
[V]
M
LEV
LEV
(10)
Looud
1
Σ ..EV
C Egy),
Ldistance
1
we MSE D.)).
FEEV
where CE denotes cross-entropy loss, and Latom and Lhond are losses for atom and bond types; Ccoordinate measures the MSE between the predicted and the ground truth coordinates; and Edistance serves as an extra constraint on coordinate predictions, to ensure correct inter-atomic distances. No- tably, we incorporate a bonded distance loss inside Luistance by adjusting the pairwise weight wy.:
{ + X, if atom and j are bonded,
Wij
(11)
11, otherwise.
This design prioritizes accurate inter-atomic distances between bonded atoms, enforcing stricter geometry in chemically important regions. À is hyperparameter that controls this constraint's 's strength. Given the components above, we define the distance between our reconstructed graph and the original graph (i.e., the reconstruction loss) as follows:
frecon EG- G G EG Y [Fratom: Lroond Levondinate, Letistance] T
(12)
where y E R. is a hyperparameter vector balancing the reconstruction terms. The final loss for UAE-3D combines Lreson with KL regularization term:
LEAE3D = Lreeon + B RE [Pkt.
(13)
where the hyperparameter B E controls the KL regularization strength.
Equia--qu Augmentations. To enforce -quivvaaia in UAE-3D, we apply random
transformations R E SE(3) to input coordinates X during training [28. 20]. Each R combines a rotation from SO(3) and a random translation from (O, 0.01Ia). Crucially, the reconstruction loss Lreson (cf. Equation (I2) remains unchanged but operates transformed inputs R G (F, E, R()). This process can be informally conceptualized as a reconstruction loss of (R G)) R. G|. where the autoencoder D(E(.) learns to preserve geometric consistency under SE(3) transformations.
3.2 Unified Latent Diffusion Modeling for 3D Molecule Generation
Diffusion Transformer (DiT). Given the unified latent space provided by our UAE-3D model, we adopt the DiT [233 as the backbone diffusion model c. for 3D molecular latent generation. Originally developed for image LDM, DiT has demonstrated strong performance in modeling latent sequences. Compared to a standard Transformer [22], DiT replaces the layernorm [29] by adaptive layernorm (adaLN) [B0], where the scale Rdd and shift b E Rd parameters conditioned on the diffusion timestep t E (0, 1]. Specifically. ly: b] = MLP maLLN (Embed, (t)), where Embed (.) is a shared linear layer and MLP is module specific to each adaLN instance. These generated shift and scale parameters act as "soft gates", enabling representations and facilitating more effective, scale-adaptive denoising.
activation of DiT's hidden
Conditional Generation with Classifier- vector we extend the adaL.N modules by combining the condition embedding Embed, with the timestep embedding Embed, (t) when computing the scale y and shift parameters. This allows the diffusion model (Z) I, c) to incorporate conditioning information during denoising:
Guidance. To enable generation with condition
yy: b] MLP (..Paaa)x(Emeed((( + Embed,(c)).
(14)
To further enforce conditioning during inference, we employ the classifier-free guidance (CFG) [II]
to find Z that has high log p(c|Z). 3. By Bayes' rule, the gradient of this objective is Vz log p(c Z) x
Vz log (gg((Z) (Z|c) Vz (Z)) Following [32]], we can interprete the diffusion model co 's output as score functions, and therefore maximize p(c Z) using the modified denoising function:
( () = (1 + w w)e (Z c) - wep (Z()) ()),
(15)
5

---

Table 1: Performance of de novo 3D molecule generation GEOM-Drugs.
indicates results
reproduced using official source codes, while other baseline results are taken from [6].
2D-Metric
FCD4
AtomStable
MolStable
V&C
V&U
V&U&N
SNN
Frag
Scaf
Train
0.251
1.000
1.000
1.000
1.000
0.000
0.585
0.999
0.584
CDGS
22.051
0.991
0.706
0.285
0.285
0.285
0.262
0,789
0.022
JODO
2.523
1.000
0.981
0.874
0,905
0.902
0.417
0.993
0.483
MiDi
7.054
0.968
0.822
0.633
0.654
0.652
0,392
0.951
0.196
EQGAT-diff
5.898
1.000
0.989
0.845
0.863
0.859
0.377
0.983
0.161
UDM-3D, ours
0.692
72.0%
1.000
0.925
0.879
0.913
0.907
0.525
0.990
0.540
3D-Metric
FCDsD4
AtomStahle
MolStable
Bond length
Bond anglel
Dihedral angle]
Train
13.73
0.861
0.028
1.56E-04
1.81E-04
1.56E-04
EDM
31.29
0.831
0.002
4.29E-01
4,96E-01
I 144B-02
JODO
19.99
0.845
0.010
8.49E-02
1.15E-02
6.68F-04
MiDi
23.14
0.750
0.003
1.17F-01
9,57F-02
4.46F-03
GcoLDM
30.68
0.843
0.008
3.91E-01
4,22E-01
1.69E-02
EQGAT-diff
26.33
0.825
0.007
1.55E-01
5,21E-02
2.10E-03
UDM-3D, ours
17.36
13.2%
0.852
0.014
9.89E-03
88.45%
5,11E-03
55.09%
1.74E-04
74.0%
where w E [0, -20) hyperparameter controlling the guidance strength; eo (Z( c) is the conditional variant of cp, incorporating the property c; (Z( t) is the unconditioned variant that ignores c. To train cg. we randomly drop condition c with certain probability Purop- allowing to learn both conditional and unconditional distributions.
Experiments
In this section, we present the experimental results of our proposed unified molecular latent space approach (UAE-3D & UDM-3D) on 3D molecule generation tasks. We comprehensively evaluate UDM-3D's performance on de novo 3D molecule generation and conditional 3D molecule genera- tion with targeted quantum properties. Our experiments demonstrate significant improvements over state-of-the- -hrt methods while maintaining computational efficiency.
4.1 Experimental Setup
Datasets. Following prior works [SL K. 3], we conduct experiments on QM9 [22]] and GEOM- Drugs [25]] datasets. QM9 [2_]] contains 130k small organic molecules (up to heavy atoms) with quantum chemical properties. We use
train/val/test splits, following [S BJ. GEOM- Drugs [2S] is a pharmaceutical-scale dataset with 450k drug-like molecules (average 44.4 atoms, max 181 atoms). We use 80%/10% 100% splits and retain the lowest-energy conformer per molecule.
Baselines. For de novo 3D molecule generation. we compare UDM-3D with CDGS 6 JODO [R]. MiDi [ G-SchNet [122], G- G-SphereNet [G4]] EDM [S], MDM
[ES). GeoLDM [ EQGAT-
diff [36] SemlaFlow [B7]] and ADiT [[8]] For conditional generation. we additionally use baselines
of BEGSDE [8] and GeoBFN [39]. We do not report MDM's performance on GEOM-DRUGS because we are unable to reproduce reasonable results using their official code. Similarly, ADIT's performance on GEOM-DRUGS is omitted, because it is not available in their paper, code, and released checkpoint.
4.2 De Novo 3D Molecule Generation
Experimental Setting. Generating a 3D molecule involves generating the structural validity (atom/bond features) and accurate spatial arrangements (3D coordinates). Therefore, we adapt the comprchensive metrics from [6. 5]. We also report all the metries for molecules from training set, serving as an approximate upper bound for performance. Our evaluation metrics can be divided into two groups: (1) 2D Metrics: Atom stability, validity & completeness (V&C), validity & uniqueness (V&U), validity& & uniqueness & novelty (V&U&N) similarity to nearest neighbor (SNN), fragment similarity (Frag), scaffold similarity (Scaf), and Fréchet ChemNet Distance (FCD) [3; (2) 3D Met- rics: Atom stability, FCD, maximum mean discrepancy (MMD) [ for the distributions of bond lengths, bond angles, and dihedral angles. More experimental details are in Appendix C.I.
Table |I and Table present the performance of UDM-3D on the de novo generation task on QM9 and GEOM- -rugg datasets, respectively, Our model demonstrates leading performances in generat-

---

Table 2: Performance of de novo 3D molecule generation QM9. indicates results reproduced us- ing official source codes, while other baseline results are taken from [R]] Some ADÎTs evaluation metrics omitted due differences in evaluation protocols. We discuss these protocol differences and their implications in Appendix C
2D-Metric
FCD1
AtomStable
MolStable
V&C
V&U
V&U&N
SNN
Frag
Scaf
Train
0.063
0,999
0.988
0,989
0.989
0.000
0.490
0.992
0.946
CDGS
0.798
0.997
0,951
0,951
0,936
0.860
0.493
0.973
0.784
JODO
0.138
0.999
0.988
0.990
0.960
0,780
0.522
0.986
0.934
MiDi
0.187
0.998
0.976
0.980
0.954
0.769
0.501
0.979
0.882
EQGAT-diff*
2.088
0.999
0,971
0.965
0.950
0.891
0.482
0.950
0.703
SemlaFlow
0.863
0.995
0.949
0.857
0.821
0.821
0.124
UDM-3D, ours
0.130 5.80%
0.999
0.988
0.983
0.973
0.950
0.508
0.987
0.898
3D-Metric
FCD
AtomStahle
MolStable
Bond length.
Bond angle]
Dihedral angle)
Train
0.877
0.994
0.953
5.44F-04
4.65F-04
1.78--04
EDM
1,285
0.986
0.817
1.30E-01
1.82E-02
6,64E-04
MDM
4.861
0.992
0.896
2.74E-01
6.60E-02
2.39E-02
JODO
0,885
0.992
0.934
1,48E-01
1.21E-02
6.29E-04
GeoLDM
1.030
0.989
0.897
2.40E-01
1.00E-02
6.59E-04
MiDi
1.100
0,983
0.842
8.96E-01
2.08E-02
8.14E-04
EQGAT-diff*
1,520
0,988
0.888
4.21E-01
1.89E-02
1.24E-03
SemlaFlow
1.127
0.971
0.787
ADiT*
2.884
0,211
9.98F-01
3,38F-02
1.46F-03
UDM-3D,
0.881
0.993
0.935
7.04E-02
9.84E-03
3,47E-04
ing chemically valid and novel molecules and achieves state-of-the-art performance in most of the metries across both datasets. Two detailed key observations emerge from the results:
Leading Geometric Accuracy. Our unified latent space enables highly accurate 3D geometric modeling, key factor in realistic molecule generation. UDM-3D significantly outperforms base- lines distributional distances of bond lengths and angles, reducing errors by an order of magnitude. These improvements hold even for complex, drug-like molecules in GEOM-Drugswhere UDM-3D achieves a 25x lower bond length error than GeoLDM (9.89E-03 v.s. 3.91E-01). This leap in ac- curacy comes from UDM-3D's ability jointly optimize chemical and geometric constraints in unified latent space, maintaining consistency across molecular modalities.
Novel and High Quality Generation. UDM-3D achieves outstanding V&U&N scores of 0.907 on GEOM- Drugs and 0.950 on QM9, while achieving ofate-of-the-art 3D stability and comparable 2D stability to strong baselines such JODO. This highlights its ability to generate novel molecules without compromising structural quality. The key lies in the unified latent space, which jointly encodes atom types, bond types, and coordinates, enabling the model to capture complex interde- pendencies among molecular features.
4.3 Conditional 3D Molecule Generation
Experimental Settings. We evaluate conditional generation of molecules with target quantum prop- erties using the protocol from [S] [Z7]. Specifically, our target properties include Co. H4, a, CHомо- ELUMO, Ac, and we evaluate the Mean Absolute Error (MAE) between target and predicted proper- ties. More details on the properties and settings are provided in Appendix CM
Bable&rreport MAE results for conditional generation on QM9. UDM-3D achieves the lowest MAE in five out of six properties, showing strong capability in generating molecules that match target properties. On average, it reduces MAE by 42.2% compared to GeoLDM, with a notable 52.7% improvement in predicting the HOMO-LUMO gap (Ae). These results highlight the importance of our unified latent space that effectively integrates all molecular modalities during generation. While GeoLDM struggles to correlate all the modalities due to its fragmented latent spaces, our unified latent representation inherently preserves the interplay between 3D geometry, bonds, and electronic molecule-property property relationships.
characteristics, enabling precise modeling over molecule-
4.4 Training& & Sampling Efficiency
We further analyze the training and sampling efficiency of UDM-3D compared to baselines on the QM9 dataset. These experiments are performed with an NVIDIA A100 GPU.
7

---

Table 3: Mean Absolute Error (MAE) for conditional 3D molecule generation QM9.
Method
μ (D)
a (Bohr")
C.
cal K)
EHоMO
(meV)
ÉLUMO (meV)
Ae (meV)
U-Bound
1.613
8.98
6.879
645
1457
1464
L-Bound
0.043
0.09
0.040
39
EDM
1.123
2.78
1.065
371
601
671
EEGSDE
0.777
2.50
0.941
302
447
487
GeoLDM
1.108
2.37
1.025
340
522
587
GeoBFN
0,998
2.34
0.949
328
516
577
JODO
0.628
1,42
0,581
226
256
335
UDM-3D,
0.603
1.54 +-8
0.553
216
247
313
Table 4: Comparing DiT with Transformer and PerceiverlO for latent diffusion modeling. Models have the same depth and hidden size. A.S. - atom stability.
Metric
Perceiver10
(a) Translation.
(b) Rotation.
(c) Rot. + Trans.
A.S.ap
0.993
0.983
0.972
Figure 4: t-SNE visualizations of UAB-3D's latents under SE(3)
A.S.p
0,999
augmentations. (a) Translations along fixed direction. (b) Rota- tions along fixed axis. (c) Sequential rotations followed by trans- lations. Color gradients show increasing distances angles.
V&C
0,983
0.938
0.933
V&U
0.973
0.922
0,931
V&U&N
0,950
Training Efficiency. As Figure Ei()) shows, JDM-3D's training costs 52 hours (14h UAE-3D + 38h UDM-3D), which 5.3 times faster than GeoLDM and 2.7 times faster than JODO. This efficiency stems from: (1) Our decoupled training paradigm: UAE-3D first learns 3D molecule compression, allowing DiT to focus solely on generative modeling the compressed latents. (2) Efficient diffu- sion with DIT: Unlike previous molecular diffusion models that require complex architectures for molti-modality and equivariance, DiT's simple and highly parallelizable design allows for faster training and sampling in the unified latent space. Crucially, speedup does not compromise per- formance, as guaranteed by UAE-3D's ncar-lossless molecular reconstruction (ef. Figure ((a;))).
Sampling Speed. Figure ((d) shows that UDM-3D generates each molecule in just 0.08 which is 7.3x faster than EDM/GeoLLDM and 9.8x faster than JODO. This speed advantage comes from our unified latent space, which simplifies DiT's modeling and avoids complex neural architectures.
4.5 Analysis and Ablation Studies
UAE-3D's Latents Present Structured Variations across Geometric Movements. Figure @shows the t-SNE visualizations of UAE-3D's latent representation under different SE(3) augmentations. Each dot represents a molecular embedding in the latent space, and the color gradient from light to dark corresponds to an increase in the translation distance (middle) the rotation angle (right) for the same molecule. We observe that the same molecule's representations remain close in the latent space and the augmentation scale increases, the representation changes along consistent direction. Figure IT shows that the representations change moving directions when the augmentation is switched from rotation translation. These observations demonstrate that UAE-3D captures meaningful and structured variations in molecular representations across SE(3) augmentations.
3D Molecule Samples by UDM-3D. We present 3D molecules generated by our method case studies in Figure E The generated molecules are chemically valid and exhibit diverse and complex structures, showcasing UDM-3D's ability to generate realistic 3D molecular conformations.
Ablating the DiT Backbone. We compare DiT with alternative transformer-basee architectures for de novo 3D molecule generation on the QM9 dataset (cf. Table g) Specifically, we evaluate against vanilla Transformer and Perceiver10[aT]] a modern architecture designed for structured inputs and outputs. The results show that DiT consistently outperforms both baselines across all metries, which we attribute to its adaptive LayerNorm layers. These layers enable DiT to effectively handle data with varying noise scales, thereby improving diffusion performance. Moreover, when paired with alternative diffusion neural architectures such as Transformer or PerceiverIO, our UAE- 3D framework still achieves meaningful and comparable performances, demonstrating its robustness

---

Table 5: The performance influence of Seuuv-equivariant augmentations the QM9 dataset.
2D- 2D-Metric
FCD4
AtomStable
MolStable
V&C
V&U
V&U&N
SNN
Frag
Scaf
No aug.
0.581
0.995
0,947
0,950
0.935
0.921
0,492
0.971
0.799
Rot.
0.315
0,995
0,948
0,951
0.943
0.927
0.493
0.979
0.875
Trans.
0.202
0,999
0,986
0.980
0.967
0,944
0.507
0.981
0.884
+ Trans. + Rot.
0.130
0.999
0.988
0.983
0.973
0.950
0.508
0.987
0.898
3D-Metric
AtomStable PCDpp4
MolStable
Bond length)
Bond angle]
Dihedral angle]
No aug.
1.065
0.993
0.879
1.01E-01
1.18E-02
1.08E-03
+ Rot.
1.007
0,993
0.876
8, 8.12E-02
9.99E-03
3,90E-04
+ Trans.
0,902
0,989
0.896
7.08E-02 -02
1.04E-02
6.81E-04
TTrans. + Rot.
0.881
0.993
0.935
7.04F.002
9.84E-03
3.47F:004
Table 6: RMSD results when trained with different SE(3) augmentations on QM9.
Train with
RMSD (x100* xx00AA) AÅ)
No aug.
1.4
Rot.
1.1
Trans.
0.9
Figure 5: Visualization of random samples generated by UDM-3D on QM9.
Trans. + Rot.
to different backbone choices. While further hyperparameter tuning could potentially improve the performance of these alternatives, the results already highlight the advantage of DiT in our setting.
Squii-Equi Augmentations. Table shows the performance influence of various SE(3)- equivariant augmentations on the QM9 dataset. The results reveal that data augmentation improves overall performance by providing diverse geometric views during UDM-3D training.
In addition to the standard performance metrics, we further investigate UAE-3D's reconstruction error on the test set when trained under different 3D transformations. Table presents the RMSD of the reconstructed molecules. By including both rotation and translation augmentations, we observe significant reduction in reconstruction error (from 1.4 x 10-3 0.2 x 10-3), achieving the lowest RMSD. This shows that our SE(3) augmentations effectively enhance the model's 's ability learn geometric consistency by exploring more comprehensive 3D molecular space.
5 Related Works
3D Molecule Generation. Early efforts for 3D molecule generation focus autoregressive ap- proaches [12, IR B2I,, constructing 3D molecules sequentially by adding atoms or molecular frag- ments. With the success of diffusion models across various domains [22, 13, 93], they also become the de facto method for 3D molecule generation [5L, 8]. However, the early diffusion works can easily generate invalid molecules because of overlooking the bond information. To bridge this gap, subse- quent works [6L 9L [II] additionally consider bond generation by building a separate diffusion process. However, generating different modalities in separate diffusion processes unnecessarily complicates the model design, reducing efficiency. It also risks compromising the consistency between these modalities. To address this, UDM-3D performs generative modeling in UAE-3D's unified latent space that integrates all molecular modalities, improving efficiency and generation quality.
LDMs for 3D Molecule Generation. For efficiency. LDMs employ a VAE [[ to compress raw data into a low-dimensional latent space, where a diffusion model performs generative mod- eling [] This method has been very popular for image generation [23]. However, existing 3D molecular LDMs still face challenges in separated latent spaces. For example, GeoLDM [ com- presses atom features and 3D coordinates separately. failing to build unified latent space. Other works [5%, П4] process different molecular components (atoms/subbrapphs or atom types/coordinates) in separate channels, increasing modeling complexity. In 3D protein generation, PepGLAD [[S] compresses 1D sequences and 31) structures into two separate latent spaces. Our UDM-3D addresses these issues by employing unified latent space, significantly improved generation efficiency.

---

6 Conclusion and Future Works
In this paper, we propose UAE-3D to compress the multi-modal features of 3D molecules into unified latent space, and demonstrate the effectiveness of latent diffusion modeling on this space by introducing UDM-3D. By integrating atom types, chemical bonds, and 3D coordinates into single latent space, our model effectively addresses the inherent challenges of multi-modality and SE(3) equivariance for 3D molecule generation. Extensive experiments on GEOM-Drugs and QM9 confirm leading performances in both de novo and conditional 3D molecule generation, setting new benchmarks for both quality and efficiency. UAE-3D's success highlights the benefits of building unified latent space for molecular design. Moving forward, we will transfer this unified latent space to new modalities, including proteins and RNAs, and extend it to broader molecular design tasks, including structure-based drug design. While UAE-3D's latents are near-lossless compressions of the original molecule, it inspires the exploration of jointly modeling UAE-3D's latents with text sequences for text-guided 3D molecule generation, following [95--11], as well as joint modeling of molecular and biological text representations across broader domains, including proteins [52]] and single-cell transcriptomics [5:3].
Limitations
Scope. In this work, explore 3D molecule generation under the settings of unconditional gener- ation and conditional generation targetting at quantum chemical properties. Other molecule genera- tion tasks, such structure-based drug design [B4]], of the scope of this work. We leave this exploration to future works.
Molecule Size. Our method relies separate module decide each generated molecule's ber of atom, following prior works [R [Sl M. Specifically, we use the molecule size distribution measured the training dataset to sample new molecule size. While this method shows decent performance on the tested tasks, we conjecture this will be one of the bottlenecks when the tested task become based method for molecule generation, in which the generated molecule size is automatically con- trolled by the auto-regressive process.
complex. Resolving this issue might demand the incorporation of auto-regressive
Acknowledgements
This research/project is supported by the National Research Foundation, Singapore under its Na- tional Large Language Models Funding Initiative (AISG Award No: AISG-NMLL-22224000) the Ministry of Education (MOE T125 R2523993309 and MOE T2EP20125-0039) the Agency for Science, Technology and Research (A*STAR) (H25J6a0034), and the National Natural Science Foundation of China (62572449).
Any opinions, findings and conclusions recommendations expressed in this material are those of the author(s) and do not reflect the views of the National Research Foundation, Singapore.
10

---

References
[1] Wengong Jin, Regina Barzilay, and Tommi S. Jaakkola. Junction tree variational autoencoder for molecular graph generation. ICML, volume 80 of Proceedings of Machine Learning Research, pages 2328-2332 PMLR, 2018.
[2] Ross Irwin, Spyridon Dimitriadis, Jiazhen He, and Esben Jannik Bjerrum. Chemformer: pre-trained transformer for computational chemistry. Mach. Learn. Sci. Technol. 3(1):15022. 2022.
[3] Daniil Polykovskiy, Alexander Zhebrak, Benjamin Sanchez-Lengelingg Sergey Golovanov, Ok- tai Tatanov, Stanislay Belyaev, Rauf Kurbanov, Aleksey Artamonov, Vladimir Aladinskiy, Mark Veselov, Artur Kadurin, Simon Johansson, Hongming Chen, Sergey Nikolenko, Alan Aspuru-Guzik, and Alex Zhavoronkov. Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models. Frontiers in Pharmacology, 2020.
[4] Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: flow-based autoregressive model for molecular graph generation. In ICLR. OpenReview.net, 2020.
[5] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In ICML, volume 162 of Proceedings of Machine Learning Research, pages 8867-8887. PMLR, 2022.
[6] Han Huang. Leilei Sun, Bowen Du, and Weifeng Lv. Learning joint 2-d and 3-d graph diffu- sion models for complete molecule generation. IEEE Transactions on Neural Networks and Learning Systems, 2024.
[7] Zaixi Zhang, Yaosen Min, Shuxin Zheng, and Qi Liu. Molecule generation for target protein binding with structural motifs. In ICLR. OpenReview.ner, 2023.
[8] Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy-guided SDE for inverse molecular design. ICLR. OpenReview.net, 2023.
[9] Clément Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In ECML/PKDD (2), volume 14170 of Lecture Notes in Computer Science, pages 560-576. Springer, 2023.
[10] Tuan Le, Julian Cremer, Frank Noé, Djork-Amé Clevert, and Kristof T. Schütt. Navigating the design space of equivariant diffusion-based generative models for de novo 3d molecule generation. ICLR. OpenReview.net, 2024.
[11] Chenqing Hua, Sitao Luan, Minkai Xu, Zhitao Ying, Jie Fu, Stefano Ermon, and Doina Pre- cup. Mudiff: Unified diffusion for complete molecule generation. In LoG, volume 231 of Proceedings of Machine Learning Research, page 33. PMLR, 2023.
[12] Niklas W. A. Gebauer, Michael Gastegger, and Kristof Schütt. admaetry-adapted generation of 3d point sets for the targeted discovery of molecules. NeurIPS, pages 7564 7576,757,, 2019.
[13] Niklas W. A. Gebauer, Michael Gastegger, Stefaan S. P. Hessmann, Kristof Schütt. Inverse design 3d molecular structures with conditional generative neural networks. CoRR, abs/2109.0453 04824, 2021.
Müller, and
[14] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, pages 10674-10685. IEEE, 2022.
[15] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
[16] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In ICLR, 2014.

---

[17] Minkai Xu, Alexander S. Powers, Ron O. Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In ICMI, volume 202 of Proceedings of Machine Learning Research, pages 38592-38610. PMLR, 2023.
[18) Xiangzhe Kong, Yinjun Jia, Wenbing Huang, and Yang Liu. Full-atom peptide design with geometric latent diffusion. In NeurIPS, 2024.
[19] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In ICML, volume 139 of Proceedings of Machine Learning Research, pages 9323- 9332. PMLR, 2021.
[20] Yuyang Wang, Ahmed A. A. Elhag, Navdeep Jaitly, Joshua M. Susskind, and Miguel Ángel Bautista. Swallowing the bitter pill: Simplified scalable conformer generation. In ICML. OpenReview.net, 2024.
[21] Cameron Diao and Ricky Loynd. Relational attention: Generalizing transformers for graph- structured tasks. In ICLR. OpenReview.ne,, 2023.
[22] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz. Kaiser, and Illia Polosukhin. Attention all you need. In NIPS, pages 5998-6008 2017.
[23] William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, pages 4172-4182. IEEE, 2023.
[24] Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 2014.
[25] Simon Axelrod and Rafael Gómez-Bombarelli. GEOM: energy-annotated molecular confor- mations for property prediction and molecular generation. CoRR, abs/2006.05531, 2020.
[26] Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. CoRR, abs/2107.00630 2021.
[27] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2022.
[28] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pages 1-3, 2024.
[29] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E, Hinton. Layer normalization. CoRR, abs/1607.0644 2016.
[30] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoning with general conditioning layer. In AAAI, pages 3942-3951. AAAI Press, 2018.
[31] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. CoRR, abs/2207.12599 2022.
[32] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR. OpenReview.net, 2021.
[33] Han Huang, Leilei Sun, Bowen Du, and Weifeng Lv. Conditional diffusion based on discrete graph structures for molecular graph generation. In AAAI, pages 4302-4311. AAAI Press, 2023.
[34] Youzhi Luo and Shuiwang Ji. autoregressive flow model for 3d molecular geometry gener- ation from scratch. In International conference learning representations (ICLR), 2022.
12

---

[35] Lei Huang, Hengtong Zhang, Tingyang Xu, and Ka-Chun Wong. MDM: model for 3d molecule generation. AAAI, pages 5105-5112. AAAI Press, 2023.
molecular diffusion
[36] Tuan Le, Frank Noé, and Djork-Aréé Clevert. Representation learning on biomolecular struc- tures using equivariant graph attention. In LoG, volume 198 of Proceedings of Machine Learn- ing Research, page 30. PMLR, 2022.
[37] Ross Irwin, Alessandro Tibo, Jon Paul Janet, and Simon Olsson. Semlaflow efficient 3d molecular generation with latent attention and equivariant flow matching. In AISTATS, volume 258 of Proceedings of Machine Learning Research, pages 3772-3780. PMLR, 2025.
[38] Chaitanya K Joshi, Xiang Fu, Yi-Lun Liao, Vahe Gharakhanyan, Benjamin Kurt Miller, Anuroop Sriram, and Zachary W Ulissi. All-atom diffusion transformers: Unified generative modelling of molecules and materials. arXiv preprint arXiv:2503.03965 2025.
[39] Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei-Ying Ying Ma. Unified generative modeling of 3d molecules with bayesian flow networks. In The Twelfth International Conference on Learning Representations, 2023.
[40] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander J. Smola. A kernel two-sample test. J. Mach. Learn. Res., 13:723-7733 2012.
[41] Andrew Jaegle, Sebastian Borgeaud, Jean- cann-Bapistte Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier J. Hé- naff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver IO: A general architecture for structured inputs & outputs. In ICLR. OpenReview.net, 2022.
[42] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynane namiis.. In ICML, volume 37 of JMLR Workshop and Conference Proceedings, pages
JMLR.org, 2015.
[43] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans image synthe- sis. In NeurlPS, pages 8780-8794, 2021.
[44] Tian Bian, Yifan Niu, Heng Chang. Divin Yan, Junzhou Huang. Yu Rong, Tingyang Xu, Jia Li, and Hong Cheng- Hierarchical graph latent diffusion model for conditional molecule genera- tion. In Proceedings of the 33rd ACM International Conference on Information and Knowledge Management, pages 130-140. ACM, 2024.
[45] Carl Edwards, Tuan Manh Lai, Kevin Ros, Garrett Honke, Kyunghyun Cho, and Heng Ji. Translation between molecules and natural language. In EMNLP, pages 375-413. Association for Computational Linguistics, 2022.
[46] Zhiyuan Liu, Sihang Li, Yanchen Luo, Hao Fei, Yixin Cao, Kenji Kawaguchi, Xiang Wang, and Tat-Seng Chua. Molca: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. In EMNLP, 2023.
[47] Zhiyuan Liu, Yanchen Luo, Han Huang, Enzhi Zhang, Sihang Li, Junfeng Fang, Yaorui Shi, Xiang Wang, Kenji Kawaguchi, and Tat-Seng Chua. NEXT-MOL: 3d diffusion meets Id language modeling for 3d molecule generation. In The Thirteenth International Conference on Learning Representations, 2025.
[48] Sihang Li, Zhiyuan Liu, Yanchen Luo, Xiang Wang, Xiangnan He, Kenji Kawaguchi, Tat-Seng Chua, and Qi Tian. 3d-molm: Towards 3d molecule-text interpretation in language models. In ICLR, 2024.
[49] Yanchen I Luo, Junfeng Fang, Sihang Li, Zhiyuan Liu, Jiancan Wu, An Zhang, Wenjie Du, and Xiang Wang. Text-guided small molecule generation via diffusion model. iScience, 27(11), 2024.
13

---

[50] Zhiyuan Liu, Yaorui Shi, An Zhang, Sihang Li, Enzhi Zhang, Xiang Wang, Kenji Kawaguchi, and Tat-Seng Chua. Reactxr: Understanding molecular reaction-ship via reaction- contextualized molccule-text pretraining. In Findings of the Association for Computational Linguistics: ACL 2024, Association for Computational Linguistics, 2024.
[51] Liang Wang. Chao Song, Zhiyuan Liu, Yu Rong. Qiang Liu, Shu Wu, and Liang Wang. Diffu- sion models for molecules: A survey of methods and tasks. arXiv, abs/2502.09511, 2025.
[52] Zhiyuan Liu, An Zhang. Hao Fei, Enzhi Zhang, Xiang Wang, Kenji Kawaguchi. and Tat- Seng Chua. Prott3: Protein-to-text generation for text-based protein understanding. In ACL. Association for Computational Linguistics, 2024.
[53] Yaorui Shi, Jiaqi Yang, Changhao Nai, Sihang Li, Junfeng Fang, Xiang Wang, Zhiyuan Liu, and Yang Zhang. Language-enhaneee representation learning for single-cell transcriptomics. arXiv preprint arXiv:2503. 09427, 2025.
[54] Yanru Qu, Keyue Qiu, Yuxuan Song, Jingjing Gong, Jiawei Han, Mingyue Zheng, Hao Zhou, and Wei- Ying Ma. Molcraft: Structure-based drug design in continuous parameter space, arXiv preprint arXiv:2404. 12141, 2024.
[55] Martin Simonovsky and Nikos Komodakis. Graphvae: Towards generation of small graphs using variational autoencoders. In ICANN (I), volume 11139 of Lecture Notes in Computer Science, pages 412-422. Springer, 2018.
[56] Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In NIPS, pages 133666315,, 2017.
[57] Ali Razavi, Aäron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VQ- 2-E22. In NeurlPS, pages 14837-14847, 2019.
[58] Myeonghun Lee and Kyoungmin Min. MGCVAE: multi-objective inverse design via molecu- lar graph conditional variational autoencoder. J. Chem. Inf. Model., 62(12))2943-299) 2950, 2022.
[59] Seokho Kang and Kyunghyun Cho. Conditional molecular design with deep generative models.
J. Chem. Inf. Model., 59(1):43-335 52, 2019.
[60] Qi Liu, Miltiadis Allamanis, Mare Brockschmidt, and Alexander L. Gaunt. Constrained graph variational autoencoders for molecule design. In NeurIPS, pages 7806-7815, 2018.
[61] Yuanqi Du, Tianfan Fu, Jimeng Sun, and Shengchao Liu. Molgensurvey: A systematic survey in machine learning models for molecule design. CoRR, abs/2203. 14500, 2022.
[62] Bin Dai and David P. Wipf. Diagnosing and enhancing VAE models. In ICLR (Poster). Open- Review.net, 2019.
[63] Jyoti Aneja, Alexander G. Schwing, Jan Kautz, and Arash Vahdat. A contrastive learning approach for training variational autoencoder priors. In NeurIPS, pages 480-493, 2021.
[64] Zhiyuan Liu, Yaorui Shi, An Zhang, Enzhi Zhang, Kenji Kawaguchi, Xiang Wang, and Tat- Seng Chua. Rethinking tokenizer and decoder in masked graph modeling for molecules. In NeurlPS, 2023.
[65] Arash Vahdat, Karsten Kreis, and Jan Kautz, Score-based generative modeling in latent space. In NeurlPS, pages 1 11287-11302, 2021.
[66] Xiaohui Zeng, Arash Vahdat, Francis Williams, Zan Gojeic, Or Litany, Sanja Fidler, and Karsten Kreis. LION: latent point diffusion models for 3d shape generation. In NewrIPS, 2022.
[67] Kristina Preuer, Philipp Renz, Thomas Unterthiner, Sepp Hochreiter, and Gunter Klambauer. Fréchet chemnet distance: metric for generative models for molecules in drug discovery. Journal of chemical information and modeling. 58(9):1736-1741. 2018.
14

---

[68] Jörg Degen, Christof Wegscheid-GGelachh Andrea Zaliani, and Matthias Rarey. On the art of compiling and using 'drug-like" chemical fragment spaces. ChemMedChem, 1(10):503-1157,, 2008.
[69] Guy W. Bemis and Mark A. Murcko. The properties of known drugs. 1. molecular frameworks. Journal Medicinal Chemistry, 39(15))2887-2893, 1996. PMID: 8709122.
[70] Jan Jensen. Xyz2mol.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims in abstract and introduction are well supported by results in the experimental section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. 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 it 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]
Justification:
The limitation section Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
16

---

Justification:
We do not include theoretical results.
Guidelines:
The answer NA means that the paper does 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 in the statement of any theo- rems.
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 to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix 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 experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The detailed needed for reproducibility are included in Section 9 Appendix C and our released code at https://githiii Cot.:////o///y11///AEU-----
Guidelines:
The answer NA means that the paper does 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 make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution a specific model and empirical evaluation, may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way 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 sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
17

---

Does the paper provide open the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Our code and used dataset is released at https: wtt.ulu..mmmyyy00030
UAE-3D4
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ https//ccccbbliiguuudeeeeee /oodeddee//oddSbmbiiiaaa ssionoolicy 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. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The settings are detailed in Section @ and Appendix O
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, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [No]
Justification: The experiments are mostly computationally expensive and we did not report error bars, following the baseline methods.
Guidelines:
The answer NA means that the paper does not 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 with given experimental conditions).
18

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
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 execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
These information are included in Section B and Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums 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 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We adhere to 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: [NA]
Justification: This work focuses on generating small molecules. It is pure computational research for applications in chemistry, material, and biology without possible harm to the human society
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
19

---

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 spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 or unintentional) misuse the technology.
there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our method does not involve such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions 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 we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The datasets are cited in Section A and Appendix C
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, snperswithcodd www.coocccaa.oee datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
20

---

For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification: Our codes released at htthps///gtabbb.c////yo//////// include readme file documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 pa- per include the full text instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: We have no crowdsourcing and no involvement of human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) 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 the NeurIPS Code of Ethics and the guidelines for their institution.
21

---

For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it 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, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy lttess///eerrissccc///ofeerreeee for what should or should not be described.
22

---

A More Experimental Results
A.1 More Ablation Studies and Analysis
Table 7:
The influence of latent space dimensionality for 3D molecule generation the QM9
dataset.
Latent Dim Atom Acc
Bond Acc RMSD) AtomSttall MalStable V&C
V&lI
VAII& Band Length
Band Angle] Dihedral Angle!
0.8932
0.325
3.27E-01
2.20E-01
4.19E-03
8
0.9998
0,9746
0.0006
0.987
0.882
0.942
0.923
0.923
1.29E-01
L.77E-02
8.14E-04
16
1,0000
1.0000
0.0012
0,999
0.988
0.983
0,973
0.950
7.01E-02
9.84E-05
3.47E-04
32
1,0000
1,0000
0.0003
0.986
0.867
0.943
0.918
0.918
9.98E-02
1442-02
7,77E-04
Table 8: The influence of R-Trans layer depth for reconstruction on the QM9 dataset.
R- Trans layers
Atom Ace
Bond Acc
RMSD
3
1,0000
0.9999
0.000326
6
1.0000
1.0000
0.000203
9
1.0000
1.0000
0.000229
12
1,0000
1.0000
0,000209
Table 9: The influence of different bonded distance loss weight À on the QM9 dataset.
A/ VAE Recon.
Atom Accuracy (%)
Bond Accuracy (%)
Coordinate RMSD (À)
(w/o X) 1
1.000
1.000
0.0034
1.000
1.000
0.0010
1.000
1.000
0.0007
10
1.000
1.000
0.0002
20
1.000
1.000
0.0002
/ 2D-Metric
FCD4
AtomStable
MolStable
V&C
V&U
V&U&N
SNN
Frag
Scaf
1
0.283
0.996
0.959
0.962
0,940
0,940
0.494
0.978
0.874
5
0.255
0,997
0.965
0.970
0.947
0,947
0.500
0.983
0.888
10
0.130
0,999
0.988
0.983
0.973
0,950
0.508
0.987
0.898
20
0.172
0,999
0.988
0.980
0.969
0.947
0.507
0.987
0.898
// 3D-Metric
FCD
AtomStable
MolStable
Bond length
Bond angle
Dihedral angle
1 0.901
0.985
0.857
1.25F-01
1880-002
1.06F-03
5
0.892
0.989
0.909
1.12F-01
1.19F-02
7.18F-04
10
0.881
0.993
0.935
7.04E-02
9.84E-03
3.47B-04
20
0.889
0.993
0.929
7.05E-02
9.84E-03
3.45E-04
Reconstruction Errors. To assess the reconstruction fidelity of UAE-3D, we evaluate perfor- mance the held-out test sets of QM9 and GEOM-Drugs We report atom-type and bond-type accuracies as well as the coordinates RMSD, which jointly reflect the VAE's ability to compress and reconstruct both invariant and equivariant molecular features. Table [ presents the reconstruction results. On the QM9 dataset, UDM-3D achieves perfect reconstruction accuracy (100% for both atoms and bonds) with a near-zero coordinate RMSD of 0.0002 A, substantially overperforming Ge- oLDM and ADiT. Notably, ADiT does not retain and reconstruct chemical bond information, which a limitation when reconstructing 3D molecules. On the more complex GEOM- -rugs dataset, UDM-3D again achieves perfect reconstruction (100.0% atom and bond accuracy, 0.0008 À RMSD), whereas GcoLDM exhibits noticeable degradation. This demonstrates that UDM-3D's unified latent space not only faithfully compresses all molecular modalities but also generalizes robustly to large and diverse molecular datasets. These results show that UAE-3D effectively learns unified latent representation that preserves both the chemical identity and geometric precision of molecules. These findings align with the visual comparison in Figure ab),,), and reinforce claim that UAE-3D achieves near-lossless reconstructionan essential prerequisite for accurate latent diffusion modeling.
Unified vs. Separated Latent Space. To validate the effectiveness of our unified latent space design, we conduct an ablation study comparing UDM-3D with a variant that employs separate latent spaces for invariant (2D) and equivariant (3D) molecular features. Specifically, we split the unified VAE
23

---

Table 10: Comparison between unified and separated latent space the QM9 dataset.
2D-Metric
FCD]
AtomStable
MolStable
V&C
V&U
V&U&N
SNN
Frag
Scaf
Unified
0.130
0.999
0.988
0.983
0.973
0.950
0.508
0.987
0.898
Separated
0.351
0.995
0.952
0.943
0.920
0.913
0.341
0.940
0.682
3D-Metric
FCD]
AtomStable
MolStable
Bond length.
Bond angle)
Dihedral angle]
Unified Separated
0.881
0.993
0.935
7.04F-02 18.4F-02
9.84F-03 123F-03
3.47F-04
2.356
0.982
0.872
7.03F-04
Table 11: Reconstruction accuracy (%) for atom and bond types, and coordinate RMSD on the test
sets.
Atom Accuracy (%)
Bond Accuracy (%)
Coordinate
Dataset
Method
RMSD (À)
GeoL.DM
96.9
93.7
0.2526
GFOM-Drugs
UDM-3D, ours
100.0
100.0
0.0008
GeoL.M
98.6
96.2
0.1830
QM9
ADIT
85.7
0.3598
UDM-3D, ours
100.0
100.0
0.0002
into two independent neural networks: one maintaining the R- R-Trans architecture and taking 2D information (atom types, bond types) as input, while the other uses a vanilla transformer structure and takes 3D information (atom types and coordinates) input. Each network generates separate latent space, with dimensions set to half the original hyperparameter. The LDM then trained on the concatenated latent space of these two networks. Table III summarizes the results on the QM9 dataset. We observe that UDM-3D with a unified latent space significantly outperforms the separated variant across all 2D and 3D generation metrics. For instance, the FCD score improves from 0.351 to 0.130, and the MolStable metric increases from 0.952 to 0.988. This demonstrates that jointly encoding invariant and equivariant features into a single latent representation enables more effective diffusion modeling, as the LDM can learn coherent distribution over unified modality rather than needing to capture complex interactions between two separate modalities. These findings validate our design choice of a unified latent space as crucial for achieving state-of-the-an performance in 3D molecule generation.
Influenece of Latent Dimensions. To investigate the impact of latent dimensionality on 3D molecule generation, we conduct ablation study on the QM9 dataset. Table @ summarizes the reconstruction errors and the performance metrics across different latent space dimensions (4, 8, 16, and 32). The results indicate that too-low latent dimension (e.g.. 4) fails capture the necessary information, leading to inferior performance. Increasing the dimension to 16 results in substan- tial performance improvements, while further increasing to 32 does not yield significant additional gains. This demonstrates that a latent dimension of 16 provides an optimal tradc-off between model capacity and reconstruction fidelity.
Influence of R-Trans Layer Depth. To assess the effect of depth in the Relational Transformer encoder in UAE-3D, we conduct an ablation study by varying the number of R-Trans layers (3, 6. 9, and 12) on the QM9 dataset regarding reconstruction error. Table 8 summarizes the results. We observe that reducing the number of transformer layers to 3 leads to noticeable increase in reconstruction error, while increasing it to 9 or 12 does not yield further improvements, as the recon- struction error is already near-zero. This indicates that a moderate depth of 6 layers is sufficient for capturing the necessary relational information among atoms and bonds for effective reconstruction.
Bonded Distance Loss. Table g presents ablation study the weight À for the bonded distance loss on the QM9 dataset. We vary À to analyze its effect on reconstruction accuracy and the preser- vation of bond geometries. The results indicate that setting À 10 achieves an optimal balance between compressing coordinate information and maintaining geometric fidelity. This experiment validates the necessity of tuning the bonded distance loss weight enhance model performance.
Ring Distribution. Following the evaluation protocol of (54]], we computed the Percentage (%) of molecular modes in terms of ring distribution for our generated molecules and compared them
24

---

Table 12: Comparison of ring distributions between molecules in the QM9 training set and those generated by our model.
3-Ring
4-Ring
5-Ring
6-Ring
(7+)-Rings
QM9 train set
43.14%
39 20.20%%
39.31%
12.72% 12.90%
2.92%
Ours
41.69%
38.20%
36.37%
3.13%
against the QM9 training set. The results in Table 12 provide quantitative assessment of how well our model preserves the structural diversity of the training dataset.
30
Same Molsculc
20
Potaloni
Translation -
Rotation
10
20
0-
Transiation -Z
-20
Rotation
A Coneeeeee
440
Rotation
-20
-60
-80
800
-60
-40
20
20
60
-30
-20
-10
0
20
30
Figure 6: Additional t-SNE visualizations of UAE-3D's latents under SE(3) augmentations. (a) 360- degree rotation along the z-axis. (b) same molecule under different SE(3) augmentations.
More Visualizations of Latent Space. In addition to Figure H. conducted more t-SNE visual- izations of UAE-3D's molecular latent representations further examine its geometric sensitivity and consistency. In Figure E(a), we apply 20-step rotations along the same axis (7-axis) to different molccules, completing a full 360-degree rotation. Remarkably, the latent trajectories all molecule form loops and return to their original positions, reflecting the geometric sensitivity rotational consis- tency in the learned representations. In Figure BKb), each color now corresponds to a different SE(3) augmentation applied the same molecule, rather than to different molecules in previous plots. We observe clear separation among representations generated by translations and rotations along different axes (x/y/z), indicating that UAE-3D effectively distinguishes between types of geometric transformations in its latent space. These results further validate the structured and disentangled nature of UAE-3D's latent geometry.
A.2 More Visualizations of Generated Molecules
To further demonstrate the chemical validity and structural diversity molecules generated by UDM-3D, we provide extended visual comparisons across the GEOM-Drugs and QM9 datasets in Figure and Figure B
B More Related Works
Molecular Variational Autoencoders. Variational Autoencoders (VAEs) have been widely ex- plored for generative modeling (166, 555-57], providing a framework to encode data into a structured latent space while enabling both generation and reconstruction [16]] In molecular generation, VAEs have been applied
molecular graphs, as seen in JT-VAE [I]), MGCVAE [58], SSVAE [59], and CGVAE [60]. However, these methods have several limitations. First, they primarily focus 2D
25

---

molecular representations and often rely non-transformer architectures, limiting their scalabil- ity for more complex molecular generation tasks [61].. Additionally, VAF-based methods typically assume simple Gaussian prior, which may fail to accurately model the complex posterior distribu- tions required for effective molecular generation [62. 63 Notably, masked graph modeling [[2]] also builds molecular graph auto-encoder, but aiming for representation learning, instead of molecule generation. With recent advancements in Latent Diffusion Models (LDMs) demonstrating success across multiple domains tion, as seen in GeoLDM [ and PepGLAD]IR) These approaches leverage diffusion modeling in latent space, offering improved generative capacity and efficiency.
[65] b6], LDMs have been introduced for 3D molecule and protein genera-
Comparison to Other Unified Generative Molecule Models. Concurrent with our work, ADIT [38] from FAIR also explores latent diffusion for 3D molecules using variational autoen- coder to compress structures into a unified latent space. However, our approach differs in three key ways: (1) ADIT does not achieve near-lossless compression. In contrast, we explicitly evaluate reconstruction quality and show significantly lower errors than ADiT (ef. Table DD), which trans- lates better 3D generation performance (cf. Table D). (2) ADIT omits chemical bonds as part of the inputs and reconstruction targets, leading to information loss. Although bond types can be inferred from 3D coordinates, such predictions are not 100% accurate. For example, the QM9 dataset [224] reports 3% inconsistency across inference methods. To address this, we explicitly model bonds using relational transformer [21] and include bond reconstruction as an objective, offering a more robust and chemically faithful modeling approach. (3) ADTT provides insufficient evaluation for small molecule generation. ADIT reports only two metrics (validity and uniqueness) for the QM9 dataset, while our evaluation includes 15 metrics, covering 2D structure, 3D geometry, stabil- ity, validity, uniqueness, and other critical aspects (cf. Table D. We conjecture that this is because ADiT is more focused crystal material generation, instead of small molecule generation.
MoICRAFT [[24] also aims to unify multiple molecular modalities for generation. However, it pro- cesses atom types and coordinates separatelyconcateeatiig them while use special design to pre- serve SE(3) equivariance for coordinates. In contrast, our model encodes atom types, coordinates, and bond types into single unified latent representation, enabling effortless multi-modal model- ing using a uni-modal latent diffusion model Diffusion Transformer [23]]. Moreover, MoICRAFT focuses on the task of cencttrr-based drug design, which is different than our major benchmarks.
Comparison to Other Equivariant 3D Molecule Generation Methods. We adopt SE(3) data augmentations to train a variational autoencoder without explicitly incorporating equivariant archi- tectures, allowing the model to learn geometric symmetries directly from data. This strategy is motivated by AlphaFold3 [28], which demonstrated that SE(3) augmentations are sufficient for general-purpose purpose networks learn equivariance and invariance. As shown in Table our model achieves near-zero reconstruction errors on 3D molecules subjected to random rotations and transla- tions, showing successful learning of SE(3) equivariance. Compared to models with built-in SE(3) equivariance, such as EDM [I9]] and JODO [6],, our method avoids the architectural complexities of handling multi-modality and enforeing SE(3) equivariant constraints. This design choice allows us to use the highly optimized and parallelizable diffusion model, Diffusion Transformer [[4]] for latent diffuusio modeling. Consequently, we achieve significant efficiency gains, with up to 8.6x faster training and 9.8x faster inference (cf. Figure ((c:)).
C More Experimental Settings
C.1 Experimental Settings for De Novo 3D Molecule Generation
Setup. For the GEOM-Drugs dataset, we use only the ground-state conformer with the lowest energy for each molecule, following the setup in JODO [6]. We adopt the same train/validation/test split ratio of 8:1:1, ensuring a fair comparison. For the QM9 dataset, we use the WOK//1K//1KK train//alidaaiion/esst split. For both datasets, all atoms are shifted zero center-of- of-ass before input. All baseline results are either directly borrowed from JODO [X]] or reproduced using their released codebase under identical settings.
Hyperparameters. [T shows the key hyperparameters used for training the UAF-3D and UDM-3D models.
26

---

Table 13: Hyperparameters of the UAE-3D and UDM-3D models.
Parameter
UAE
UDM
epochs
2000
10000
atom loss weight (7imom)
1.0
bond loss weight (?bund)
1.0
coordinate weight (Ycvnddintee)
1.0
distance loss weight (7ddssmece)
1.0
bonded distance loss weight (A)
10.0
KLD loss weight batch size optimizer
1e-8
512
Adam
Icarning rate
lc-4
weight decay
lc-5
translation augmentation scale
0.1
encoder hidden size
encoder #heads
8
encoder #blocks
6
latent dimension
16
decoder hidden size
decoder #heads
decoder #blocks
diffusion hidden size diffusson #heads diffuusio #layers diffusion mlp ratio condition drop (Parop)
512
8
4.0
0.1
Evaluation Metrics. Our evaluation framework for de novo 3D molecule generation encompasses complementary metrics to assess molccular validity, diversity, and geometric fidelity. These metrics operate at two levels:
2D Structural Analysis:
Validity & Stability: Measures adherence to chemical rules:
Atom Stability: Percentage atoms with chemically valid valency, determined by cross- referencing the bond counts with allowed configurations for each atom type (e.g., carbon with 4 bonds). This is evaluated using RDKit's atom valency parser. Formal charges are taken into account in the valency specification.
Molecule Stability: Percentage of molecules with valid valency for all atoms pass the above valency check.
Validity & Completeness (V&C): Fraction of fully connected, syntactically correct (SMILES- parsable) molecules, excluding fragmented or hypervalent structures.
Diversity: Quantifies exploration of chemical space:
Validity & Uniqueness (V &U:; Percentage of unique molecules among valid ones, calculated as the ratio of unique SMILES strings the total number of valid molecules.
Validity & Uniqueness & Novelty (V&U&N): Percentage of unique molecules among valid ones that are also novel, calculated the ratio of unique SMILES strings the total number of valid molecules, excluding those present in the training set.
Distribution Alignment: Convert all the valid molecules into SMILES strings and compare gener- ated/test distributions on the MOSES benchmark [G]]
Fréchet ChemNet Distance (FCD): Similarity between generated and reference distributions using activations from the penultimate layer of ChemNet [[67], pretrained molecular prop- erty predictor. Lower scores indicate better alignment.
Similaritt to the nearest neighbor (SNN): Tanimoto similarity between the generated molecules and their nearest neighbors in the test set, calculated using Morgan fingerprints. Fragment Similarity (Frag): Tanimoto similarity between the generated molecules and their nearest neighbors in the test set, calculated using BRICS [68] fragments.
27

---

Scaffold Similarity (Scaf): Frequencies of Bemis-Murcko scaffolds [6%] in the generated molecules compared to the test set.
3D Geometric Analysis:
Conformer Quality: The 3D coordinates of generated molecules are converted into bond graphs using a predefined lookup of atomic distance thresholds (per atom pair). Based on these bonds, 2D molecular graphs are reconstructed and then evaluated for atom and molecule stability using the same criteria as in 2D Structural Analysis.
Atom Stability: Consistency of valency 3D-derived 2D graphs. Reconstructs 2D bond orders from 3D atomic distances (via tabulated bond-length thresholds) and checks valency compliance.
Molecule Stability: Percentage of molecules with valid valency for all atoms pass the above valency check.
FCD3D: Same computation as standard FCD, but applied to the reconstructed molecules from 3D coordinates. Captures distributional similarity under geometric constraints.
Geometric Fidelity: Measures spatial arrangement accuracy via Maximum Mean Discrepancy (MMD) for:
Bond lengths: MMD between distributions of bond lengths for eight frequent bond types
(e.g., C-C, C-N, C-O). Computed using Gaussian kernel with bandwidth tuned via the median heuristic.
-Bond angles: MMD between the bond angle distributions of generated vs. reference molecules. Triplets of bonded atoms (e.g., C-C-C) define each angle.
Dihhedral angles: MMD between distributions of dihedral (torsional) angles over four-atom sequences (e.g., C-C-C-C). Captures conformational realism.
Comparison to Other Evaluation Protocols. Discrepancies in evaluation methodologies can sig- nificantly impact reported performance metrics. For instance, ADiT's validity scores are obtained by reconstructing molecules from atomic coordinates using pymatgen. core. .Molecul, followed by bond inference via RDKit (likely using the xyz2mo1 heuristic [+]] This pipeline ensures topologi- cal consistency (i.e.. SMILES validity) but does not explicitly enforce chemical stability constraints such correct valency or bond geometry. In contrast, our evaluation employs RDKit's Chem. Mol for molecule reconstruction, which enforces stricter valency checks and chemical correctness. And our broad evaluation criteria comprehensively measure the validity, diversity, and geometric fidelity of 3D molecule generation.
C.2 Experimental Setting for Conditional 3D Molecule Generation
Setup. We conduct conditional molecule generation the QM9 dataset following the protocol in EDM [S] and EEGSDE [8]. The properties are evaluated by the QM9 property predictor [[4]] pre-trained on half of the QM9 dataset. We use the property classifier pre-trained by EDM [[]]
Specially, we optimize molecules toward six key electronic properties from quantum chemistry for
conditional generation:
Electronic Response:
Dipole Moment (μ)): Molecular polarity measure
Polarizability (a:): Induced dipole response electric fields
Thermodynamic Properties:
Heat Capacity (C.):: Energy absorption at constant volume
Electronic Structure:
НОМО YLUOO (ELUMO): Frontier orbital energies
HOMO-LUMO Gap (Ae): Critical for reactivity and conductivity
The evaluation protocol follows rigorous split-and-validate strategy [5]]
QM9 training data divided into disjoint subsets D, (50k) and D3 (50k)
Property predictor de trained exclusively on D.
28

---

Generated molecules compared against o.'s predictions D, (lower-bound baseline) and ran- dom predictions without any relation between molecule and property (upper-hound baseline)
Performance quantified via Mean Absolute Error (MAE) across all properties
This approach ensures fair assessment of conditional generation without data leakage, with de 's Db. performance establishing the theoretical minimum achievable error.
Algorithm Training Algorithm for UAE-3D and UDM-3D
Require: Molecular dataset D, encoder &, decoder D, diffusion model 1: Initialize encoder &, decoder D, and diffusion model cs
2: Stage 1: Train UAE-3D
3: while not converged do
4: Sample batch of 3D molecules G D
5: Z {G) (Encode)
6: G D [ecode)}
7: Lrcca (Eq. I |Reconstruction loss)
8: LUAE-3D (Eq. (T3) (VAE loss)
9: Update & and 'D using LUAE3D
10: end while
Stage Train UDM-3D
12: while not converged do
13: Sample batch of latent sequences (0) ~ E(D)
Sample diffusion timestep ~ Uniform(O,1)
15:
O (() + v1 a where ~ N(O,,)) I) (Diffusion)
16:
é = c0 (Z
17:
Laimusion é el (Noise prediction loss)
18: Update Eg using Laiffusion
19:
end while
Algorithm 2 UDM-3D Sampling with CFG
Require: Trained UAE-3D encoder E, DiT Eg
Require: Guidance strength w, timesteps {
Require: Target property (optional)
1: ( ~ (O, I) (Initial noise)
2: for k K downto do
3: f - to
4: Compute Guidance:
5: Econd f- (e) c)
6: funcond <- Eg ( (Z .2,, 0)
7:
é
(1+ w) Econd W'nccond (CFG blending)
8: Denoise:
9:
1) DDPM DDPMDDMtePZ ē,tt [TS]]]
10: end for
Decode:
12: G [F. F,E,X) + ()) (UAE-3D decoder)
C.3 Pseudo Code
We provide the training and sampling algorithms for UDM-3D in Algorithms 1 and A respectively. The training algorithm for UDM-3D involves two stages: training the VAE (UAE-3D) and training the DiT (UDM-3D). The sampling algorithm for UDM-3D with CFG involves iteratively denoising latent sequence Z using the DiT and the CFG guidance.
29

---

Figure 7: Visualization of random samples generated by UDM-3D on GEOM-Druse
30

---

Figure 8: More visualization of random samples generated by UDM-3D on QM9.
31


Reviews:
[{"review_id": "n4d6tA18AD", "reviewer": "Reviewer_cDUj", "summary": "This paper tackles the challenge of multi-modal 3D molecular generation by proposing UAE-3D, a unified variational autoencoder capable of compressing atom types, bond types, and 3D coordinates into a single latent space while maintaining near lossless reconstruction. By leveraging SE(3)-equivariant data augmentations in training, the method learns equivariant representations without relying on 3D-specific architectural biases. The authors demonstrate that their unified latent space allows for the application of latent diffusion models (specifically, Diffusion Transformer, DiT) to generate realistic molecules efficiently and accurately.", "questions": "- Could the authors provide more formal justification (or empirical evidence across much more diverse datasets) for the claim that SE(3)-equivariant augmentation paired with the unified VAE suffices to ensure geometric and chemical preservation, as opposed to specialized architectures?\n- What is the effect of varying the latent space dimension, depth of R-Trans encoder, DiT architecture (vs other transformer-based or non-transformer LDM backbones), or hyperparameters $\\gamma$, $\\lambda$? Is the model robust to such design choices or sensitive in particular to one?\n- Are there patterns to the errors or invalid/implausible generations outside of what is summarized quantitatively in Tables 1/2? Could the authors expand with more systematic failure analysis (qualitative or quantitative)?", "limitations": "The limitations section is in the appendix, but should be summarized more explicitly in the main paper.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- The paper is methodologically well-grounded. The approach replaces separate modal latent spaces for molecular generation with a unified latent space, validated by strong empirical results. The SE(3)-equivariant augmentations for learning equivariant representations are elegantly simple and effective, and the use of a Relational Transformer allows seamless integration of node and edge features.\n\n- As visualized in Figure 2 , the model achieves 100% atom/bond accuracy and extremely low coordinate RMSD on reconstruction—substantially better than prior work (GeoLDM) and critical for downstream generative modeling. Demonstrates that training and inference times are significantly reduced compared to strong baselines (GeoLDM and JODO), largely due to unified latent modeling and use of a general diffusion backbone.\n\n- The model and pipeline architecture are clearly depicted in Figure 3, providing an accessible overview of feature extraction, latent encoding, and generative modeling.\n\n- This paper offers insightful ablations, showing the value of each key component (unified latent space, DiT backbone, data augmentation) and the effects on RMSD and geometry metrics. Figures 4(a-c) (t-SNE plots) demonstrate latent representation consistency under SE(3) transformations, aligning with the claim of geometric structure preservation.\n\n**Weakness**\n- The paper relies heavily on empirical demonstrations, lacking formal theoretical analysis of why the unified latent space retains all necessary information for 3D generation, or convergence guarantees for the learning of equivariance solely via data augmentation. While the gains are empirically solid, more discussion or theoretical intuition would improve the scientific contribution.\n\n- The ablation studies in Section 4.5, Table 5, and Table 6 are thoughtful but could be further strengthened. For example, the effect of latent dimension size, number of R-Trans layers, or replacing the DiT backbone with other modern diffusion architectures could be explored to validate robustness and architectural choices beyond what is already discussed. It would also be instructive to see whether the benefits of unified latent space and data augmentation persist across substantially larger or more complex molecules.\n\n- The paper claims broad applicability of the unified latent approach to diverse 3D molecular and multimodal datasets, but empirical validation is limited to QM9 and GEOM-Drugs. Explicit discussion of the limitations in generalization, scalability to much larger chemical spaces, or application to proteins/RNAs is somewhat deferred to future work.", "was_revised": true, "final_justification": "I will keep the rating unchanged", "rebuttal": "> (W1) The paper relies heavily on empirical demonstrations, lacking formal theoretical analysis of why the unified latent space retains all necessary information for 3D generation, or convergence guarantees for the learning of equivariance solely via data augmentation. While the gains are empirically solid, more discussion or theoretical intuition would improve the scientific contribution.\n\nThank you for the constructive feedback. 3D molecular generation is based on data consisting of atomic types, chemical bonds, and 3D coordinates. Our method integrates this information into a unified latent space, allowing the model to learn all necessary geometric and chemical features in one space. While we currently lack rigorous theoretical analysis, we provide the following intuitive explanations:\n1. **Near-lossless Latent Space**: The unified latent space is designed to be near-lossless, meaning it can perfectly reconstruct the original 3D coordinates and chemical information. As shown in Figure 2, our VAE achieves 100% atom/bond accuracy and near-zero coordinate RMSD on the test set. This ensures that all necessary information is preserved during encoding and decoding. \n2. **SE(3)-Equivariant Augmentation**: By applying SE(3) augmentations during training, the model learns to be equivariant to rotations and translations. This means that the model can generate 3D coordinates that are consistent with the original molecular structure. To validate this, we conducted experiments on 1000 molecules in the test set, applying random SE(3) transformations (rotations and/or translations) and measuring the RMSD of the molecules reconstructed by the model. As shown in the following Table, The results show that the model maintains low RMSD under different SE(3) transformations, confirming its SE(3) equivariance.\n\n| Transformation | RMSD (Å) |\n|---|---|\n| No transform. | 0.002 |\n| Random rot. | 0.002 |\n| Random trans. | 0.002 |\n| Random rot. & trans. | 0.002 |\n\n> (W2) The ablation studies in Section 4.5, Table 5, and Table 6 are thoughtful but could be further strengthened. For example, the effect of latent dimension size, number of R-Trans layers, or replacing the DiT backbone with other modern diffusion architectures could be explored to validate robustness and architectural choices beyond what is already discussed. \n> \n> (Q2) What is the effect of varying the latent space dimension, depth of R-Trans encoder, DiT architecture (vs other transformer-based or non-transformer LDM backbones), or hyperparameters? Is the model robust to such design choices or sensitive in particular to one?\n\nWe thank the reviewer for the insightful questions. We provide the following studies and new results:\n\n1. **Latent Dimension Ablation**\n\n    As already shown in Appendix Table 7, we observe that too low a latent dimension fails to capture the necessary information, leading to poor performance. Increasing the dimension to 16 significantly improves performance, while further increasing it to 32 does not provide significant additional improvement. And in the updated Table 7 below, we also provide the reconstruction error for QM9 dataset for more clarity:\n\n    | Latent Dim | Atom Acc | Bond Acc | RMSD | AtomStable | MolStable | V\\&U\\&N | Bond Len$\\downarrow$ | Angle$\\downarrow$ | Dihedral$\\downarrow$ |\n    | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n    | 4 | 0.9201 | 0.8932 | 0.080 | 0.914 | 0.325 | 0.566 | 3.27E-01 | 2.20E-01 | 4.19E-03 |\n    | 8 | 0.9998 | 0.9746 | 0.006 | 0.987 | 0.882 | 0.923 | 1.29E-01 | 1.77E-02 | 8.14E-04 |\n    | 16 | 1.0000 | 1.0000 | 0.002 | 0.999 | 0.988 | 0.950 | 7.04E-02 | 9.84E-03 | 3.47E-04 |\n    | 32 | 1.0000 | 1.0000 | 0.003 | 0.986 | 0.867 | 0.918 | 9.98E-02 | 1.42E-02 | 7.77E-04 |\n\n    Likewise, dimensions 4 and 8 yield higher reconstruction error, indicating they are insufficient to capture molecular complexity. The latent dimension of 16 achieves near-zero reconstruction error, while 32 does not significantly improve performance and even slightly increases the error, indicating that too large a latent dimension may not be beneficial for reconstruction.\n\n2. **R-Trans Layer Depth Ablation (VAE reconstruction)**\n\n    | R-Trans layers | Atom Acc | Bond Acc | RMSD |\n    | --- | --- | --- | --- |\n    | 3 | 1.0000 | 0.9999 | 0.00326 |\n    | 6 | 1.0000 | 1.0000 | 0.00203 |\n    | 9 | 1.0000 | 1.0000 | 0.00229 |\n    | 12 | 1.0000 | 1.0000 | 0.00209 |\n\n    We observe that decreasing the number of our transformer layers to 3 leads to a significant increase in reconstruction error, while increasing it to 9/12 can not further improve the reconstruction quality because the reconstruction error is already near-zero. \n\n3. **Backbone Ablation**\n\n   As shown in Table 4, we already compared our DiT backbone with a standard Transformer architecture. In the updated Table 4 below, we add a comparison with PerceiverIO, a modern transformer-based architecture designed for structured inputs and outputs. The results show that DiT consistently outperforms PerceiverIO across all metrics, which we attribute to its adaptive LayerNorm layers. These layers enable DiT to effectively handle data with different noise scales, thereby improving diffusion performance.\n\n   | Structure | 3D-AtomStable | 2D-AtomStable | V\\&C | V\\&U | V\\&U\\&N | \n    |---|---|---|---|---|---|\n   | Transformer | 0.983 | 0.997 | 0.938 | 0.922 | 0.922 |\n   | PerceiverIO | 0.972 | 0.990 | 0.933 | 0.931 | 0.931 |\n   | DiT | 0.993 | 0.999 | 0.983 | 0.973 | 0.950 |\n   \n   Moreover, when paired with other diffusion neural architectures, like Transformer and PerceiverIO, our model still achieves meaningful and comparable performances. The performance can also be potentially improved given more hyperparameter tuning, which is unfortunately not available given rebuttal's limited time. This observation demonstrate that our UAE-3D model is robust to different diffusion neural architectures.\n\n> (W2) It would also be instructive to see whether the benefits of unified latent space and data augmentation persist across substantially larger or more complex molecules.\n> (W3) The paper claims broad applicability of the unified latent approach to diverse 3D molecular and multimodal datasets, but empirical validation is limited to QM9 and GEOM-Drugs. Explicit discussion of the limitations in generalization, scalability to much larger chemical spaces, or application to proteins/RNAs is somewhat deferred to future work.\n\nThank you for the suggestion. We agree that broader empirical validation is important. While QM9 and GEOM-Drugs are the standard benchmarks, we provide the following evidence to support scalability:\n\n* **GEOM-Drugs** includes molecules up to **181 atoms**, which are comparable in size to small peptides. Our method achieves **perfect reconstruction** on this dataset (see Table 9), showing robustness to large-scale molecular inputs.\n* **CrossDocked2020** is a large-scale dataset of protein-ligand complexes, containing over 100,000 protein-ligand pairs and up to over **400 atoms** per complex. We trained our VAE on molecules from this dataset, including amino acids and ligands, achieving:\n\n  | Atom Acc  | RMSD(Å) |\n  | ---  | --- |\n  | 1.000  | 0.011 |\n\nThese results show that UAE-3D can scale to **larger, protein-related molecules** and maintain high fidelity. We did not test bond type prediction accuracy yet because the bond type information is not available in the dataset for protein's amino acids. Adapting to full protein/RNA generation will be addressed in future work, as discussed in our Appendix A.\n\n> (Q1) Could the authors provide more formal justification (or empirical evidence across much more diverse datasets) for the claim that SE(3)-equivariant augmentation paired with the unified VAE suffices to ensure geometric and chemical preservation, as opposed to specialized architectures?\n\nThank you for the suggestion. We are happy to provide empirical results on more diverse datasets to support our claims. As shown in the above table, we have conducted experiments on the CrossDocked2020 dataset, which contains large protein-ligand complexes. The results indicate that our model maintains low reconstruction loss and can preserve complex geometric and chemical information even on large-scale protein-ligand molecules. We are also conducting further validation experiments on larger PDB datasets to strengthen our claims.\n\n> (Q3) Are there patterns to the errors or invalid/implausible generations outside of what is summarized quantitatively in Tables 1/2? Could the authors expand with more systematic failure analysis (qualitative or quantitative)?\n\nWe appreciate the suggestion and offer the following observations:\n\n- **Error Types**: During generation, failure cases often involve incorrect **aromatic bonds**, which are particularly sensitive to small coordinate perturbations. This occurs rarely but can lead to **invalid valence assignments**. Specifically, because the interatomic distance (i.e., bond length) in the generated molecules determines the bond type, for example, the carbon-carbon bond length in the benzene ring is 1.40Å, which is between a single bond (1.54Å) and a double bond (1.34Å). Therefore, during the generation process, these bonds are likely mismatched due to small coordinate shifts, resulting in incorrect bond assignments for some carbon atoms.\n- **Quantitative Insights**: Across 10k samples, invalid molecule rate remains around 1%, and ring-distribution (Table 10) closely matches that of the training set, indicating no mode collapse.\n\n> (Limitations) The limitations section is in the appendix, but should be summarized more explicitly in the main paper.\n\nThe limitation section is now in Appendix due to our limited space in main paper. We appreciate the advice, and will revise to include it in the main paper in the camera-ready version.\n\n[1] Perceiver IO: A General Architecture for Structured Inputs & Outputs, Jaegle et al., ICLR 2022."}, {"review_id": "hSyZbmJaJM", "reviewer": "Reviewer_uzP6", "summary": "In this work, the authors present a VAE-based framework that compresses the multi-modal features (positions, atom types, and bonds) of 3D molecules into a latent space. Using a diffusion model, they show that sampling on this latent space is effective and demonstrate this by conditional generation on QM9 and unconditional generation on GEOM DRUGs and the QM9 dataset. This approach, due to its disjoint training pipeline, is faster than a few existing frameworks for generation.", "questions": "Q1. How does the model truly ensure that the equivariance to SE(3) is truly present in the generated samples? In both the augmented and non-augmented cases.\n\nQ2. Given that this method relies on a lossless latent space, how can to ensure that this is not overfitting on training data? Could you also report (only) uniqueness scores in Tables 2 and 5?\n\nQ3. Could you provide additional details on how augmentation is done and added to the model? \n\nQ4. How to update this model if one of the modalities says bonds have missing data. Would this approach overfit on the bonds present and not generate the bonds not seen before? Is this a good model for molecule discovery, and if so, could you elaborate on this?\n\nQ5. The tsne plot is unclear to me. If this is in a Center of mass frame, how do I interpret translations? What do these plots contribute? \n\n\nOther suggestions:\n- Add other joint modeling papers and other missing citations in the Table for comparison.", "limitations": "Limitations section is not present in the main paper.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n1. This paper presents joint modeling of continuous and discrete features for both molecular generation and property prediction. Due to its decoupled nature of 1. training a VAE to have a latent space 2. performing diffusion on latents it is faster than a few previous methods.\n\n\n\n### Weakness\n1. The results presented in all the table do not have error bars. \n2. This work only compares with a few discrete graph generation models, and leaves out a few joint modeling papers.  \n3. The paper misses important citations in the text/table\n- PONITA, Bekkers et al: https://openreview.net/forum?id=dPHLbUqGbr \n- FlowMol, Dunn et al, https://arxiv.org/abs/2404.19739\n- MolFM, https://arxiv.org/pdf/2307.09484", "was_revised": true, "final_justification": "The authors respond to all my questions in the rebuttal. The questions about the lack of generality for the lossless latent space still remain for me. Due to the fair response to the rest of the questions, I have updated my score for the paper.", "rebuttal": "> (W1) The results presented in all the table do not have error bars.\n\nThank you for the suggestion. We have now added standard deviation as error bars to the reported results in Tables 1, 2, and 3. For clarity, we summarize the results with error bars below (based on 5 independent inference runs):\n\nTable 1 (GEOM-Drugs)\n| 2D-Metric | FCD$\\downarrow$ | AtomStable | MolStable | V&C | V&U | V&U&N | SNN | Frag | Scaf |\n|---|---|---|---|---|---|---|---|---|---|\n| UDM-3D | 0.692 $\\pm$ 0.012 | 1.000 $\\pm$ 0.001 | 0.925 $\\pm$ 0.002 | 0.879 $\\pm$ 0.001 | 0.913 $\\pm$ 0.005 | 0.907 $\\pm$ 0.004 | 0.525 $\\pm$ 0.08 | 0.990 $\\pm$ 0.007 | 0.540 $\\pm$ 0.001 |\n\n| 3D-Metric | FCD$\\downarrow$ | AtomStable | MolStable | Bond length$\\downarrow$ | Bond angle$\\downarrow$ | Dihedral angle$\\downarrow$ |\n|---|---|---|---|---|---|---|\n| UDM-3D | 17.36 $\\pm$ 0.610 | 0.852 $\\pm$ 0.002 | 0.014 $\\pm$ 0.000 | 9.89E-03 $\\pm$ 0.02E-03 | 5.11E-03 $\\pm$ 0.05E-03 | 1.74E-04 $\\pm$ 0.02E-04 |\n\nTable 2 (QM9)\n| 2D-Metric | FCD$\\downarrow$ | AtomStable | MolStable | V&C | V&U | V&U&N | SNN | Frag | Scaf |\n|---|---|---|---|---|---|---|---|---|---|\n| UDM-3D | 0.130 $\\pm$ 0.003 | 0.999 $\\pm$ 0.001 | 0.988 $\\pm$ 0.002 | 0.983 $\\pm$ 0.002 | 0.973 $\\pm$ 0.001 | 0.950 $\\pm$ 0.001 | 0.508 $\\pm$ 0.05 | 0.987 $\\pm$ 0.003 | 0.898 $\\pm$ 0.005 |\n\n| 3D-Metric | FCD$\\downarrow$ | AtomStable | MolStable | Bond length$\\downarrow$ | Bond angle$\\downarrow$ | Dihedral angle$\\downarrow$ |\n|---|---|---|---|---|---|---|\n| UDM-3D | 0.881 $\\pm$ 0.002 | 0.993 $\\pm$ 0.001 | 0.935 $\\pm$ 0.001 | 7.04E-02 $\\pm$ 0.05E-02 | 9.84E-03 $\\pm$ 0.04E-03 | 3.47E-04 $\\pm$ 0.03E-04 |\n\n> (W2) This work only compares with a few discrete graph generation models, and leaves out a few joint modeling papers.\n\nWe want to clarify that JODO[1], MiDi[2], and EQGAT[3] are all strong baselines with joint modeling capability.\n\n- JODO propose a joint 2D and 3D graph diffusion model that generates geometric graphs representing complete molecules with atom types, formal charges, bond information, and 3D coordinates.\n- MiDi jointly generates molecular graphs and their corresponding 3D arrangement of atoms.\n- EQGAT takes continuous atom positions, while chemical elements and bond types are categorical variables, and uses a diffusion model to generate molecules.\n\n> (W3) The paper misses important citations in the text/table\n> \n> - PONITA, Bekkers et al\n> - FlowMol, Dunn et al,\n> - MolFM,\n\nWe thank the reviewer for pointing these related works, and are happy to include results for PONITA and FlowMol in the Table below for comparison. However, MolFM is not designed for the same benchmark, thus cannot be compared directly, we have included it in the related work session for comparison.\n\n| Method | Dataset | 2D-AtomStable (%) | 2D-MolStable (%) |\n|---|---|---|---|\n| FlowMol | GEOM-DRUGS | 99.0 | 67.5 |\n| Ours | GEOM-DRUGS | 100.0 | 92.5 |\n| PONITA | QM9 | 98.9 | 87.8 |\n| FlowMol | QM9 | 99.7 | 96.2 |\n| Ours | QM9 | 99.9 | 98.8 |\n\nNote that all results are copied from the original papers, and the AtomStable and MolStable reported in PONITA and FlowMol actually correspond to our 2D metrics (following the practice of EDM[4] and JODO[1]). We did not compare other evaluation metrics because they are not reported in these baselines' original paper.\n\n> (Limitations) Limitations section is not present in the main paper.\n\nThe limitation section is now in Appendix due to our limited space in main paper. We appreciate the advice, and will revise to include it in the main paper in the camera-ready version.\n\n> (Formatting)\n> \n> - EQGAT model is not cited or referred to in the text.\n> - The citations listed with a '-', like 1-4, 9-11, do not have access to all the numbers.\n\nWe thank the reviewer for catching this. In fact, EQGAT has been correctly cited in the Introduction section, with citation number 10, but we will also cite it in more appropriate places.\nWe note that similar citation ranges are commonly used in NeurIPS papers [5,6]. That said, we agree that explicit citations improve clarity and have revised accordingly. We will correct the ambiguous dash-style citation ranges (e.g., “1-4”) to list individual references.\n\n> (Q1) How does the model truly ensure that the equivariance to SE(3) is truly present in the generated samples? In both the augmented and non-augmented cases.\n\nThank you for the question. You can gently refer to Table 6 in our main paper. We provide RMSD (Root Mean Square Deviation) to measure equivariance in generated samples under different augmented and non-augmented scenarios. The low error validates that our latent diffusion process retains SE(3)-equivariance, as the generated samples closely match the original 3D coordinates.\n\n> (Q2) Given that this method relies on a lossless latent space, how can to ensure that this is not overfitting on training data? Could you also report (only) uniqueness scores in Tables 2 and 5?\n\nWe appreciate your question. We report reconstruction error on the **test sets**, in Figure 1 (a;b) and Table 9 in Appendix. These molecules are different from those in the training dataset, eliminating the possibility that our model is overfitting the training dataset.\n\nThe uniqueness scores are already reported in Tables 1 and 2, where “V&U” denotes molecules that are both valid and unique. A molecule is defined as unique only if it is also valid (following EDM[4], JODO[1])—otherwise, the notion of uniqueness becomes ill-defined. Our model achieves uniqueness scores of 0.913 on GEOM-DRUGS and 0.973 on QM9, demonstrating its ability to generate a diverse set of valid molecular structures.\n\n> (Q3) Could you provide additional details on how augmentation is done and added to the model?\n\nYes. As described in Section 3.1 of the paper, given a molecule’s 3D coordinates $\\mathbf{x}$, we apply a random rigid transformation by sampling a rotation $\\mathbf{R} \\sim \\mathrm{SO}(3)$ and a translation $\\mathbf{t} \\sim \\mathcal{N}(\\mathbf{0}, 0.01\\mathbf{I})$, and computing the augmented coordinates as $\\tilde{\\mathbf{x}} = \\mathbf{R}\\mathbf{x} + \\mathbf{t}$. This transformation alters the molecule’s global orientation and position in space but preserves its internal structure. The augmented coordinates $\\tilde{\\mathbf{x}}$, along with other graph features (i.e., atom types and bonds), are then fed into UAE-3D, which is trained to reconstruct $\\tilde{\\mathbf{x}}$, thereby encouraging SE(3)-equivariance through augmentation.\n\n> (Q4) How to update this model if one of the modalities says bonds have missing data. Would this approach overfit on the bonds present and not generate the bonds not seen before? Is this a good model for molecule discovery, and if so, could you elaborate on this?\n\nWe appreciate the insightful question and would like to clarify the following:\n- **Missing Bond Information can be mostly inferred using xyz2mol or OpenBabel.** Bond data can be inferred from the 3D coordinates. If we know the bonding connection is imperfect in the training dataset, we can use tools like xyz2mol and OpenBabel to pre-process the training data to verify and fill-in the missing chemical bonds. This process can handle most cases (>98%) according to the GEOM-DRUGS paper[7] These results should be good enough for a lot of cases and the little error is tolerable, given that the GEOM-DRUGs' ground truth annotation has 1.2% error of unstable molecules.\n- **Our model can function without bond information completely.** If the bond information is missing, we can train a new model without bond information as input and reconstruction target. The VAE can reconstruct the 3D coordinates without explicit bond information. \nFor example, in the CrossDocked2020 dataset which contains protein-ligand complexes, the bond type information is not available for protein's amino acids, but our model can still reconstruct the 3D coordinates of the protein-ligand complexes with high fidelity as shown in the Table below:\n\n| Atom Acc  | RMSD(Å) |\n| ---  | --- |\n| 1.000  | 0.011 |\n\n> (Q5) The tsne plot is unclear to me. If this is in a Center of mass frame, how do I interpret translations? What do these plots contribute?\n\nThank you for your comments. The t-SNE plot is intended to show the relative positions of real molecules in **latent space** (not the physical 3D coordinate system) under different coordinate transformations. Before these transformations, the molecules remain in a center of mass frame, but after the translations, they no longer remain centered.\n\nThe key insight of this plot is that the latent code of UAE-3D exhibits structural changes with the real geometric motion. This illustrates the meaningful structural changes in the molecular representation in SE(3)-equivariant augmentation.\n\nReferences:\n\n[1] Learning Joint 2-D and 3-D Graph Diffusion Models for Complete Molecule Generation, Han et al., TNNLS 2024.\n\n[2] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation, Vignac et al., ECML 2024.\n\n[3] Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation, Le et al., ICLR 2024.\n\n[4] Equivariant Diffusion for Molecule Generation in 3D, Hoogeboom et al., ICML 2022.\n\n[5] Convolutional Differentiable Logic Gate Networks, Petersen et al., NeurIPS 2024.\n\n[6] CSPG: Crossing Sparse Proximity Graphs for Approximate Nearest Neighbor Search, Yang et al., NeurIPS 2024.\n\n[7] GEOM, energy-annotated molecular conformations for property prediction and molecular generation, Axelrod et al., Scientific Data 2022."}, {"review_id": "Wm0wlSZsAV", "reviewer": "Reviewer_78B5", "summary": "The paper describes a two step approach to molecular generation where first a VAE (built using transformers) is trained to learn a unified per-atom latent space of 3D molecules and then a diffusion model is trained to generate such spaces. The model demonstrates success in both unconditional and conditional evaluations and is significantly more efficient than the compared approaches.", "questions": "How do you achieve near zero reconstruction with a VAE? The latent space is variational so one would expect sampling from the generated embedding to produce different outputs. Do you only evaluate the mean? What are sampled outputs like?\n\nTo what extent is this model permutation invariant?\n\nCan you design an ablation study or other experiment that supports your claim the accuracy of your model is due to the unified latent space instead of any of the other architectural and training procedure choices you made?\n\nWhat advantages does your approach have over recently described flow matching approaches?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths.  The described approach is highly efficient as diffusion can be done in the unified latent space using standard architectures.  The paper is clearly written and the model makes interesting use of transformers.\n\nWeaknesses. \nThis is a fast moving field and the SoA has advanced beyond the presented baselines (e.g., SemlaFlow outperforms the presented work both in terms of accuracy and speed).\n\nThe performance of the model is repeatedly attributed to the unified latent space, but no evidence is provided supporting this assertion.", "was_revised": true, "final_justification": "The additional evaluations greatly strengthen the paper - they show they outperform a more relevant SotA baseline and perform an ablation study demonstrating the effect of their choice of latent space.", "rebuttal": "> (W1) This is a fast moving field and the SoA has advanced beyond the presented baselines (e.g., SemlaFlow outperforms the presented work both in terms of accuracy and speed).\n> \n> (Q4) What advantages does your approach have over recently described flow matching approaches?\n\nWe thank the reviewer for pointing out recent advances. In response, we have added a comparison to SemlaFlow, a recent baseline. As shown in the following Table, our model outperforms SemlaFlow in both sampling speed and generation quality. The results are reproduced using the official code and pre-trained model checkpoints and evaluated with the same evaluation metrics as our model. \nNote that the missing metrics in SemlaFlow are due to the frequent NaN results in the calculation of related metrics caused by the existing errors in the generated molecules, making it difficult to compute reasonable reporting numbers. \nWe will further organize and analyze the results of SemlaFlow, and include this comparison and analysis in the revised version.\n\nTable: 2D Evaluation metrics on the QM9 dataset.\n| Method | FCD$\\downarrow$ | AtomStable | MolStable | V&C | V&U | V&U&N | SNN | Frag | Scaf |\n|---|---|---|---|---|---|---|---|---|---|\n| SemlaFlow | 0.863 | 0.995 | 0.949 | 0.857 | 0.821 | 0.821 | 0.124 | - | - |\n| Ours | 0.130 | 0.999 | 0.988 | 0.983 | 0.973 | 0.950 | 0.508 | 0.987 | 0.898 |\n\nTable: 3D evaluation metrics on the QM9 dataset.\n| Method | FCD$\\downarrow$ | AtomStable | MolStable | Bond length$\\downarrow$ | Bond angle$\\downarrow$ | Dihedral angle$\\downarrow$ |\n|---|---|---|---|---|---|---|\n| SemlaFlow | 1.127 | 0.971 | 0.787 | - | - | - |\n| Ours | 0.881 | 0.993 | 0.935 | 7.04E-02 | 9.84E-03 | 3.47E-04 |\n\nTable: Sampling speed comparison on the QM9 dataset.\n| Method | Sampling Time (s/molecule) |\n|---|---|\n| SemlaFlow | 0.993 |\n| Ours | 0.081 |\n\nOur approach offers two key advantages over recent flow-matching models such as SemlaFlow:\n\n- **Better Generation Quality**: Our model achieves significantly better generation quality across multiple metrics, including FCD, AtomStable, MolStable, and others. This is due to our unified latent space design that captures both 2D and 3D molecular features effectively.\n- **Improved sampling speed**: Our model achieves **10x faster sampling** compared to SemlaFlow, as shown in the above table. This is due to our latent diffusion model (LDM) design, which allows for efficient sampling in a unified latent space. In contrast, flow-matching models typically require more complex sampling procedures that are slower and less efficient.\n\n> (W2) The performance of the model is repeatedly attributed to the unified latent space, but no evidence is provided supporting this assertion.\n> \n> (Q3) Can you design an ablation study or other experiment that supports your claim the accuracy of your model is due to the unified latent space instead of any of the other architectural and training procedure choices you made?\n\nThank you for this valuable suggestion. To better validate the contribution of the unified latent space, we have conducted an additional ablation experiment to compare our model with its variant which have separate latent space for equivariant (3D) and invariant (2D) modalities. Specifically, we split the unified VAE into two independent neural networks: one maintaining the R-Trans architecture and taking 2D information (atom types, bond types) as input, while the other uses a vanilla transformer structure and takes 3D information (atom types and coordinates) as input. Each network generates a separate latent space, with dimensions set to half of the original hyperparameter. The LDM is then trained on the concatenated latent space of these two networks.\n\nTable: 2D Evaluation metrics on the QM9 dataset.\n| Latent Space | FCD$\\downarrow$ | AtomStable | MolStable | V&C | V&U | V&U&N | SNN | Frag | Scaf |\n|---|---|---|---|---|---|---|---|---|---|\n| Unified | 0.130 | 0.999 | 0.988 | 0.983 | 0.973 | 0.950 | 0.508 | 0.987 | 0.898 |\n| Separated | 0.351 | 0.995 | 0.952 | 0.943 | 0.920 | 0.913 | 0.341 | 0.940 | 0.682 |\n\nTable: 3D evaluation metrics on the QM9 dataset.\n| Latent Space | FCD$\\downarrow$ | AtomStable | MolStable | Bond length$\\downarrow$ | Bond angle$\\downarrow$ | Dihedral angle$\\downarrow$ |\n|---|---|---|---|---|---|---|\n| Unified | 0.881 | 0.993 | 0.935 | 7.04E-02 | 9.84E-03 | 3.47E-04 |\n| Separated | 2.356 | 0.982 | 0.872 | 18.4E-02 | 123E-03 | 7.03E-04 |\n\nThese results support our claim that the unified latent space plays a key role in our improved generation quality. We attribute this improvement to the early fusion of 2D and 3D information in the VAE encoder, which allows the LDM to model a single coherent modality. In contrast, for the separated latent spaces, the LDM needs to learn over two distinct modalities and their interactions, which increases the complexity of generation.\n\nWe will include this ablation and discussion in the revised manuscript.\n\n> (Q1) How do you achieve near zero reconstruction with a VAE? The latent space is variational so one would expect sampling from the generated embedding to produce different outputs. Do you only evaluate the mean? What are sampled outputs like?\n\nThank you for the insightful question. \n\n**We evaluate the VAE posterior's sampled latents instead of mean.** Our VAE setup follows the design in Stable Diffusion [1], where we use a very small KL regularization coefficient (1e-8; see Appendix Table 11) to prioritize reconstruction fidelity over latent regularization. This results in a low-variance posterior distribution ($\\mathrm{e}^{-8.1786}$), where samples from the posterior remain close to the mean. Consequently, the decoder can produce stable and high-fidelity reconstructions from sampled latents.\n\nImportantly, we do not remove the KL term entirely. While doing so still allows for accurate reconstruction, it significantly degrades the generation performance of the LDM, as shown in our ablation below. We attribute this to a key advantage of VAEs over AEs for generative modeling: the variational posterior ensures a structured and smooth latent space, which is critical for generation. \n\nAlthough we did not explicitly evaluate using the posterior mean as input to the LDM, we expect this to yield similarly poor results based on the observed impact of removing KL regularization.\n\nTable: Reconstruction error on the QM9 dataset.\n| KL weight | Atom Acc | Bond Acc | RMSD |\n|---|---|---|---|\n| 1E-8 | 1.0000 | 1.0000 | 0.002 |\n| 0 | 1.0000 | 1.0000 | 0.002 |\n\nTable: 2D Evaluation metrics on the QM9 dataset.\n| KL weight | AtomStable | MolStable | V&C | V&U | V&U&N |\n|---|---|---|---|---|---|\n| 1E-8 | 0.999 | 0.988 | 0.983 | 0.973 | 0.950 |\n| 0 | 0.995 | 0.961 | 0.982 | 0.962 | 0.950 |\n\nTable: 3D evaluation metrics on the QM9 dataset.\n| KL weight | AtomStable | MolStable |\n|---|---|---|\n| 1E-8 | 0.993 | 0.935 |\n| 0 | 0.982 | 0.833 |\n\n> (Q2) To what extent is this model permutation invariant?\n\nOur model is permutation equivariant, not invariant, in both the VAE and LDM components. This is because both modules produce atom-wise outputs and are implemented using Transformer-based architectures. Since the Transformer is inherently permutation equivariant—meaning that permuting the input atoms results in a corresponding permutation of the output—the overall model preserves this property.\n\nReferences:\n\n[1] High-Resolution Image Synthesis with Latent Diffusion Models. Rombach et al., CVPR 2022."}, {"review_id": "2OAbCJaETu", "reviewer": "Reviewer_eGEZ", "summary": "Molecules possess multiple modalities, such as atom types, inter-atomic bonds, and 3D spatial positions. Latent-based generative molecular models typically use distinct latent spaces to capture different properties like equivariance and invariance. In this manuscript, the authors propose a generative model that utilizes a single latent space—embedding all modalities within the same representation used for a diffusion process. The proposed method is faster than existing approaches while maintaining great accuracy.", "questions": "- In the section SE(3)-Equivariant Augmentations, the authors could cite other works that use similar data augmentation techniques—for instance, AlphaFold 3.\n\n- The manuscript refers to both Lossless and near-Lossless reconstruction. I slightly prefer the term near-Lossless, as any trained method may result in minor reconstruction losses.\n\n- For the sampling time comparisons: did you re-run each model (e.g., GeoLDM), or are the reported sampling times taken from their original papers?\n\n- I would be interested in further exploration of the learned latent space. For example, does it preserve molecular clusters?\n\n- What variant of Maximum Mean Discrepancy (MMD) do you use? The result depends on the kernel choice.\n\n*Minor Comments:*\n\n> Equivariance through our tailor-made SE(3) augmentations\n\nUnless the authors introduce novel augmentations, this phrasing may be overstated, as these data augmentation strategies are fairly common in small molecule and protein design.\n\n> ** indicates results reproduced\n\nThis sentence should not begin with a symbol.", "limitations": "The limitations of the work are not highlighted—either through experiments or in an explicit discussion. The authors should clarify and discuss limitations more clearly.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nThe authors provide extensive comparisons against many baseline methods. The proposed model appears faster while achieving comparable or superior results to existing approaches.\n\n**Weaknesses**\n\nOn GEOM-Drugs results:\n\n- JODO should be bold for the Frag metric.\n\nOn QM9 results:\n\n- JODO should be bold in Table 2 for the SNN metric.\n- In Table 2, EQGAT-diff should be bold for the AtomStable metric (2D).\n\nOverall, most results in Table 2 are very close, especially for the 2D metrics.\n\nIt is difficult to assess whether the model truly performs better for de novo design, as many scores are quite similar to those of JODO. If the authors are aware of another dataset that better differentiates between methods, additional experiments would be appreciated. On the other hand, since most studies use these datasets, the benchmark might be saturated. Still, the results are interesting, as they suggest that a different modeling paradigm can achieve comparable or better outcomes. I would suggest revising phrasing such as \"Breakthrough in Geometric Accuracy,\" as this claim is not clearly supported by the results in Tables 1 and 2. It is perfectly acceptable that the method does not outperform across all metrics—what is compelling is that a different architecture and training paradigm can achieve competitive results while being more efficient. This, in itself, can be a valuable contribution.", "was_revised": true, "final_justification": "I would like to thank the authors for engaging in the discussion. I think the new experiments improve the quality of the manuscript.", "rebuttal": "> (W1) On GEOM-Drugs results: JODO should be bold for the Frag metric.\n\nThank you for the gentle reminder. We have fixed these in the revision.\n\n> (W2) On QM9 results: JODO should be bold in Table 2 for the SNN metric. In Table 2, EQGAT-diff should be bold for the AtomStable metric (2D). Overall, most results in Table 2 are very close, especially for the 2D metrics.\n\nThank you for pointing this out. We have corrected the boldface formatting of Table 2 accordingly: JODO is now bolded for the SNN metric, and EQGAT-diff is bolded for AtomStable (2D). We also agree that most 2D metrics are close, which further emphasizes the necessity to explore beyond 2D geometry and evaluate holistic 3D generative fidelity. \nWe also want to emphasize that our method shows significantly better performance on 3D metrics, especially Bond Length, Bond Angle, and Dihedral Angle on the both datasets, indicating that our method achieves higher accuracy and reliability in 3D molecular generation.\n\n> (W3) It is difficult to assess whether the model truly performs better for de novo design, as many scores are quite similar to those of JODO. If the authors are aware of another dataset that better differentiates between methods, additional experiments would be appreciated. On the other hand, since most studies use these datasets, the benchmark might be saturated. Still, the results are interesting, as they suggest that a different modeling paradigm can achieve comparable or better outcomes. I would suggest revising phrasing such as \"Breakthrough in Geometric Accuracy,\" as this claim is not clearly supported by the results in Tables 1 and 2. It is perfectly acceptable that the method does not outperform across all metrics—what is compelling is that a different architecture and training paradigm can achieve competitive results while being more efficient. This, in itself, can be a valuable contribution.\n\nThank you for the appreciation. We agree and have toned down the statement of \"Breakthrough in Geometric Accuracy\" to \"Leading Geometric Accuracy\" to better reflect the empirical findings. While our method does not always dominate the benchmarks, it demonstrates improved or competitive  performance under a novel modeling paradigm. We believe the introduction of a unified, lossless latent space—combined with efficient training—constitutes a meaningful contribution to the field.\n\nWe also acknowledge the potential saturation of current benchmarks and are actively exploring more challenging and discriminative datasets for future work.\n\n> (Limitations) The limitations of the work are not highlighted—either through experiments or in an explicit discussion. The authors should clarify and discuss limitations more clearly.\n\nThe limitation section is now in Appendix A due to our limited space in main paper. We appreciate the advice, and will revise to include it in the main paper once have more space in the camera-ready version.\n\n> (Q1) In the section SE(3)-Equivariant Augmentations, the authors could cite other works that use similar data augmentation techniques—for instance, AlphaFold 3.\n\nThank you for the suggestion. We have included references to prior works such as AlphaFold 3 [2] and MCF [1] that also employ SE(3)-equivariant augmentations in Section 3.1 of the revised manuscript. \n\n> (Q2) The manuscript refers to both Lossless and near-Lossless reconstruction. I slightly prefer the term near-Lossless, as any trained method may result in minor reconstruction losses.\n\nWe appreciate this semantic clarification. In the revised version, we will consistently use the term *near-lossless* to acknowledge the minimal but nonzero reconstruction loss.\n\n> (Q3) For the sampling time comparisons: did you re-run each model (e.g., GeoLDM), or are the reported sampling times taken from their original papers?\n\nAll reported sampling times are obtained by re-running the official implementations of baselines (e.g., GeoLDM, EDM) under the same hardware and sampling settings to ensure fair comparison.\n\n> (Q4) I would be interested in further exploration of the learned latent space. For example, does it preserve molecular clusters?\n\nWe appreciate your insightful suggestion. We have clustered the molecules in the test set based on their Scaffold and embedded all molecules into the unified latent space, resulting in a t-sne visualization. We observe that molecules with the same scaffolds are relatively close in the latent space, showing that our unified latent space effectively captures molecular structure. Unfortunately, we are unable to provide this visualization in image form at this rebuttal stage, but we will include it in the final version.\n\n> (Q5) What variant of Maximum Mean Discrepancy (MMD) do you use? The result depends on the kernel choice.\n\nWe follow the same kernel and implementation as in baseline works (e.g., EDM, JODO), using an RBF kernel with a fixed bandwidth. For completeness, we will provide these details in the revised Appendix.\n\n> Equivariance through our tailor-made SE(3) augmentations: unless the authors introduce novel augmentations, this phrasing may be overstated, as these data augmentation strategies are fairly common in small molecule and protein design.\n\nWe appreciate this note. While our SE(3)-equivariant augmentations follow standard practice, we emphasize that their integration into both encoding and diffusion stages contributes to the equivariance of the unified latent space. We have revised the phrasing to avoid overstatement and better reflect the novelty in usage rather than in augmentation design per se.\n\n> \\* indicates results reproduced – this sentence should not begin with a symbol.\n\nThank you for pointing this out. We will rewrite these sentences as *The results marked with an asterisk (\\*) are reproduced using official source codes*.\n\nReference:\n\n[1] Swallowing the Bitter Pill: Simplified Scalable Conformer Generation, Wang et al., ICML 2024.\n\n[2] Accurate structure prediction of biomolecular interactions with AlphaFold 3, Abramson et al., Nature 2024."}]

Author rebuttal:

Meta-review: This paper tackles the challenging problem of unified latent encoding for 3D molecules, given their inherently multimodal nature. The proposed autoencoder uses R-Trans to encode atom types, bonds, and coordinates into a single latent space and a generic transformer for decoding, trained with well-motivated loss functions and SE(3)-equivariant augmentations. The approach achieves low reconstruction error and enables high-quality DiT-based generation. Reviewers highlighted the significance of the unified latent space, the efficiency of the framework, and the strong empirical performance. Concerns centered mainly on missing baselines and the limited scope of evaluation, but many of these were addressed convincingly in the rebuttal. Overall, I find the contribution meaningful and well-supported, and I recommend acceptance.
Decision: Accept (poster)