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

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.

💡 Research Summary

Elign introduces a novel post‑training framework that aligns E(3)‑equivariant diffusion models with high‑fidelity physical potentials while preserving the fast inference speed of the original generative model. The authors identify two major computational bottlenecks in existing physics‑guided diffusion approaches: (1) the high cost of quantum‑chemical evaluations such as density functional theory (DFT), and (2) the need to query these expensive oracles at every diffusion step during sampling. To address the first bottleneck, Elign leverages a pretrained “foundational” machine‑learning force field (MLFF) that has been trained on large, diverse quantum‑mechanical datasets. This MLFF provides rapid approximations of both potential energy and atomic forces, effectively replacing on‑the‑fly DFT calculations with cheap neural‑network evaluations.

The second bottleneck is tackled by moving the physical steering from inference time to a dedicated post‑training phase. The reverse diffusion process is cast as a finite‑horizon Markov decision process (MDP) where each denoising step corresponds to a stochastic policy. Starting from a pretrained diffusion model (the base policy), the authors fine‑tune the policy using reinforcement learning (RL) to maximize preferences defined by the MLFF. The key contribution is the Force‑Energy Disentangled Group Relative Policy Optimization (FED‑GRPO) algorithm. FED‑GRPO defines two separate reward components: an energy reward (negative of the MLFF‑predicted energy) and a force reward (negative of the squared norm of the MLFF‑predicted forces). Each reward is independently z‑score normalized per diffusion timestep, thereby disentangling their scales and preventing one from dominating the other. The normalized rewards are then combined to compute an advantage estimate used in a PPO‑style trust‑region update, which constrains the updated policy to stay close to the original diffusion model (KL‑regularization). This design mirrors the “post‑training alignment” used in large language models, ensuring that the model retains its learned data distribution while incorporating physical preferences.

Experiments are conducted on QM9 (small organic molecules with DFT reference data) and GEOM‑Drugs (a diverse set of drug‑like conformers). The baseline is an E(3)‑equivariant diffusion model trained by score matching on semi‑empirical structures, without any physics guidance. Elign is then applied as a post‑training step. Evaluation metrics include mean absolute error (MAE) of DFT energies, root‑mean‑square (RMS) forces, and structural diversity. Results show that Elign reduces the energy MAE by roughly 15–20 % and the RMS force by a comparable margin compared to the unguided baseline, while matching the baseline’s sampling speed because no energy or force evaluations are required at inference time. In contrast, methods that perform runtime DFT or semi‑empirical guidance incur orders‑of‑magnitude slower generation.

The paper’s contributions can be summarized as follows:

1. Cost‑effective physical signals – By substituting DFT with a pretrained, broad‑coverage MLFF, Elign obtains high‑quality energy and force information at negligible computational cost.
2. Amortized inference – Physical steering is shifted entirely to a post‑training RL phase, eliminating per‑step oracle calls during generation.
3. FED‑GRPO algorithm – A novel RL objective that disentangles energy and force rewards, applies per‑timestep group normalization, and uses trust‑region updates to preserve the pretrained diffusion distribution.

Overall, Elign demonstrates that high‑fidelity, physically stable 3‑D molecular conformations can be generated at the speed of standard diffusion samplers, opening the door for large‑scale applications in drug discovery, materials design, and computational chemistry where both accuracy and throughput are critical. Future directions include scaling to even larger foundational force fields, incorporating multi‑objective physical constraints (e.g., dipole moments, vibrational spectra), and extending the alignment framework to other equivariant generative paradigms such as normalizing flows or equivariant VAEs.