Crys-JEPA: Joint Embedding for Crystal Generation

• Crys-JEPA is a joint-embedding predictive architecture that constructs an energy-aware latent space for de novo crystal generation.
• It addresses the stability–novelty trade-off by ranking candidate crystals using surrogate energy metrics and a refined screening pipeline.
• Empirical evaluations on MP-20 and Alex-MP-20 datasets demonstrate significant improvements in generating thermodynamically stable and novel inorganic materials.

Crys-JEPA is a joint-embedding predictive architecture designed for de novo crystal generation, targeting the accelerated discovery of novel and stable inorganic materials. Addressing the pronounced stability–novelty trade-off inherent in likelihood-based crystal generative models, Crys-JEPA constructs an energy-aware latent space that enables surrogate stability assessment and an efficient screening-and-refinement pipeline. This framework enables significant improvements in the generation of materials that are both thermodynamically feasible and distinct from known structures, leveraging embedding-based comparison rather than expensive energy evaluations (Liu et al., 14 May 2026).

1. Objectives and the Stability–Novelty Trade-off

De novo crystal generation seeks to identify new periodic materials that are both realistic—meaning physically valid and thermodynamically stable—and novel, i.e., not simple reproductions of known database entries. Stability is operationally defined by an energy-above-hull threshold (ΔE < ε, e.g., ε = 0.1 eV/atom), signifying proximity to the convex hull of formation energy. Novelty requires generated samples to populate low-density regions of the empirical distribution.

Empirical analyses of state-of-the-art likelihood-based generators—including diffusion models, normalizing flows, and VAEs—demonstrate a strong negative correlation between stability and novelty. Structures generated close to the training distribution (high-density regions) are stable but not novel, while those farther away exhibit low overlap with known structures but quickly become thermodynamically unstable. This observation reveals that the intersection of stability and novelty is a narrow manifold in configuration space, poorly explored by standard generative training paradigms (Liu et al., 14 May 2026).

2. JEPA: Energy-Aware Embedding Construction

The core of Crys-JEPA is a joint-embedding predictive architecture (JEPA) that induces an energy-aware latent space for crystals. JEPA comprises a shared transformer encoder and an MLP-based predictor, leveraging a contrastive InfoNCE-like objective weighted by formation energy differences to learn representations where Euclidean or cosine distances reflect relative energetic stability.

Architecture and Loss

• Input Representation: Each crystal $C_t$ is encoded as atom-wise feature vectors $V_t \in \mathbb{R}^{n \times (3+100+6)}$, capturing atomic types, lattice, and positional information.
• Augmentations: To construct positive pairs for contrastive learning, energy-preserving random translations in fractional coordinates and SO(3) lattice rotations are applied, producing a context crystal $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$.
• Embedding Process: The encoder $E(\cdot)$ (an 8-layer transformer, $d=512$) maps both $V_c$ and $V_t$ to embeddings $H_c, H_t \in \mathbb{R}^{n \times d}$, with a global representation extracted via a [CLS] token.
• Prediction Head: The predictor $P(\cdot)$ (MLP) combines $H_c$ and the augmentation parameters $V_t \in \mathbb{R}^{n \times (3+100+6)}$0, outputting a predicted embedding $V_t \in \mathbb{R}^{n \times (3+100+6)}$1:

$V_t \in \mathbb{R}^{n \times (3+100+6)}$2

• Energy-Aware InfoNCE Loss: Alignment between $V_t \in \mathbb{R}^{n \times (3+100+6)}$3 and its positive $V_t \in \mathbb{R}^{n \times (3+100+6)}$4 is enforced, while repulsion between negatives is scaled by their formation energy difference:

$V_t \in \mathbb{R}^{n \times (3+100+6)}$5

with

$V_t \in \mathbb{R}^{n \times (3+100+6)}$6

After training, embeddings $V_t \in \mathbb{R}^{n \times (3+100+6)}$7 are arranged such that $V_t \in \mathbb{R}^{n \times (3+100+6)}$8 correlates with $V_t \in \mathbb{R}^{n \times (3+100+6)}$9, providing a surrogate for energy-based comparison in latent space.

3. Surrogate Stability Assessment via Embedding Distance

Crys-JEPA enables efficient stability assessment by reference-based embedding distance metrics:

• Reference Set: For a candidate crystal $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$0, collect all training crystals $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$1 sharing the same chemical system or any subsystem (for A2B, one also includes A, B, AB2, etc.).
• Distance Metric: Compute the average squared distance in embedding space:

$C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$2

where $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$3. A small $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$4 indicates embedding proximity to energetically similar references, suggesting higher likelihood of stability.

• Screening: Generated crystals are ranked by ascending $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$5, with the lowest values selected as the most promising for further consideration.

4. Screening-and-Refinement Pipeline

Crys-JEPA institutes a closed-loop generative refinement process structured as follows:

1. Sampling: Draw $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$6 candidate crystals from the base generative model $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$7.
2. Pre-screening: In silico relax each candidate with a machine-learned force field (MLFF) and retain those satisfying validity (V), uniqueness (U), and novelty (N) relative to the training set.
3. JEPA Screening: For each surviving candidate, construct its reference set and compute $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$8. Candidates are then sorted by $C_c = \mathcal{T}_t \circ \mathcal{R}_r(C_t)$9.
4. Selection: Select the top $E(\cdot)$0 of candidates (by $E(\cdot)$1) for fine-tuning.
5. Refinement: Fine-tune the base model $E(\cdot)$2 on this set for 1–5 epochs using the diffusion objective, yielding a refined model $E(\cdot)$3.

One iteration suffices to achieve substantial performance improvements, though additional cycles could be applied if required (Liu et al., 14 May 2026).

Pipeline Algorithm (Structured Summary)

Step	Operation	Output/Effect
1. Generate	Sample $E(\cdot)$4 crystals from $E(\cdot)$5	Candidate set
2. Pre-screen	MLFF relax; VU.N filter	Valid, unique, novel subset $E(\cdot)$6
3. JEPA screen	Reference set $E(\cdot)$7 compute $E(\cdot)$8 for each	Rank $E(\cdot)$9 by $d=512$0
4. Select top $d=512$1	Select $d=512$2 from $d=512$3	High-quality pool
5. Fine-tune	Train $d=512$4 on $d=512$5 (1–5 epochs)	Updated generator $d=512$6

5. Empirical Performance and Metrics

Crys-JEPA was benchmarked on the MP-20 and Alex-MP-20 datasets. The principal evaluation metrics are S.U.N (stability, uniqueness, novelty) and V.S.U.N (validity, stability, uniqueness, novelty), computed over $d=512$7 generated structures as follows:

$d=512$8

where $d=512$9 indicate satisfaction of the respective criteria for sample $V_c$0.

Performance improvement (DFT-evaluated) relative to MatterGen baseline:

Dataset	Baseline V.S.U.N	Crys-JEPA Refined	Relative Improvement
MP-20	26.4%	47.9%	+81.4%
Alex-MP-20	35.1%	64.1%	+82.6%

The approach demonstrates consistent and substantial gains in the joint satisfaction of realistic, stable, unique, and novel structure criteria (Liu et al., 14 May 2026).

6. Training and Implementation Details

• JEPA Encoder: 8-layer transformer ($V_c$1, 16 heads, no dropout), trained on ~840k crystals with formation energy labels, batch size $V_c$2, $V_c$3, Adam optimizer ($V_c$4), for 50 epochs.
• Generative Backbone: Diffusion (DDPM, $V_c$5 timesteps, cosine $V_c$6-schedule), denoiser is a 12-layer transformer ($V_c$7, dropout 0.01), trained with AdamW on both MP-20 and Alex-MP-20.
• Screening Efficiency: Embedding inference for 10k crystals requires ~48 seconds on a single NVIDIA L40S GPU, substantially faster than MLFF-based surrogates.
• Refinement Hyperparameters: Candidate pool size $V_c$8–10,000, selection rate $V_c$9–10%. Larger $V_t$0 improves performance; excessive $V_t$1 admits unstable outliers.
• Hardware and Environment: AMD EPYC 9554 CPU, NVIDIA L40S GPU, PyTorch, Linux platform.

7. Significance and Broader Applications

Crys-JEPA delineates a paradigm where energy-based surrogate screening in an embedding space obviates the need for costly quantum mechanical calculations during candidate selection. This approach enables scalable crystal generation embodying both physical plausibility and structural novelty, and the embedding-based screening method is agnostic to specific generative architectures. The joint-embedding principle underlying Crys-JEPA has also been repurposed in other domains, such as CryoLVM for cryo-EM density map analysis, where JEPA-based pretraining facilitates rapid adaptation to diverse downstream tasks via rich structural latent representations (Fu et al., 2 Feb 2026).

A plausible implication is that this architecture template, coupling contrastive, property-aware embedding learning with downstream surrogate-based screening, provides a generalizable solution to problems characterized by a narrow overlap between realism (feasibility) and innovation (distance from training distribution). Future extensions may explore alternative property surrogates, reference class construction, or generative model backbones.