- The paper introduces CrystalREPA, a plug-and-play framework that transfers stability-aware priors from universal MLIPs to crystal generative models.
- It employs an element-aware contrastive loss to align atom-wise representations, enhancing energy accuracy, validity, and structural fidelity with minimal overhead.
- Empirical results across multiple models and datasets show improved metastable rates and reduced energy above hull, underscoring enhanced physical plausibility.
CrystalREPA: Transferring Stability-Aware Atomistic Priors from Universal MLIPs to Crystal Generative Models
Context and Motivation
The inverse design of crystalline solids demands generative models capable of efficiently exploring atomic configuration spaces subject to the periodic and physical constraints characteristic of crystals. State-of-the-art generative frameworks (diffusion and flow-based models) trained on equilibrium-structure datasets lack direct exposure to the full underlying potential energy surface (PES), with generation thereby guided primarily by geometric priors rather than principled stability supervision. In contrast, universal machine learning interatomic potentials (MLIPs) are extensively trained on large-scale, explicitly labeled datasets and yield robust atom-wise representations highly sensitive to structure, energy, and forces across an immense chemical space. This paper investigates the representation gap between contemporary crystal generative models and universal MLIPs, and proposes a framework for explicitly transferring stability-aware priors from the latter to the former.
CrystalREPA: Methodology and Architectural Innovations
The proposed framework, CrystalREPA, is a generic, plug-and-play training-time augmentation for crystal generative models. The core mechanism is atom-wise representation alignment via an element-aware contrastive loss between the internal hidden states of the generator's GNN encoder on intermediate (noisy) structures and the frozen atom-wise representations extracted from a universal MLIP on the corresponding clean crystal. Key architectural aspects are as follows:
- Representation Alignment: Instead of architectural or sampling modifications, CrystalREPA focuses strictly on the encoder's atom-wise embeddings, projecting the generative model's hidden states via an MLP head into the MLIP's feature space. This enables alignment of geometric and energetic priors without interference with generative flexibility or downstream sampling.
- Element-Aware Contrastive Loss (EA-NCE): The alignment loss leverages a symmetric InfoNCE objective with an element-aware mask, removing negatives corresponding to pairs of the same element to avoid representation collapse and false-negative signals due to intrinsic chemical similarities.
- Teacher Choice and Efficiency: Universal MLIPs are kept frozen and clean-structure representations are precomputed, resulting in marginal training overhead and no extra inference cost. The method is agnostic to the choice of generative model or MLIP architecture.
Empirical Results: Quality Gains and Robustness
CrystalREPA is comprehensively evaluated across three generative frameworks (CrystalFlow, MatterGen, DiffCSP), ten universal MLIP teachers (including DPA, MACE, ORB, SevenNet, and MatterSim), and two large benchmark datasets (MP-20, Alex-MP-20). The key findings are:
- Consistent Improvements: Across models and datasets, CrystalREPA delivers robust gains in Metastable Rate, MSUN Rate (Metastable, Stable, Unique, Novel), energy above hull (Ehull), RMSD (structural fidelity), and validity. For example, in MatterGen, Metastable% increases from 57.4 to 66.2 and MSUN% from 25.3 to 29.2, with the lowest Ehull and improved fidelity.
- Representation Gap Closure: Probing experiments demonstrate that atom-wise representations from vanilla generative models are substantially less predictive of formation energy than MLIP-derived representations. CrystalREPA narrows this gap substantially, especially in early GNN layers, indicating enhanced stability-awareness in generative encodings.
- Minimal Overhead: Training overhead is ≈1%, and inference remains unchanged—no additional computational cost or complexity is incurred during crystal generation.
- Complementarity with Conditional Generation: Unlike energy-conditioned generation, which may hurt structural validity and incurs inference overhead, CrystalREPA improves stability, fidelity, and validity in a balanced manner and does not preclude simultaneous use with other conditional methodologies.
Theoretical and Practical Implications
A pronounced and counterintuitive finding is that MLIP leaderboard accuracy (e.g., Matbench Discovery scores) does not reliably predict transfer effectiveness as CrystalREPA teachers. Instead, performance gain is highly correlated with the distinguishability of the teacher's atom-wise representation space, as quantified by new metrics: Chemical, Environment, and Total Resolution Indices (CRI, ERI, TRI). This has the following implications:
- Transfer Criteria: Practitioners should prioritize MLIP teachers with highly resolved, well-separated atom-wise representations (across element and local environment) rather than exclusively relying on energy/force benchmark scores when seeking generative transfer.
- Representation Normalization: Representation normalization can enhance transfer by increasing representation space distinguishability, as evidenced by improvements in both metrics and Metastable Rate.
From a practical standpoint, CrystalREPA enables generative models to natively favor physically plausible—rather than merely high-likelihood—crystals, significantly advancing inverse design capabilities for computational materials discovery. It achieves this without trading off diversity, validity, or unique/novel structure yield.
Limitations and Future Directions
While CrystalREPA addresses the crucial geometry-vs-energy representation gap, some limitations remain:
- Generality: The framework aligns only atom-level representations and ignores higher-order (e.g., bond, motif) features, which may further enhance transfer.
- Applicability: Experiments focus on crystalline materials; extension to molecular and biomolecular generation would require investigation.
- Data Overlap: Data overlap between the generator and MLIP training sets is not necessary for observed gains; improvements persist when overlap is minimized.
Further research is warranted along several axes:
- Representation alignment at higher abstraction levels (edges/motifs/periodicity).
- Investigation of transfer in non-crystalline domains (e.g., molecules, proteins).
- Combination strategies with strong conditional and guided generation paradigms.
Conclusion
CrystalREPA introduces a lightweight, general framework to bridge the gap between crystal generative models and physical stability priors learned by universal MLIPs. By enforcing atom-wise representation alignment with chemically resolved, stability-aware teacher models, CrystalREPA significantly enhances the thermodynamic quality, validity, and fidelity of generated crystals while preserving sampling efficiency and diversity. Its effectiveness is robust across architectures and datasets, and its transferability is controlled by teacher representation structure rather than conventional MLIP benchmark accuracy. This methodology constitutes a substantial step toward integrating generative and atomistic modeling for accelerated computational materials discovery.
Reference: "CrystalREPA: Transferring Physical Priors from Universal MLIPs to Crystal Generative Models" (2605.08960)
