World Models as Neural Surrogates
for DFT-Scale Materials Property Simulation

Abstract

DFT simulation is the primary bottleneck in high-throughput materials discovery. We propose a two-stage neural surrogate: a Joint Embedding Predictive Architecture (JEPA) encodes crystal graph structure into a physics-informed latent representation, and a lightweight transformer reads out formation energy, bandgap, and stability. The JEPA learns physical regularities via masked embedding prediction without explicit quantum-mechanical labels. This decomposition is well-suited to replacing repetitive DFT operations in screening loops. Validation on Materials Project and OQMD is ongoing.

Keywords: world models · joint embedding predictive architecture · neural surrogates · materials discovery · crystal graphs

1 · Introduction

High-throughput computational materials discovery relies on evaluating large candidate pools against property targets before committing to laboratory synthesis. Density functional theory provides the accuracy required for this evaluation, but its O(N³) scaling with electron count makes it prohibitively expensive at screening scale: a single DFT calculation for a moderately complex structure takes minutes to hours on a computing cluster, and a full screening campaign may involve tens of thousands of candidates [Jain et al., 2013]. The dominant cost in modern materials foundry pipelines is therefore not synthesis or characterisation but simulation.

Neural surrogate models have addressed this bottleneck with increasing success. Graph neural networks such as MEGNet [Chen et al., 2019] and ALIGNN [Choudhary & DeCost, 2021] operate directly on crystal graph representations and achieve strong performance on standard benchmarks, approaching DFT-level accuracy on formation energy prediction. However, these architectures are task-specific regressors: each property requires a separate model, and the learned representations do not capture the joint latent structure underlying multiple properties simultaneously. More fundamentally, they do not model the spatial predictive structure of the simulation itself — the fact that a material's properties emerge from the physics of its local atomic environments and their interactions across the lattice.

We propose grounding the surrogate in a Joint Embedding Predictive Architecture [Assran et al., 2023], or more commonly called a World Model: a self-supervised framework that learns latent representations by predicting masked regions of a structured input in embedding space rather than in raw feature space. Applied to materials, the JEPA encoder learns to predict the latent representation of missing or masked structural components from visible context, forcing the model to internalise the spatial and physical regularities of crystal structures without supervision. The resulting latent state z is then consumed by a lightweight transformer readout head that predicts material properties autoregressively. The two stages have complementary roles: the JEPA encoder does the hard work of structural understanding; the readout head does the cheap work of property regression from an already-informative representation.

2 · Proposed architecture

Stage 1 · JEPA structural encoder

The input to the encoder is a crystal graph G = (V, E), where nodes v ∈ V represent atomic sites with features (atomic number, oxidation state, electronegativity) and edges e ∈ E encode interatomic distances and bond angles derived from the lattice geometry. A fraction of nodes is masked, and the encoder E_θ is trained to predict the latent representations of masked nodes from the visible context:

L_JEPA = 𝔼[ ‖ E_θ(G_visible)_pred − sg(E_φ(G_full)) ‖² ],

where sg(·) denotes stop-gradient and E_φ is an exponential moving average (target) encoder following the I-JEPA formulation [Assran et al., 2023]. This objective operates entirely in latent space: the model never reconstructs raw atomic positions, only their embeddings. The result is a global latent vector z ∈ ℝ^d_z summarising the material's structural state.

Stage 2 · Transformer readout head

A lightweight transformer D_ψ receives z as a conditioning token and autoregressively predicts a short sequence of property values:

ŷ_t = D_ψ(z, ŷ_<t),    t ∈ {E_f, E_g, s},

where E_f is formation energy per atom (eV/atom), E_g is electronic bandgap (eV), and s ∈ {0, 1} is a binary thermodynamic stability flag. The autoregressive structure allows the model to condition bandgap prediction on the formation energy estimate, reflecting the physical correlation between thermodynamic stability and electronic structure. The readout head is intentionally small — 4–6 transformer layers, d = 256 — as the representational burden is carried by the JEPA encoder. The full system is quantisable to 4-bit GGUF format for deployment on resource-constrained hardware, with total inference cost per candidate in the millisecond range.

Training and data

The JEPA encoder is pre-trained self-supervised on crystal graphs from the Open Quantum Materials Database [Saal et al., 2013] (~1M structures). The readout head is then fine-tuned supervised on Materials Project [Jain et al., 2013] property labels (~150K structures with DFT-computed E_f, E_g, and stability). The composite fine-tuning loss is:

L = λ₁ ‖ Ê_f − E_f ‖² + λ₂ ‖ Ê_g − E_g ‖² + λ₃ L_BCE(ŝ, s).

No DFT calculations are performed at any stage: the surrogate learns entirely from pre-computed simulation records, making it deployable without quantum chemistry software.

Surrogate scope

The surrogate targets small-scale and repetitive DFT operations: rapid property screening of candidate structures prior to high-fidelity validation runs. The intended use is reducing a candidate pool of O(10⁴) structures to O(10²) before expensive simulation is committed. It is not intended to replace full DFT for final structural confirmation. This bounded scope makes the surrogate's accuracy requirements tractable: a MAE of ~0.05 eV/atom on formation energy is sufficient for screening-level ranking, even if insufficient for publication-quality property reporting.

3 · Discussion

The World Model formulation offers a specific advantage over direct graph regression: by learning to predict masked structural regions in latent space, the encoder is forced to model the physical dependencies between atomic environments — the same dependencies that DFT resolves quantum-mechanically. This is qualitatively different from supervised graph regression, where the model learns a mapping from full graph to property label without necessarily internalising the spatial structure of that mapping. The hypothesis is that a JEPA-pretrained encoder generalises better to novel chemistries, where atomic environments may be unseen in training but their physical regularities are not.

The two-stage decomposition also has a practical benefit for active discovery loops: the JEPA encoder can be frozen after pre-training, and the readout head can be fine-tuned cheaply on new property targets or new chemical spaces without retraining the full model. This makes the architecture compatible with the iterative, campaign-structured workflow of a materials foundry, where simulation objectives shift as the campaign progresses.

Status

This paper is an architectural proposal submitted for review. It is published here as a preprint. Validation on Materials Project and OQMD is ongoing. The authors welcome feedback and discussion on the approach and its intended scope.

Cite

@misc{srivastava2026worldmodelsdft,
  title  = {World Models as Neural Surrogates for DFT-Scale
            Materials Property Simulation},
  author = {Srivastava, Arjun and Karakoti, Gaurav},
  year   = {2026},
  note   = {Preprint. Submitted for review.},
  url    = {https://deuxvoir.example/research/world-models-neural-surrogates-dft}
}