MiAD: Mirage Atom Diffusion in Crystals

• MiAD is a generative diffusion framework that introduces a mirage atom type to enable dynamic atom counts in crystal synthesis.
• It employs an equivariant joint Markov diffusion model combining Gaussian, Wrapped Normal, and categorical diffusions for lattice, coordinates, and atom types.
• MiAD achieves a state-of-the-art 8.2% S.U.N. rate, outperforming baselines in stability and demonstrating improved uniqueness and novelty.

Mirage Atom Diffusion (MiAD) is a diffusion-based generative framework for de novo crystal generation that explicitly models the appearance and disappearance of atoms during the generative process. Introduced as a solution to the rigidity of fixed atom count in prior diffusion models, MiAD achieves state-of-the-art results on metrics of crystal stability, uniqueness, and novelty (S.U.N.) by employing a simple mirage infusion technique. This approach introduces an auxiliary “mirage” atom type, expands the generative domain to permit dynamic atom counts, and integrates these modifications within an equivariant joint Markov diffusion model (Okhotin et al., 18 Nov 2025).

1. Joint Diffusion Modeling for Crystal Structures

MiAD treats a crystal unit cell as a triplet $M = (L, F, A)$, where $L \in \mathbb{R}^{3 \times 3}$ is the lattice matrix, $F \in [0,1)^{N \times 3}$ denotes fractional coordinates of $N$ atoms, and $A \in \{1, \dots, K\}^N$ the atom types. The model employs a Markov noising process $q(M_{0:T}) = q(M_0) \prod_{t=1}^T q(M_t|M_{t-1})$ to corrupt data into pure noise, and a learnable reverse denoising process $$p_\theta(M_{0:T}) = p(M_T)\prod_{t=1}^T p_\theta(M_{t-1}|M_t)$$ to synthesize samples.

The forward kernel factorizes as

$$q(M_t|M_{t-1}) = q(L_t|L_{t-1}) \cdot q(F_t|F_{t-1}) \cdot q(A_t|A_{t-1})$$

with component-wise treatment as follows:

• Lattice: A Gaussian DDPM diffusion: $$q(L_t|L_{t-1}) = \mathcal{N}(L_t; \sqrt{1-\beta_t} L_{t-1}, \beta_t I)$$.
• Fractional coordinates: Modeled on the torus via Wrapped Normal kernels $q(F_t|F_{t-1})$, with score-matching loss leveraging Riemannian gradients.
• Atom types: Handled by categorical D3PM diffusion, $$q(A_{t,i}|A_{t-1,i}) = \mathrm{Cat}(A_{t,i}; Q_t \cdot \mathrm{onehot}(A_{t-1,i}))$$.

The total loss combines lattice, coordinate, and atom-type terms as $L = \kappa_L L_L + \kappa_F L_F + \kappa_A L_A$. The neural network is an equivariant GNN (CSPNet), respecting atom permutation, O(3) rotation, and periodicity.

2. Mirage Infusion: Dynamic Atom Count Handling

Traditional diffusion models for crystal generation fix $N$ (atom count) at the outset, constraining trajectory variability. MiAD's mirage infusion addresses this by:

• Introducing an extra “mirage” atom type (type 0).
• Fixing a large maximum atom count $N_m \geq \max$ training atom count; working instead with

$$\overline{M} = (L,\,\overline{F},\,\overline{A}), \qquad \overline{A}_i \in \{0,1,\dots,K\},\quad i=1\dots N_m$$

• Infusion operator: Extends a genuine crystal $M$ with $N$ real atoms to $\overline{M}$ by appending $N_m-N$ mirage atoms (type 0) with uniformly random coordinates.
• Reduction operator: After sampling, genuine atoms are identified by type $>0$, while mirage atoms (type 0) are removed.

All model components and losses are computed in the expanded $N_m$-atom domain, but only non-mirage atoms contribute to the coordinate loss, via a masked objective:

$$L_F^{mask} = \mathbb{E}_{t,\overline{M}_0,\overline{M}_t}\sum_{i\in\mathrm{Real}(\overline{M}_0)} \left\|\nabla_{\overline{F}_{t,i}} \log q(\overline{F}_t|\overline{F}_0) - s_{\theta,i}(\overline{M}_t, t) \right\|^2$$

This enables the model to generate/deactivate atoms dynamically during denoising.

3. Mathematical Model and Equivariant Constraints

Both forward and reverse processes operate over the expanded joint space: $$\begin{align*} \text{Forward:} & \qquad q(\overline{M}_{0:T}) = q(\overline{M}_0)\prod_{t=1}^T q(\overline{M}_t|\overline{M}_{t-1}) \ \text{Reverse:} & \qquad p_\theta(\overline{M}_{0:T}) = p(\overline{M}_T)\prod_{t=1}^T p_\theta(\overline{M}_{t-1}|\overline{M}_t) \end{align*}$$ The full training loss is

$$L = \kappa_L L_L + \kappa_F L_F^{mask} + \kappa_A L_A$$

with $L_L$ the KL loss for the lattice, $L_F^{mask}$ for masked coordinates, and $L_A$ for atom-type transitions across all $N_m$ atoms.

Equivariance under group actions is enforced so that, for any group element $g$ (e.g., permutation, rotation, translation),

$$\begin{align*} p_\theta(g \cdot L_T) &= p(L_T) \ p_\theta(g \cdot M_{t-1}|g \cdot M_t) &= p_\theta(M_{t-1}|M_t) \ q(g \cdot M_t|g \cdot M_{t-1}) &= q(M_t|M_{t-1}) \end{align*}$$

4. Neural Network Architecture and Symmetry Handling

MiAD employs the CSPNet GNN architecture from DiffCSP. Key features include:

• Atom embeddings updated via message-passing layers, pooling over neighbors within a fixed cutoff.
• Lattice representation processed by a feed-forward network on $L_t^T L_t$ for O(3) equivariance.
• Fractional-coordinate updates use pairwise Fourier features to impose periodicity.

The fixed-size $N_m$ representation guarantees the network is agnostic to the actual count of non-mirage atoms, handling dynamic atom populations transparently.

5. Sampling Algorithm

MiAD's generative process operates as follows:

1. Initialization:
   • $L_T \sim \mathcal{N}(0, I)$ (lattice)
   • $$\overline{F}_T \sim \mathrm{Uniform}(0,1)^{N_m\times 3}$$ (fractional coordinates)
   • $\overline{A}_T \sim \mathrm{Cat}(1/(K+1))^{N_m}$ (atom types)
2. Denoising:
   • For $t = T$ down to $1$: predict and sample from $p_\theta(\overline{M}_{t-1}|\overline{M}_t)$.
3. Reduction:
   • Discard all atoms $i$ with $\overline{A}_{0,i} = 0$ to obtain a valid sample $M_0$.

This routine allows for emergence or disappearance of atoms at any reverse diffusion step.

6. Experimental Evaluation

Dataset and Training

Experiments are conducted on the MP-20 subset of the Materials Project, containing 45,231 stable crystals (train/val/test split 60/20/20%). MiAD and baseline models are evaluated using 10,000 generated samples each. Hyperparameters follow prior work (DiffCSP/CSPNet): hidden dimension 512, 6 GNN layers, maximum 100 atoms, maximum 20 neighbors, and cutoff 7.0 Å. For MiAD, $N_m=25$ is used (≈2.7× the mean atom count). Training uses batch size 256, Adam optimizer with learning rate $10^{-3}$.

Quantitative Results

Key performance is measured by the S.U.N. rate (fraction of generated crystals that are stable by DFT [E$_{\text{hull}}<$0 eV], unique via structure matching, and novel/unseen in training). Select results:

Model	Stable %	Unique %	S.U.N. %
DiffCSP	5.0	66.0	3.3
FlowMM	4.6	60.9	2.8
FlowLLM	13.9	33.8	4.7
WyFormer	5.2	92.3	4.8
ADiT	12.8	50.8	6.5
MiAD	12.5	65.2	8.2

MiAD achieves an 8.2% S.U.N. rate, a 2.5× improvement over the unmodified DiffCSP baseline and a 25% gain over the previous state of the art. Similar improvements are observed on ML-potential-based stability metrics and larger datasets.

7. Limitations and Prospective Research

Several limitations are identified:

• Computational Overhead: Mirage infusion increases the average atom count (by ≈2.7×), resulting in roughly 4× slower training and sampling and elevated memory requirements.
• Hyperparameter Sensitivity: Performance depends on $N_m$ and loss coefficient $\kappa_A$; ablations suggest optimal $N_m=25$ and balanced loss weights ($36\%$ lattice, $46\%$ coord, $18\%$ atom type).
• Failure Modes: Inadequate masking or improper initialization of mirage atoms may lead to degenerate behaviors such as trivial atom insertions or removals, or reduced generative diversity.
• Metrics Coverage: The S.U.N. metric, while comprehensive for stability and novelty, does not capture space-group diversity or structural variety, suggesting the need for additional evaluation criteria.
• Generalization: Mirage infusion is directly compatible with DiffCSP, MatterGen, FlowMM, CrystalFlow. Extending the approach to architectures such as transformers or to new domains (e.g., drug design, nanoporous materials) is proposed as a promising future direction.

A plausible implication is that the mirage infusion technique could be generalized beyond crystals to other structured generative domains that require variable count handling, provided equivariance and masking are appropriately enforced.

---

MiAD demonstrates that introducing a mirage atom type to joint diffusion models for crystal structure generation enables flexible atom count and substantially improves de novo inorganic materials discovery performance, as measured by the S.U.N. metric. The technique's conceptual simplicity and strong empirical results indicate its potential as a new standard for diffusion-based generative modeling in materials science (Okhotin et al., 18 Nov 2025).