Toward Better Geometric Representations for Molecule Generative Models

Technical Summary: LENSEs – Latent Enhancement for Non-smooth Structural Encodings

Problem Statement
Geometric representation-conditioned generation (GeoRCG) has emerged as an efficient paradigm for 3D molecule generation, decoupling the modeling of molecular representations from the generation of 3D structures. In this framework, a pretrained molecular encoder (e.g., UniMol, Frad) defines a latent space, a Representation Diffusion Model (RDM) samples a molecule-level vector from this space, and a generator reconstructs the 3D structure conditioned on this vector.

However, the effectiveness of GeoRCG is fundamentally limited by the quality of the latent space provided by pretrained encoders. The authors identify two critical bottlenecks:

1. Non-smooth Representation Manifolds: Pretrained encoders, optimized primarily for property prediction via coordinate denoising, produce representations with high Lipschitz constants. Small perturbations in molecular conformation lead to disproportionately large shifts in the latent space. This non-smoothness makes the distribution difficult for the RDM to model, often necessitating compensatory mechanisms like hand-tuned noise injection.
2. Underexploited Multi-level Semantics: Pretrained encoders exhibit hierarchical features where shallow layers capture local structural motifs (e.g., functional groups) and deeper layers encode global semantics. Existing GeoRCG pipelines typically utilize only the final-layer representation, discarding the rich semantic information preserved in earlier layers. Consequently, the generator must relearn these semantics from scratch, increasing optimization difficulty.

Methodology: LENSEs
To address these issues, the authors propose LENSEs (Latent Enhancement for Non-smooth Structural Encodings), a framework that actively refines the conditioning latent space during generative training. Rather than relying passively on a frozen encoder, LENSEs introduces three complementary mechanisms:

1. Representation Head with Learnable Layer Pooling:
   Instead of using a single deterministic output from the encoder, LENSEs aggregates multi-layer features. It employs a learnable layer-pooling mechanism to weight representations from different encoder depths ($H^{(l)}$). These pooled features are mapped to the parameters of a diagonal Gaussian posterior ($\mu, \sigma$) via projection heads. The latent code $z$ is sampled using the reparameterization trick. This head is trained jointly with the generator to extract representations that are both semantically rich and generation-friendly, regularized by a KL-divergence loss to ensure a compact latent space.

2. Molecule Perceptual Loss ($L_{PERC}$):
   Inspired by perceptual losses in computer vision, this loss aligns the high-level semantic representation of the denoised (generated) molecule with that of the clean ground-truth molecule within the frozen encoder's feature space. By passing both the generated and clean molecules through the encoder and minimizing the distance between their pooled graph-level representations, the generator is encouraged to preserve meaningful structural patterns (e.g., functional groups) rather than just minimizing low-level coordinate errors. A stop-gradient operation is applied to the ground-truth representation.

3. Representation Alignment (REPA) Loss ($L_{REPA}$):
   Motivated by the observation that generator hidden states naturally align with encoder features, this loss explicitly aligns the node-level features of the generator with those of the pretrained encoder. A learnable projector maps generator features into the encoder's feature space, and cosine similarity is maximized between corresponding nodes. This reduces the semantic gap between pretraining and generation, accelerating training and improving the acquisition of chemically meaningful knowledge.

The total training objective combines the base generation loss ($L_{gen}$) with the KL loss, perceptual loss, and REPA loss.

Key Results
The authors evaluate LENSEs on the GEOM-DRUG and QM9 datasets, comparing it against state-of-the-art baselines including GeoRCG, CanonFlow, and SemlaFlow.

• Generation Quality: On the challenging GEOM-DRUG dataset, LENSEs achieves state-of-the-art performance. It attains 97.28% validity and 98.51% molecule stability, surpassing the previous best GeoRCG (95.3% validity, 97.20% stability) and CanonFlow (95.9% validity). It also achieves the lowest energy and strain metrics, indicating superior physical plausibility.
• Representation Smoothness: Analysis of the Lipschitz constant reveals that the refined latent space is significantly smoother than the raw encoder space. On GEOM-DRUG, the empirical Lipschitz constant is reduced by 4.6×, confirming that the representation manifold is more amenable to diffusion modeling.
• Representation Richness: The effective rank of the representation increases from 65.32 (raw encoder) to 88.42 (LENSEs), indicating a reduction in dimensional collapse and a richer, more structured latent space.
• Downstream Utility: When used for downstream regression tasks on QM9 properties with a frozen encoder and representation head, LENSEs outperforms the raw encoder across all properties, demonstrating that the generative training process itself acts as an effective pretraining paradigm for refining molecular representations.
• Sampling Efficiency: LENSEs demonstrates improved sampling efficiency, achieving high validity and stability with fewer generation steps compared to the baseline.

Significance and Claims
The paper claims that high-quality 3D molecule generation depends not only on the generator architecture but critically on the quality of the conditioning representation. LENSEs demonstrates that the latent space of pretrained encoders can be systematically improved to be smoother, more semantically informative, and better aligned with generative objectives.

The authors posit that their work establishes a bidirectional relationship between generation quality and representation quality. Specifically, they suggest that generative training with alignment objectives can serve as a potential pretraining paradigm for molecular encoders, refining representations to be more predictive and useful for downstream tasks. The work highlights that useful molecular representations are not inherently generation-friendly and require explicit refinement to unlock their full potential in generative modeling.