MolAlign3D: Enhancing Fixed-Dimensional E(3)-Equivariant Latent Space for High-Fidelity 3D Molecular Reconstruction and Editing

Zitao Chen, Jiatong Ji, Yinjun Jia, Wei-Ying Ma, Yanyan Lan
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:14721-14746, 2026.

Abstract

Recent advances in 3D molecular modeling have achieved high-fidelity structural synthesis, yet these models often lack an explicit and manipulable representation space. To address this, MolFLAE introduced a fixed-dimensional, E(3)-equivariant latent space, providing a novel framework for molecular editing independent of atom counts. However, because its latent space was primarily optimized for geometric reconstruction, it remains semantically shallow and inadequate for comprehensive representation learning. In this work, we propose MolAlign3D, which evolves this architecture into a unified semantic-generative engine. By anchoring MolFLAE’s manipulable latents with embeddings from a pre-trained molecular encoder, we yield a manifold that is both semantically dense and geometrically precise. Experiments show that MolAlign3D achieves high-fidelity molecular reconstruction and attains comparable performance on molecular property prediction benchmarks. Notably, the integration of rich semantic priors significantly enhances zero-shot molecular manipulation, including atom-number editing and latent-space interpolation, outperforming prior fixed-dimensional equivariant latent baseline.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-chen26ay, title = {{M}ol{A}lign3{D}: Enhancing Fixed-Dimensional E(3)-Equivariant Latent Space for High-Fidelity 3{D} Molecular Reconstruction and Editing}, author = {Chen, Zitao and Ji, Jiatong and Jia, Yinjun and Ma, Wei-Ying and Lan, Yanyan}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {14721--14746}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/chen26ay/chen26ay.pdf}, url = {https://proceedings.mlr.press/v306/chen26ay.html}, abstract = {Recent advances in 3D molecular modeling have achieved high-fidelity structural synthesis, yet these models often lack an explicit and manipulable representation space. To address this, MolFLAE introduced a fixed-dimensional, E(3)-equivariant latent space, providing a novel framework for molecular editing independent of atom counts. However, because its latent space was primarily optimized for geometric reconstruction, it remains semantically shallow and inadequate for comprehensive representation learning. In this work, we propose MolAlign3D, which evolves this architecture into a unified semantic-generative engine. By anchoring MolFLAE’s manipulable latents with embeddings from a pre-trained molecular encoder, we yield a manifold that is both semantically dense and geometrically precise. Experiments show that MolAlign3D achieves high-fidelity molecular reconstruction and attains comparable performance on molecular property prediction benchmarks. Notably, the integration of rich semantic priors significantly enhances zero-shot molecular manipulation, including atom-number editing and latent-space interpolation, outperforming prior fixed-dimensional equivariant latent baseline.} }
Endnote
%0 Conference Paper %T MolAlign3D: Enhancing Fixed-Dimensional E(3)-Equivariant Latent Space for High-Fidelity 3D Molecular Reconstruction and Editing %A Zitao Chen %A Jiatong Ji %A Yinjun Jia %A Wei-Ying Ma %A Yanyan Lan %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-chen26ay %I PMLR %P 14721--14746 %U https://proceedings.mlr.press/v306/chen26ay.html %V 306 %X Recent advances in 3D molecular modeling have achieved high-fidelity structural synthesis, yet these models often lack an explicit and manipulable representation space. To address this, MolFLAE introduced a fixed-dimensional, E(3)-equivariant latent space, providing a novel framework for molecular editing independent of atom counts. However, because its latent space was primarily optimized for geometric reconstruction, it remains semantically shallow and inadequate for comprehensive representation learning. In this work, we propose MolAlign3D, which evolves this architecture into a unified semantic-generative engine. By anchoring MolFLAE’s manipulable latents with embeddings from a pre-trained molecular encoder, we yield a manifold that is both semantically dense and geometrically precise. Experiments show that MolAlign3D achieves high-fidelity molecular reconstruction and attains comparable performance on molecular property prediction benchmarks. Notably, the integration of rich semantic priors significantly enhances zero-shot molecular manipulation, including atom-number editing and latent-space interpolation, outperforming prior fixed-dimensional equivariant latent baseline.
APA
Chen, Z., Ji, J., Jia, Y., Ma, W. & Lan, Y.. (2026). MolAlign3D: Enhancing Fixed-Dimensional E(3)-Equivariant Latent Space for High-Fidelity 3D Molecular Reconstruction and Editing. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:14721-14746 Available from https://proceedings.mlr.press/v306/chen26ay.html.

Related Material