Chapter 3: Pre-trained Models and Transfer Learning

This chapter organizes the practical uses of materials- and chemistry-specialized pre-trained models such as MatBERT and ChemBERTa. You will learn the key points of fine-tuning with small datasets.

💡 Supplement: We compare freezing, partial freezing, and full-layer training to find the optimal balance between computational resources and accuracy.

Study Time: 25-30 minutes | Difficulty: Intermediate to Advanced

📋 What You Will Learn in This Chapter

• The importance and principles of pre-training
• Pre-trained models for materials science such as MatBERT and MolBERT
• Fine-tuning strategies
• Few-shot learning and prompt engineering
• Domain adaptation

3.1 The Importance of Pre-training

Why Pre-training Is Necessary

Challenges in materials science: - ❌ Labeled data is scarce (experimental data is expensive) - ❌ Domain-specific knowledge is required - ❌ Training from scratch is time-consuming and costly

Advantages of pre-training: - ✅ Acquire general knowledge from large-scale unlabeled data - ✅ Achieve high accuracy with a small amount of labeled data - ✅ Dramatically shorten development time (weeks → hours)

Pre-training Tasks

Examples from natural language processing: - Masked Language Model (MLM): Mask some words and predict them - Next Sentence Prediction (NSP): Predict whether two sentences are consecutive

Applications in materials science: - Masked Atom Prediction: Mask some atoms and predict them - Property Prediction: Simultaneously predict multiple material properties - Contrastive Learning: Place similar materials close together and different materials far apart

3.2 MatBERT: Materials BERT

Overview

MatBERT is a model that learns material composition formulas with BERT.

Features: - Pre-trained on composition formulas of 500k materials - Masked atom prediction task - Applicable to various property predictions via transfer learning

Tokenizing Composition Formulas

import torch
import torch.nn as nn
from transformers import BertTokenizer, BertModel

class CompositionTokenizer:
    def __init__(self):
        # Custom vocabulary (elements of the periodic table)
        self.vocab = ['[PAD]', '[CLS]', '[SEP]', '[MASK]'] + [
            'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne',
            'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca',
            # ... all elements
        ]
        self.token_to_id = {token: i for i, token in enumerate(self.vocab)}
        self.id_to_token = {i: token for i, token in enumerate(self.vocab)}

    def tokenize(self, composition):
        """
        Tokenize a composition formula

        Args:
            composition: a composition formula such as 'Fe2O3'
        Returns:
            tokens: list of tokens
        """
        import re
        # Split into elements and numbers
        pattern = r'([A-Z][a-z]?)(\d*\.?\d*)'
        matches = re.findall(pattern, composition)

        tokens = ['[CLS]']
        for element, count in matches:
            if element in self.vocab:
                # Add the element
                tokens.append(element)
                # If the count is greater than 1, repeat it that many times (simplified)
                if count and float(count) > 1:
                    for _ in range(int(float(count)) - 1):
                        tokens.append(element)
        tokens.append('[SEP]')

        return tokens

    def encode(self, compositions, max_length=32):
        """
        Convert composition formulas to IDs

        Args:
            compositions: list of composition formulas
            max_length: maximum length
        Returns:
            input_ids: (batch_size, max_length)
            attention_mask: (batch_size, max_length)
        """
        batch_input_ids = []
        batch_attention_mask = []

        for comp in compositions:
            tokens = self.tokenize(comp)
            ids = [self.token_to_id.get(token, 0) for token in tokens]

            # Padding
            attention_mask = [1] * len(ids)
            while len(ids) < max_length:
                ids.append(0)  # [PAD]
                attention_mask.append(0)

            # Truncation
            ids = ids[:max_length]
            attention_mask = attention_mask[:max_length]

            batch_input_ids.append(ids)
            batch_attention_mask.append(attention_mask)

        return torch.tensor(batch_input_ids), torch.tensor(batch_attention_mask)

# Usage example
tokenizer = CompositionTokenizer()

compositions = [
    'Fe2O3',     # iron oxide
    'LiCoO2',    # lithium cobalt oxide (battery material)
    'BaTiO3'     # barium titanate (dielectric)
]

input_ids, attention_mask = tokenizer.encode(compositions)
print(f"Input IDs shape: {input_ids.shape}")
print(f"First composition tokens: {input_ids[0][:10]}")

The MatBERT Model

class MatBERT(nn.Module):
    def __init__(self, vocab_size, d_model=512, num_layers=6, num_heads=8):
        super(MatBERT, self).__init__()

        # Embedding
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(512, d_model)

        # Transformer Encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=2048,
            batch_first=True
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)

        self.d_model = d_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, seq_len, d_model)
        """
        batch_size, seq_len = input_ids.shape

        # Token embedding
        token_embeddings = self.embedding(input_ids)

        # Positional embedding
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
        position_embeddings = self.position_embedding(positions)

        # Sum
        embeddings = token_embeddings + position_embeddings

        # Transformer
        # Convert attention_mask for the Transformer (0→-inf, 1→0)
        transformer_mask = (1 - attention_mask).bool()
        output = self.transformer_encoder(embeddings, src_key_padding_mask=transformer_mask)

        return output

# Usage example
vocab_size = len(tokenizer.vocab)
model = MatBERT(vocab_size, d_model=512, num_layers=6, num_heads=8)

embeddings = model(input_ids, attention_mask)
print(f"Embeddings shape: {embeddings.shape}")  # (3, 32, 512)

Pre-training: Masked Atom Prediction

def masked_atom_prediction_loss(model, input_ids, attention_mask, mask_prob=0.15):
    """
    Pre-training via masked atom prediction

    Args:
        model: the MatBERT model
        input_ids: (batch_size, seq_len)
        attention_mask: (batch_size, seq_len)
        mask_prob: masking probability
    Returns:
        loss: the loss
    """
    batch_size, seq_len = input_ids.shape

    # Mask randomly
    mask_token_id = tokenizer.token_to_id['[MASK]']
    mask = torch.rand(batch_size, seq_len) < mask_prob
    mask = mask & (attention_mask == 1)  # Exclude padding positions

    # Save the original tokens
    original_input_ids = input_ids.clone()

    # Apply the mask
    input_ids[mask] = mask_token_id

    # Forward
    embeddings = model(input_ids, attention_mask)

    # Prediction head
    prediction_head = nn.Linear(model.d_model, vocab_size)
    logits = prediction_head(embeddings)

    # Loss computation (only for masked positions)
    criterion = nn.CrossEntropyLoss(ignore_index=-100)
    labels = original_input_ids.clone()
    labels[~mask] = -100  # Ignore non-masked positions

    loss = criterion(logits.view(-1, vocab_size), labels.view(-1))

    return loss

# Pre-training loop (simplified version)
def pretrain_matbert(model, dataloader, epochs=10):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for input_ids, attention_mask in dataloader:
            loss = masked_atom_prediction_loss(model, input_ids, attention_mask)

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

            total_loss += loss.item()

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

    return model

3.3 Fine-tuning Strategies

What Is Fine-tuning?

Definition: Additional training to adapt a pre-trained model to a specific task

Strategies: 1. Full Fine-tuning: Update all parameters 2. Feature Extraction: Use only the embedding layers and train only the prediction head 3. Partial Fine-tuning: Update only some layers

Implementation: Band Gap Prediction

class MatBERTForBandgap(nn.Module):
    def __init__(self, matbert_model, d_model=512):
        super(MatBERTForBandgap, self).__init__()
        self.matbert = matbert_model

        # Prediction head
        self.bandgap_predictor = nn.Sequential(
            nn.Linear(d_model, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 1)
        )

    def forward(self, input_ids, attention_mask):
        # MatBERT embeddings
        embeddings = self.matbert(input_ids, attention_mask)

        # Use the [CLS] token embedding
        cls_embedding = embeddings[:, 0, :]

        # Band gap prediction
        bandgap = self.bandgap_predictor(cls_embedding)
        return bandgap

# Fine-tuning
def finetune_for_bandgap(pretrained_model, train_loader, val_loader, strategy='full'):
    """
    Fine-tuning for band gap prediction

    Args:
        pretrained_model: pre-trained MatBERT
        train_loader: training data loader
        val_loader: validation data loader
        strategy: 'full', 'feature', 'partial'
    """
    model = MatBERTForBandgap(pretrained_model)

    # Freeze parameters according to the strategy
    if strategy == 'feature':
        # Freeze MatBERT
        for param in model.matbert.parameters():
            param.requires_grad = False
    elif strategy == 'partial':
        # Freeze lower layers, update only upper layers
        for i, layer in enumerate(model.matbert.transformer_encoder.layers):
            if i < 3:  # Freeze the lower 3 layers
                for param in layer.parameters():
                    param.requires_grad = False

    # Optimization
    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)
    criterion = nn.MSELoss()

    # Training loop
    best_val_loss = float('inf')
    for epoch in range(20):
        model.train()
        train_loss = 0
        for input_ids, attention_mask, bandgaps in train_loader:
            predictions = model(input_ids, attention_mask)
            loss = criterion(predictions, bandgaps)

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

            train_loss += loss.item()

        # Validation
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for input_ids, attention_mask, bandgaps in val_loader:
                predictions = model(input_ids, attention_mask)
                loss = criterion(predictions, bandgaps)
                val_loss += loss.item()

        train_loss /= len(train_loader)
        val_loss /= len(val_loader)

        print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_matbert_bandgap.pt')

    return model

3.4 Few-shot Learning

Overview

Few-shot learning: Learning a new task from a small number of samples (a few to a few dozen)

Importance in materials science: - Data for novel materials is extremely scarce - Acquiring experimental data is expensive - Rapid prototyping is required

Prototypical Networks

class PrototypicalNetwork(nn.Module):
    def __init__(self, matbert_model, d_model=512):
        super(PrototypicalNetwork, self).__init__()
        self.encoder = matbert_model

    def forward(self, support_ids, support_mask, query_ids, query_mask, support_labels):
        """
        Classification with Prototypical Networks

        Args:
            support_ids: support set inputs (n_support, seq_len)
            support_mask: support set mask
            query_ids: query inputs (n_query, seq_len)
            query_mask: query mask
            support_labels: support set labels (n_support,)
        Returns:
            predictions: predicted labels for the queries
        """
        # Embeddings of the support set and queries
        support_embeddings = self.encoder(support_ids, support_mask)[:, 0, :]  # [CLS]
        query_embeddings = self.encoder(query_ids, query_mask)[:, 0, :]

        # Compute each class's prototype (mean embedding)
        unique_labels = torch.unique(support_labels)
        prototypes = []
        for label in unique_labels:
            mask = (support_labels == label)
            prototype = support_embeddings[mask].mean(dim=0)
            prototypes.append(prototype)

        prototypes = torch.stack(prototypes)  # (num_classes, d_model)

        # Distances between queries and prototypes
        distances = torch.cdist(query_embeddings, prototypes)  # (n_query, num_classes)

        # Predict the class of the nearest prototype
        predictions = torch.argmin(distances, dim=1)

        return predictions

# Usage example: 3-way 5-shot classification
# 3 classes, 5 samples per class
n_classes = 3
n_support_per_class = 5
n_query = 10

support_ids = torch.randint(0, vocab_size, (n_classes * n_support_per_class, 32))
support_mask = torch.ones_like(support_ids)
support_labels = torch.arange(n_classes).repeat_interleave(n_support_per_class)

query_ids = torch.randint(0, vocab_size, (n_query, 32))
query_mask = torch.ones_like(query_ids)

proto_net = PrototypicalNetwork(model)
predictions = proto_net(support_ids, support_mask, query_ids, query_mask, support_labels)
print(f"Predictions: {predictions}")

3.5 Prompt Engineering

Prompts in Materials Science

Prompt: Providing additional information to the model to improve performance

Example:

# Normal: 'Fe2O3'
# With prompt: '[OXIDE] Fe2O3 [BANDGAP]'

Implementation

class PromptedMatBERT(nn.Module):
    def __init__(self, matbert_model, d_model=512):
        super(PromptedMatBERT, self).__init__()
        self.matbert = matbert_model

        # Task-specific prompt embeddings (learnable)
        self.task_prompts = nn.Parameter(torch.randn(10, d_model))  # 10 types of tasks

    def forward(self, input_ids, attention_mask, task_id=0):
        """
        Args:
            input_ids: (batch_size, seq_len)
            attention_mask: (batch_size, seq_len)
            task_id: task ID (0-9)
        """
        batch_size = input_ids.size(0)

        # Standard embeddings
        embeddings = self.matbert(input_ids, attention_mask)

        # Prepend the task prompt
        task_prompt = self.task_prompts[task_id].unsqueeze(0).expand(batch_size, -1, -1)
        embeddings = torch.cat([task_prompt, embeddings], dim=1)

        return embeddings

# Usage example
prompted_model = PromptedMatBERT(model)

# Task 0: band gap prediction
embeddings_task0 = prompted_model(input_ids, attention_mask, task_id=0)

# Task 1: formation energy prediction
embeddings_task1 = prompted_model(input_ids, attention_mask, task_id=1)

print(f"Embeddings with prompt shape: {embeddings_task0.shape}")

3.6 Domain Adaptation

Overview

Domain adaptation: Adapting a model trained on a source domain to a target domain

Example: - Source: inorganic materials data - Target: organic molecules data

Adversarial Domain Adaptation

class DomainClassifier(nn.Module):
    def __init__(self, d_model=512):
        super(DomainClassifier, self).__init__()
        self.classifier = nn.Sequential(
            nn.Linear(d_model, 256),
            nn.ReLU(),
            nn.Linear(256, 2)  # source or target
        )

    def forward(self, embeddings):
        return self.classifier(embeddings)

class DomainAdaptiveMatBERT(nn.Module):
    def __init__(self, matbert_model):
        super(DomainAdaptiveMatBERT, self).__init__()
        self.matbert = matbert_model
        self.domain_classifier = DomainClassifier()
        self.task_predictor = nn.Linear(512, 1)  # e.g., band gap prediction

    def forward(self, input_ids, attention_mask, alpha=1.0):
        """
        Args:
            alpha: strength of domain adaptation
        """
        embeddings = self.matbert(input_ids, attention_mask)[:, 0, :]

        # Task prediction
        task_output = self.task_predictor(embeddings)

        # Domain prediction (using a gradient reversal layer)
        # Omitted here for simplicity
        domain_output = self.domain_classifier(embeddings)

        return task_output, domain_output

# Training loop (simplified version)
def train_domain_adaptive(model, source_loader, target_loader, epochs=20):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
    task_criterion = nn.MSELoss()
    domain_criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        for (source_ids, source_mask, source_labels), (target_ids, target_mask, _) in zip(source_loader, target_loader):
            # Source domain
            source_task, source_domain = model(source_ids, source_mask)
            source_domain_labels = torch.zeros(source_ids.size(0), dtype=torch.long)  # source = 0

            # Target domain
            target_task, target_domain = model(target_ids, target_mask)
            target_domain_labels = torch.ones(target_ids.size(0), dtype=torch.long)  # target = 1

            # Losses
            task_loss = task_criterion(source_task, source_labels)
            domain_loss = domain_criterion(source_domain, source_domain_labels) + \
                          domain_criterion(target_domain, target_domain_labels)

            total_loss = task_loss + 0.1 * domain_loss

            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

        print(f"Epoch {epoch+1}, Task Loss: {task_loss.item():.4f}, Domain Loss: {domain_loss.item():.4f}")

3.7 Summary

Key Points

1. Pre-training: Acquire general knowledge from large-scale unlabeled data
2. Fine-tuning: Specialize to a task with a small amount of data
3. Few-shot learning: Learn a new task from a few samples
4. Prompt engineering: Represent task information in embeddings
5. Domain adaptation: Transfer knowledge across different domains

Preparing for the Next Chapter

In Chapter 4, we will learn molecule generation and inverse materials design using diffusion models.

📝 Exercises

Exercise 1: Conceptual Understanding

For the three fine-tuning strategies (Full, Feature Extraction, Partial), explain the situations in which each is best suited.

Exercise 2: Implementation

Fill in the blanks in the following code to complete a function that loads a pre-trained model and fine-tunes it.

def load_and_finetune(pretrained_path, train_loader, val_loader):
    # Load the pre-trained model
    matbert = MatBERT(vocab_size=______, d_model=512)
    matbert.load_state_dict(torch.load(______))

    # Build the fine-tuning model
    model = MatBERTForBandgap(______)

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

    # Training loop
    for epoch in range(10):
        model.train()
        for input_ids, attention_mask, targets in train_loader:
            predictions = model(______, ______)
            loss = ______(predictions, targets)

            optimizer.zero_grad()
            ______.backward()
            optimizer.step()

    return model

def load_and_finetune(pretrained_path, train_loader, val_loader):
    # Load the pre-trained model
    matbert = MatBERT(vocab_size=len(tokenizer.vocab), d_model=512)
    matbert.load_state_dict(torch.load(pretrained_path))

    # Build the fine-tuning model
    model = MatBERTForBandgap(matbert)

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

    # Training loop
    for epoch in range(10):
        model.train()
        for input_ids, attention_mask, targets in train_loader:
            predictions = model(input_ids, attention_mask)
            loss = criterion(predictions, targets)

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

    return model

Exercise 3: Application

List three scenarios in which few-shot learning is particularly useful in materials science, and explain the reason for each.

🚀 Implementation Exercises: Transformer for Materials

Exercise 1: Implementing MatBERT (BERT for Materials)

import torch
import torch.nn as nn
from transformers import BertConfig, BertModel

class MaterialsBERT(nn.Module):
    def __init__(self, vocab_size=120, d_model=768, num_layers=12, num_heads=12):
        """
        Materials BERT implementation

        Args:
            vocab_size: number of atom types + special tokens
            d_model: hidden layer dimension
            num_layers: number of Transformer layers
            num_heads: number of attention heads
        """
        super().__init__()

        # BERT configuration
        config = BertConfig(
            vocab_size=vocab_size,
            hidden_size=d_model,
            num_hidden_layers=num_layers,
            num_attention_heads=num_heads,
            intermediate_size=d_model * 4,
            hidden_dropout_prob=0.1,
            attention_probs_dropout_prob=0.1,
            max_position_embeddings=512
        )

        self.bert = BertModel(config)

    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        """
        Args:
            input_ids: (batch_size, seq_len) atomic number sequence
            attention_mask: (batch_size, seq_len)
            token_type_ids: (batch_size, seq_len)
        Returns:
            outputs: BERT outputs with pooler_output
        """
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids
        )

        return outputs

# Usage example
mat_bert = MaterialsBERT(vocab_size=120, d_model=768)

# Dummy data: Fe2O3 (iron oxide)
# [CLS] Fe Fe O O O [SEP]
input_ids = torch.tensor([[101, 26, 26, 8, 8, 8, 102]])  # 101=[CLS], 102=[SEP]
attention_mask = torch.ones_like(input_ids)

outputs = mat_bert(input_ids, attention_mask)
print(f"Last hidden state shape: {outputs.last_hidden_state.shape}")  # (1, 7, 768)
print(f"Pooler output shape: {outputs.pooler_output.shape}")  # (1, 768)

Exercise 2: Implementing MatGPT (GPT for Materials Generation)

from transformers import GPT2Config, GPT2LMHeadModel

class MaterialsGPT(nn.Module):
    def __init__(self, vocab_size=120, d_model=768, num_layers=12, num_heads=12):
        """
        Materials GPT for generative tasks

        Args:
            vocab_size: number of atom types + special tokens
            d_model: hidden layer dimension
            num_layers: number of Transformer layers
            num_heads: number of attention heads
        """
        super().__init__()

        config = GPT2Config(
            vocab_size=vocab_size,
            n_positions=512,
            n_embd=d_model,
            n_layer=num_layers,
            n_head=num_heads,
            resid_pdrop=0.1,
            embd_pdrop=0.1,
            attn_pdrop=0.1
        )

        self.gpt = GPT2LMHeadModel(config)

    def forward(self, input_ids, labels=None):
        """
        Args:
            input_ids: (batch_size, seq_len)
            labels: (batch_size, seq_len) for training
        """
        outputs = self.gpt(input_ids=input_ids, labels=labels)
        return outputs

    def generate_composition(self, start_tokens, max_length=50, temperature=1.0):
        """
        Generate a composition formula

        Args:
            start_tokens: (1, start_len) starting tokens
            max_length: maximum generation length
            temperature: sampling temperature
        """
        self.eval()
        with torch.no_grad():
            for _ in range(max_length - start_tokens.size(1)):
                outputs = self.gpt(start_tokens)
                logits = outputs.logits[:, -1, :] / temperature

                probs = torch.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)

                start_tokens = torch.cat([start_tokens, next_token], dim=1)

                # Stop at the [SEP] token
                if next_token.item() == 102:
                    break

        return start_tokens

# Usage example
mat_gpt = MaterialsGPT(vocab_size=120, d_model=768)

# Generate: [CLS] Fe ... (generate an oxide)
start = torch.tensor([[101, 26]])  # [CLS] Fe
generated = mat_gpt.generate_composition(start, max_length=20)
print(f"Generated sequence: {generated}")

Exercise 3: Implementing MatT5 (T5 for Materials Seq2Seq)

from transformers import T5Config, T5ForConditionalGeneration

class MaterialsT5(nn.Module):
    def __init__(self, vocab_size=120, d_model=512, num_layers=6):
        """
        Materials T5 for sequence-to-sequence tasks
        (e.g., composition → properties description)

        Args:
            vocab_size: vocabulary size
            d_model: model dimension
            num_layers: number of encoder and decoder layers
        """
        super().__init__()

        config = T5Config(
            vocab_size=vocab_size,
            d_model=d_model,
            d_kv=64,
            d_ff=d_model * 4,
            num_layers=num_layers,
            num_decoder_layers=num_layers,
            num_heads=8,
            dropout_rate=0.1
        )

        self.t5 = T5ForConditionalGeneration(config)

    def forward(self, input_ids, attention_mask=None, labels=None):
        """
        Args:
            input_ids: (batch_size, src_len) input sequence
            labels: (batch_size, tgt_len) target sequence
        """
        outputs = self.t5(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )
        return outputs

    def predict_properties(self, composition_ids, max_length=50):
        """
        Generate a property description from a composition formula

        Args:
            composition_ids: (batch_size, seq_len) composition formula
            max_length: maximum generation length
        """
        self.eval()
        with torch.no_grad():
            outputs = self.t5.generate(
                composition_ids,
                max_length=max_length,
                num_beams=4,
                early_stopping=True
            )
        return outputs

# Usage example
mat_t5 = MaterialsT5(vocab_size=120, d_model=512)

# Input: Fe2O3 → Output: "semiconductor bandgap 2.0 eV"
input_ids = torch.tensor([[26, 26, 8, 8, 8]])  # Fe Fe O O O
outputs = mat_t5.predict_properties(input_ids, max_length=20)
print(f"Predicted properties: {outputs}")

🧪 Implementing SMILES/SELFIES Tokenization

SMILES Tokenizer

import re
from typing import List, Dict

class SMILESTokenizer:
    """
    Complete tokenization of SMILES strings

    Supports:
    - aromaticity (c, n, o, s)
    - stereochemistry (@, @@, /, \\)
    - branches ((, ))
    - bonds (-, =, #, :)
    - rings (digits)
    """

    def __init__(self):
        # Regular expression pattern (in priority order)
        self.pattern = r'(\[[^\]]+\]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])'

        # Special tokens
        self.special_tokens = {
            '[PAD]': 0,
            '[CLS]': 1,
            '[SEP]': 2,
            '[MASK]': 3,
            '[UNK]': 4
        }

        # Build the vocabulary
        self.vocab = self._build_vocab()
        self.token_to_id = {token: i for i, token in enumerate(self.vocab)}
        self.id_to_token = {i: token for token, i in self.token_to_id.items()}

    def _build_vocab(self) -> List[str]:
        """Build the vocabulary"""
        vocab = list(self.special_tokens.keys())

        # Element symbols
        elements = ['C', 'N', 'O', 'S', 'P', 'F', 'Cl', 'Br', 'I',
                   'c', 'n', 'o', 's', 'p']  # aromatic

        # Symbols
        symbols = ['(', ')', '[', ']', '=', '#', '-', '+', '\\', '/',
                  ':', '.', '@', '@@']

        # Digits
        numbers = [str(i) for i in range(10)]

        vocab.extend(elements + symbols + numbers)

        return vocab

    def tokenize(self, smiles: str) -> List[str]:
        """
        Tokenize a SMILES string

        Args:
            smiles: SMILES string

        Returns:
            tokens: list of tokens

        Examples:
            >>> tokenizer = SMILESTokenizer()
            >>> tokenizer.tokenize("CC(C)Cc1ccc(cc1)C(C)C(=O)O")
            ['C', 'C', '(', 'C', ')', 'C', 'c', '1', 'c', 'c', 'c', '(', ...]
        """
        tokens = re.findall(self.pattern, smiles)
        return ['[CLS]'] + tokens + ['[SEP]']

    def encode(self, smiles: str, max_length: int = 128) -> Dict[str, torch.Tensor]:
        """
        Convert a SMILES string to IDs

        Args:
            smiles: SMILES string
            max_length: maximum length

        Returns:
            encoding: input_ids, attention_mask
        """
        tokens = self.tokenize(smiles)

        # Convert tokens to IDs
        ids = [self.token_to_id.get(token, self.token_to_id['[UNK]'])
               for token in tokens]

        # Padding
        attention_mask = [1] * len(ids)
        while len(ids) < max_length:
            ids.append(self.token_to_id['[PAD]'])
            attention_mask.append(0)

        # Truncation
        ids = ids[:max_length]
        attention_mask = attention_mask[:max_length]

        return {
            'input_ids': torch.tensor([ids]),
            'attention_mask': torch.tensor([attention_mask])
        }

    def decode(self, ids: List[int]) -> str:
        """Reconstruct a SMILES string from IDs"""
        tokens = [self.id_to_token.get(id, '[UNK]') for id in ids]
        # Remove special tokens
        tokens = [t for t in tokens if t not in self.special_tokens]
        return ''.join(tokens)

# Usage example
tokenizer = SMILESTokenizer()

# Ibuprofen
smiles = "CC(C)Cc1ccc(cc1)C(C)C(=O)O"
tokens = tokenizer.tokenize(smiles)
print(f"Tokens: {tokens[:10]}...")

encoding = tokenizer.encode(smiles)
print(f"Input IDs shape: {encoding['input_ids'].shape}")
print(f"First 10 IDs: {encoding['input_ids'][0][:10]}")

# Decode
decoded = tokenizer.decode(encoding['input_ids'][0].tolist())
print(f"Decoded: {decoded}")

SELFIES Tokenizer

try:
    import selfies as sf
except ImportError:
    print("Install selfies: pip install selfies")

class SELFIESTokenizer:
    """
    SELFIES (SELF-referencIng Embedded Strings) Tokenizer

    Advantages:
    - Generates 100% valid molecules
    - Grammatically correct
    - More robust than SMILES
    """

    def __init__(self):
        self.special_tokens = {
            '[PAD]': 0,
            '[CLS]': 1,
            '[SEP]': 2,
            '[MASK]': 3
        }

        # Common SELFIES tokens
        self.vocab = self._build_vocab()
        self.token_to_id = {token: i for i, token in enumerate(self.vocab)}
        self.id_to_token = {i: token for token, i in self.token_to_id.items()}

    def _build_vocab(self) -> List[str]:
        """
        Build the SELFIES vocabulary

        Common tokens:
        [C], [N], [O], [=C], [=N], [Ring1], [Branch1], etc.
        """
        vocab = list(self.special_tokens.keys())

        # Basic tokens
        common_tokens = [
            '[C]', '[N]', '[O]', '[S]', '[P]', '[F]', '[Cl]', '[Br]', '[I]',
            '[=C]', '[=N]', '[=O]', '[#C]', '[#N]',
            '[Ring1]', '[Ring2]', '[Branch1]', '[Branch2]',
            '[O-1]', '[N+1]', '[nop]'
        ]

        vocab.extend(common_tokens)
        return vocab

    def smiles_to_selfies(self, smiles: str) -> str:
        """Convert SMILES to SELFIES"""
        try:
            selfies = sf.encoder(smiles)
            return selfies
        except Exception as e:
            print(f"Encoding error: {e}")
            return ""

    def selfies_to_smiles(self, selfies: str) -> str:
        """Convert SELFIES to SMILES"""
        try:
            smiles = sf.decoder(selfies)
            return smiles
        except Exception as e:
            print(f"Decoding error: {e}")
            return ""

    def tokenize(self, selfies: str) -> List[str]:
        """
        Tokenize a SELFIES string

        Args:
            selfies: SELFIES string

        Returns:
            tokens: list of tokens

        Examples:
            >>> tokenizer = SELFIESTokenizer()
            >>> tokenizer.tokenize("[C][C][Branch1][C][C][C]")
            ['[CLS]', '[C]', '[C]', '[Branch1]', '[C]', '[C]', '[C]', '[SEP]']
        """
        tokens = list(sf.split_selfies(selfies))
        return ['[CLS]'] + tokens + ['[SEP]']

    def encode(self, selfies: str, max_length: int = 128) -> Dict[str, torch.Tensor]:
        """Convert a SELFIES string to IDs"""
        tokens = self.tokenize(selfies)

        # Convert tokens to IDs (unknown tokens are added dynamically)
        ids = []
        for token in tokens:
            if token not in self.token_to_id:
                new_id = len(self.vocab)
                self.vocab.append(token)
                self.token_to_id[token] = new_id
                self.id_to_token[new_id] = token
            ids.append(self.token_to_id[token])

        # Padding
        attention_mask = [1] * len(ids)
        while len(ids) < max_length:
            ids.append(self.token_to_id['[PAD]'])
            attention_mask.append(0)

        # Truncation
        ids = ids[:max_length]
        attention_mask = attention_mask[:max_length]

        return {
            'input_ids': torch.tensor([ids]),
            'attention_mask': torch.tensor([attention_mask])
        }

# Usage example
if 'sf' in dir():
    tokenizer_selfies = SELFIESTokenizer()

    # Convert from SMILES to SELFIES
    smiles = "CC(C)Cc1ccc(cc1)C(C)C(=O)O"
    selfies = tokenizer_selfies.smiles_to_selfies(smiles)
    print(f"SELFIES: {selfies}")

    # Tokenize
    tokens = tokenizer_selfies.tokenize(selfies)
    print(f"Tokens: {tokens[:10]}...")

    # Encode
    encoding = tokenizer_selfies.encode(selfies)
    print(f"Encoded shape: {encoding['input_ids'].shape}")

⚠️ Practical Pitfalls and How to Address Them

1. Overfitting in Fine-tuning

Problem: Validation loss diverges when training on a small dataset

# ❌ Problem: Updating all parameters with a large learning rate
def wrong_finetuning():
    model = MatBERTForBandgap(pretrained_matbert)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)  # Too large!

    for epoch in range(100):  # Too many epochs, too
        for batch in train_loader:
            loss = compute_loss(batch)
            loss.backward()
            optimizer.step()

# ✅ Solution: Layer-wise learning rate decay + Early stopping
def correct_finetuning():
    model = MatBERTForBandgap(pretrained_matbert)

    # Layer-wise learning rate
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {
            'params': [p for n, p in model.matbert.named_parameters()
                      if not any(nd in n for nd in no_decay)],
            'weight_decay': 0.01,
            'lr': 2e-5  # Keep the pre-trained part small
        },
        {
            'params': [p for n, p in model.matbert.named_parameters()
                      if any(nd in n for nd in no_decay)],
            'weight_decay': 0.0,
            'lr': 2e-5
        },
        {
            'params': model.bandgap_predictor.parameters(),
            'lr': 1e-4  # Keep the prediction head large
        }
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_parameters)

    # Early stopping
    best_val_loss = float('inf')
    patience = 5
    patience_counter = 0

    for epoch in range(100):
        train_loss = train_epoch(model, train_loader, optimizer)
        val_loss = validate(model, val_loader)

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_model.pt')
            patience_counter = 0
        else:
            patience_counter += 1

        if patience_counter >= patience:
            print(f"Early stopping at epoch {epoch}")
            break

    # Restore the best model
    model.load_state_dict(torch.load('best_model.pt'))
    return model

2. The Domain Shift Problem

Problem: Applying a model pre-trained on inorganic materials to organic molecules

# ❌ Problem: Applying directly despite different domains
def wrong_domain_adaptation():
    # Pre-train on inorganic materials
    matbert = pretrained_on_inorganic_materials()

    # Fine-tune directly on organic molecule data
    # → Poor performance!
    finetune_on_organic_molecules(matbert)

# ✅ Solution: Intermediate task transfer
def correct_domain_adaptation():
    # Step 1: Pre-train on inorganic materials
    matbert = pretrained_on_inorganic_materials()

    # Step 2: Continue training on an intermediate task (between inorganic and organic)
    # Example: metal-organic framework (MOF) data
    matbert = continual_pretrain_on_mof(matbert)

    # Step 3: Fine-tune on organic molecule data
    model = finetune_on_organic_molecules(matbert)

    return model

# Or: Domain-adversarial training
class DomainAdversarialTraining:
    def train(self, source_data, target_data):
        for source_batch, target_batch in zip(source_data, target_data):
            # Source domain: task loss
            source_output = model(source_batch)
            task_loss = compute_task_loss(source_output, source_batch.labels)

            # Both domains: domain classification loss (gradient reversal)
            source_domain_pred = domain_classifier(source_output, reverse_gradient=True)
            target_domain_pred = domain_classifier(target_output, reverse_gradient=True)

            domain_loss = compute_domain_loss(source_domain_pred, target_domain_pred)

            total_loss = task_loss + 0.1 * domain_loss
            total_loss.backward()
            optimizer.step()

3. Mistakes in the Masking Strategy for Masked Language Modeling

Problem: The masking pattern is biased

# ❌ Problem: Random masking (chemically meaningless)
def wrong_masking(composition_ids):
    mask_prob = 0.15
    mask = torch.rand(composition_ids.shape) < mask_prob
    composition_ids[mask] = MASK_TOKEN_ID
    return composition_ids

# ✅ Solution: Chemically meaningful masking
def chemically_aware_masking(composition_ids, element_groups):
    """
    Masking that considers element groups

    Args:
        composition_ids: (batch, seq_len)
        element_groups: {group_id: [element_ids]}
            Example: {0: [26, 27, 28], 1: [8, 16]}  # transition metals, chalcogens
    """
    mask_prob = 0.15
    masked_ids = composition_ids.clone()

    for i in range(composition_ids.size(0)):
        # Mask by chemical group
        for group_id, element_ids in element_groups.items():
            group_positions = torch.isin(composition_ids[i], torch.tensor(element_ids))
            if group_positions.sum() > 0:
                # Mask a portion within the group
                mask_within_group = torch.rand(group_positions.sum()) < mask_prob
                group_indices = torch.where(group_positions)[0]
                masked_positions = group_indices[mask_within_group]
                masked_ids[i, masked_positions] = MASK_TOKEN_ID

    return masked_ids

# Usage example
element_groups = {
    0: [26, 27, 28, 29],  # Fe, Co, Ni, Cu (transition metals)
    1: [8, 16, 34],       # O, S, Se (chalcogens)
    2: [3, 11, 19]        # Li, Na, K (alkali metals)
}

masked_composition = chemically_aware_masking(composition_ids, element_groups)

4. Mistakes in Support Set Selection for Few-shot Learning

Problem: The support set is biased

# ❌ Problem: Selecting the support set randomly
def wrong_support_selection(dataset, k=5):
    indices = torch.randperm(len(dataset))[:k]
    return dataset[indices]

# ✅ Solution: Support set selection that considers diversity
def diverse_support_selection(dataset, embeddings, k=5):
    """
    Select diverse samples with K-means

    Args:
        dataset: the dataset
        embeddings: (N, d) sample embeddings
        k: support set size
    """
    from sklearn.cluster import KMeans

    # Cluster with K-means
    kmeans = KMeans(n_clusters=k, random_state=42)
    labels = kmeans.fit_predict(embeddings.numpy())

    # Select the sample closest to each cluster center
    support_indices = []
    for i in range(k):
        cluster_indices = torch.where(torch.tensor(labels) == i)[0]
        cluster_embeddings = embeddings[cluster_indices]
        cluster_center = kmeans.cluster_centers_[i]

        # The sample closest to the center
        distances = torch.norm(cluster_embeddings - torch.tensor(cluster_center), dim=1)
        closest_idx = cluster_indices[torch.argmin(distances)]
        support_indices.append(closest_idx.item())

    return dataset[support_indices]

# Usage example
# Precompute the dataset embeddings
embeddings = compute_embeddings(dataset, matbert)
support_set = diverse_support_selection(dataset, embeddings, k=10)

5. Insufficient Optimization in Prompt Engineering

Problem: Poor performance with fixed prompts

# ❌ Problem: A manually designed fixed prompt
class FixedPromptModel(nn.Module):
    def __init__(self, matbert):
        super().__init__()
        self.matbert = matbert
        # Fixed prompt
        self.prompt = nn.Parameter(torch.randn(1, 10, 768), requires_grad=False)

# ✅ Solution: A learnable prompt (Prefix-Tuning)
class LearnablePromptModel(nn.Module):
    def __init__(self, matbert, prompt_length=10, num_tasks=5):
        super().__init__()
        self.matbert = matbert
        self.prompt_length = prompt_length

        # Learnable prompts per task
        self.task_prompts = nn.Parameter(torch.randn(num_tasks, prompt_length, 768))

        # Freeze the MatBERT parameters
        for param in self.matbert.parameters():
            param.requires_grad = False

    def forward(self, input_ids, task_id=0):
        batch_size = input_ids.size(0)

        # Input embeddings
        input_embeddings = self.matbert.embeddings(input_ids)

        # Prepend the task-specific prompt
        prompt = self.task_prompts[task_id].unsqueeze(0).expand(batch_size, -1, -1)
        embeddings = torch.cat([prompt, input_embeddings], dim=1)

        # Pass through the Transformer
        outputs = self.matbert.encoder(embeddings)

        return outputs

# Training
model = LearnablePromptModel(pretrained_matbert, prompt_length=10, num_tasks=5)

# Optimize only the prompts (drastically reduces the number of parameters)
optimizer = torch.optim.Adam([model.task_prompts], lr=1e-3)

✅ Chapter 3 Completion Checklist

Conceptual Understanding (10 items)

• [ ] Can explain the importance and advantages of pre-training
• [ ] Understand the principle of Masked Language Modeling
• [ ] Can explain the differences between Full/Feature Extraction/Partial Fine-tuning
• [ ] Understand the principle of few-shot learning (Prototypical Networks)
• [ ] Understand the concept of prompt engineering
• [ ] Can explain the need for domain adaptation
• [ ] Understand the relationship between pre-training tasks and downstream tasks
• [ ] Can quantitatively evaluate the effect of transfer learning
• [ ] Understand the characteristics of materials-specialized models such as MatBERT and MolBERT
• [ ] Can explain the differences and use cases of BERT/GPT/T5

Implementation Skills (15 items)

• [ ] Can implement MatBERT
• [ ] Can implement MatGPT
• [ ] Can implement MatT5
• [ ] Can implement a SMILES tokenizer
• [ ] Can implement a SELFIES tokenizer
• [ ] Can implement Masked Atom Prediction
• [ ] Can implement fine-tuning strategies (Full/Feature/Partial)
• [ ] Can implement Prototypical Networks
• [ ] Can implement learnable prompts
• [ ] Can implement domain-adversarial training
• [ ] Can implement early stopping
• [ ] Can configure a layer-wise learning rate
• [ ] Can save and load pre-trained models
• [ ] Can leverage the Hugging Face Transformers library
• [ ] Can integrate a custom tokenizer into Transformers

Debugging Skills (5 items)

• [ ] Can detect overfitting and address it with regularization
• [ ] Can detect domain shift and apply adaptation methods
• [ ] Can evaluate the validity of a masking strategy
• [ ] Can evaluate the quality of a few-shot support set
• [ ] Can visualize and analyze the effect of prompts

Application Ability (5 items)

• [ ] Can apply a pre-trained model to a new material property prediction task
• [ ] Can combine multiple pre-training tasks to improve performance
• [ ] Can design a domain adaptation strategy
• [ ] Can combine few-shot learning with data augmentation
• [ ] Can optimize performance through prompt engineering

Data Processing (5 items)

• [ ] Can preprocess SMILES data
• [ ] Can convert to SELFIES
• [ ] Can implement data augmentation (SMILES enumeration)
• [ ] Can split data by domain
• [ ] Can generate episodes for few-shot learning

Evaluation Skills (5 items)

• [ ] Can quantitatively evaluate the effect of pre-training (vs from scratch)
• [ ] Can compare and evaluate fine-tuning strategies
• [ ] Can properly evaluate few-shot performance (N-way K-shot)
• [ ] Can measure the effect of domain adaptation
• [ ] Can analyze the impact of prompts

Theoretical Background (5 items)

• [ ] Read the MatBERT/MolBERT papers
• [ ] Read the BERT paper (Devlin et al., 2019)
• [ ] Read the GPT paper
• [ ] Read at least one paper on few-shot learning
• [ ] Understand transfer learning theory

Completion Criteria

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

Next Chapter: Chapter 4: Generative Models and Inverse Design

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