27.3 Materials and Genomics Models

"I screened 2.2 million crystal structures before breakfast. The hard part was deciding which 380,000 to call 'stable.' The experimentalists are still catching up."

A Graph Neural Network With a Materials Project Account

1. GNoME: Scaling Materials Discovery

In 2023, a single model screened more candidate crystals in weeks than human crystallographers had validated in over a century, and 380,000 of those candidates turned out to be stable new materials.

That model is GNoME (Graph Networks for Materials Exploration; Merchant et al., 2023), a GNN developed by Google DeepMind that predicted 2.2 million crystal structures to be stable, of which 380,000 were subsequently confirmed, expanding known stable inorganic materials by roughly an order of magnitude. The model uses an active learning loop: GNoME proposes candidate crystal structures, density functional theory (DFT, a quantum-mechanical method that approximates the electronic structure of atoms to compute energies and forces from first principles) calculations evaluate their stability, and the results are fed back to improve the model. Figure 27.4 illustrates this closed-loop pipeline.

Formation energy is the energy released or absorbed when a compound forms from its constituent elements in their standard states. A negative value means the compound is more stable than its separated elements, making it a candidate for real-world synthesis. The GNN ingests the crystal graph and runs several rounds of message passing (an iterative process in which each node in the graph collects and aggregates information from its neighbors, progressively building representations that encode the wider structural context), building atom-level representations that encode local chemical environments. It then pools those representations into a single scalar prediction of formation energy per atom. Use formation energy screening to triage thousands of hypothetical compositions quickly. Switch to full DFT or experimental validation only for top candidates that pass the screening threshold.

Architecture and Usage

What. A GNN that takes a crystal structure (atoms and their 3D positions in a periodic unit cell) as input and predicts the formation energy per atom.

Why. DFT calculations are the gold standard for predicting material stability, but each calculation takes minutes to hours on a compute cluster. GNoME predicts formation energy in milliseconds, enabling screening of millions of candidate structures that would be computationally infeasible with DFT alone.

How. The model represents each crystal as a graph where nodes are atoms and edges connect atoms within a cutoff distance. Node features encode element identity; edge features encode distance and direction. Multiple rounds of message passing aggregate information from neighboring atoms, and a final readout layer predicts the total formation energy. Figure 27.3.1 illustrates crystal-to-graph conversion and GNN message passing for formation energy prediction.

When. Use GNoME when you need rapid stability screening of hypothetical crystal structures. It is particularly valuable for compositional exploration: testing whether a new combination of elements could form a stable material before committing to expensive DFT calculations or synthesis experiments.

The formation energy \(E_f\) determines thermodynamic stability:

where \(E_{\text{total}}\) is the total energy of the compound, \(n_i\) is the number of atoms of element \(i\), and \(\mu_i\) is the chemical potential of element \(i\) in its reference state. A material is thermodynamically stable if it lies on the convex hull (the lower boundary of the lowest-energy phases plotted against composition) of formation energies for its composition space.

from pymatgen.core import Structure, Lattice
import numpy as np

# Build a crystal structure programmatically: NaCl (rock salt)
lattice = Lattice.cubic(5.64)  # cubic lattice, a = 5.64 Angstroms
nacl = Structure(
    lattice,
    species=["Na", "Cl", "Na", "Cl", "Na", "Cl", "Na", "Cl"],
    coords=[
        [0.0, 0.0, 0.0], [0.5, 0.5, 0.5],  # Na at origin, Cl at body center
        [0.5, 0.5, 0.0], [0.0, 0.0, 0.5],  # alternating pattern
        [0.5, 0.0, 0.5], [0.0, 0.5, 0.0],
        [0.0, 0.5, 0.5], [0.5, 0.0, 0.0],
    ]
)

print(f"Formula: {nacl.composition.reduced_formula}")
print(f"Space group: {nacl.get_space_group_info()[0]}")
print(f"Number of atoms: {nacl.num_sites}")
print(f"Lattice parameter: {nacl.lattice.a:.2f} Angstroms")
print(f"Volume: {nacl.volume:.2f} Angstrom^3")

# Graph representation for GNN input
from pymatgen.analysis.graphs import StructureGraph
from pymatgen.analysis.local_env import CrystalNN

crystal_nn = CrystalNN()
struct_graph = StructureGraph.from_local_env_strategy(nacl, crystal_nn)

print(f"\nCrystal graph:")
print(f"  Nodes (atoms): {len(struct_graph.graph.nodes)}")
print(f"  Edges (bonds): {len(struct_graph.graph.edges)}")
print(f"  Avg coordination: "
      f"{np.mean([struct_graph.graph.degree(n) for n in struct_graph.graph.nodes]):.1f}")

2. MACE-MP-0: A Universal Interatomic Potential

MACE-MP-0 (Batatia et al., 2024) is a universal interatomic potential. It predicts the energy and forces of any atomic configuration across the entire periodic table. Unlike GNoME, which predicts formation energies of complete crystal structures, MACE-MP-0 predicts the potential energy surface (the function that maps every possible arrangement of atoms to its total energy, so that the slope at any point gives the force on each atom). Given any arrangement of atoms, it computes the energy and the force on each atom. This enables molecular dynamics simulations, geometry optimization, and phonon calculations (where phonons are quantized lattice vibrations that govern thermal and mechanical properties of solids) at near-DFT accuracy with orders-of-magnitude speedup. (As of 2024, the team released an updated checkpoint called MACE-MP-0b with improved accuracy across the periodic table, particularly for transition metals and rare earths; the architecture and API remain the same.)

A force prediction that points the wrong direction can send a molecular dynamics simulation careening into a physically impossible state within nanoseconds. Getting the geometry of forces right is not merely a nice property; it is the difference between a simulation that matches experiment and one that explodes.

MACE uses an equivariant message-passing architecture. "Equivariant" means that if you rotate the atomic configuration, the predicted forces rotate accordingly. This is a hard physics constraint: forces are vectors, and their direction must transform correctly under rotation. Standard GNNs do not enforce this constraint and must learn it from data, which requires more parameters and more training examples.

The MACE architecture builds higher-order equivariant features through tensor products of atomic embeddings with spherical harmonics (mathematical functions defined on the surface of a sphere that encode directional information about interatomic vectors):

where \(h_i^{(t)}\) is the feature vector of atom \(i\) at layer \(t\), \(R_{ij}^{(t)}\) encodes the relative position of atoms \(i\) and \(j\) as spherical harmonics, and \(\otimes\) denotes the tensor product that couples geometric and chemical information. The learnable weights \(W^{(t)}\) determine how this coupled information is mixed.

# MACE-MP-0 usage through the mace-torch package
# pip install mace-torch

from mace.calculators import mace_mp
from ase import Atoms
from ase.build import bulk
import numpy as np

# Create a silicon crystal
si = bulk("Si", "diamond", a=5.43)
print(f"Silicon unit cell: {si.get_chemical_formula()}")
print(f"Positions:\n{si.get_positions()}")

# Load the pretrained MACE-MP-0 calculator
calc = mace_mp(model="medium", default_dtype="float64")
si.calc = calc

# Compute energy and forces
energy = si.get_potential_energy()
forces = si.get_forces()
stress = si.get_stress()

print(f"\nEnergy: {energy:.4f} eV")
print(f"Energy per atom: {energy / len(si):.4f} eV/atom")
print(f"Max force: {np.abs(forces).max():.6f} eV/Angstrom")
print(f"Stress tensor (Voigt): {stress.round(4)}")

# Compare: perturb one atom and recompute
si_perturbed = si.copy()
si_perturbed.calc = calc
si_perturbed.positions[0] += [0.1, 0.0, 0.0]  # displace by 0.1 Angstrom

energy_pert = si_perturbed.get_potential_energy()
forces_pert = si_perturbed.get_forces()

print(f"\nAfter 0.1 A displacement:")
print(f"Energy change: {energy_pert - energy:.4f} eV")
print(f"Restoring force on atom 0: {forces_pert[0]} eV/Angstrom")

3. CHGNet: Charge-Informed Potentials

CHGNet (Deng et al., 2023) extends the universal potential concept by explicitly tracking atomic charges and magnetic moments. Standard interatomic potentials (including MACE-MP-0) predict energy as a function of atomic positions only, ignoring the electronic degrees of freedom. CHGNet adds charge and spin as additional variables, enabling accurate simulation of electrochemical reactions, magnetic phase transitions, and redox processes (reactions involving electron transfer between atoms, which change their oxidation states).

CHGNet was pretrained on 1.6 million structures from the Materials Project, with energies, forces, stresses, and magnetic moments computed by DFT. On its benchmark test sets, the model achieves energy prediction errors below 30 meV/atom and force errors below 60 meV/Angstrom across the periodic table (circa 2023). (As of 2024, CHGNet v0.3 further reduced these errors and added broader coverage of magnetic and electrochemical systems.)

from chgnet.model.model import CHGNet
from pymatgen.core import Structure, Lattice

# Load pretrained CHGNet
chgnet = CHGNet.load()

# Predict properties for LiFePO4: a battery cathode material
# where charge tracking matters
lifepo4 = Structure.from_spacegroup(
    "Pnma",  # olivine space group
    Lattice.from_parameters(10.33, 6.01, 4.69, 90, 90, 90),
    species=["Li", "Fe", "P", "O", "O", "O", "O"],
    coords=[
        [0.0, 0.0, 0.0],
        [0.282, 0.25, 0.975],
        [0.095, 0.25, 0.418],
        [0.097, 0.25, 0.743],
        [0.457, 0.25, 0.206],
        [0.166, 0.046, 0.285],
        [0.166, 0.454, 0.285],
    ]
)

prediction = chgnet.predict_structure(lifepo4)
print(f"Material: LiFePO4 (olivine)")
print(f"Energy: {prediction['e']:.4f} eV/atom")
print(f"Max force: {abs(prediction['f']).max():.4f} eV/Angstrom")
print(f"Magnetic moments: {prediction['m'].round(2)}")

Where GNoME, MACE-MP-0, and CHGNet learn the physics of atomic arrangements in crystals, a parallel revolution applies the same self-supervised recipe to a very different substrate: the four-letter sequences that encode biological information in DNA.

4. Nucleotide Transformer: Genomic Language at Scale

The Nucleotide Transformer (Dalla-Torre et al., 2023) brings the protein language modeling paradigm to DNA. The model is a transformer encoder pretrained on nucleotide sequences from 3,200 diverse genomes using masked nucleotide prediction. The model family spans four sizes (50M to 2.5B parameters) and comes in two variants: one pretrained on human genome reference sequences, and one pretrained on multi-species sequences.

DNA has a four-letter alphabet (A, C, G, T), much smaller than the 20-letter amino acid alphabet. To compensate, the Nucleotide Transformer uses a tokenization scheme based on 6-mers (contiguous sequences of 6 nucleotides), producing a vocabulary of \(4^6 = 4{,}096\) tokens. This captures local sequence context within each token, similar to how byte-pair encoding captures subword structure in NLP.

from transformers import AutoTokenizer, AutoModel
import torch

# Load Nucleotide Transformer (500M variant)
model_name = "InstaDeepAI/nucleotide-transformer-v2-500m-multi-species"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
model.eval()

# A DNA sequence containing a known promoter region
# (simplified TATA box + downstream sequence)
dna_sequence = (
    "GCGCGCATAAAAGGCGCGCGCATGCTAGCTAGCTAGCTAGCGATCGATCG"
    "ATCGATCGATCGATCGTAGCTAGCTAGCAATTGCAATTGCTAGCTAGCTAG"
)

inputs = tokenizer(dna_sequence, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

embeddings = outputs.last_hidden_state
print(f"Input sequence length: {len(dna_sequence)} nucleotides")
print(f"Token count: {inputs['input_ids'].shape[1]}")
print(f"Embedding shape: {embeddings.shape}")
print(f"Embedding dimension: {embeddings.shape[2]}")

# Mean-pool for a sequence-level embedding
seq_embedding = embeddings[0, 1:-1].mean(dim=0)  # exclude special tokens
print(f"Sequence embedding: {seq_embedding.shape}")
print(f"Embedding norm: {seq_embedding.norm():.4f}")

In practice, these per-token and sequence-level embeddings serve as input features for downstream classifiers. By adding a lightweight prediction head (typically a linear layer or a small MLP) on top of the frozen or fine-tuned embeddings, researchers have achieved state-of-the-art results on tasks such as promoter detection, splice site prediction, enhancer identification, and variant effect scoring, all without task-specific architectures.

5. DNABERT-2: Efficient Genomic Understanding

DNABERT-2 (Zhou et al., 2024) addresses a limitation of k-mer tokenization. Fixed k-mers create an arbitrary segmentation boundary: the sequence ACGTAG tokenized as 6-mers starting at position 0 produces different tokens than the same sequence starting at position 1. DNABERT-2 replaces k-mers with byte pair encoding (BPE), the same tokenization strategy used by GPT models. BPE learns a data-driven vocabulary that captures variable-length motifs, from single nucleotides to longer functional elements.

DNABERT-2 achieves strong performance on 28 genomic prediction benchmarks with only 117M parameters, roughly 20x smaller than the largest Nucleotide Transformer. This efficiency comes from three design choices: BPE tokenization (better coverage per token), ALiBi positional encoding (Attention with Linear Biases, which adds a learned distance penalty to attention scores instead of fixed position embeddings, letting the model handle variable sequence lengths without retraining), and flash attention (a memory-efficient attention algorithm that tiles the computation to minimize GPU memory reads, reducing both memory footprint and wall-clock time).

from transformers import AutoTokenizer, AutoModel
import torch

# Load DNABERT-2
tokenizer = AutoTokenizer.from_pretrained(
    "zhihan1996/DNABERT-2-117M", trust_remote_code=True
)
model = AutoModel.from_pretrained(
    "zhihan1996/DNABERT-2-117M", trust_remote_code=True
)
model.eval()

# Compare tokenization: fixed k-mer vs BPE
dna = "ATCGATCGATCGATCG"

# DNABERT-2 uses BPE
bpe_tokens = tokenizer.tokenize(dna)
print(f"Sequence: {dna}")
print(f"BPE tokens: {bpe_tokens}")
print(f"Number of BPE tokens: {len(bpe_tokens)}")

# Encode and extract embeddings
inputs = tokenizer(dna, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

print(f"\nEmbedding shape: {outputs.last_hidden_state.shape}")
print(f"Embedding dim: {outputs.last_hidden_state.shape[-1]}")

# Demonstrate: encode two promoter sequences and compare
promoter_active = "GCGCGCTATAAAAGGCGCGCATGCTAGCTAGCTA"    # contains TATA box
promoter_silent = "GCGCGCNNNNNNNNGGCGCGCATGCTAGCTAGCTA"   # disrupted TATA box

for name, seq in [("Active", promoter_active), ("Silent", promoter_silent)]:
    inp = tokenizer(seq, return_tensors="pt")
    with torch.no_grad():
        out = model(**inp)
    emb = out.last_hidden_state[0, 1:-1].mean(dim=0)
    print(f"{name} promoter embedding norm: {emb.norm():.4f}")

DNABERT-2's success with a compact architecture raises a natural question: across all five models in this section, what design principles separate the approaches that scale efficiently from those that rely on brute-force parameter counts?

6. Cross-Domain Patterns

Looking across these five models and the five from Section 27.2, several patterns emerge that generalize beyond any single domain:

Architecture follows symmetry. Sequences (proteins, DNA) have translational symmetry: a motif's function does not depend on its position. Crystals have rotational and translational symmetry in 3D. Transformers with positional encoding match the first; equivariant message-passing networks match the second. Violating a domain's symmetry forces the model to learn invariance from data, wasting capacity and training compute.

Pretraining data defines the ceiling. No model outperforms the information content of its pretraining corpus. ESM-2 captures evolutionary patterns present in UniRef but cannot predict properties absent from protein sequences (e.g., post-translational modifications). GNoME captures thermodynamic stability as computed by DFT but inherits DFT's known errors (e.g., underestimation of band gaps). Understanding the pretraining data's coverage and biases is essential for knowing when a foundation model's predictions are trustworthy.

Small models can win with better design. DNABERT-2 (117M) outperforms models 20x larger through better tokenization and attention mechanisms. CHGNet, with relatively few parameters, can achieve accuracy comparable to larger universal potentials for battery materials by incorporating charge information as an additional inductive bias. Scale matters, but architecture and training objective matter more per parameter. In short: match your architecture to your domain's symmetries, your tokenizer to its natural units, and your pretraining data to the questions you actually need answered.