Research map / Probabilistic and scientific modeling / Molecules and proteins
WGFormer: An SE(3)-Transformer Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes WGFormer, an SE(3)-Transformer whose Sinkhorn-based attention layers are interpreted as Wasserstein gradient flows minimizing a latent energy, used in an auto-encoder to refine RDKit conformations into DFT-level ground-state conformations.
Problem
Predicting molecular ground-state (energy-minimized) 3D conformations accurately and efficiently; existing learning methods are empirically designed and lack interpretability and physically meaningful energy minimization.
Approach
Designs a Transformer layer whose attention uses Sinkhorn scaling with QKV adjustments so that stacked layers correspond to Wasserstein gradient flow steps minimizing an entropic optimal transport energy over latent atom mixtures; the decoder predicts coordinate residuals as weighted displacements, keeping SE(3)-equivariance.
Key results
- On Molecule3D random split test, WGFormer achieves D-MAE 0.392 and C-RMSD 0.664, beating the best baseline ConfOpt-ThreeAtom (0.430 D-MAE, 0.736 C-RMSD).
- On QM9 test, WGFormer reduces C-RMSD to 0.206 versus 0.246 for ConfOpt-ThreeAtom, roughly 15% lower.
- WGFormer runs over 50% faster per molecule than the top baselines while using 30 layers, reported as 0.150 s/mol versus 18.058 s/mol for GeoDiff on QM9.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use WGFormer instead of ConfOpt-ThreeAtom | Molecular ground-state conformation prediction · Molecule3D (random) test | C-RMSD: 0.736 -> 0.664 | Helped |
| Use WGFormer instead of ConfOpt-ThreeAtom | Molecular ground-state conformation prediction · Molecule3D (scaffold) test | C-RMSD: 0.691 -> 0.627 | Helped |
| Use WGFormer instead of ConfOpt-ThreeAtom | Molecular ground-state conformation prediction · QM9 test | C-RMSD: 0.246 -> 0.206 | Helped |
| Use WGFormer instead of ConfOpt on DRUGS | Molecular ground-state conformation prediction · DRUGS | D-RMSE: 1.272 -> 1.281 | Hurt |
| Add MAE loss to RMSD-only training (ablation) | WGFormer · Molecule3D (random) validation | C-RMSD: 0.716 -> 0.662 | Helped |
| Add RMSD loss to MAE-only training (ablation) | WGFormer · Molecule3D (random) validation | C-RMSD: 0.688 -> 0.662 | Helped |
| Add Sinkhorn module and QKV adjustment together (ablation) | WGFormer attention head · Molecule3D (scaffold) validation | C-RMSD: 0.661 -> 0.618 | Helped |
| Add Sinkhorn module without QKV adjustment (ablation) | WGFormer attention head · Molecule3D (scaffold) validation | C-RMSD: 0.661 -> 0.656 | Helped |
| Add QKV adjustment without Sinkhorn module (ablation) | WGFormer attention head · Molecule3D (scaffold) validation | C-RMSD: 0.661 -> 0.658 | Helped |
| Replace 30 distinct layers with 6 layers repeated 5 times (ablation) | WGFormer · QM9 test | C-RMSD: 0.206 -> 0.219 | Hurt |
Surprising findings
- Latent energy and xTB potential energy decrease across layers with Pearson correlation 0.885 ± 0.033, supporting the interpretation as energy minimization.
- Using 6 repeated layers (each applied 5 times) yields performance comparable to 30 distinct layers on QM9.
Limitations
- For non-Dirac atom measures, the model does not correspond to Wasserstein gradient flows when M goes to infinity.
- Experiments are limited to Molecule3D and QM9, with DRUGS only briefly reported in the appendix.
Experimental setup
Datasets: Molecule3D QM9
Benchmarks: Molecule3D (random split) Molecule3D (scaffold split) QM9 DRUGS
Metrics: D-MAE D-RMSE C-RMSD inference time
Baselines: GTMGC ConfOpt-TwoAtom ConfOpt-ThreeAtom SE(3)-Transformer EGNN GeoDiff TorsionDiff
Builds on: Sinkformer ConfOpt Uni-Mol SE(3)-Transformer
Closest papers by meaning
- DiSCO: Diffusion Schrödinger Bridge for Molecular Conformer Optimization (AAAI 2024 · similarity 0.89)
- Swallowing the Bitter Pill: Simplified Scalable Conformer Generation (ICML 2024 · similarity 0.89)
- Learning Over Molecular Conformer Ensembles: Datasets and Benchmarks (ICLR 2024 · similarity 0.87)
- SE(3)-Stochastic Flow Matching for Protein Backbone Generation (ICLR 2024 · similarity 0.86)
- Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling (ICML 2025 · similarity 0.86)
- Structure-Aware E(3)-Invariant Molecular Conformer Aggregation Networks (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.