Source-linked AI summary
OxiGen: Oxidation-State-Aware Crystal Generation
Dylan John, Kim E. Jelfs, Alex M. Ganose, Eleonora Giunchiglia
TL;DR
Crystal generators can produce charge-neutral materials yet poorly reproduce oxidation-state distributions observed in synthesised materials, limiting a key chemical validity signal. OxiGen explicitly generates oxidation states and enforces neutrality through exact structured inference. It improves oxidation-state fidelity, leads evaluated methods in stable, unique, and novel crystal rate, and maintains high compositional validity under band-gap conditioning.
Problem
Existing crystal generators can produce charge-neutral materials but poorly reproduce the oxidation-state distributions observed in synthesised materials, while stability metrics may miss improbable assignments.
Source: §1 · p. 1Method
OxiGen explicitly represents oxidation states during crystal diffusion and uses a structured output layer with exact finite-state-automaton inference to enforce charge neutrality.
Source: §1 · p. 2, §4.3 · p. 4Results
Takeaways & Limitations
The species representation and structured output layer require minimal architectural modification and could be incorporated into other crystal generative models.
Source: §6 · p. 9OxiGen may mischaracterise or exclude materials without chemically meaningful integer oxidation states, and charge neutrality alone does not enforce other aspects of chemical validity.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Stable and compositionally valid crystals may still rely on uncommon oxidation states that are less likely to be experimentally realisable. OxiGen addresses this gap by explicitly generating oxidation states while enforcing charge neutrality.
1 INTRODUCTION
- Oxidation-state frequencies vary widely, and compositional validity checks for a charge-neutral assignment without accounting for how common its oxidation states are.Uncommon assignments may require particular bonding environments, external conditions, or specialised synthesis routes.Source: §1 · p. 1
- Existing generative models can produce crystals that are charge balanced only by rare oxidation states, which stability filtering does not necessarily identify.Some generated structures predicted to lie on the convex hull nevertheless have improbable oxidation-state assignments.Source: §1 · p. 1
- OxiGen explicitly represents oxidation states during generation so it can learn their frequencies from training data and enforce global charge neutrality.Its structured output layer uses dynamic programming over a finite-state automaton for exact inference.Source: §1 · p. 2
2 RELATED WORK
Prior crystal-generation work addresses structure, composition, or chemical constraints in different ways. OxiGen places oxidation states inside a crystal diffusion model that jointly generates composition and structure.
Crystal Generative Models
- Crystal-GFN enforces charge neutrality but generates composition, space group, and lattice parameters rather than atomic coordinates; Mastej et al. (2026) use oxidation-state validity as a reward.CrysVCD represents element–oxidation-state pairs but relies on a separate diffusion model for structures, while ChargeDIFF models charge density rather than compositional validity (Park et al., 2026b).Source: §2 · p. 2
- OxiGen incorporates oxidation states directly in crystal diffusion, jointly generating composition and structure while enforcing charge neutrality by construction.This contrasts with approaches that use oxidation states only as a reward or generate formulas separately from structures.Source: §2 · p. 2
Structured Prediction
- OxiGen represents charge neutrality with a finite-state automaton and uses dynamic programming for exact MAP decoding and exact constrained sampling.DINGO uses finite automata and dynamic programming for exact constrained MAP decoding in diffusion language models (Suresh et al., 2025).Source: §2 · p. 2
3 PRELIMINARIES
The preliminaries define a crystal through its unit-cell representation, formal oxidation-state assignments, and the discrete diffusion process used to generate categorical species.
Crystal representation
- A crystal is represented by a unit cell whose translations tile 3D space, with elements, fractional coordinates, and a lattice matrix specifying its sites and geometry.Cartesian coordinates are obtained from the lattice matrix and fractional coordinates.Source: §3 · pp. 2–3
Oxidation states
- Formal oxidation states are integer charges, and charge neutrality requires their sum across sites to equal zero.The work assumes every crystal has a charge-neutral integer oxidation-state assignment, while noting this assumption can be problematic for some materials.Source: §3 · p. 3
Discrete diffusion
- Discrete diffusion corrupts categorical species assignments through a forward process and learns a reverse process that progressively denoises them.In absorbing-state diffusion, each species is independently replaced by a mask token according to a schedule.Source: §3 · p. 3
4 PROPOSED METHOD: OXIGEN
OxiGen generates element–oxidation-state species alongside crystal coordinates and lattice parameters. Its structured inference layer couples site-wise predictions to enforce charge neutrality during generation.
4.1 SPECIES REPRESENTATION
- Each site is represented by a species token pairing an element with an oxidation state, so OxiGen generates species assignments rather than element identities alone.The species vocabulary comprises element–oxidation-state pairs observed in synthesised materials.Source: §4.1 · p. 3
4.2 DIFFUSION MODEL
- OxiGen jointly generates species assignments, fractional coordinates, and the lattice matrix through a diffusion process over the noisy crystal state.The forward process independently corrupts each component.Source: §4.2 · p. 3
- The model retains MatterGen’s network architecture and continuous diffusion for coordinates and lattice, modifying the discrete representation and diffusion over species.The denoiser predicts site-wise species probabilities, but independent sampling can violate global charge neutrality.Source: §4.2 · pp. 3–4
4.3 STRUCTURED INFERENCE
- OxiGen conditions factorised site predictions on charge neutrality and uses dynamic programming over a finite-state automaton for exact inference.This restricts the structured distribution to charge-neutral assignments without summing over the exponentially large space of complete assignments.Source: §4.3 · p. 4
- The automaton tracks cumulative charge across sites, with charge-neutral assignments corresponding to paths that start and end at zero.Forward and backward dynamic-programming weights enable tractable inference, and the resulting inference is invariant to site ordering.Source: §4.3 · pp. 4–5
- A complete forward or backward dynamic-programming pass has time complexity O(N^2|S|(zmax − zmin)).The terminal forward weight F_N(0) equals the total probability mass of charge-neutral assignments and is the normalising constant.Source: §4.3 · p. 5
- Viterbi decoding with backpointers returns the exact maximum-probability species assignment under the charge-neutral structured distribution.The recurrence tracks the maximum probability for prefixes with each cumulative charge.Source: §4.3 · p. 5
- Sequential sampling from backward-weighted site conditionals yields an exact joint sample that preserves dependencies among sites unmasked in the same reverse step.At each reverse step, OxiGen samples a complete charge-neutral clean species assignment before sampling the intermediate reverse state; training uses the standard factorised masked-diffusion objective.Source: §4.3 · pp. 5–6
5 EXPERIMENTS
The experiments evaluate oxidation-state prediction and OxiGen’s de novo and band-gap-conditioned crystal generation. OxiGen improves oxidation-state fidelity while performing strongly on stability-related measures and maintaining compositional validity under conditioning.
5.1 OXIDATION-STATE PREDICTION
- The oxidation-state predictor achieves higher site and crystal accuracy than both baselines, and structured Viterbi decoding enforces neutrality while further improving both accuracies over argmax decoding.Evaluation uses 28,492 ICSD structures and a composition-disjoint train/validation/test split of 22,879/2,831/2,782.Source: §5.1 · p. 6
- The predictor labels MP-20 to create MP-20-OS; filtering removes 5.9% of structures, including 3.7% without a charge-neutral assignment.The resulting dataset contains 25,536/8,495/8,545 train/validation/test structures.Source: §5.1 · p. 6
5.2 CRYSTAL GENERATION
- Generation comparisons retrain five baselines and OxiGen on MP-20-OS for five seeds, sampling 10,240 crystals per method per seed.The baselines are DiffCSP, MatterGen, FlowMM, OMatG, and Crystalite.Source: §5.2 · p. 7
- Compositional validity omits electronegativity balance and mixed valence, while oxidation-state fidelity metrics exclude metallic alloys and elemental structures.The frequency score targets the test-set value; distribution distance compares generated and test-set oxidation-state distributions.Source: §5.2 · p. 7
- OxiGen closely matches test-set compositional validity and oxidation-state frequency score while achieving the lowest oxidation-state distribution distance.Figure 2 shows closer reproduction of the test-set validity curve and selected-element oxidation-state distributions than the baselines.Source: §5.2 · p. 7, Figure 2 · p. 8
- OxiGen achieves the highest S.U.N. rate, lowest relaxation RMSD, and second-lowest mean Ehull among evaluated methods.It also substantially improves oxidation-state fidelity over baselines.Source: §5.2 · pp. 7–8, Table 2 · p. 7
- Structured sampling enforces charge neutrality by construction and substantially improves all three oxidation-state metrics; species representation alone yields modest gains while most assignments violate neutrality.Training with the structured objective adds overhead without consistent additional improvement, so it is not used for OxiGen.Source: §5.2 · p. 8, Table 3 · p. 8
- At band-gap targets of 1, 3, 5, and 7 eV, both models shift predicted distributions toward targets, while OxiGen maintains high validity and higher S.U.N. rates than MatterGen.MatterGen’s compositional validity decreases sharply as target band gap increases; Figure 3 reports one fine-tuning seed.Source: §5.2 · p. 9, Figure 3 · p. 9
6 CONCLUSIONS
OxiGen explicitly represents oxidation states during crystal generation and enforces charge neutrality through finite-state-automaton dynamic programming. It improves oxidation-state fidelity and generation outcomes, but its formal oxidation-state framework and predictor have important limitations.
6 CONCLUSIONS
- OxiGen enforces charge neutrality during generation while explicitly representing oxidation states.The structured inference uses dynamic programming over a finite-state automaton.Source: §6 · p. 9
- Across experiments, OxiGen better reproduces synthesised-material oxidation-state distributions, has the highest stable, unique, and novel crystal rate, and maintains high validity under band-gap conditioning.The authors expect the species representation and structured output layer to be readily incorporated into other crystal generative models.Source: §6 · p. 9
- Formal integer oxidation states can be ambiguous or undefined for substantial covalency, charge delocalisation, or metallic bonding, risking mischaracterisation or exclusion of valid materials.Predicted oxidation-state assignments may also carry errors or biases, and charge neutrality alone does not enforce electronegativity balance or local coordination and bonding consistency.Source: §6 · p. 9
AI DISCLOSURE
The section discloses the authors’ use of generative AI and describes exact inference for oxidation-state assignments. It also notes that independent site-wise argmax predictions do not ensure chemical admissibility or charge neutrality.
AI DISCLOSURE
- The authors used generative AI to assist research, coding, baselines, literature identification, manuscript editing, and figure creation, and take responsibility for the final work.They state that the authors proposed the question, developed the core methodology and experimental design, and analysed the results.Source: AI DISCLOSURE · p. 10
A STRUCTURED OUTPUT LAYER
- Dynamic programming over cumulative charge identifies the maximum-weight charge-neutral assignment, and backtracking recovers its exact MAP solution.The proof establishes the forward weights by induction and uses the assignment-independent normalising constant to equate maximising unnormalised weight with maximising the model probability.Source: A STRUCTURED OUTPUT LAYER · p. 15
- The sequential procedure samples from the structured distribution, assigning zero probability to non-neutral assignments.For a neutral assignment, the backward-weight factors telescope to its structured-distribution probability.Source: A STRUCTURED OUTPUT LAYER · p. 15
B OXIDATION-STATE PREDICTION DETAILS
- Site-wise argmax decoding does not ensure element-specific chemical admissibility or charge neutrality for the complete assignment.The predictor assigns oxidation states to sites in a crystal with known composition and structure, using a full oxidation-state vocabulary and categorical cross-entropy training.Source: B OXIDATION-STATE PREDICTION DETAILS · p. 15
B.1 VOCABULARY AND DATASET CONSTRUCTION
The oxidation-state vocabulary is built from element-specific admissible states, and the ICSD structures are curated to retain charge-neutral, consistently labelled examples. The retained structures are split by reduced composition.
B.1 VOCABULARY AND DATASET CONSTRUCTION
- The oxidation-state vocabulary uses the SMACT v4.0.0 ICSD-24 filter with consensus 6, low commonality, and zero states excluded.This stricter-than-default consensus threshold is intended to exclude extremely unusual oxidation states; each element has an admissible subset of the vocabulary.Source: §B.1 · p. 16
- ICSD structures are retained only when every site has an integer, element-admissible oxidation state and the complete assignment is charge neutral.Curation also removes duplicates and structures over 200 atoms, reduces cells to primitive Niggli form, and restricts mixed valence to designated elements.Source: §B.1 · p. 16
- Structures are grouped by reduced composition into disjoint splits containing 22,879 training, 2,831 validation, and 2,782 test structures.No reduced composition occurs in more than one split.Source: §B.1 · p. 16
B.2 MODEL ARCHITECTURES
This material describes oxidation-state predictor architectures and structured inference, then details the MP-20-OS dataset and OxiGen’s crystal and species diffusion. The species representation makes oxidation states explicit, while structured inference enforces charge neutrality.
B.2 MODEL ARCHITECTURES
- BERTOS predicts oxidation states from composition tokens, whereas CrystaliteOS uses a geometry-enhanced crystal transformer with a per-site oxidation-state classifier.The two predictors use separate best-validation checkpoints for argmax and Viterbi decoding.Source: §B.2 · p. 16
B.3 STRUCTURED INFERENCE
- Structured inference restricts assignments to element-admissible, charge-neutral states and limits mixed valence to designated elements.For other elements, sites are grouped into one transition sharing an oxidation state; designated mixed-valence elements retain site-level transitions, and Viterbi decoding returns the MAP assignment.Source: §B.3 · pp. 16–17
B.4 MP-20 OXIDATION-STATE LABELLING
- 42,576 structures remain in MP-20-OS from 45,229 MP-20 structures after filtering.The final dataset comprises 25,536 training, 8,495 validation, and 8,545 test structures.Source: §B.4 · p. 17, Table 4 · p. 17
C OXIGEN DETAILS
- OxiGen adapts MatterGen by jointly denoising species, fractional coordinates, and the lattice, with a single GemNet-dT network coupling their reverse predictions.The forward corruption processes factorise, and OxiGen retains MatterGen’s coordinate and lattice diffusion and network architecture.Source: C OXIGEN DETAILS · pp. 18–19
C.1 SPECIES DIFFUSION
- OxiGen represents each site as an element–oxidation-state species and replaces MatterGen’s atomic-identity diffusion with continuous-time absorbing-state diffusion.The model adds a learned oxidation-state embedding to the element embedding; its categorical output head predicts over the joint species vocabulary.Source: §C.1 · p. 19
- Species sites are masked independently during forward diffusion, and the denoiser predicts clean species from the full noisy crystal while only masked sites contribute to its objective.The implementation uses a numerically stabilised masking schedule approximating the fully masked terminal condition.Source: §C.1 · pp. 19–20
- The default OxiGen model applies the structured output layer during sampling, while charge-neutral structured training is used only in an ablation.The structured training objective computes exact one-site marginals; in the continuous-time limit, only single-site reveal events contribute to the NELBO.Source: §C.1 · p. 20
C.2 TRAINING AND SAMPLING
OxiGen retains MatterGen’s coordinate and lattice procedures and generates species with charge-neutral sampling. The section also describes structured-inference overhead, band-gap conditioning, and the relaxed-structure evaluation protocol.
C.2 TRAINING AND SAMPLING
- OxiGen retains MatterGen’s network, coordinate-diffusion, and lattice-diffusion procedures, and generates species using charge-neutral sampling.Generation excludes noble gases, Tc, Pm, and elements with atomic number Z ≥84; training for 900 epochs takes approximately 3.5 hours on four GH200 GPUs.Source: §C.2 · p. 20, §C.3 · p. 21
C.3 COMPUTATIONAL OVERHEAD
- Each structured-inference dynamic-programming pass costs O(N 2|S|(zmax − zmin)), with one backward pass drawing a charge-neutral sample per diffusion step, while structured training uses one forward and one backward pass to compute marginals.Training time is measured per epoch on MP-20-OS, and sampling time per batch of 128 structures.Source: §C.3 · p. 20
D BAND-GAP CONDITIONAL GENERATION
- Band-gap conditioning fine-tunes MatterGen and OxiGen on MP-20-OS band-gap labels and uses classifier-free guidance for continuous scores and discrete predictions.Generated-crystal band gaps are predicted from MatterSim-relaxed structures using iComFormer, whose reported test MAE is 0.193 eV.Source: D BAND-GAP CONDITIONAL GENERATION · p. 21
E.1 STABILITY, UNIQUENESS, AND NOVELTY
- Stability, uniqueness, novelty, and relaxation metrics are evaluated on relaxed structures, with a crystal counted as S.U.N. only if it is stable, unique, and novel.Stability is defined as Ehull ≤0.1 eV/atom; uniqueness and novelty use StructureMatcher, with novelty assessed against the Alex-MP reference dataset.Source: §E.1 · p. 21
E.2 OXIDATION-STATE METRICS
The section defines oxidation-state fidelity metrics and their reference distributions, then reports de novo-generation comparisons, including OxiGen’s balance of stability, uniqueness, and novelty.
E.2 OXIDATION-STATE METRICS
- The oxidation-state frequency score is the highest geometric-mean species frequency over valid assignments, counting each distinct species once and averaging over structures with an assignment.Valid assignments are charge neutral and disallow mixed valence; species frequencies come from ICSD-24.Source: §E.2 · pp. 22–23
- The minimum-frequency metric counts structures admitting a valid assignment whose every species frequency meets threshold τ.It reports the fraction of evaluated structures satisfying that condition.Source: §E.2 · p. 23
- Oxidation-state distributions are compared with macro-averaged Jensen–Shannon distance after weighting valid assignments by their species-frequency products and marginalising by element.Distributions are normalised separately for each element; the MP-20-OS test split is the reference, and compositions without valid assignments are excluded.Source: §E.2 · p. 23
F.1 DE NOVO GENERATION
- OxiGen achieves the highest S.U.N. rate and lowest relaxation RMSD among the compared methods, balancing stability, uniqueness, and novelty.The comparison uses the de novo-generation evaluation.Source: §F.1 · p. 23, Table 9 · p. 24
- OxiGen is comparable to MatterGen on structural validity and coverage, and slightly worse on element-count distribution distance.These proxy metrics use unrelaxed structures; distribution metrics are computed only on structures that are compositionally and structurally valid.Source: §F.1 · pp. 23–24, Table 10 · p. 24
- Under LeMat-GenBench, OxiGen has the highest validity, uniqueness, and M.S.U.N. rate among evaluated methods, and is competitive on other metrics.OxiGen uses 2,500 generated structures and is trained on MP-20-OS, while leaderboard baselines are trained on MP-20; the benchmark’s stability thresholds differ from the primary evaluation.Source: §F.1 · p. 24
F.2 ABLATION STUDY
The section compares band-gap-conditioned generation, ablates structured sampling and training, and illustrates generated crystals.
F.2 ABLATION STUDY
- Structured sampling raises stability and slightly raises S.U.N. while modestly reducing uniqueness and novelty; structured training provides no further improvement.Relative to the element-representation baseline, OxiGen slightly increases stability and S.U.N.Source: §F.2 · pp. 24–25, Table 12 · p. 26
- Structured sampling substantially improves agreement with CrystaliteOS oxidation-state assignments, while structured training adds no further improvement.Site agreement measures matching atomic sites; crystal agreement requires matching assignments at every site.Source: §F.2 · p. 25, Table 13 · p. 26
F.3 BAND-GAP CONDITIONAL GENERATION
- OxiGen has a higher S.U.N. rate than MatterGen at all four band-gap targets; its stability is higher despite lower novelty, and RMSE is comparable at 1 eV and better at 3, 5, and 7 eV.The band-gap-conditioned results use one fine-tuning seed and 1,024 generated structures per target.Source: §F.3 · p. 25, Table 14 · p. 26
F.4 EXAMPLE CRYSTALS
- Figure 4 presents examples of crystals generated by OxiGen that are stable, unique, and novel.Source: Figure 4 · p. 27