Elign: Equivariant Diffusion Model Alignment from Foundational Machine Learned 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) that estimates molecular energies and forces. Second, we eliminate repeated run-time queries by shifting physical steering to the post-training phase. To achieve the second amortization, we formulate reverse diffusion as a reinforcement learning problem and propose to use Group Relative Policy Optimization (GRPO) to fine-tune the denoising policy. Our objective combines a potential-based energy reward and a force-based stability reward, which are optimized in a disentangled fashion. 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.

Computers can now generate three-dimensional shapes of molecules — a capability useful for designing drugs, materials, and new chemical compounds. The catch is that today's AI generators often produce molecules that look plausible but aren't physically stable: real molecules in nature wouldn't actually settle into the shapes the AI proposes. Researchers can fix this by checking each generated shape against physics during generation — but the most accurate physics method (called density functional theory, or DFT) is extremely slow, and it has to be re-run at every step of the AI's process. This makes high-quality molecule generation prohibitively expensive. We introduce Elign, a method that removes both bottlenecks. First, instead of slow DFT, we use a much faster, AI-based physics approximation known as a machine-learning force field, which has been pretrained on large quantities of chemistry data. Second, instead of consulting physics at every generation step, we use reinforcement learning to bake the physics knowledge directly into the generator during a one-time fine-tuning phase. After fine-tuning, the generator produces physically realistic molecules at exactly the same speed as the original AI — no further physics calculations are needed when generating new molecules. In our experiments, molecules produced by Elign are measurably closer to physical equilibrium — lower energies and smaller residual forces when checked against gold-standard DFT calculations — while generation remains as fast as the original, unguided AI. In short: Elign teaches an AI molecule generator to respect the laws of physics, without making it any slower at its job.