Score-based diffusion models for accurate crystal-structure inpainting and reconstruction of hydrogen positions

This paper presents a score-based diffusion model that combines materials science and computer vision techniques to accurately and efficiently reconstruct missing hydrogen atom positions in crystal structures, achieving a success rate exceeding 97% compared to traditional DFT or unconditioned approaches.

The Gist

Imagine you have a 3D puzzle of a crystal, but someone has taken out all the tiny pieces representing hydrogen atoms. In the real world, finding where these hydrogen pieces belong is incredibly hard. It's like trying to spot a ghost in a room using a flashlight that barely sees anything; standard X-ray cameras often miss them entirely, forcing scientists to use expensive, massive machines (neutron diffraction) to find them. Because of this, many crystal databases are full of "ghosts"—structures where the hydrogen pieces are either missing entirely or just guessed at based on a hunch.

This paper introduces a new "AI detective" that can fill in those missing pieces with high accuracy. Here is how they did it, explained simply:

The Problem: The "Missing Piece" Puzzle

Scientists have a powerful AI tool called MatterGen that is great at creating new crystal puzzles from scratch. However, the researchers wanted to use it for a different job: inpainting.

Think of "inpainting" like the "Healing Brush" tool in photo editing software. If you have an old photo with a scratch or a missing piece, the tool looks at the surrounding pixels and intelligently fills in the gap. The researchers wanted to do this for crystals: look at the known parts of the structure and intelligently fill in the missing hydrogen spots.

The Solution: Borrowing from Art Restoration

The team realized that the best tools for "filling in the blanks" weren't actually in chemistry; they were in computer vision (how computers "see" images). They took a technique called TD-Paint, originally designed to restore damaged images, and taught it to understand crystal structures.

They trained a new version of the AI (called pos-only-TD) with a specific rule:

• The Knowns: The AI is told, "These atoms are real and fixed; do not touch them."
• The Unknowns: The AI is told, "These spots are empty; guess what goes here."

Unlike older methods that would guess the whole picture from scratch (which is like trying to redraw a whole painting because you lost one brushstroke), this AI only focuses on filling the holes while respecting the existing structure.

How It Works: The "Denoising" Dance

The AI works like a sculptor starting with a block of noisy, random clay.

1. Start: It starts with a crystal where the hydrogen spots are just random noise (like static on an old TV).
2. The Process: Step-by-step, the AI "cleans" the noise. It asks, "Based on the atoms I do know are there, what should the hydrogen atoms look like?"
3. The Twist: The new method (TD-Paint) is smarter about this. It knows that the "known" atoms are already perfect and shouldn't be "noisy." It only adds noise to the missing spots and cleans those up, making the process much faster and more accurate.

The Results: A 99% Success Rate

The team tested this on a huge library of crystals.

• The Test: They took crystals with known hydrogen positions, hid the hydrogens, and asked the AI to find them again.
• The Score: The AI succeeded in finding the exact original structure 97% of the time.
• The Bonus: In the other cases, the AI didn't just fail; it often found a better version of the crystal—one that is more stable and energetic than the original "guess" found in the database.

When they filtered out the crystals that were known to be just theoretical guesses (not real experiments), the success rate jumped to 99%.

Why This Matters

• Speed and Cost: Instead of needing expensive neutron machines to find hydrogens, scientists can now use this AI to predict them instantly on a standard computer.
• Fixing Bad Data: The AI can act as a quality control check. If the AI predicts a hydrogen position that is very different from what's in the database, it might mean the database entry was wrong or a "bad guess" all along.
• Beyond Hydrogen: While they tested it on hydrogen, the method is "hydrogen-agnostic." This means the same tool could be used to fill in missing pieces for other atoms, like lithium in battery materials, without needing to be retrained from scratch.

In short, the researchers took a tool designed for art restoration, taught it the language of crystals, and gave it the ability to "see" the invisible hydrogen atoms that have been missing from our scientific maps for decades.

Technical Summary

Technical Summary: Score-based diffusion models for accurate crystal-structure inpainting and reconstruction of hydrogen positions

Problem Statement
The accurate determination of atomic positions for hydrogen atoms in crystalline materials remains a significant challenge in computational materials science. While hydrogen is abundant and critical for applications ranging from energy storage to semiconductor industry, its low X-ray and electron scattering power makes it difficult to locate via standard X-ray diffraction. Consequently, historical crystallographic databases (such as MC3D, derived from COD, ICSD, and MPDS) often contain structures where hydrogen positions are either omitted, inserted via heuristics, or determined with significant uncertainty. Although experimental advances (e.g., quantum crystallography) and expensive neutron scattering can resolve these positions, existing databases still rely heavily on incomplete or approximate data. The core problem addressed is the reconstruction of these missing hydrogen positions given a known "host" crystal structure (the non-hydrogen framework) to generate accurate, stable, and complete crystal structures.

Methodology
The authors propose a generative approach combining score-based diffusion models with techniques adapted from computer vision image inpainting.

1. Model Architecture: The study utilizes Microsoft's MatterGen, a pre-trained score-based diffusion model originally designed to generate novel crystal structures. The authors retrain and adapt this model for the specific task of inpainting (filling in missing data) rather than full generation.
2. Inpainting Strategies: Several approaches are benchmarked:
   • MatterGen Baseline: The original model applied to the task with fixed lattice and atomic types.
   • pos-only: A model retrained to denoise only atomic positions, keeping lattice and types fixed.
   • pos-only-RePaint: The retrained model combined with the RePaint algorithm, which iteratively resamples known regions to harmonize them with the inpainted regions.
   • pos-only-TD (Time-Dependent): A novel adaptation of the TD-Paint concept from computer vision. This model is trained to handle variable noise levels per site within a single structure. It conditions the denoising process on "clean" known positions (zero noise) while iteratively predicting the missing hydrogen sites (high noise). This allows the model to directly utilize the known host structure without adding noise to it.
3. Workflow:
   • Start with a host structure (hydrogen sites removed).
   • Randomly initialize $N_H$ hydrogen sites in the unit cell.
   • Apply the pos-only-TD diffusion model for $N_{steps}$ iterations to predict hydrogen positions.
   • Perform a constrained relaxation of the hydrogen sites using a Machine Learning Interatomic Potential (MLIP), specifically NequIP, while keeping the host atoms and cell fixed.
   • Perform a full relaxation of all positions (including host atoms) using NequIP.
4. Evaluation Metrics: Success is defined in two ways:
   • Structural Match: The predicted structure matches the reference (from the database) within pymatgen's StructureMatcher tolerances.
   • Lower Energy or Structural (LES) Match: The prediction either matches the structure or finds a configuration that is energetically more stable than the reference (according to DFT or MLIP).

Key Results
The method was evaluated on two datasets derived from the MC3D database:

• DFT Dataset: Structures with known hydrogen positions that have been relaxed using Density Functional Theory (DFT).

• EXP Dataset: The original experimental structures before DFT curation.

• Performance Comparison: The pos-only-TD model significantly outperformed the baseline MatterGen, the pos-only model, and the pos-only-RePaint approach.

  • RePaint improved variance but increased computational cost (requiring ~1000 steps vs. 300) without significantly boosting the matching rate.
  • The pos-only-TD model achieved a structural matching rate of 88% on a subset where a DFT-based "pinball" method reached ~87%, but with a more efficient global optimization approach.

• Success Rates:

  • When generating a single sample per structure, the method achieved high matching rates.
  • By generating multiple samples (e.g., $k=10$ or $k=30$) and selecting the most stable configuration, the LES matching rate exceeded 97% for the DFT+MLIP dataset and 97.2% for the EXP+MLIP dataset.
  • When excluding structures flagged as "theoretical" in the source databases, the success rate on the EXP dataset exceeded 99%.

• Stability and Generalization:

  • Many "non-matching" predictions were found to be energetically more stable than the original reference structures, suggesting the experimental references in the databases were suboptimal or incorrect.
  • The model demonstrated zero-shot generalization to structures with up to 40 atoms per unit cell, despite being trained only on structures with up to 20 atoms. While performance degraded slightly for larger systems, the LES matching rate remained above 90%.

• Computational Efficiency: The diffusion approach is significantly faster than the DFT-based pinball method, which suffers from combinatorial explosion and high computational costs.

Significance and Claims
The paper claims that transferring image inpainting techniques (specifically TD-Paint) to materials science enables a faster and more accurate completion of host structures compared to unconditioned diffusion models or traditional DFT-based approaches.

• Database Curation: The method provides a robust tool for curating and extending materials databases (like MC3D) by reliably reconstructing missing hydrogen positions, thereby enabling high-throughput screening of hydrogen-containing materials.
• Discovery of Stable Configurations: The approach does not merely reconstruct known structures; it often identifies configurations that are more stable than those reported in experimental databases, potentially correcting historical errors or heuristic insertions.
• Hydrogen-Agnostic Transferability: The models are trained in a "hydrogen-agnostic" manner. The authors claim this architecture can be easily transferred to other inpainting tasks, such as intercalating ions (e.g., Li, Na) into cathode materials or predicting muon stopping sites, without the need for retraining.
• Foundation for Future Work: The paper positions the method as a foundation that can be fine-tuned for larger structures or specific applications, rather than requiring full retraining.

The authors emphasize that their success rate definition includes both structural recovery and the discovery of more stable configurations, framing the latter not as a failure of reconstruction but as a successful prediction of a more physically accurate state.