OxiGen: Oxidation-State-Aware Crystal Generation

arXiv:2610.08296 · cs.LG, cond-mat.mtrl-sci · Submitted 2026-10-06 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "OxiGen: Oxidation-State-Aware Crystal Generation".

Jane: The gist The proposed OxiGen model, an oxidation-state-aware crystal diffusion model, substantially improves oxidation-state fidelity and generates the highest rate of stable, unique,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, OxiGen is basically taking crystal generation—which usually just guesses atoms—and adding this explicit layer where every site has not just an element, but also an oxidation state, s equals a z.

Jane: That means they aren't just generating a structure; they're generating a chemical recipe that specifies how oxidized every single atom is.

Lu: Exactly. The core idea is that existing models only check if the total charge adds up to zero at the end, but OxiGen enforces that constraint *during* the generation process itself using something called structured inference.

Meng: Structured inference sounds complicated when you're trying to generate millions of possibilities; how does that actually help us run this thing without it taking forever?

Tom: Well, they’ve set up a finite-state automaton for charge neutrality. They use dynamic programming over that automaton, with forward and backward weights— F i(c) and B i(c) —to do an exact constrained sampling.

Jane: So instead of throwing away all the "wrong" structures because they’re not neutral, the model uses this DP to guide the sampling path to only those that are chemically possible in terms of charge balance.

Lu: Right. It replaces a standard element assignment with this species assignment s because it handles both identity and state at once. This is how they achieve that higher fidelity score f OS.

Meng: I'm still worried about the computational cost you mentioned earlier; if we have to do that DP pass at every step of the diffusion process, how long does a full generation take?

Tom: They estimated the overhead for sampling alone is around three hundred ninety-nine point five two seven plus or minus four point nine two three seconds per batch on a single NVIDIA GH200 with twelve gigabytes of memory. That's significant inference time.

Jane: So, they’re trading that extra time for a much more chemically accurate result, which is what the results show—it gets better oxidation-state fidelity and lower distribution distance d OS compared to models like MatterGen.

Lu: It means they've shown that just representing the states isn't enough; you have to enforce that global charge neutrality during sampling for the metrics to actually improve. The ablation study backs that up, confirming training with just the neutral objective isn't adding much on its own.

Meng: That’s a key point for deployment—if we can’t easily run that structured inference in production without a massive slowdown, we have to rethink how we scale this kind of constraint enforcement.

Tom: And in terms of the big picture, they’ve shown this approach works across different property conditions too. When you condition generation on band-gap labels, OxiGen keeps high compositional validity even when the target band gap is really wide—like three, five, or seven electron volts.

Jane: That’s interesting because many models lose their reliability when you ask them for something that's really rare in the training data. This suggests OxiGen is robust in inverse design scenarios where the properties are sparse.

Lu: It also successfully generates examples of stable, unique materials like CaAuF five and Tl two InBiS four proving it can explore that chemical space effectively. But the limitation they flag is that formal integer oxidation states can still be ambiguous for systems with a lot of charge delocalization.

Meng: So, the method works well when those states are relatively clear, but we need to keep in mind that this isn't a universal fix for every kind of complex chemistry.

Tom: The implication is that incorporating prior chemical knowledge through simple constraints like charge neutrality can guide generative models toward more chemically plausible materials, even with these known limitations.

Jane: It really shows that adding explicit chemical rules, even simple ones like charge balance, makes a huge difference in the quality of the generated output. What we have to keep in mind is that this only handles charge neutrality; it doesn't enforce other constraints like electronegativity balance.

Lu: That's true. Charge neutrality is just one aspect of chemical validity; it doesn’t cover everything needed for a perfectly realistic material description.

Meng: For production systems, we have to build in mechanisms like this structured sampling if we want our outputs to be reliable and not just statistically plausible but chemically sound.

The paper's summary: Tom: So, OxiGen isn't just about fixing the charge neutrality; they showed that explicitly representing oxidation states in a species assignment s equals a, z is fundamentally better than just using element assignments.

Jane: It means the model learns those actual oxidation state frequencies directly from the data instead of having to guess them based on some abstract rule.

Lu: Exactly. The ablation study was pretty telling; training with just the neutral objective didn't give you much improvement, but enforcing charge neutrality during sampling substantially improved all three metrics— f OS, d OS, and compositional validity.

Meng: So the main implication is that representation learning needs to be coupled with a constraint satisfaction mechanism like this structured output layer if we want chemically meaningful results.

Tom: It also makes the model much more robust when you're doing inverse design under property constraints. When they conditioned generation on band-gap labels, OxiGen maintained high compositional validity across the whole range of targets.

Jane: That’s significant because other models, like MatterGen, saw their validity drop sharply as the target band gap increased. So OxiGen keeps things chemically plausible even when you ask it for very specific, rare properties.

Lu: It also improved site-specific prediction accuracy through this structured Viterbi decoding in the CrystaliteOS predictor; it enforces charge neutrality for all predictions instead of just using a standard argmax decoding.

Meng: But what about the computational trade-off again? We talked about that inference time, and that's what we have to deal with when we try to deploy this kind of quality boost at scale.

Tom: That overhead is around three hundred ninety-nine point five two seven plus or minus four point nine two three seconds per batch on a single NVIDIA GH200 with twelve gigabytes of memory. It’s definitely a factor for running inference on production hardware.

Jane: So, the big change here for someone just listening is this: we move from models that generate structures that are just technically neutral to models that generate structures whose oxidation state distributions look like they came from actual lab synthesis.

Lu: The future work they hint at is incorporating more complex chemical constraints beyond just charge neutrality because formal integer states have their limits with highly covalent or delocalized systems.

Meng: I wonder if we could simplify the structured inference layer for production; maybe a less complex automaton structure would save us some of that overhead while keeping the fidelity boost.

Tom: The paper confirms that this method works very well across different types of targets, including de novo generation and band-gap conditioned generation, showing it’s versatile.

Jane: It really shows that incorporating prior chemical knowledge through simple constraints can guide generative models towards more chemically plausible materials overall.

Lu: I'm sure the next step is exploring how to integrate more complex physical chemistry rules into this diffusion framework, moving beyond just charge balance.

Meng: So for now, we have a high-fidelity method that works well but requires significant computational resources and doesn't solve every single chemical ambiguity.

The paper's improvements: Tom: So, to wrap up "OxiGen: Oxidation-State-Aware Crystal Generation," they’ve shown that by explicitly representing oxidation states as species s equals a, z and enforcing charge neutrality during sampling, you get much better fidelity over existing crystal generative models.

Jane: It really boils down to this: the model learns the distributions of oxidation states you see in actual synthesised materials instead of just guessing based on abstract rules.

Lu: I think the main implication is that incorporating prior chemical knowledge through simple constraints can guide generative models toward more chemically plausible materials, even if it means dealing with the ambiguity of formal integer oxidation states for some systems.

Meng: From a production standpoint, we have to remember that this method requires significant computational resources; training takes about three point five hours over nine hundred epochs on four NVIDIA GH200 GPUs.

Tom: And inference has that overhead we discussed earlier, around three hundred ninety-nine point five two seven plus or minus four point nine two three seconds per batch on a single NVIDIA GH200 with twelve gigabytes of memory. That’s the cost we have to consider when deploying this kind of quality improvement.

Jane: It changes how we think about inverse design because it gives us a way to generate structures that are not just technically neutral, but chemically sensible and stable, even when you condition them on very specific property targets.

Lu: The future work they mentioned is looking at incorporating more complex physical chemistry rules because formal integer oxidation states still have limitations for highly covalent or delocalized systems.

Meng: If we want to deploy this kind of quality boost, I wonder if we could simplify the structured inference layer for production; maybe a less complex automaton structure would save us some of that overhead while keeping the fidelity boost.

Tom: The paper "OxiGen: Oxidation-State-Aware Crystal Generation" proves that by adding these explicit chemical constraints—the species representation and the structured enforcement—they manage to generate stable, unique crystals with a fidelity that matches what you'd expect from materials actually made in a lab.

Lalam: I am certain of this answer based on the literature review; OxiGen proves that coupling explicit species representation with dynamic programming for constraint satisfaction is necessary to move beyond mere charge neutrality and achieve high-fidelity oxidation-state predictions in crystal generation.

Conclusion: Tom: So we've been looking at "OxiGen: Oxidation-State-Aware Crystal Generation," and basically, they showed that explicitly representing oxidation states as species s equals a, z and enforcing charge neutrality during sampling gives you much better fidelity over existing crystal generation models.

Jane: It really boils down to this: the model learns the actual oxidation state frequencies you see in synthesized materials instead of just guessing based on some abstract rule.

Lu: The paper shows that enforcing global charge neutrality during sampling is needed because just representing oxidation states on its own only provides a modest improvement over just using element representation, and the structured sampling is what really boosts those metrics.

Jane: It really shows that by adding these explicit chemical constraints—the species representation and the structured enforcement—they manage to generate stable, unique crystals with a fidelity that matches what you'd expect from materials actually made in a lab.

Dylan John, Kim E. Jelfs, Alex M. Ganose, Eleonora Giunchiglia

Department of Chemistry, Imperial College London · Department of Electrical and Electronic Engineering, Imperial College London

cs.LG, cond-mat.mtrl-sci

Submitted: 2026-10-06

Updated: 2026-10-06

Comments: 27 pages, 4 figures

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 88/100

The gist: The gist The proposed OxiGen model, an oxidation-state-aware crystal diffusion model, substantially improves oxidation-state fidelity and generates the highest rate of stable, unique, and novel

Key concepts

Oxidation-State-Aware Crystal Diffusion Model: This is a type of AI model designed to generate crystal structures. Unlike simpler models, OxiGen specifically tracks the oxidation state (how many electrons an atom has lost or gained) for every atom during generation. It treats the creation process like a diffusion, where it refines the structure while ensuring that all atoms in the final crystal have correct charges.
Species Assignment s = (a, z): Instead of just assigning an element (like Carbon), OxiGen assigns a 'species' to each site. This species is a pair: 'a' represents the element identity, and 'z' represents its oxidation state. For example, instead of just saying there is one Oxygen atom, it specifies that the oxygen is in an -2 oxidation state. This detailed assignment allows the model to capture more complex chemical realities.
Structured Inference for Charge Neutrality: This is a mathematical technique used during generation to guarantee that the final crystal has a net charge of zero. OxiGen uses a finite-state automaton and dynamic programming over this structure. This allows the model to calculate all possible valid charge assignments efficiently, ensuring that every generated material adheres to the fundamental rule of electrical neutrality.
Diffusion Process and Sampling: The model generates a crystal by starting with noise and iteratively refining it through a diffusion process. The key is that the generation steps are factored across three parts: species assignment, fractional coordinates, and the lattice matrix. During sampling, OxiGen first finds a valid charge-neutral species assignment before generating the physical structure, ensuring chemical consistency from the start.

Summary

The gist The proposed OxiGen model, an oxidation-state-aware crystal diffusion model, substantially improves oxidation-state fidelity and generates the highest rate of stable, unique, and novel crystals among evaluated methods

OxiGen's Core Innovation

OxiGen is an oxidation-state-aware crystal diffusion model that explicitly represents oxidation states during generation It enforces global charge neutrality by construction using a structured output layer with exact inference over a finite-state automaton This approach addresses the limitation of existing models that only generate materials with charge-neutral oxidation-state assignments, which poorly reproduce the distributions observed in synthesised materials The model achieves this by replacing the standard element assignment with a species assignment s = (a, z), where s specifies both the element identity and its oxidation state

Structured Inference for Charge Neutrality

To enforce charge neutrality, OxiGen introduces a structured output layer that exploits the simple additive structure of the charge-neutrality constraint for tractable structured inference This layer represents charge neutrality as a finite-state automaton, where transitions are governed by the cumulative charge ci = X j=1 z(s0,j) Dynamic programming over this automaton enables exact constrained sampling and MAP decoding without summing over an exponentially large space of species assignments Specifically, the forward weight Fi(c) gives the total probability mass of assignments to the first i sites with cumulative charge c, and the backward weight Bi(c) gives the probability mass of assignments to the remaining sites that complete charge c to zero

Diffusion Process and Sampling

OxiGen is a crystal diffusion model that jointly generates species assignment, fractional coordinates, and lattice matrix The forward diffusion process factorises across these components: q(Mt Mr) = q(St Sr) q(Ft Fr) q(Lt Lr) The reverse process is parameterised by a shared SE(3)-equivariant GemNet-dT message-passing network that takes the complete noisy crystal as input During sampling, the model first samples a complete charge-neutral clean species assignment Sˆ0 ∼ peθ(S0 Mt, t) and then samples Sr ∼ q(Sr St, Sˆ0) from the tractable posterior of the forward absorbing process This procedure guarantees that the final species assignment is charge neutral

Oxidation-State Prediction and Fidelity

The model utilizes a predictor to assign an oxidation state to each atomic site, defined by a factorised distribution pϕ(Z M) = Y N i=1 pϕ,i(z i M) For the species component, OxiGen uses an absorbing-state discrete diffusion process based on MDLM (Sahoo et al., 2024), where the denoiser predicts pθ,i(s0,i Mt, t) The structured training objective Lneutral is used in the ablation study to evaluate the contribution of enforcing charge neutrality during sampling

Experimental Evaluation and Results

Empirically, OxiGen substantially improves oxidation-state fidelity over existing crystal generative models In de novo generation using MP-20-OS, OxiGen achieves a higher oxidation-state frequency score fOS and lower oxidation-state distribution distance dOS compared to baselines like MatterGen Furthermore, in band-gap conditioned generation, OxiGen maintains consistently high compositional validity across the full range of targets, whereas MatterGen's compositional validity decreases sharply as the target band gap increases OxiGen also achieves a higher S.U.N. rate than MatterGen at all four conditioning targets

Conclusion and Limitations

OxiGen demonstrates that incorporating prior chemical knowledge through simple heuristics and constraints can guide generative models towards more chemically plausible materials A limitation of OxiGen is that it inherits the limitations of formal integer oxidation states, which can be ambiguous or undefined for systems with substantial covalency or charge delocalisation Moreover, charge neutrality is only one aspect of chemical validity; OxiGen does not enforce other constraints such as electronegativity balance Future work could incorporate additional chemical constraints and explore more flexible inductive biases

Computational Overhead

Structured inference requires a single backward dynamic-programming pass to draw a charge-neutral sample at each diffusion step, with a time complexity of O(N 2S(zmax − zmin)) Structured training requires one forward and one backward dynamic-programming pass to compute the marginals of the structured distribution, also with complexity O(N 2S(zmax − zmin)) The computational overhead of the structured output layer during sampling is approximately 309.527±4.923 s/batch on a single NVIDIA GH200 120GB The structured training time is approximately 3.5 hours for OxiGen over 900 epochs on four GH200 GPUs

Example Crystals

OxiGen successfully generates examples of stable, unique, and novel crystals including CaAuF5, CsKCl2, KNaAlH5, Rb2SbOF5, and Tl2InBiS4 The generated structures are shown in Figure 4 The model's ability to generate these structures highlights its effectiveness in exploring the chemical space

Ablation Study Insights

Ablation studies confirm that the structured output layer substantially improves all three oxidation-state metrics, suggesting that explicitly representing oxidation states is insufficient on its own Training with the structured objective does not provide any consistent additional improvement while adding training overhead The species representation alone provides only a modest improvement over the element representation, with the majority of generated assignments violating charge neutrality Structured sampling enforces charge neutrality by construction and substantially improves all three oxidation-state metrics This suggests that explicitly representing oxidation states is insufficient on its own, and that enforcing global charge neutrality during sampling is needed to ensure chemically meaningful oxidation-state assignments

Conditional Generation Performance

When conditioned on band-gap labels from MP-20-OS, OxiGen achieves a higher S.U.N. rate than MatterGen at all four conditioning targets Band-gap RMSE is comparable at 1 eV and better for OxiGen at the 3, 5, and 7 eV targets This demonstrates the value of structured sampling for inverse design, where conditioning on sparsely represented targets can otherwise lead to chemically invalid compositions The model maintains high validity even under property conditioning

Evaluation Metrics Summary

The evaluation metrics include compositional validity, the oxidation-state frequency score fOS, and the oxidation-state distribution distance dOS OxiGen substantially improves oxidation-state fidelity, closely matching the MP-20-OS test-set values for compositional validity and oxidation-state frequency score while achieving the lowest oxidation-state distribution distance The results show that OxiGen closely reproduces both the test-set compositional validity curve and oxidation-state distributions

Further Results

In the LeMat-GenBench evaluation, OxiGen achieves the highest validity, uniqueness, and M.S.U.N.

Improvements for AI systems

Improved oxidation-state fidelity in crystal generation by utilizing OxiGen's structured output layer to enforce charge neutrality during generation, which substantially improves oxidation-state fidelity over existing crystal generative models. This allows the system to generate crystals that reproduce the distributions of oxidation states observed in synthesised materials rather than just being charge-neutral.
Enhanced stability and novelty in de novo crystal design by employing OxiGen, which achieves the highest rate of stable, unique, and novel crystals among evaluated methods. This is achieved by using a species representation where each site is characterized by a species s = (a, z), explicitly representing oxidation states during generation.
Improved property-conditioned inverse design by maintaining high compositional validity even under constraints, as OxiGen maintains high compositional validity even under property conditioning. This capability allows the system to generate crystals with desired properties while ensuring they are chemically plausible, preventing the generation of structures that are only charge-balanced using rare oxidation states.
More robust oxidation-state prediction by integrating structured Viterbi decoding into the CrystaliteOS predictor, which enforces charge neutrality for all predictions and improves site accuracy over standard argmax decoding. This ensures that predicted oxidation states are not just locally plausible but also globally consistent with the required charge balance.
More accurate structural representation of chemical validity by moving from representing a crystal as M = (A,F, L) to M = (S,F, L), where S is the species assignment, ensuring that the overall crystal is represented as M = (S,F, L) and that S must satisfy charge neutrality.

Abstract

Generative models have the potential to accelerate inorganic materials discovery by enabling inverse design, but generating experimentally realisable crystals remains challenging. Oxidation states are widely used to assess the compositional validity of crystals and guide inorganic materials discovery. While existing generative models for crystals can generate materials with charge-neutral oxidation-state assignments, they poorly reproduce the distributions of oxidation states observed in synthesised materials. To address this limitation, we propose OxiGen, an oxidation-state-aware crystal diffusion model that explicitly represents oxidation states during generation. OxiGen enforces global charge neutrality by construction using a structured output layer with exact inference over a finite-state automaton. Empirically, OxiGen substantially improves oxidation-state fidelity, generates the highest rate of stable, unique, and novel crystals among evaluated methods, and maintains high compositional validity even under property conditioning.