Elign: Equivariant Diffusion Model Alignment from Foundational Machine Learning Force Fields

TL;DR

Elign tackles the gap between diffusion-based molecular generation and thermodynamic realism by performing post-training alignment that amortizes expensive quantum evaluations. It introduces a foundation MLFF to provide energy and force signals and formulates reverse diffusion as an RL problem, optimizing a Force–Energy Disentangled GRPO objective that uses potential-based shaping to bias toward low-energy configurations. Theoretical analysis connects the KL-regularized alignment to a Gibbs-like tilt of the terminal distribution, and empirical results show significant gains in stability (e.g., QM9 and GEOM-Drugs) while preserving unguided inference speed. The approach yields conformations with lower MLFF energies and forces, and DFT oracle checks confirm improved physical fidelity, making it a scalable pathway to physically guided 3D molecular generation.

Abstract

Generative models for 3D molecular conformations must respect Euclidean symmetries and concentrate probability mass on thermodynamically favorable, mechanically stable structures. However, E(3)-equivariant diffusion models often reproduce biases from semi-empirical training data rather than capturing the equilibrium distribution of a high-fidelity Hamiltonian. While physics-based guidance can correct this, it faces two computational bottlenecks: expensive quantum-chemical evaluations (e.g., DFT) and the need to repeat such queries at every sampling step. We present Elign, a post-training framework that amortizes both costs. First, we replace expensive DFT evaluations with a faster, pretrained foundational machine-learning force field (MLFF) to provide physical signals. Second, we eliminate repeated run-time queries by shifting physical steering to the training phase. To achieve the second amortization, we formulate reverse diffusion as a reinforcement learning problem and introduce Force--Energy Disentangled Group Relative Policy Optimization (FED-GRPO) to fine-tune the denoising policy. FED-GRPO includes a potential-based energy reward and a force-based stability reward, which are optimized and group-normalized independently. Experiments show that Elign generates conformations with lower gold-standard DFT energies and forces, while improving stability. Crucially, inference remains as fast as unguided sampling, since no energy evaluations are required during generation.

Figures (10)

• Figure 1: Staged pipeline for equilibrium molecular generation. (i) Pretraining: score-matching trains an E(3)-equivariant diffusion model $\pi_\theta^{\text{pre}}$ on approximate structures. (ii) Optional SFT: conditional fine-tuning improves adherence to discrete specifications. (iii) Preference model: a foundation MLFF $\phi$ provides energy/force signals. (iv) Post-training: RL fine-tunes the sampler to improve stability without run-time oracle calls.
• Figure 2: Overview of Elign. (a) Rollout branching with shared prefix: starting from a CoM-free Gaussian prior at $t=T$, an EGNN policy denoises to $t=0$. At $t=T_{\mathrm{prefix}}$, we cache a shared prefix state and branch into $K$ rollouts with independent noise. Each trajectory is propagated to its terminal state $\bm{z}_0$ and scored by a foundational MLFF. (b) Energy-based reward shaping: intermediate predicted clean geometries $\hat{\bm{z}}_{0|t}$ are evaluated by the MLFF, and local energy differences $E_\phi(\hat{\bm{z}}_{0|t}) - E_\phi(\hat{\bm{z}}_{0|t-1})$ provide dense shaping signals that bias the policy toward lower-energy conformations. (c) FED-GRPO: terminal energy and RMS force rewards are z-score normalized separately per timestep, then combined to compute the final advantage for policy updates.
• Figure 3: Computational efficiency comparison. (Left) Inference time per molecule for post-trained models versus guidance-based methods. (Right) Training epoch time comparing different reward oracles on a log scale. For a fair comparison, all methods use terminal-only rewards (no PBRS). Evaluation was conducted on a node with an NVIDIA H100 GPU and Intel Xeon Sapphire Rapids CPUs (Xeon Platinum 6458Q, AVX-512 enabled).
• Figure 4: DFT oracle evaluation of generated molecules from QM9. Cumulative distribution functions comparing baseline and post-trained models on force RMS (left) and formation energy per atom (right). The results are obtained from 1024 samples.
• Figure 5: Geometry overlays on GEOM-Drugs. Top row: baseline diffusion model (green). Bottom row: post-trained model (pink). In each panel, the generated conformer (dark) is overlaid on the reference structure (light gray). Post-training improves connectivity and overall geometric fidelity to the reference conformer.

Theorems & Definitions (10)

• Theorem 1: Energy-aligned terminal distribution
• Theorem 2: Energy error implies distribution error
• Theorem 3: Energy-aligned terminal distribution (restated)
• proof
• Theorem 4: Energy error implies distribution error (restated)
• Lemma 1: Likelihood-ratio bound implies TV bound
• proof
• proof : Proof of Theorem \ref{['thm:energy_to_dist_app']}
• Theorem 5: PBRS induces an approximate (alchemical) force in the reverse drift
• proof