Reproduce “A standard transformer and attention with linear biases for molecular conformer generation” on OpenResearch

About this paper

Viatcheslav Gurev, Timothy Rumbell

IBM Research introduces S23D, a diffusion model utilizing a standard transformer with an attention mechanism incorporating linear biases for molecular conformer generation. The model achieves state-of-the-art recall on the GEOM-DRUGS benchmark with significantly fewer parameters (e.g., 10% the size of previous large models) and demonstrates improved precision through a post-hoc chirality correction step.

Key insights

• Standard transformer architectures can achieve state-of-the-art recall in molecular conformer generation with significantly fewer parameters than previous non-equivariant models by using effective positional encodings.
• Implementing an ALiBi-like relative positional encoding, which linearly biases attention based on shortest path distances, is computationally efficient and improves performance over graph Laplacian eigenvectors.
• A two-stage training protocol that initially removes hydrogens and then finetunes with them substantially accelerates training for large molecular datasets.
• Post-hoc chirality correction is crucial for improving precision metrics in non-equivariant models and for fair comparison with equivariant models.

Summary by alphaXiv.