Crystal structure diffusion

t = —

structure

All three crystal components are denoised jointly by one E(3)-equivariant graph network:

Fractional coordinates — diffused on the periodic 3-torus via wrapped-normal score matching, sampled with a predictor–corrector (Langevin) scheme that respects translational symmetry.
Atom types — each atom is relaxed to a continuous 100-dimensional one-hot vector, diffused with Gaussian DDPM, then decoded back to a discrete element by arg-max.
Lattice (3×3) — the lattice matrix is diffused with Gaussian DDPM, equivariant to global rotations of the unit cell.

DPO — aligning the generator with stability

The pre-trained generator is aligned with a stability objective using Direct Preference Optimization. Unlike online RL (which queries a reward model inside the training loop), DPO is offline: sample a pool of crystals once, score each by CHGNet stability, pair a more-stable winner against a less-stable loser, and train the generator to prefer the winner — a supervised-style loss with no reward model in the loop and no policy-gradient instability.

1Sample a poolcrystals from the SFT generator

→

2Score stabilityCHGNet energy / atom

→

3Build pairsstable winner vs loser

→

4Prefer winnerDPO update (β-regularized)

How DPO works on one preference pair

One real winner/loser pair; the gauge plays back the real σ(z) = P(prefer winner) during training.

Mathematical formulation

Standard DPO turns a preference (winner x_w ≽ loser x_l) into a logistic loss on the policy-vs-reference log-likelihood ratio:

𝓛 = − 𝔼 log σ ( β [ log π_θ(x_w)π_ref(x_w) − log π_θ(x_l)π_ref(x_l) ] )

A diffusion model's likelihood is intractable, so each crystal's term becomes its denoising error, decomposed over the three DiffCSP components — lattice (L), fractional coordinates (F), atom types (A) — under a shared noise level t:

Δ_φ(x) = w_L ‖ε_L − ε_φ^L‖² + w_F ‖s_F − s_φ^F‖² + w_A ‖ε_A − ε_φ^A‖²φ ∈ {θ, ref}

Lower Δ means the model denoises that crystal better (higher implied likelihood). The implemented loss compares policy to reference on the winner and the loser:

z = −βT [ (Δ_θ(x_w) − Δ_ref(x_w)) − (Δ_θ(x_l) − Δ_ref(x_l)) ], 𝓛 = −𝔼 log σ(z)

x_w is the more thermodynamically stable crystal (lower CHGNet energy/atom), x_l the less stable — the only supervision DPO needs.
β controls how far the policy may drift from the frozen reference π_ref: small β gives a stronger shift but risks chemistry collapse (reward hacking); large β stays close to π_ref.
Only a lightweight adapter + the T_c-conditioning pathway are updated; the equivariant backbone and π_ref's component errors are frozen/pre-cached, so each step forwards only π_θ.
σ(z) = P(prefer winner); over the β=0.5 run it rose 0.51 → 0.67.

Result: the generator shifts toward more stable crystals

Energy-per-atom distribution shift

Smaller β pushes the generator toward lower (more stable) energies. But β=0.1's dramatic shift is reward hacking — it collapses onto a handful of stable transition-metal chemistries (158 of 826 chemical systems), so its gain is diversity loss, not better generation. β=0.5 (headline) gives a genuine, diversity-preserving shift; larger β stay near the baseline.