# Toward Better Geometric Representations for Molecule Generative Models

arXiv: https://arxiv.org/abs/2605.07693 · HTML: https://mlsanity.com/paper/2605.07693/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 58 references that are in the map.
- First posted on arXiv 8 May 2026.

## Builds on (most cited first, 30 of 40)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Feature Pyramid Networks for Object Detection](https://mlsanity.com/paper/1612.03144/) (2016)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking](https://mlsanity.com/paper/2210.01776/) (2022)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [Structure-based Drug Design with Equivariant Diffusion Models](https://mlsanity.com/paper/2210.13695/) (2022)
- [Geometric Latent Diffusion Models for 3D Molecule Generation](https://mlsanity.com/paper/2305.01140/) (2023)
- [3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction](https://mlsanity.com/paper/2303.03543/) (2023)
- [SPIn-NeRF: Multiview Segmentation and Perceptual Inpainting with Neural Radiance Fields](https://mlsanity.com/paper/2211.12254/) (2022)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation](https://mlsanity.com/paper/2302.09048/) (2023)
- [Return of Unconditional Generation: A Self-supervised Representation Generation Method](https://mlsanity.com/paper/2312.03701/) (2023)
- [Diffusion Model as Representation Learner](https://mlsanity.com/paper/2308.10916/) (2023)
- [Mixed Continuous and Categorical Flow Matching for 3D De Novo Molecule Generation](https://mlsanity.com/paper/2404.19739/) (2024)
- [Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation](https://mlsanity.com/paper/2309.17296/) (2023)
- [SemlaFlow -- Efficient 3D Molecular Generation with Latent Attention and Equivariant Flow Matching](https://mlsanity.com/paper/2406.07266/) (2024)
- [Geometry-Complete Diffusion for 3D Molecule Generation and Optimization](https://mlsanity.com/paper/2302.04313/) (2023)
- [Group SELFIES: A Robust Fragment-Based Molecular String Representation](https://mlsanity.com/paper/2211.13322/) (2022)
- [Fractional Denoising for 3D Molecular Pre-training](https://mlsanity.com/paper/2307.10683/) (2023)
- [A Survey of Generative AI for de novo Drug Design: New Frontiers in Molecule and Protein Generation](https://mlsanity.com/paper/2402.08703/) (2024)
- [Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation](https://mlsanity.com/paper/2305.12347/) (2023)
- [FlowMol3: Flow Matching for 3D De Novo Small-Molecule Generation](https://mlsanity.com/paper/2508.12629/) (2025)

## Built on by (most cited first, 1 of 1)

- [Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains](https://mlsanity.com/paper/2610.11454/) (2026)

## Abstract (arXiv)

Geometric representation-conditioned molecule generation provides an effective paradigm that decouples molecule representation modeling from structure generation. By decoupling molecule generation into two stages-first generating a meaningful molecule representation, and then generating a 3D molecule conditioned on this representation-the efficiency and quality of the generation process can be significantly enhanced. However, its effectiveness is fundamentally limited by the quality of the representation space: pretrained molecular encoders, such as UniMol, produce representations that are non-smooth and not fully exploited during the generative training process. In this work, we propose LENSEs, a framework that better exploits the potential of molecule representations in representation-conditioned generation methods. In particular, LENSEs introduces three complementary mechanisms: (1) a representation head, simultaneously trained during generative tasks, that extracts multi-level representations from the pretrained encoder; (2) a molecule perceptual loss that optimizes the generator in a semantic-informative representation space; and (3) a node-level representation alignment (REPA) loss that explicitly aligns the generator's hidden states with encoder representations, reducing the semantic gap between pretraining and generation. We demonstrate the effectiveness of these improvements through extensive molecule generation tasks. Specifically, on the challenging molecule generation dataset GEOM-DRUG, LENSEs achieves 97.28% validity and 98.51% molecule stability, surpassing existing advanced methods. Further analyses through Lipschitz constant reduction (4.6x) and QM9 probing tasks also demonstrate the smoother, more informative refined representations, establishing generative training with alignment objectives as a potential pretraining paradigm for molecular encoders.