Reinforcement Learning Guides Diffusion to New Crystals

The dilemma at the heart of generative discovery

Discovering a new functional material means searching a design space so vast that almost all of it has never been examined. Generative models arrived with real promise here, learning to produce chemically plausible crystals far faster than any hand built search. The catch is philosophical as much as technical. These models are trained to maximize the likelihood of their data, which is a polite way of saying they learn to make more of what already exists.

That goal quietly works against discovery. A compound that closely resembles known materials tends to be stable, because it sits in a well understood corner of chemistry. A compound that is genuinely novel, by definition, sits somewhere stranger, where stability is less certain. So a model tuned for likelihood drifts toward the safe and familiar, and a scientist hunting for the new and useful has to fight the very tool meant to help.

The authors call this the novelty and stability dilemma, and it is the pivot of the whole paper. Their insight is that you do not fix it by training on more data. You fix it by changing the objective, telling the model to optimize for what you actually want, which is chemically novel structures that are still thermodynamically viable. That shift from imitating data to pursuing an explicit goal is exactly what reinforcement learning is built for.

Diffusion in a compressed latent space

Before the reinforcement learning can happen, the generative backbone has to be the right shape. Crystals are awkward objects to generate, because a structure mixes three very different things, the discrete list of which atoms are present, the continuous fractional positions of those atoms, and the lattice that defines the box they live in. Applying a diffusion model directly to that jumble is unstable, because the discrete and continuous parts want different treatment.

The team, building on their earlier Chemeleon model, sidestep this by first squeezing a crystal into a single continuous latent vector using a variational autoencoder. The messy mix of atom types, coordinates, and lattice becomes one smooth representation, and the diffusion happens there rather than on the raw atoms. This latent approach is not just tidier. It reconstructs crystals with 99.4 percent accuracy on a standard test set, well above a comparable baseline, and it samples several times faster than working on the atoms directly. The denoiser itself is a diffusion transformer, the same architecture behind many modern image and video generators. Readers who want the groundwork on latent diffusion will find a clear neighbour in our piece on latent set diffusion for 3D shape reconstruction, which uses the same compress then denoise idea on a different kind of structure.

How the reward reshapes the model

With a clean latent space in hand, the authors treat the reverse diffusion process as a sequence of decisions and apply reinforcement learning to it. The method is Group Relative Policy Optimization, or GRPO, the same family of algorithm that has been sharpening the reasoning of large language models. Its appeal here is that it avoids training a separate value network, which would be expensive and unstable over the long horizon of a diffusion trajectory. Instead it generates a group of candidate structures for the same condition and scores each one relative to the group average.

What the model is rewarded for is the crux, and the authors are careful to make the reward verifiable rather than a learned guess. It is a weighted sum of four physically grounded terms.

Each term captures something a scientist would actually care about. Creativity rewards structures that are both unique within the batch and novel against known databases, softened with a continuous distance measure so borderline cases give a smooth gradient rather than a hard yes or no. Stability rewards low energy above the convex hull, the standard proxy for whether a crystal will hold together, estimated cheaply with a machine learning force field instead of slow first principles calculations.

The two diversity terms are the quiet heroes. Built from a kernel based distance between the compositions and structures the model produces, they push the batch to spread out across chemical space rather than piling onto a few winning motifs. The authors show that removing this diversity reward causes the model to collapse onto a narrow set of structures, exploiting the same easy wins over and over. Diversity is what keeps the exploration honest.

What the numbers show

The headline benchmark asks each model to generate ten thousand crystals from scratch, then measures how many are metastable, unique, and novel all at once, a combined score the authors abbreviate as mSUN. It is a demanding metric, because a structure has to clear every bar to count.

The jump is large. Reinforcement learning lifted the combined score from 15.9 percent to 61.3 percent, a relative improvement of nearly three hundred percent over the same model before tuning, and it comfortably cleared the strongest baseline. Novelty climbed from 62 percent to almost 98 percent, and metastability rose too, which is the striking part. The model became more adventurous and more physically sound at the same time, rather than trading one for the other. On a smaller standard dataset the same pattern held, with the tuned model reaching the highest novelty of any model tested.

The authors are upfront about the one metric that moved the wrong way. Uniqueness slipped from 99.4 percent to 88.7 percent, a mild case of the model repeating itself. This is the known signature of on policy reinforcement learning, where the policy learns to revisit its own high reward discoveries, and even the diversity reward cannot fully prevent it when you sample ten thousand structures. It is a fair price for the gains elsewhere, and the paper does not hide it.

Our approach extends the Pareto frontier, achieving both high novelty and high stability. Such improvements cannot be achieved through conventional heuristic fine tuning methods, which remain constrained by the underlying training data distribution. Park and Walsh, Nature Machine Intelligence, 2026

Escaping the tradeoff, not just balancing it

The most satisfying result is a picture rather than a number. When you plot novelty against stability for all the models, they fall along a curve where doing better on one means doing worse on the other. That curve is the dilemma made visible. The reward tuned model does not sit on the curve. It sits beyond it, further out in both directions than any competitor, which means it found a way to be more novel and more stable together rather than picking a point on the old tradeoff.

Digging into what the model actually generated explains how. The reward nudged it toward transition metals, whose flexible bonding supports a wider range of stable structures, and away from the oxygen heavy compounds the untuned model favoured. It leaned into more complex chemistries with three and four different elements, and it explored more polarizable species that the baseline rarely touched. In other words, the reward did not just make the model try harder. It changed which neighbourhoods of chemistry the model visited. This is a very different flavour of control from the reward alignment used on image generators, a topic our review of preference alignment for diffusion models covers from the visual side, and it shows the same idea paying off in a scientific domain.

Aiming at a target property

Generating novel crystals from scratch is impressive, but a materials scientist usually wants something specific, a compound with a particular property. The authors show the same framework handles this too. They asked the model to generate crystals with a band gap near three electronvolts, a value chosen precisely because it sits at the edge of what the training data covers, making it a hard target.

The reinforcement learning approach produced a sharp cluster of structures right at the target, where the two standard guidance methods it was compared against either drifted vaguely or collapsed almost entirely. Of the structures it generated, a healthy fraction landed both in the desired band gap range and passed the stability and novelty bars, and after a proper physics based relaxation those numbers improved further. Guiding generation toward a goal without wrecking chemical validity is exactly the kind of controllable design that inverse materials discovery has been reaching for, and it echoes the reward shaping seen in our look at reward aligned 3D generation.

Where it falls short

The authors are refreshingly candid about the limits, and a fair reading has to include them.

The drop in uniqueness is real and points at a deeper issue. On policy reinforcement learning is prone to a subtle mode collapse, where the model narrows onto its own successes, and the diversity reward mitigates this without curing it. Push the reward for diversity too hard and a different problem appears, reward hacking, where the model games the score by generating needlessly complex compounds with five or six elements whose predicted stability may be an artifact of a thin reference database rather than real chemistry. The sweet spot between too little exploration and too much gaming has to be tuned by hand.

There is also a gap between what the model generates and what a chemist could make. The framework produces ordered periodic crystals and does not yet distinguish a genuine ordered compound from a disordered alloy that only looks ordered when forced into a repeating cell. Some of its more exotic outputs are better understood as approximations of solid solutions than as distinct new phases. The authors flag this honestly as a limitation of the representation itself.

Finally, the whole pipeline rests on proxies. Stability is judged by a machine learning force field standing in for expensive first principles calculations, and novelty by distance measures against existing databases. Both are reasonable and both are imperfect, and the reward is only ever as trustworthy as the proxies underneath it. A structure the model believes is stable and novel still has to survive a real calculation and, eventually, a real synthesis, neither of which this work performs.

Why it matters

Step back and the contribution is a change of aim. For years the way to improve a generative model for materials was to train it on more and better data, which only ever sharpens its imitation of the known. This work shows that the more powerful lever is the objective itself. By rewarding a model for novelty, stability, and diversity through verifiable physical measures, the authors turned a passive imitator into an active explorer that pushes into the sparse regions where new materials actually live.

The design is also deliberately modular, which is what gives it reach. The reward is a set of interchangeable terms, so the same machinery that hunts for novel stable crystals can be pointed at a target band gap, or in principle at any measurable property, simply by swapping in a new reward. That flexibility, paired with an open source release, makes it less a single result than a template for controllable scientific generation. The broader current of language model and diffusion research moving into materials, which our coverage of multimodal models for materials science traces, gains a principled way to say not just make something plausible but make something new and worth making.

A reference implementation

The code below is a compact, runnable version of the core idea. It defines a small latent diffusion policy, rolls out a group of samples for one condition, scores each with a stand in reward that blends stability and novelty and diversity, computes the group relative advantage, and updates the policy with a clipped GRPO objective and a penalty that keeps it close to the pretrained model. The authors full code is linked under the block.

# Reinforcement learning guided latent diffusion for crystals
# Compact reference for the GRPO update on a latent denoising policy.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


class LatentPolicy(nn.Module):
    # Predicts the mean of the next latent state during reverse diffusion.
    # A fixed variance turns each step into a Gaussian we can score.
    def __init__(self, dim=32, cond=8, hidden=128, log_std=-2.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + cond + 1, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, dim),
        )
        self.log_std = nn.Parameter(torch.full((dim,), log_std))

    def mean(self, z, cond, t):
        tt = torch.full((z.shape[0], 1), float(t))
        return self.net(torch.cat([z, cond, tt], dim=-1))

    def step_logprob(self, z, z_next, cond, t):
        # log probability of the taken denoising step under a Gaussian policy
        mu = self.mean(z, cond, t)
        std = self.log_std.exp()
        return (-0.5 * (((z_next - mu) / std) ** 2) - self.log_std).sum(-1)


def rollout(policy, cond, dim=32, steps=6):
    # reverse diffusion from noise to a final latent, recording each step
    z = torch.randn(cond.shape[0], dim)
    traj = []
    for t in range(steps, 0, -1):
        mu = policy.mean(z, cond, t)
        std = policy.log_std.exp()
        z_next = mu + std * torch.randn_like(mu)
        traj.append((z.detach(), z_next.detach(), t))
        z = z_next
    return z, traj


def reward_fn(z0):
    # stand in for the verifiable reward, blends stability, novelty, diversity
    stability = -torch.relu(z0.norm(dim=-1) - 4.0)          # cheap energy proxy
    novelty = torch.tanh(z0.std(dim=-1))                    # spread as a novelty proxy
    return 1.0 * stability + 1.0 * novelty


def grpo_step(policy, ref_policy, cond, opt, clip=1e-3, beta=1.0):
    z0, traj = rollout(policy, cond)
    rewards = reward_fn(z0)
    # group relative advantage, normalized within the sampled group
    adv = (rewards - rewards.mean()) / (rewards.std() + 1e-6)

    total = 0.0
    for z, z_next, t in traj:
        logp = policy.step_logprob(z, z_next, cond, t)
        with torch.no_grad():
            logp_ref = ref_policy.step_logprob(z, z_next, cond, t)
        ratio = (logp - logp.detach()).exp()               # one on policy step
        unclipped = ratio * adv
        clipped = torch.clamp(ratio, 1 - clip, 1 + clip) * adv
        surrogate = -torch.min(unclipped, clipped).mean()
        kl = (logp - logp_ref).mean()                      # stay near the pretrained model
        total = total + surrogate + beta * kl

    opt.zero_grad()
    total.backward()
    opt.step()
    return rewards.mean().item()


if __name__ == "__main__":
    dim, cond_dim, group = 32, 8, 64
    policy = LatentPolicy(dim, cond_dim)
    ref_policy = LatentPolicy(dim, cond_dim)
    ref_policy.load_state_dict(policy.state_dict())        # frozen reference
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)

    cond = torch.randn(group, cond_dim)                    # one shared condition, a group of rollouts
    for update in range(30):
        r = grpo_step(policy, ref_policy, cond, opt)
        if (update + 1) % 10 == 0:
            print(f"update {update + 1}  mean reward {r:.4f}")

Conclusion

The core achievement here is a generative model for crystals that was taught to want the right things. By wrapping a latent diffusion model in reinforcement learning and rewarding it for novelty, stability, and diversity through verifiable physical measures, the authors turned a system that imitated known materials into one that hunts for new ones. The combined discovery score jumped from about 16 percent to 61 percent, and novelty rose to nearly 98 percent, without sacrificing the physical soundness that usually falls apart when a model reaches for the unfamiliar.

The conceptual shift worth remembering is that the objective, not the data, was the real lever. Generative models had been improving by learning to imitate more faithfully, which is precisely the wrong instinct for discovery. Recasting generation as reward driven exploration let the model step past the old novelty and stability tradeoff rather than merely picking a point on it, and the visible proof is a model that lands beyond the frontier every competitor is stuck on.

The approach travels because the reward is modular. The same machinery that finds novel stable crystals can chase a target band gap, or in principle any property you can measure, by swapping the reward terms. Paired with a fast latent backbone and an open release, that makes this less a single benchmark win than a reusable recipe for controllable scientific generation, and a template other domains can borrow.

The honest limits keep it grounded. The model repeats itself more than the untuned version, it can be lured into gaming an overweighted diversity reward, it does not yet tell an ordered compound from a disordered alloy, and every reward rests on proxies that stand in for slow calculations and real synthesis. None of these sink the result, but they mark the distance between a promising generator and a validated discovery engine.

Future directions follow naturally. Add regularization to curb the reward hacking, teach the representation to handle disorder, and, most importantly, close the loop by feeding real calculations and eventually real experiments back into the reward. If those steps hold up, the quiet promise here, that you can tell a generative model to value the new and the viable at once, becomes a practical way to search for the materials that do not exist yet.

Frequently asked questions