Chapter 2: Transformer Architectures for Materials

This chapter surveys the differences and scope of representative models such as Matformer. It also covers connections to crystals and graphs.

💡 Note: Performance varies greatly depending on the input representation (structure / composition / text) and the design of the pretraining task.

Study Time: 30-35 min | Difficulty: Intermediate to Advanced

📋 What You Will Learn in This Chapter

• Design principles of Transformer architectures specialized for materials science
• Matformer: Materials Transformer for Property Prediction
• CrystalFormer: Crystal Structure Representation
• ChemBERTa: Molecular SMILES Representation Learning
• Perceiver IO: Integration of Diverse Data
• Implementation exercise: Material property prediction with Matformer

2.1 The Need for Materials-Specialized Transformers

Limitations of General-Purpose Transformers

Problems with using natural language processing Transformers as-is: - ❌ 3D structural information of molecules and materials is lost - ❌ Chemical bonds and interatomic distances cannot be considered - ❌ Periodic boundary conditions (crystals) cannot be handled - ❌ Physical constraints (conservation laws, symmetry) are ignored

Characteristics of Materials-Specialized Transformers

Required extensions: - ✅ Embedding of 3D structure: atomic coordinates, distances, angles - ✅ Periodic boundary conditions: repetition of the crystal lattice - ✅ Physical constraints: symmetry, equivariance - ✅ Integration of diverse data: structure + composition + experimental data

2.2 Matformer: Materials Transformer

Overview

Matformer (Chen et al., 2022) is a Transformer model that predicts properties from the crystal structure of materials.

Characteristics: - Nested Transformer: hierarchical processing at the atom level and the crystal level - Distance-aware Attention: considers interatomic distances - Elastic Inference: dynamically adjusts computational cost and memory

Architecture

Atom Embedding

import torch
import torch.nn as nn
import numpy as np

class AtomEmbedding(nn.Module):
    def __init__(self, num_atoms=118, d_model=256):
        """
        Atom embedding layer

        Args:
            num_atoms: Number of atom types (periodic table, 118 elements)
            d_model: Embedding dimension
        """
        super(AtomEmbedding, self).__init__()
        self.embedding = nn.Embedding(num_atoms, d_model)

    def forward(self, atomic_numbers):
        """
        Args:
            atomic_numbers: (batch_size, num_atoms) atomic numbers
        Returns:
            embeddings: (batch_size, num_atoms, d_model)
        """
        return self.embedding(atomic_numbers)

# Usage example: NaCl crystal
batch_size = 2
num_atoms = 8  # Number of atoms in the unit cell

# Atomic numbers: Na(11), Cl(17)
atomic_numbers = torch.tensor([
    [11, 17, 11, 17, 11, 17, 11, 17],  # Sample 1
    [11, 17, 11, 17, 11, 17, 11, 17]   # Sample 2
])

atom_emb = AtomEmbedding(num_atoms=118, d_model=256)
embeddings = atom_emb(atomic_numbers)
print(f"Atom embeddings shape: {embeddings.shape}")  # (2, 8, 256)

Distance-aware Attention

Attention that considers interatomic distances:

class DistanceAwareAttention(nn.Module):
    def __init__(self, d_model, num_heads, max_distance=10.0):
        """
        Distance-aware Attention

        Args:
            d_model: Model dimension
            num_heads: Number of attention heads
            max_distance: Maximum distance (Å)
        """
        super(DistanceAwareAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.max_distance = max_distance

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

        # Distance embedding
        self.distance_embedding = nn.Linear(1, num_heads)

    def forward(self, x, distance_matrix):
        """
        Args:
            x: (batch_size, num_atoms, d_model)
            distance_matrix: (batch_size, num_atoms, num_atoms) interatomic distances (Å)
        """
        batch_size = x.size(0)

        # Q, K, V
        Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # Attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(self.d_k)

        # Distance bias
        # Larger value for closer distances, smaller value for farther distances
        distance_bias = self.distance_embedding(distance_matrix.unsqueeze(-1))  # (batch, num_atoms, num_atoms, num_heads)
        distance_bias = distance_bias.permute(0, 3, 1, 2)  # (batch, num_heads, num_atoms, num_atoms)

        # Transform distance with a Gaussian function (higher score for closer atoms)
        distance_factor = torch.exp(-distance_matrix.unsqueeze(1) / 2.0)  # (batch, 1, num_atoms, num_atoms)

        scores = scores + distance_bias * distance_factor

        # Softmax
        attention_weights = torch.softmax(scores, dim=-1)

        # Apply attention
        output = torch.matmul(attention_weights, V)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.W_o(output)

        return output, attention_weights

# Usage example
d_model = 256
num_heads = 8
num_atoms = 8

dist_attn = DistanceAwareAttention(d_model, num_heads)

x = torch.randn(2, num_atoms, d_model)
# Interatomic distances of the NaCl crystal (simplified)
distance_matrix = torch.tensor([
    [[0.0, 2.8, 3.9, 4.8, 3.9, 5.5, 4.8, 6.7],  # Distances from atom 1
     [2.8, 0.0, 2.8, 3.9, 5.5, 3.9, 6.7, 4.8],
     # ... omitted
     [6.7, 4.8, 5.5, 3.9, 4.8, 3.9, 2.8, 0.0]]
]).repeat(2, 1, 1)  # Replicate for batch_size

output, attn_weights = dist_attn(x, distance_matrix)
print(f"Output shape: {output.shape}")  # (2, 8, 256)

Matformer Block

class MatformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff=1024, dropout=0.1):
        """
        Basic block of Matformer

        Args:
            d_model: Model dimension
            num_heads: Number of attention heads
            d_ff: Intermediate dimension of the Feed-Forward layer
            dropout: Dropout rate
        """
        super(MatformerBlock, self).__init__()

        self.distance_attention = DistanceAwareAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)

        # Feed-Forward Network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, distance_matrix):
        # Distance-aware Attention + Residual
        attn_output, _ = self.distance_attention(x, distance_matrix)
        x = self.norm1(x + self.dropout1(attn_output))

        # Feed-Forward + Residual
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_output))

        return x

2.3 CrystalFormer: Crystal Structure Transformer

Overview

CrystalFormer is a Transformer that considers the periodic boundary conditions of crystals.

Characteristics: - Wyckoff position embedding: considers crystal symmetry - Fractional Coordinates: representation in fractional coordinates - Space Group Encoding: embedding of space group information

Fractional Coordinate Embedding

class FractionalCoordinateEncoding(nn.Module):
    def __init__(self, d_model):
        super(FractionalCoordinateEncoding, self).__init__()
        self.coord_linear = nn.Linear(3, d_model)

    def forward(self, fractional_coords):
        """
        Args:
            fractional_coords: (batch_size, num_atoms, 3) fractional coordinates [0, 1)
        Returns:
            encoding: (batch_size, num_atoms, d_model)
        """
        # Trigonometric embedding
        freqs = torch.arange(1, d_model // 6 + 1, dtype=torch.float32)
        coords_expanded = fractional_coords.unsqueeze(-1) * freqs

        encoding = torch.cat([
            torch.sin(2 * np.pi * coords_expanded),
            torch.cos(2 * np.pi * coords_expanded)
        ], dim=-1)

        # Adjust dimension with a linear transformation
        encoding = encoding.flatten(start_dim=2)
        encoding = self.coord_linear(encoding)

        return encoding

Handling Periodic Boundary Conditions

def compute_periodic_distance(coords1, coords2, lattice_matrix):
    """
    Distance calculation considering periodic boundary conditions

    Args:
        coords1: (num_atoms1, 3) fractional coordinates
        coords2: (num_atoms2, 3) fractional coordinates
        lattice_matrix: (3, 3) lattice vector matrix
    Returns:
        distances: (num_atoms1, num_atoms2) minimum distances (Å)
    """
    # Convert to Cartesian coordinates
    cart1 = torch.matmul(coords1, lattice_matrix)
    cart2 = torch.matmul(coords2, lattice_matrix)

    # Consider all periodic images (range of -1, 0, 1)
    offsets = torch.tensor([
        [i, j, k] for i in [-1, 0, 1]
                  for j in [-1, 0, 1]
                  for k in [-1, 0, 1]
    ], dtype=torch.float32)  # 27 combinations

    min_distances = []
    for offset in offsets:
        offset_cart = torch.matmul(offset, lattice_matrix)
        shifted_cart2 = cart2 + offset_cart

        # Distance calculation
        diff = cart1.unsqueeze(1) - shifted_cart2.unsqueeze(0)
        distances = torch.norm(diff, dim=-1)
        min_distances.append(distances)

    # Select the minimum distance
    min_distances = torch.stack(min_distances, dim=-1)
    min_distances, _ = torch.min(min_distances, dim=-1)

    return min_distances

# Usage example: simple cubic lattice
fractional_coords = torch.tensor([
    [0.0, 0.0, 0.0],  # Atom 1
    [0.5, 0.5, 0.5]   # Atom 2
])

lattice_matrix = torch.tensor([
    [5.0, 0.0, 0.0],
    [0.0, 5.0, 0.0],
    [0.0, 0.0, 5.0]
])  # 5 Å cubic lattice

distances = compute_periodic_distance(fractional_coords, fractional_coords, lattice_matrix)
print("Distance matrix (Å):")
print(distances)

2.4 ChemBERTa: Molecular SMILES Representation Learning

Overview

ChemBERTa is a model that learns molecular SMILES strings with BERT.

Characteristics: - Based on RoBERTa (an improved version of BERT) - Pretrained on 10M molecules - Transfer learning achieves high accuracy even with small amounts of data

SMILES Tokenization

from transformers import RobertaTokenizer

class SMILESTokenizer:
    def __init__(self):
        # Tokenizer for ChemBERTa
        self.tokenizer = RobertaTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")

    def encode(self, smiles_list):
        """
        Tokenize SMILES strings

        Args:
            smiles_list: List of SMILES
        Returns:
            input_ids: Token IDs
            attention_mask: Mask
        """
        encoded = self.tokenizer(
            smiles_list,
            padding=True,
            truncation=True,
            max_length=128,
            return_tensors='pt'
        )
        return encoded['input_ids'], encoded['attention_mask']

# Usage example
smiles_list = [
    'CCO',  # Ethanol
    'CC(C)Cc1ccc(cc1)C(C)C(=O)O',  # Ibuprofen
    'CN1C=NC2=C1C(=O)N(C(=O)N2C)C'  # Caffeine
]

tokenizer = SMILESTokenizer()
input_ids, attention_mask = tokenizer.encode(smiles_list)

print(f"Input IDs shape: {input_ids.shape}")
print(f"Attention mask shape: {attention_mask.shape}")
print(f"First molecule tokens: {input_ids[0][:10]}")

Using the ChemBERTa Model

from transformers import RobertaModel

class ChemBERTaEmbedding(nn.Module):
    def __init__(self, pretrained_model="seyonec/ChemBERTa-zinc-base-v1"):
        super(ChemBERTaEmbedding, self).__init__()
        self.bert = RobertaModel.from_pretrained(pretrained_model)

    def forward(self, input_ids, attention_mask):
        """
        Args:
            input_ids: (batch_size, seq_len)
            attention_mask: (batch_size, seq_len)
        Returns:
            embeddings: (batch_size, hidden_size)
        """
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)

        # Use the embedding of the [CLS] token
        cls_embedding = outputs.last_hidden_state[:, 0, :]

        return cls_embedding

# Molecular property prediction model
class MoleculePropertyPredictor(nn.Module):
    def __init__(self, hidden_size=768, num_properties=1):
        super(MoleculePropertyPredictor, self).__init__()
        self.chemberta = ChemBERTaEmbedding()
        self.predictor = nn.Sequential(
            nn.Linear(hidden_size, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, num_properties)
        )

    def forward(self, input_ids, attention_mask):
        embeddings = self.chemberta(input_ids, attention_mask)
        predictions = self.predictor(embeddings)
        return predictions

# Usage example
model = MoleculePropertyPredictor(num_properties=1)  # Example: logP prediction
predictions = model(input_ids, attention_mask)
print(f"Predictions shape: {predictions.shape}")  # (3, 1)

2.5 Perceiver IO: Integration of Diverse Data

Overview

Perceiver IO is a Transformer that can integrate and process different types of data.

Applications in materials science: - Structural data + compositional data - Experimental data + computational data - Images + text + numerical values

Architecture

Simple Implementation

class PerceiverBlock(nn.Module):
    def __init__(self, latent_dim, input_dim, num_latents=64):
        super(PerceiverBlock, self).__init__()
        self.num_latents = num_latents
        self.latent_dim = latent_dim

        # Latent array (learnable)
        self.latents = nn.Parameter(torch.randn(num_latents, latent_dim))

        # Cross-Attention: Latent → Input
        self.cross_attn = nn.MultiheadAttention(latent_dim, num_heads=8, batch_first=True)

        # Self-Attention: Latent → Latent
        self.self_attn = nn.MultiheadAttention(latent_dim, num_heads=8, batch_first=True)

        # Embed the input
        self.input_projection = nn.Linear(input_dim, latent_dim)

    def forward(self, x):
        """
        Args:
            x: (batch_size, seq_len, input_dim) input data
        Returns:
            latents: (batch_size, num_latents, latent_dim)
        """
        batch_size = x.size(0)

        # Embed the input
        x_embed = self.input_projection(x)

        # Replicate the latents
        latents = self.latents.unsqueeze(0).repeat(batch_size, 1, 1)

        # Cross-Attention: Latent (Query) ← Input (Key, Value)
        latents, _ = self.cross_attn(latents, x_embed, x_embed)

        # Self-Attention: within the Latent
        latents, _ = self.self_attn(latents, latents, latents)

        return latents

# Usage example: integrating structural data and compositional data
batch_size = 2
seq_len = 20
input_dim = 128
latent_dim = 256

perceiver = PerceiverBlock(latent_dim, input_dim, num_latents=32)

# Structural data (e.g., atomic coordinates)
structure_data = torch.randn(batch_size, seq_len, input_dim)

latents = perceiver(structure_data)
print(f"Latent representation shape: {latents.shape}")  # (2, 32, 256)

2.6 Implementation Exercise: Material Property Prediction with Matformer

Complete Implementation Example

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Dataset
class MaterialsDataset(Dataset):
    def __init__(self, num_samples=100):
        self.num_samples = num_samples

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # Dummy data (in practice, obtained from Materials Project, etc.)
        num_atoms = 8
        atomic_numbers = torch.randint(1, 30, (num_atoms,))  # Atomic numbers
        positions = torch.randn(num_atoms, 3)  # Atomic coordinates (Å)
        distance_matrix = torch.cdist(positions, positions)  # Distance matrix

        # Target: band gap (eV)
        target = torch.randn(1)

        return atomic_numbers, distance_matrix, target

# Matformer model (simplified version)
class SimpleMatformer(nn.Module):
    def __init__(self, d_model=256, num_heads=8, num_layers=4):
        super(SimpleMatformer, self).__init__()

        self.atom_embedding = AtomEmbedding(num_atoms=118, d_model=d_model)

        self.layers = nn.ModuleList([
            MatformerBlock(d_model, num_heads)
            for _ in range(num_layers)
        ])

        self.pooling = nn.AdaptiveAvgPool1d(1)
        self.predictor = nn.Sequential(
            nn.Linear(d_model, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

    def forward(self, atomic_numbers, distance_matrix):
        # Atom embedding
        x = self.atom_embedding(atomic_numbers)

        # Matformer blocks
        for layer in self.layers:
            x = layer(x, distance_matrix)

        # Global pooling
        x = x.transpose(1, 2)  # (batch, d_model, num_atoms)
        x = self.pooling(x).squeeze(-1)  # (batch, d_model)

        # Prediction
        output = self.predictor(x)
        return output

# Training
def train_matformer():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Data
    dataset = MaterialsDataset(num_samples=100)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

    # Model
    model = SimpleMatformer(d_model=256, num_heads=8, num_layers=4).to(device)

    # Optimization
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    # Training loop
    model.train()
    for epoch in range(5):
        total_loss = 0
        for atomic_numbers, distance_matrix, target in dataloader:
            atomic_numbers = atomic_numbers.to(device)
            distance_matrix = distance_matrix.to(device)
            target = target.to(device)

            # Forward
            predictions = model(atomic_numbers, distance_matrix)
            loss = criterion(predictions, target)

            # Backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")

    return model

# Execution
trained_model = train_matformer()

2.7 Summary

Key Points

1. Matformer: distance-aware attention, hierarchical structure
2. CrystalFormer: periodic boundary conditions, fractional coordinates, space groups
3. ChemBERTa: SMILES representation learning, transfer learning
4. Perceiver IO: integration of diverse data

Preparation for the Next Chapter

In Chapter 3, we will learn about pretrained models (MatBERT, MolBERT) and fine-tuning.

📝 Exercises

Problem 1: Conceptual Understanding

List three reasons why Distance-aware Attention is superior to ordinary Attention in materials science.

Problem 2: Implementation

Implement a simple function that computes distances without considering periodic boundary conditions.

def compute_simple_distance(coords1, coords2):
    """
    Simple distance calculation (without periodic boundary conditions)

    Args:
        coords1: (num_atoms1, 3)
        coords2: (num_atoms2, 3)
    Returns:
        distances: (num_atoms1, num_atoms2)
    """
    # Implement here
    pass

def compute_simple_distance(coords1, coords2):
    diff = coords1.unsqueeze(1) - coords2.unsqueeze(0)
    distances = torch.norm(diff, dim=-1)
    return distances

Problem 3: Application

Using ChemBERTa, design a model that predicts the aqueous solubility of molecules. Explain the required layers and configuration.

class SolubilityPredictor(nn.Module):
    def __init__(self):
        super(SolubilityPredictor, self).__init__()
        self.chemberta = ChemBERTaEmbedding()  # 768 dimensions

        self.predictor = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 1)  # Solubility (continuous value)
        )

    def forward(self, input_ids, attention_mask):
        embeddings = self.chemberta(input_ids, attention_mask)
        solubility = self.predictor(embeddings)
        return solubility

📊 Data Licenses and Terms of Use

Materials Datasets

• Materials Project: CC BY 4.0
• Computational data for over 690,000 materials
• Citation: Jain, A. et al. APL Materials 1, 011002 (2013)
• OQMD (Open Quantum Materials Database): available for academic use
• Over 1 million DFT calculation data
• Citation: Saal, J. E. et al. JOM 65, 1501-1509 (2013)
• AFLOW: CC BY 4.0
• Over 3.5 million material data

Molecular and Chemical Datasets

• ZINC15: free for academic use, commercial use requires inquiry
• Over 1 billion commercially available compounds
• ChEMBL: CC BY-SA 3.0
• Over 2 million bioactive compounds
• QM9: CC0 1.0
• Quantum chemistry calculation data for 130,000 molecules

Crystal Structure Datasets

• COD (Crystallography Open Database): public domain
• Over 500,000 crystal structures
• ICSD (Inorganic Crystal Structure Database): license purchase required
• Over 250,000 inorganic crystal structures

🔧 Code Reproducibility Guidelines

Environment Setup (for Materials Science)

# requirements.txt
torch==2.0.1
transformers==4.30.2
pymatgen==2023.8.10  # Materials science library
rdkit==2023.3.2  # Chemical computation
ase==3.22.1  # Atomic Simulation Environment
numpy==1.24.3
scipy==1.11.1
matplotlib==3.7.1

Reproducibility of Distance Matrix Calculation

import numpy as np
from pymatgen.core import Structure
from pymatgen.analysis.local_env import CrystalNN

def reproducible_distance_matrix(structure, cutoff=10.0, seed=42):
    """
    Reproducible distance matrix calculation

    Args:
        structure: pymatgen Structure object
        cutoff: Distance cutoff (Å)
        seed: Random seed
    """
    np.random.seed(seed)

    # Retrieve atomic coordinates
    positions = structure.cart_coords

    # Distance matrix calculation (considering periodic boundary conditions)
    distance_matrix = structure.distance_matrix

    # Apply cutoff
    distance_matrix[distance_matrix > cutoff] = cutoff

    return distance_matrix

# Usage example
from pymatgen.core import Lattice, Structure

# NaCl structure
lattice = Lattice.cubic(5.64)
species = ["Na", "Cl"]
coords = [[0, 0, 0], [0.5, 0.5, 0.5]]
structure = Structure(lattice, species, coords)

dist_matrix = reproducible_distance_matrix(structure)
print(f"Distance matrix shape: {dist_matrix.shape}")

Matformer Configuration Parameters

# Matformer configuration
matformer_config = {
    'model': {
        'd_model': 256,
        'num_heads': 8,
        'num_layers': 4,
        'd_ff': 1024,
        'dropout': 0.1,
        'max_atoms': 50,  # Maximum number of atoms
        'max_distance': 10.0  # Distance cutoff (Å)
    },
    'distance_embedding': {
        'num_gaussians': 50,
        'gaussian_start': 0.0,
        'gaussian_end': 10.0
    },
    'training': {
        'batch_size': 32,
        'learning_rate': 1e-4,
        'num_epochs': 100,
        'scheduler': 'cosine_with_warmup',
        'warmup_epochs': 10
    },
    'seed': 42
}

⚠️ Practical Pitfalls and Solutions

1. Miscalculation of Periodic Boundary Conditions

Problem: Interatomic distances across lattice boundaries are not calculated correctly

# ❌ Wrong: simple Euclidean distance
def wrong_periodic_distance(pos1, pos2, lattice):
    return np.linalg.norm(pos1 - pos2)

# ✅ Correct: minimum image convention
def correct_periodic_distance(frac1, frac2, lattice_matrix):
    """
    Distance calculation considering periodic boundary conditions

    Args:
        frac1, frac2: fractional coordinates (3,)
        lattice_matrix: lattice vectors (3, 3)
    """
    # Minimum image convention: -0.5 <= delta < 0.5
    delta_frac = frac1 - frac2
    delta_frac = delta_frac - np.floor(delta_frac + 0.5)

    # Convert to Cartesian coordinates
    delta_cart = np.dot(delta_frac, lattice_matrix)

    return np.linalg.norm(delta_cart)

# Debugging method
print("Fractional coordinates:", frac1, frac2)
print("Delta (before wrapping):", frac1 - frac2)
print("Delta (after wrapping):", delta_frac)
print("Minimum distance:", correct_periodic_distance(frac1, frac2, lattice_matrix))

2. Handling Missing Values in Atomic Number Embedding

Problem: Handling unknown atom types or empty sites

# ❌ Problem: error on unknown atoms
class NaiveAtomEmbedding(nn.Module):
    def __init__(self):
        self.embedding = nn.Embedding(118, 256)  # 118 element types

    def forward(self, atomic_numbers):
        return self.embedding(atomic_numbers)  # Error for values >= 118!

# ✅ Solution: add an unknown token
class RobustAtomEmbedding(nn.Module):
    def __init__(self, num_atoms=118, d_model=256):
        super().__init__()
        # +1 is for unknown atoms
        self.embedding = nn.Embedding(num_atoms + 1, d_model, padding_idx=0)
        self.num_atoms = num_atoms

    def forward(self, atomic_numbers):
        # Clip out-of-range atomic numbers
        atomic_numbers = torch.clamp(atomic_numbers, 0, self.num_atoms)
        return self.embedding(atomic_numbers)

# Test
embedding = RobustAtomEmbedding()
test_atoms = torch.tensor([1, 6, 8, 200])  # 200 is out of range
output = embedding(test_atoms)
print(f"Output shape: {output.shape}")  # No error

3. Memory Efficiency of Distance-aware Attention

Problem: The distance matrix is too large, causing OOM

# ❌ Problem: keeping distances for all atom pairs
def memory_intensive_distance_attention(x, all_distances):
    # all_distances: (batch, max_atoms, max_atoms)
    # Memory: batch * max_atoms^2 * 4 bytes
    pass

# ✅ Solution: sparse distance matrix
def sparse_distance_attention(x, positions, cutoff=8.0):
    """
    Compute only within the cutoff distance

    Args:
        x: atomic features (batch, num_atoms, d_model)
        positions: atomic coordinates (batch, num_atoms, 3)
        cutoff: distance cutoff (Å)
    """
    batch_size, num_atoms, _ = positions.shape

    # Distance matrix calculation
    diff = positions.unsqueeze(2) - positions.unsqueeze(1)
    distances = torch.norm(diff, dim=-1)

    # Cutoff mask
    mask = distances < cutoff

    # Attention scores (-inf outside the cutoff)
    scores = compute_attention_scores(x)
    scores = scores.masked_fill(~mask, float('-inf'))

    return scores

4. Handling Special Characters in SMILES Tokenization

Problem: Symbols for aromaticity and stereochemistry are lost

# ❌ Wrong: only alphabetic characters and digits
def wrong_smiles_tokenize(smiles):
    return list(filter(str.isalnum, smiles))

# ✅ Correct: preserve special characters
def correct_smiles_tokenize(smiles):
    """
    Complete SMILES tokenization

    Aromaticity: c, n, o, s (lowercase)
    Stereochemistry: @, @@
    Branching: (), []
    Bonds: -, =, #, :, /,
    """
    import re
    pattern = r'(\[[^\]]+\]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])'
    tokens = re.findall(pattern, smiles)

    return tokens

# Test: complex SMILES
smiles = "C[C@H](N)C(=O)O"  # L-alanine (with stereochemistry)
tokens = correct_smiles_tokenize(smiles)
print(f"Tokens: {tokens}")
# ['C', '[C@H]', '(', 'N', ')', 'C', '(', '=', 'O', ')', 'O']

5. Space Group Encoding in CrystalFormer

Problem: Efficient representation of the 230 space groups

# ❌ Inefficient: one-hot encoding (230 dimensions)
def inefficient_space_group_encoding(space_group_number):
    encoding = torch.zeros(230)
    encoding[space_group_number - 1] = 1
    return encoding

# ✅ Efficient: learnable embedding
class SpaceGroupEmbedding(nn.Module):
    def __init__(self, d_model=64):
        super().__init__()
        # 230 space groups + one for unknown
        self.embedding = nn.Embedding(231, d_model)

    def forward(self, space_group_numbers):
        """
        Args:
            space_group_numbers: (batch,) space group numbers from 1 to 230
        """
        return self.embedding(space_group_numbers)

# Further: exploit the hierarchical structure of space groups
class HierarchicalSpaceGroupEmbedding(nn.Module):
    def __init__(self, d_model=64):
        super().__init__()
        # Crystal systems (7 types)
        self.crystal_system_emb = nn.Embedding(8, d_model // 2)
        # Point groups (32 types)
        self.point_group_emb = nn.Embedding(33, d_model // 2)

    def forward(self, crystal_system, point_group):
        cs_emb = self.crystal_system_emb(crystal_system)
        pg_emb = self.point_group_emb(point_group)
        return torch.cat([cs_emb, pg_emb], dim=-1)

✅ Chapter 2 Completion Checklist

Conceptual Understanding (10 items)

• [ ] Can explain the difference between general-purpose Transformers and materials-specialized Transformers
• [ ] Understand the hierarchical architecture of Matformer
• [ ] Understand the principle of Distance-aware Attention
• [ ] Understand the importance of periodic boundary conditions
• [ ] Understand the conversion between fractional coordinates and Cartesian coordinates
• [ ] Understand the space group encoding of CrystalFormer
• [ ] Understand the relationship between ChemBERTa and RoBERTa
• [ ] Understand the cross-attention mechanism of Perceiver IO
• [ ] Can list three or more reasons why materials-specialized Transformers are needed
• [ ] Can explain the application scenarios of each model (Matformer, CrystalFormer, ChemBERTa)

Mathematical and Physical Understanding (5 items)

• [ ] Can derive the distance calculation formula under periodic boundary conditions
• [ ] Can write the transformation matrix from fractional coordinates to Cartesian coordinates
• [ ] Understand the formula for distance embedding using Gaussian basis functions
• [ ] Understand the relationship between lattice vectors and reciprocal lattice vectors
• [ ] Understand the concept of Wyckoff positions

Implementation Skills (15 items)

• [ ] Can implement the AtomEmbedding class
• [ ] Can implement DistanceAwareAttention
• [ ] Can implement MatformerBlock
• [ ] Can implement distance calculation considering periodic boundary conditions
• [ ] Can implement fractional coordinate encoding
• [ ] Can use ChemBERTaEmbedding
• [ ] Can implement and use a SMILES tokenizer
• [ ] Can implement PerceiverBlock
• [ ] Can load crystal structures using PyMatGen
• [ ] Can retrieve Materials Project data
• [ ] Can visualize a distance matrix
• [ ] Can interpret attention weights in a materials-science sense
• [ ] Can handle varying numbers of atoms in batch processing (padding)
• [ ] Can implement a model training loop
• [ ] Can evaluate prediction results (MAE, RMSE)

Debugging Skills (5 items)

• [ ] Can detect miscalculations of periodic boundary conditions
• [ ] Can handle out-of-range atomic number errors
• [ ] Can identify and resolve memory efficiency problems
• [ ] Can debug errors in SMILES tokenization
• [ ] Can verify the symmetry of a distance matrix

Application Ability (5 items)

• [ ] Can apply Matformer to a new material property prediction task
• [ ] Can apply ChemBERTa to molecular property prediction
• [ ] Can integrate multiple data sources (structure + composition)
• [ ] Can extend an existing model to add new features
• [ ] Can propose ways to utilize prediction results in experimental design

Data Processing (5 items)

• [ ] Can preprocess Materials Project data
• [ ] Can canonicalize SMILES
• [ ] Can standardize crystal structures (primitive cell)
• [ ] Can implement data augmentation
• [ ] Can appropriately split training/validation/test data

Theoretical Background (5 items)

• [ ] Have read the Matformer paper (Chen et al., 2022)
• [ ] Have read the ChemBERTa paper
• [ ] Have read the Perceiver paper
• [ ] Understand the fundamentals of crystallography (Bravais lattices, space groups)
• [ ] Understand the fundamentals of materials science (bonding, lattice defects)

Completion Criteria

• Minimum criterion: Achieve 40 or more items (80%)
• Recommended criterion: Achieve 45 or more items (90%)
• Excellent criterion: Achieve all 50 items (100%)

Next Chapter: Chapter 3: Pretrained Models and Transfer Learning

Author: Yusuke Hashimoto (Tohoku University) Last Updated: October 19, 2025