Constrained Diffusion for Accelerated Structure Relaxation of Inorganic Solids with Point Defects

1. Problem Statement

The paper addresses the computational bottleneck associated with high-throughput first-principles simulations (specifically Density Functional Theory, or DFT) for point defects in inorganic solids, using Bismuth Telluride ($Bi_2Te_3$) as a case study.

• The Challenge: Point defects (vacancies, anti-sites) drastically alter material properties (electronic, thermal, structural). However, simulating them requires large supercells and complex energy landscapes. DFT is computationally prohibitive for high-throughput exploration of the vast configuration space created by defect permutations.
• The Limitation of Existing AI: While generative models (like Diffusion Models) can generate material structures, they struggle to satisfy strict physical constraints (e.g., minimum interatomic distances, periodic boundary conditions, force minimization).
  • Conditional Diffusion: Biases training but often fails to guarantee hard constraints, leading to physically impossible structures (e.g., overlapping atoms).
  • Projected Diffusion (State-of-the-Art): Enforces constraints at every denoising step. However, this relies on the assumption that constraints can be meaningfully evaluated on noisy intermediate states. In materials science, constraint evaluators (surrogates) are often unreliable on high-noise samples, leading to "misaligned gradients" and compounding errors.
  • Post-Processing: Projects the final output onto the constraint set. This often fails because the final output may be too far from the feasible region to be corrected in a single step, or it pushes the sample away from the learned data distribution.

2. Methodology

The authors propose a Constrained Diffusion Framework that integrates a Primal-Dual Algorithm to handle complex constraints without relying on intermediate feasibility assumptions.

A. Core Framework

The method utilizes a Score-based Diffusion Model to generate atomic structures. The generation process is treated as a multi-stage optimization problem where constraints are enforced only on the terminal state ($x_0$), rather than at every intermediate timestep ($x_t$).

B. Constraint Types

The model targets three specific physical constraints for $Bi_2Te_3$:

1. Geometric Constraints: Enforce minimum interatomic distances ($d_{min}$) to prevent atomic collapse and ensure periodic boundary conditions (fractional coordinates in $[0,1]^3$).
2. Distributional Constraints: Ensure the Radial Distribution Function (RDF) of generated structures matches reference data within a 5Å cutoff, preserving local coordination shells.
3. Force Minimization: Minimize total atomic forces to ensure structural stability. This is handled via a soft constraint using a neural surrogate (MACE model) to predict forces, added to the objective function.

C. The Primal-Dual Projection Algorithm

To solve the constrained optimization problem at the final step ($t=0$), the authors employ an Augmented Lagrangian Method (ALM):

• Formulation: Instead of a hard projection, they define a Lagrangian loss function $L(y, x_0; \lambda, \mu)$ that combines:
  • The distance to the diffusion output ($\|y - x_0\|^2$).
  • Linear penalty terms for constraint violations ($\lambda_j c_j(y)$).
  • Quadratic penalty terms for constraint violations ($\frac{\mu_k}{2} c_k(y)^2$).
• Optimization: The algorithm iteratively updates the projected structure $y$ using gradient descent while simultaneously updating the Lagrange multipliers ($\lambda$) and quadratic penalties ($\mu$) to drive constraint violations toward zero.
• Key Innovation: By deferring constraint enforcement to the final "clean" sample ($x_0$), the method avoids the gradient inaccuracies associated with evaluating complex surrogates on noisy, intermediate diffusion states.

3. Key Contributions

1. Novel Constraint Handling: Extends constrained generation to handle complex, non-convex constraints modeled by tractable functions and neural surrogates using a primal-dual approach.
2. Terminal Set Constraints: Introduces a "terminal set" perspective (inspired by Model Predictive Control) where constraints are only enforced on the final output. This bypasses the "misaligned gradient" problem inherent in projected diffusion models that enforce constraints on noisy intermediate states.
3. First Application to Thermoelectrics: Provides the first study applying constrained diffusion to supplement DFT-based point defect simulation for thermoelectric materials ($Bi_2Te_3$).
4. State-of-the-Art Performance: Demonstrates superior performance across six distinct defect configurations compared to existing baselines.

4. Experimental Results

The method was evaluated on six stoichiometric configurations of $Bi_2Te_3$ representing four defect types (Te vacancies, Bi vacancies, Bi-on-Te anti-sites, Te-on-Bi anti-sites).

Baselines Compared:

• Conditional Diffusion Models (CDM)
• Projected Diffusion Models (Projected DM)
• Post-processing Optimization (Post-proc DM)

Performance Metrics:

• RMSD (Root Mean Square Deviation): Measures structural deviation from reference DFT-relaxed structures.
• RDF (Radial Distribution Function) Error: Measures global geometric consistency.
• Force (eV/Å): Total force magnitude (lower is better; indicates stability).

Key Findings (Table 1 & 2):

• RMSD: The proposed method achieved the lowest RMSD (e.g., 0.90 ± 0.38 for 16Bi+21Te) compared to Projected DM (2.54) and Post-proc DM (2.30).
• RDF Similarity: The proposed method showed significantly higher RDF similarity (error 0.30 ± 0.11) compared to baselines, which were often 2x–100x higher.
• Force Minimization: The proposed method reduced total forces by several orders of magnitude (e.g., $7.88 \times 10^{-2}$ eV/Å) compared to Conditional DM ($10^9$) and Post-proc DM ($10^6$).
• Physical Realism: Unlike Conditional DM (which produced overlapping atoms) and Projected DM (which suffered from gradient misalignment), the proposed method consistently generated physically plausible, stable structures.

5. Significance

• Accelerated Discovery: This framework enables the rapid generation of physically realistic defect structures, acting as a high-fidelity surrogate for expensive DFT calculations. This is crucial for high-throughput screening of thermoelectric materials.
• Robustness to Noise: By relaxing the assumption that constraints must hold at every diffusion step, the method overcomes a major limitation of current physics-informed generative models in materials science.
• Generalizability: The primal-dual projection approach is applicable to other scientific domains where constraints are complex, non-convex, and evaluated by expensive or noisy surrogates (e.g., protein folding, fluid dynamics).

In summary, the paper presents a robust, physics-aware generative framework that successfully bridges the gap between the flexibility of diffusion models and the rigorous constraints required for accurate materials simulation.