Applications of Modular Co-Design for De Novo 3D Molecule Generation

TL;DR

<3-5 sentence high-level summary> Megalodon presents a scalable, multi-modal transformer for de novo 3D molecule generation that jointly models 3D structure and discrete chemical features, enabling effective diffusion and flow-matching generation. By incorporating equivariant layers and self-conditioning, it achieves state-of-the-art results on unconditional and conditional generation, and introduces energy-aware benchmarks that assess conformer realism. The model scales to 40M parameters and demonstrates substantially improved valid-molecule counts and lower relaxation energies, approaching thermodynamic minima on the GEOM energy landscape. This work highlights the practical impact of combining discrete and continuous data modeling with geometry-aware transformers for accelerated, energetically feasible molecular discovery.

Abstract

De novo 3D molecule generation is a pivotal task in drug discovery. However, many recent geometric generative models struggle to produce high-quality 3D structures, even if they maintain 2D validity and topological stability. To tackle this issue and enhance the learning of effective molecular generation dynamics, we present Megalodon-a family of scalable transformer models. These models are enhanced with basic equivariant layers and trained using a joint continuous and discrete denoising co-design objective. We assess Megalodon's performance on established molecule generation benchmarks and introduce new 3D structure benchmarks that evaluate a model's capability to generate realistic molecular structures, particularly focusing on energetics. We show that Megalodon achieves state-of-the-art results in 3D molecule generation, conditional structure generation, and structure energy benchmarks using diffusion and flow matching. Furthermore, doubling the number of parameters in Megalodon to 40M significantly enhances its performance, generating up to 49x more valid large molecules and achieving energy levels that are 2-10x lower than those of the best prior generative models.

Figures (6)

• Figure 1: Megalodon Architecture: molecules are separated into 3D structures and discrete atom types, bond types, and atom charge features. All features are embedded separately, passed through a feed-forward layer, and aggregated to produce the input tokens for the fused Invariant Transformer blocks. The embedded structure features and transformer outputs for the discrete features are passed to a single EGNN layer for structure updates. The output heads consist of standard MLPs and an EGNN layer for bond refinement.
• Figure 2: Time and interpolation comparison between Megalodon and Megalodon-flow
• Figure 3: Diffusion model performance as a function of molecule size. Note the ability for Megalodon to generate valid and stable molecules with little training data support.
• Figure 4: Megalodon molecule generation dynamics generated with Imagen 2
• Figure 5: Distribution of molecule sizes