Chapter 4: Generative Models and Inverse Design

This chapter covers the basic concepts and caveats of inverse design using diffusion models and VAEs. It also addresses evaluation metrics and challenges in practical deployment.

💡 Note: Diversity and feasibility of generated results are a trade-off. Embedding physical and chemical constraints is the key.

Learning Time: 20-25 min | Difficulty: Advanced

📋 What You Will Learn in This Chapter

Principles of Diffusion Models
Conditional Generation
Molecular Generation and SMILES Generation
Materials Inverse Design
Industrial Applications and Career Paths

4.1 What Are Generative Models?

The Importance of Generative Models in Materials Science

Conventional Approach (Forward Problem):

Material Structure → Property Prediction

Inverse Design (Inverse Problem):

Desired Properties → Material Structure Generation

Advantages of Generative Models: - ✅ Automatically generate candidates from a vast search space - ✅ Multi-objective optimization (satisfying multiple properties simultaneously) - ✅ Generation that accounts for synthesizability - ✅ Discovery of novel structures beyond human intuition

flowchart LR A[Target Properties] --> B[Generative Model] C[Constraints] --> B B --> D[Candidate Materials] D --> E[Property Prediction] E --> F{Target Achieved?} F -->|No| B F -->|Yes| G[Experimental Validation] style B fill:#e1f5ff style G fill:#ffe1e1

4.2 Principles of Diffusion Models

What Are Diffusion Models?

Core Idea: Reverse the noise-adding process to generate data from noise

Forward Process (Adding Noise): $$ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t} x_{t-1}, \beta_t I) $$

Reverse Process (Removing Noise): $$ p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) $$

Visual Understanding

Simple Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class SimpleDiffusionModel(nn.Module):
    def __init__(self, input_dim, hidden_dim=256, num_timesteps=1000):
        super(SimpleDiffusionModel, self).__init__()
        self.num_timesteps = num_timesteps

        # Noise schedule
        self.betas = torch.linspace(1e-4, 0.02, num_timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

        # Noise prediction network
        self.noise_predictor = nn.Sequential(
            nn.Linear(input_dim + 1, hidden_dim),  # +1 for the timestep
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )

    def forward_process(self, x0, t):
        """
        Forward process: adding noise

        Args:
            x0: original data (batch_size, input_dim)
            t: timestep (batch_size,)
        Returns:
            xt: data with noise added
            noise: the added noise
        """
        batch_size = x0.size(0)

        # Noise level per timestep
        alpha_t = self.alphas_cumprod[t].view(-1, 1)
        sqrt_alpha_t = torch.sqrt(alpha_t)
        sqrt_one_minus_alpha_t = torch.sqrt(1 - alpha_t)

        # Sample noise
        noise = torch.randn_like(x0)

        # Add noise
        xt = sqrt_alpha_t * x0 + sqrt_one_minus_alpha_t * noise

        return xt, noise

    def predict_noise(self, xt, t):
        """
        Predict noise

        Args:
            xt: data with noise added
            t: timestep
        Returns:
            predicted_noise: the predicted noise
        """
        # Embed the timestep
        t_embed = t.float().unsqueeze(1) / self.num_timesteps

        # Predict noise
        x_with_t = torch.cat([xt, t_embed], dim=1)
        predicted_noise = self.noise_predictor(x_with_t)

        return predicted_noise

    def reverse_process(self, xt, t):
        """
        Reverse process: removing noise (one step)

        Args:
            xt: current data
            t: timestep
        Returns:
            x_prev: data one step earlier
        """
        # Predict noise
        predicted_noise = self.predict_noise(xt, t)

        # Parameters
        alpha_t = self.alphas[t].view(-1, 1)
        alpha_t_cumprod = self.alphas_cumprod[t].view(-1, 1)
        beta_t = self.betas[t].view(-1, 1)

        # Compute the previous step
        x_prev = (1 / torch.sqrt(alpha_t)) * (
            xt - (beta_t / torch.sqrt(1 - alpha_t_cumprod)) * predicted_noise
        )

        # Add noise (when t > 0)
        if t[0] > 0:
            noise = torch.randn_like(xt)
            x_prev = x_prev + torch.sqrt(beta_t) * noise

        return x_prev

    def generate(self, batch_size, input_dim):
        """
        Generate data

        Args:
            batch_size: batch size
            input_dim: data dimension
        Returns:
            x0: generated data
        """
        # Start from pure noise
        xt = torch.randn(batch_size, input_dim)

        # Run the reverse process
        for t in reversed(range(self.num_timesteps)):
            t_batch = torch.full((batch_size,), t, dtype=torch.long)
            xt = self.reverse_process(xt, t_batch)

        return xt

# Usage example: generating molecular descriptors
input_dim = 128  # descriptor dimension
diffusion_model = SimpleDiffusionModel(input_dim, hidden_dim=256, num_timesteps=100)

# Training data (dummy)
x0 = torch.randn(64, input_dim)  # descriptors for 64 molecules

# Forward process (adding noise)
t = torch.randint(0, 100, (64,))
xt, noise = diffusion_model.forward_process(x0, t)

# Noise prediction
predicted_noise = diffusion_model.predict_noise(xt, t)

# Loss
loss = F.mse_loss(predicted_noise, noise)
print(f"Training loss: {loss.item():.4f}")

# Generation
generated_data = diffusion_model.generate(batch_size=10, input_dim=input_dim)
print(f"Generated data shape: {generated_data.shape}")

4.3 Conditional Generation

Overview

Conditional Generation: Generate by providing the target properties as conditions

Example:

# Condition: band gap = 2.0 eV, formation energy < 0
# Generate: material structures that satisfy the conditions

Implementation: Conditional Diffusion

class ConditionalDiffusionModel(nn.Module):
    def __init__(self, input_dim, condition_dim, hidden_dim=256, num_timesteps=1000):
        super(ConditionalDiffusionModel, self).__init__()
        self.num_timesteps = num_timesteps

        # Noise schedule
        self.betas = torch.linspace(1e-4, 0.02, num_timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

        # Condition encoder
        self.condition_encoder = nn.Sequential(
            nn.Linear(condition_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )

        # Noise prediction network (conditional)
        self.noise_predictor = nn.Sequential(
            nn.Linear(input_dim + hidden_dim + 1, hidden_dim),  # +1 for the timestep
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )

    def predict_noise(self, xt, t, condition):
        """
        Conditional noise prediction

        Args:
            xt: data with noise added (batch_size, input_dim)
            t: timestep (batch_size,)
            condition: condition (target properties) (batch_size, condition_dim)
        Returns:
            predicted_noise: the predicted noise
        """
        # Embed the condition
        condition_embed = self.condition_encoder(condition)

        # Embed the timestep
        t_embed = t.float().unsqueeze(1) / self.num_timesteps

        # Concatenate
        x_with_condition = torch.cat([xt, condition_embed, t_embed], dim=1)

        # Predict noise
        predicted_noise = self.noise_predictor(x_with_condition)

        return predicted_noise

    def generate_conditional(self, condition, input_dim):
        """
        Conditional data generation

        Args:
            condition: condition (batch_size, condition_dim)
            input_dim: data dimension
        Returns:
            x0: generated data
        """
        batch_size = condition.size(0)

        # Start from pure noise
        xt = torch.randn(batch_size, input_dim)

        # Reverse process
        for t in reversed(range(self.num_timesteps)):
            t_batch = torch.full((batch_size,), t, dtype=torch.long)

            # Predict noise
            predicted_noise = self.predict_noise(xt, t_batch, condition)

            # Parameters
            alpha_t = self.alphas[t]
            alpha_t_cumprod = self.alphas_cumprod[t]
            beta_t = self.betas[t]

            # Compute the previous step
            xt = (1 / torch.sqrt(alpha_t)) * (
                xt - (beta_t / torch.sqrt(1 - alpha_t_cumprod)) * predicted_noise
            )

            # Add noise (when t > 0)
            if t > 0:
                noise = torch.randn_like(xt)
                xt = xt + torch.sqrt(beta_t) * noise

        return xt

# Usage example
input_dim = 128
condition_dim = 3  # band gap, formation energy, magnetic moment

conditional_model = ConditionalDiffusionModel(input_dim, condition_dim, hidden_dim=256, num_timesteps=100)

# Target properties
target_properties = torch.tensor([
    [2.0, -0.5, 0.0],  # band gap 2.0 eV, formation energy -0.5 eV, non-magnetic
    [3.5, -1.0, 2.0],  # band gap 3.5 eV, formation energy -1.0 eV, magnetic
])

# Conditional generation
generated_materials = conditional_model.generate_conditional(target_properties, input_dim)
print(f"Generated materials shape: {generated_materials.shape}")  # (2, 128)

4.4 Molecular Generation: SMILES Generation

Overview

SMILES (Simplified Molecular Input Line Entry System): Represent molecules as strings

Examples: - Ethanol: CCO - Benzene: c1ccccc1 - Aspirin: CC(=O)Oc1ccccc1C(=O)O

Transformer-based SMILES Generation

from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer

class SMILESGenerator(nn.Module):
    def __init__(self, vocab_size=1000, d_model=512, num_layers=6):
        super(SMILESGenerator, self).__init__()

        # GPT-2 config
        config = GPT2Config(
            vocab_size=vocab_size,
            n_positions=512,
            n_embd=d_model,
            n_layer=num_layers,
            n_head=8
        )

        self.gpt = GPT2LMHeadModel(config)

    def forward(self, input_ids, labels=None):
        """
        Args:
            input_ids: (batch_size, seq_len)
            labels: (batch_size, seq_len) targets for next-token prediction
        """
        outputs = self.gpt(input_ids, labels=labels)
        return outputs

    def generate_smiles(self, start_token_id, max_length=100, temperature=1.0):
        """
        Generate a SMILES string

        Args:
            start_token_id: start token ID
            max_length: maximum length
            temperature: sampling temperature (higher is more random)
        Returns:
            generated_ids: the generated token IDs
        """
        generated = [start_token_id]

        for _ in range(max_length):
            input_ids = torch.tensor([generated])
            outputs = self.gpt(input_ids)
            logits = outputs.logits[:, -1, :] / temperature

            # Sampling
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1).item()

            generated.append(next_token)

            # Stop at the end token
            if next_token == 2:  # [EOS]
                break

        return generated

# Conditional SMILES generation
class ConditionalSMILESGenerator(nn.Module):
    def __init__(self, vocab_size=1000, condition_dim=10, d_model=512):
        super(ConditionalSMILESGenerator, self).__init__()

        # Condition encoder
        self.condition_encoder = nn.Linear(condition_dim, d_model)

        # GPT-2 config
        config = GPT2Config(
            vocab_size=vocab_size,
            n_positions=512,
            n_embd=d_model,
            n_layer=6,
            n_head=8
        )
        self.gpt = GPT2LMHeadModel(config)

    def forward(self, input_ids, condition):
        """
        Args:
            input_ids: (batch_size, seq_len)
            condition: (batch_size, condition_dim) target properties
        """
        batch_size, seq_len = input_ids.shape

        # Embed the condition
        condition_embed = self.condition_encoder(condition).unsqueeze(1)  # (batch, 1, d_model)

        # Token embeddings
        token_embeddings = self.gpt.transformer.wte(input_ids)

        # Prepend the condition
        embeddings = torch.cat([condition_embed, token_embeddings], dim=1)

        # GPT-2 forward (directly from embeddings)
        outputs = self.gpt(inputs_embeds=embeddings)

        return outputs

# Usage example: generate a molecule with high solubility
condition_dim = 5  # logP, solubility, molecular weight, number of HB donors, number of HB acceptors
target_properties = torch.tensor([[1.5, 10.0, 250.0, 2.0, 3.0]])  # high solubility

conditional_smiles_gen = ConditionalSMILESGenerator(vocab_size=1000, condition_dim=condition_dim)

4.5 Materials Inverse Design Workflow

Complete Workflow

flowchart TB A[Define Target Properties] --> B[Conditional Generative Model] B --> C[Generate Candidate Materials] C --> D[Property Prediction Model] D --> E{Target Achieved?} E -->|No| F[Reject Candidate] F --> B E -->|Yes| G[Synthesizability Check] G --> H{Synthesizable?} H -->|No| F H -->|Yes| I[Stability Calculation] I --> J{Stable?} J -->|No| F J -->|Yes| K[Experimental Candidate List] style A fill:#e1f5ff style K fill:#e1ffe1

Implementation Example

class MaterialsInverseDesign:
    def __init__(self, generator, predictor, synthesizability_checker):
        """
        Materials inverse design system

        Args:
            generator: conditional generative model
            predictor: property prediction model
            synthesizability_checker: synthesizability checker
        """
        self.generator = generator
        self.predictor = predictor
        self.synthesizability_checker = synthesizability_checker

    def design_materials(self, target_properties, num_candidates=100, threshold=0.1):
        """
        Inverse-design materials

        Args:
            target_properties: target properties (condition_dim,)
            num_candidates: number of candidates to generate
            threshold: tolerance
        Returns:
            valid_materials: list of materials that passed validation
        """
        valid_materials = []

        for i in range(num_candidates):
            # 1. Generate a candidate
            candidate = self.generator.generate_conditional(
                target_properties.unsqueeze(0),
                input_dim=128
            )

            # 2. Predict properties
            predicted_properties = self.predictor(candidate)

            # 3. Compare with the target
            error = torch.abs(predicted_properties - target_properties).mean()
            if error > threshold:
                continue

            # 4. Synthesizability check
            if not self.synthesizability_checker(candidate):
                continue

            # 5. Stability check (omitted)

            # Passed
            valid_materials.append({
                'structure': candidate,
                'predicted_properties': predicted_properties,
                'error': error.item()
            })

        # Sort by error
        valid_materials.sort(key=lambda x: x['error'])

        return valid_materials

# Usage example
def simple_synthesizability_checker(structure):
    """
    Simple synthesizability check (in practice, more complex)
    """
    # Here it always returns True (in practice, use tools such as Retrosyn)
    return True

# Build the system
inverse_design_system = MaterialsInverseDesign(
    generator=conditional_model,
    predictor=lambda x: torch.randn(x.size(0), 3),  # dummy predictor
    synthesizability_checker=simple_synthesizability_checker
)

# Target properties
target = torch.tensor([2.5, -0.8, 0.0])  # band gap, formation energy, magnetic moment

# Run inverse design
designed_materials = inverse_design_system.design_materials(target, num_candidates=50)
print(f"Found {len(designed_materials)} valid materials")

# Show the top 3
for i, material in enumerate(designed_materials[:3]):
    print(f"\nMaterial {i+1}:")
    print(f"  Predicted properties: {material['predicted_properties']}")
    print(f"  Error: {material['error']:.4f}")

4.6 Industrial Applications and Careers

Real-World Success Stories

1. Drug Discovery: Finding a Novel Antibiotic

MIT (2020): - Method: Molecular generation with a diffusion model - Result: Discovery of halicin (a novel antibiotic) - Impact: 100x faster than conventional methods

2. Battery Materials: High Energy Density Electrolytes

Stanford/Toyota (2022): - Method: Transformer + reinforcement learning - Result: Solid electrolyte with 1.5x lithium conductivity - Impact: Accelerating the commercialization of all-solid-state batteries

3. Catalysts: CO₂ Reduction Catalysts

CMU (2023): - Method: Conditional generation + DFT calculation - Result: Discovery of a catalyst with 10x efficiency - Impact: Contribution to achieving carbon neutrality

Career Paths

AI Materials Design Engineer: - Role: R&D at pharmaceutical, chemical, and materials manufacturers - Salary: 8-15 million yen (Japan), $120k-$250k (US) - Required Skills: Transformers, generative models, materials science

Researcher (Academia): - Role: PI at universities and research institutes - Research Areas: AI materials science, computational materials science - Competitiveness: Nature/Science-level publications are expected

Startup Founder: - Examples: Insilico Medicine (drug discovery AI), Citrine Informatics (materials AI) - Fundraising: Series A-C, hundreds of millions to billions of yen - Success Cases: IPO, acquisition by major companies

4.7 Summary

Key Points

Diffusion Models: Generate high-quality data from noise
Conditional Generation: Design materials by specifying target properties
SMILES Generation: Generate molecular structures with a Transformer
Inverse Design: Backward search from properties to structures
Industrial Applications: Practical use is advancing in drug discovery, batteries, and catalysts

Summary of the Series

Chapter 1: Transformer fundamentals, the attention mechanism Chapter 2: Materials-specific architectures (Matformer, ChemBERTa) Chapter 3: Pretrained models, transfer learning Chapter 4: Generative models, inverse design

Next Steps: 1. Gain experience through hands-on projects 2. Read the latest papers to update your knowledge 3. Participate in Kaggle competitions to test your skills 4. Join communities to exchange information

📝 Exercises

Problem 1: Conceptual Understanding

List three advantages of diffusion models compared to conventional generative models (VAE, GAN).

Sample Answer

1. **Training stability**: Mode collapse, as seen in GANs, is less likely to occur 2. **Sample quality**: Can generate high-quality and diverse samples 3. **Flexible conditioning**: Various conditions (properties, constraints) can be easily incorporated Additional: - **Interpretability**: The generation process is step-by-step and easy to understand - **Scalability**: Efficient training even on large-scale data

Problem 2: Implementation

Write code for conditional generation that produces materials satisfying multiple target properties (band gap, formation energy) simultaneously.

def multi_objective_generation(generator, target_bandgap, target_formation_energy, num_samples=10):
    """
    Generate materials via multi-objective optimization

    Args:
        generator: conditional generative model
        target_bandgap: target band gap (eV)
        target_formation_energy: target formation energy (eV/atom)
        num_samples: number to generate
    Returns:
        generated_materials: list of generated materials
    """
    # Implement here
    pass

Sample Answer

def multi_objective_generation(generator, target_bandgap, target_formation_energy, num_samples=10):
    # Create the condition
    condition = torch.tensor([[target_bandgap, target_formation_energy]])
    condition = condition.repeat(num_samples, 1)

    # Generate
    generated_materials = generator.generate_conditional(condition, input_dim=128)

    return generated_materials

# Usage example
target_bg = 2.0  # 2.0 eV
target_fe = -0.5  # -0.5 eV/atom

materials = multi_objective_generation(conditional_model, target_bg, target_fe, num_samples=20)
print(f"Generated {materials.shape[0]} materials")

Problem 3: Application

In materials inverse design, list five important criteria for evaluating generated candidate materials, and explain each.

Sample Answer

1. **Degree of target property achievement**: - How close the predicted properties are to the target values - Pareto optimality in the case of multiple properties 2. **Synthesizability**: - Whether it can be fabricated by known synthesis methods - Availability of precursors - Feasibility of synthesis conditions (temperature, pressure) 3. **Thermodynamic stability**: - Negative formation energy (stable phase) - Most stable compared to other crystal structures - Stability against decomposition reactions 4. **Chemical validity**: - Satisfies valence rules - Reasonable bond lengths and angles - Consistent with known chemical systems 5. **Cost and environmental impact**: - Price and abundance of constituent elements - Use of hazardous elements (Cd, Pb, etc.) - Recyclability

🎓 Congratulations on Completing the Series!

By completing this series, you have mastered how to apply Transformers and generative models in materials science, from fundamentals to applications.

Next Steps

Hands-on Projects: - Material property prediction with Materials Project data - Molecular generation with the QM9 dataset - Fine-tuning on your own data
Paper Implementation: - Read and implement the Matformer paper - Take on the latest generative model papers
Competitions: - Open Catalyst Challenge - Kaggle molecular prediction competitions
Community Participation: - Hugging Face Forum - Materials Project Community - Materials science conferences (MRS, APS)

🎯 Details of Materials-Specific Transformers

ChemBERTa: Chemistry BERT

from transformers import RobertaTokenizer, RobertaModel, RobertaConfig

class ChemBERTa(nn.Module):
    """
    ChemBERTa: RoBERTa trained on 10M SMILES strings

    Features:
    - Pretrained on PubChem, ZINC, ChEMBL
    - Dedicated SMILES tokenizer
    - Optimized for molecular property prediction
    """

    def __init__(self, pretrained_model="seyonec/ChemBERTa-zinc-base-v1"):
        super().__init__()
        self.tokenizer = RobertaTokenizer.from_pretrained(pretrained_model)
        self.model = RobertaModel.from_pretrained(pretrained_model)

    def forward(self, smiles_list):
        """
        Args:
            smiles_list: List of SMILES strings

        Returns:
            embeddings: (batch_size, 768) molecular embeddings
        """
        # Tokenize
        encoded = self.tokenizer(
            smiles_list,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors='pt'
        )

        # Forward
        outputs = self.model(**encoded)

        # [CLS] token embedding
        embeddings = outputs.last_hidden_state[:, 0, :]

        return embeddings

# Usage example
chemberta = ChemBERTa()

smiles_list = [
    "CC(C)Cc1ccc(cc1)C(C)C(=O)O",  # ibuprofen
    "CN1C=NC2=C1C(=O)N(C(=O)N2C)C"  # caffeine
]

embeddings = chemberta(smiles_list)
print(f"Molecular embeddings: {embeddings.shape}")  # (2, 768)

MatBERT: Materials Composition BERT

class MatBERT(nn.Module):
    """
    MatBERT: BERT for materials composition

    Pretraining:
    - Materials Project (500k+ compositions)
    - OQMD, AFLOW datasets
    - Masked composition prediction
    """

    def __init__(self, vocab_size=120, d_model=768, num_layers=12):
        super().__init__()

        config = BertConfig(
            vocab_size=vocab_size,
            hidden_size=d_model,
            num_hidden_layers=num_layers,
            num_attention_heads=12,
            intermediate_size=3072,
            max_position_embeddings=50  # maximum number of atoms in a material
        )

        self.bert = BertModel(config)

    def forward(self, composition_ids, attention_mask=None):
        """
        Args:
            composition_ids: (batch, seq_len) sequence of atomic numbers
                             e.g., [CLS] Fe Fe O O O [SEP]

        Returns:
            outputs: BERT outputs
        """
        outputs = self.bert(
            input_ids=composition_ids,
            attention_mask=attention_mask
        )

        return outputs

# Fine-tuning example: band gap prediction
class MatBERTForBandgap(nn.Module):
    def __init__(self, matbert):
        super().__init__()
        self.matbert = matbert

        # Prediction head
        self.regressor = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, 1)
        )

    def forward(self, composition_ids, attention_mask=None):
        outputs = self.matbert(composition_ids, attention_mask)
        cls_embedding = outputs.pooler_output

        bandgap = self.regressor(cls_embedding)
        return bandgap

MatGPT: Materials Generation GPT

from transformers import GPT2LMHeadModel, GPT2Config

class MatGPT(nn.Module):
    """
    MatGPT: GPT for materials composition generation

    Applications:
    - Generating novel material compositions
    - Conditional generation (target properties → composition)
    - Automating materials design
    """

    def __init__(self, vocab_size=120, d_model=768, num_layers=12):
        super().__init__()

        config = GPT2Config(
            vocab_size=vocab_size,
            n_positions=50,
            n_embd=d_model,
            n_layer=num_layers,
            n_head=12
        )

        self.gpt = GPT2LMHeadModel(config)

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

        Args:
            start_tokens: (1, start_len) start tokens
                         e.g., [CLS] Li
            max_length: maximum generation length
            temperature: sampling temperature (low → deterministic, high → random)
            top_k: Top-k sampling

        Returns:
            generated: (1, gen_len) the generated composition formula
        """
        self.eval()

        with torch.no_grad():
            generated = self.gpt.generate(
                start_tokens,
                max_length=max_length,
                temperature=temperature,
                top_k=top_k,
                do_sample=True,
                pad_token_id=0
            )

        return generated

# Conditional generation
class ConditionalMatGPT(nn.Module):
    """
    Conditional materials generation

    Conditions: band gap, formation energy, magnetic moment
    """

    def __init__(self, matgpt, condition_dim=3):
        super().__init__()
        self.matgpt = matgpt

        # Condition encoder
        self.condition_encoder = nn.Sequential(
            nn.Linear(condition_dim, 768),
            nn.ReLU(),
            nn.Linear(768, 768)
        )

    def forward(self, input_ids, conditions):
        """
        Args:
            input_ids: (batch, seq_len)
            conditions: (batch, condition_dim) target properties

        Returns:
            logits: (batch, seq_len, vocab_size)
        """
        # Embed the conditions
        condition_embed = self.condition_encoder(conditions)
        condition_embed = condition_embed.unsqueeze(1)  # (batch, 1, 768)

        # Input embeddings
        input_embeddings = self.matgpt.gpt.transformer.wte(input_ids)

        # Prepend the conditions
        embeddings = torch.cat([condition_embed, input_embeddings], dim=1)

        # GPT forward
        outputs = self.matgpt.gpt(inputs_embeds=embeddings)

        return outputs.logits

# Usage example
matgpt = MatGPT(vocab_size=120)
cond_matgpt = ConditionalMatGPT(matgpt, condition_dim=3)

# Target: band gap 2.5 eV, formation energy -1.0 eV, non-magnetic
target_conditions = torch.tensor([[2.5, -1.0, 0.0]])

# Generation start token
start = torch.tensor([[101]])  # [CLS]

# Generate
with torch.no_grad():
    logits = cond_matgpt(start, target_conditions)
    # Generate the next token by sampling
    probs = torch.softmax(logits[:, -1, :], dim=-1)
    next_token = torch.multinomial(probs, num_samples=1)

print(f"Next token: {next_token}")

🔬 Details of Transfer Learning Strategies

Strategy 1: Full Fine-tuning

def full_finetuning(pretrained_model, train_loader, val_loader):
    """
    Update all parameters

    When to apply:
    - Sufficient target data (thousands of samples or more)
    - Similar domain
    - When aiming for the highest accuracy
    """
    model = pretrained_model

    # Update all parameters
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)

    # Learning rate scheduler
    num_training_steps = len(train_loader) * epochs
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_training_steps)

    best_val_loss = float('inf')

    for epoch in range(epochs):
        model.train()
        for batch in train_loader:
            optimizer.zero_grad()

            outputs = model(**batch)
            loss = outputs.loss

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            scheduler.step()

        # Validation
        model.eval()
        val_loss = evaluate(model, val_loader)

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

    return model

Strategy 2: Adapter Tuning

class AdapterLayer(nn.Module):
    """
    Adapter: high performance with few parameters

    Idea: Insert an Adapter (a small bottleneck NN) into each Transformer layer
    """

    def __init__(self, d_model, adapter_size=64):
        super().__init__()

        self.adapter = nn.Sequential(
            nn.Linear(d_model, adapter_size),  # Down-project
            nn.ReLU(),
            nn.Linear(adapter_size, d_model)   # Up-project
        )

        # Residual connection
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, x):
        """
        Args:
            x: (batch, seq_len, d_model)

        Returns:
            x + adapter(x): Residual connection
        """
        residual = x
        x = self.layer_norm(x)
        x = self.adapter(x)
        return residual + x

class MatBERTWithAdapters(nn.Module):
    """
    MatBERT + Adapters

    Advantages:
    - Number of updated parameters: 1-2% of full model
    - Performance: 95-98% of full fine-tuning
    - Adapters can be swapped across multiple tasks
    """

    def __init__(self, pretrained_matbert, adapter_size=64):
        super().__init__()
        self.matbert = pretrained_matbert

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

        # Add an adapter to each Transformer layer
        self.adapters = nn.ModuleList([
            AdapterLayer(768, adapter_size)
            for _ in range(12)  # 12 layers
        ])

    def forward(self, input_ids, attention_mask=None):
        # MatBERT forward (frozen)
        outputs = self.matbert(input_ids, attention_mask, output_hidden_states=True)

        hidden_states = outputs.hidden_states

        # Apply an Adapter to each layer
        for i, adapter in enumerate(self.adapters):
            hidden_states[i+1] = adapter(hidden_states[i+1])

        # Output of the final layer
        final_hidden = hidden_states[-1]

        return final_hidden

# Usage example
pretrained = MatBERT(vocab_size=120)
model_with_adapters = MatBERTWithAdapters(pretrained, adapter_size=64)

# Train only the Adapters
trainable_params = sum(p.numel() for p in model_with_adapters.adapters.parameters())
total_params = sum(p.numel() for p in model_with_adapters.parameters())

print(f"Trainable params: {trainable_params} ({trainable_params/total_params*100:.2f}%)")

Strategy 3: LoRA (Low-Rank Adaptation)

class LoRALayer(nn.Module):
    """
    LoRA: Low-Rank Adaptation of Large Language Models

    Idea: Decompose the weight matrix update into low rank
    W_new = W_frozen + BA (B: m×r, A: r×n, r << m,n)
    """

    def __init__(self, in_features, out_features, rank=8):
        super().__init__()

        self.rank = rank

        # Low-rank matrices (trainable)
        self.lora_A = nn.Parameter(torch.randn(rank, in_features) / rank)
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))

    def forward(self, x, frozen_weight):
        """
        Args:
            x: (batch, seq_len, in_features)
            frozen_weight: (out_features, in_features) frozen weight

        Returns:
            output: (batch, seq_len, out_features)
        """
        # Frozen part
        output = torch.matmul(x, frozen_weight.T)

        # LoRA part
        lora_output = torch.matmul(x, self.lora_A.T)
        lora_output = torch.matmul(lora_output, self.lora_B.T)

        return output + lora_output

class MatBERTWithLoRA(nn.Module):
    """
    MatBERT + LoRA

    Advantages:
    - Number of updated parameters: 0.1-1% of full model
    - Performance: on par with full fine-tuning
    - LoRA can be merged at inference time (no slowdown)
    """

    def __init__(self, pretrained_matbert, rank=8):
        super().__init__()
        self.matbert = pretrained_matbert

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

        # Add LoRA to the attention QKV
        self.lora_layers = nn.ModuleDict()
        for layer_idx in range(12):
            self.lora_layers[f'layer_{layer_idx}_q'] = LoRALayer(768, 768, rank)
            self.lora_layers[f'layer_{layer_idx}_v'] = LoRALayer(768, 768, rank)

    def forward(self, input_ids, attention_mask=None):
        # Omitted: integrate LoRA into the attention computation
        pass

# Usage example
model_with_lora = MatBERTWithLoRA(pretrained, rank=8)

trainable_params = sum(p.numel() for p in model_with_lora.lora_layers.parameters())
total_params = sum(p.numel() for p in model_with_lora.parameters())

print(f"Trainable params: {trainable_params} ({trainable_params/total_params*100:.3f}%)")

🎓 Implementing Pretraining for Materials

Pretraining Task 1: Masked Atom Prediction

def pretrain_masked_atom_prediction(model, dataloader, epochs=100):
    """
    Masked Atom Prediction (MAP)

    Task: predict masked atoms
    Example: Fe [MASK] O → Fe Fe O (Fe2O3)
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss(ignore_index=0)  # Pad token

    model.train()

    for epoch in range(epochs):
        total_loss = 0

        for batch in dataloader:
            composition_ids = batch['composition_ids']  # (batch, seq_len)

            # Mask 15% of the atoms
            mask_prob = 0.15
            masked_composition, labels = mask_atoms(composition_ids, mask_prob)

            # Forward
            outputs = model(masked_composition)
            logits = outputs.logits  # (batch, seq_len, vocab_size)

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

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

            total_loss += loss.item()

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

    return model

def mask_atoms(composition_ids, mask_prob=0.15):
    """
    Mask atoms

    Strategy:
    - 80%: replace with [MASK]
    - 10%: replace with a random atom
    - 10%: leave unchanged
    """
    labels = composition_ids.clone()
    masked_composition = composition_ids.clone()

    # Select masking targets
    mask = torch.rand(composition_ids.shape) < mask_prob
    mask[:, 0] = False  # exclude [CLS]
    mask[:, -1] = False  # exclude [SEP]

    # 80% to [MASK]
    mask_token_mask = torch.rand(composition_ids.shape) < 0.8
    masked_composition[mask & mask_token_mask] = MASK_TOKEN_ID

    # 10% to random atoms
    random_mask = torch.rand(composition_ids.shape) < 0.1
    random_atoms = torch.randint(1, 119, composition_ids.shape)
    masked_composition[mask & random_mask] = random_atoms[mask & random_mask]

    # 10% left as is

    # Ignore labels at unmasked positions
    labels[~mask] = -100

    return masked_composition, labels

Pretraining Task 2: Contrastive Learning

class ContrastiveLearning(nn.Module):
    """
    Contrastive Learning for Materials

    Idea: place similar materials close together and different materials far apart
    """

    def __init__(self, matbert, temperature=0.07):
        super().__init__()
        self.matbert = matbert
        self.temperature = temperature

    def forward(self, compositions1, compositions2, labels):
        """
        Args:
            compositions1: (batch, seq_len) Augmented sample 1
            compositions2: (batch, seq_len) Augmented sample 2
            labels: (batch,) 1 if similar, 0 if dissimilar

        Returns:
            loss: Contrastive loss
        """
        # Embeddings
        emb1 = self.matbert(compositions1).pooler_output  # (batch, 768)
        emb2 = self.matbert(compositions2).pooler_output

        # Normalize
        emb1 = F.normalize(emb1, dim=-1)
        emb2 = F.normalize(emb2, dim=-1)

        # Cosine similarity
        similarity = torch.matmul(emb1, emb2.T) / self.temperature  # (batch, batch)

        # Loss: InfoNCE
        loss = F.cross_entropy(similarity, torch.arange(emb1.size(0), device=emb1.device))

        return loss

# Data augmentation
def augment_composition(composition_ids):
    """
    Data augmentation for composition formulas

    Methods:
    - Shuffle atom order (Fe2O3 → O3Fe2)
    - Substitute same-group elements (LiCoO2 → NaCoO2)
    """
    # Implementation omitted
    pass

✅ Chapter 4 Completion Checklist

Conceptual Understanding (10 items)

[ ] I understand the principles of diffusion models (forward/reverse process)
[ ] I understand the mechanism of conditional generation
[ ] I can explain the differences and advantages of SMILES vs SELFIES
[ ] I understand the materials inverse design workflow
[ ] I can explain the difference between ChemBERTa and MatBERT
[ ] I understand the differences and use cases of full fine-tuning / Adapter / LoRA
[ ] I understand the principle of Masked Atom Prediction
[ ] I understand the application of Contrastive Learning to materials science
[ ] I can explain the evaluation metrics of generative models (validity, diversity, novelty)
[ ] I understand the constraints in materials inverse design (synthesizability, stability)

Implementation Skills (15 items)

[ ] I can implement SimpleDiffusionModel
[ ] I can implement ConditionalDiffusionModel
[ ] I can implement SMILESGenerator
[ ] I can implement ConditionalSMILESGenerator
[ ] I can use ChemBERTa
[ ] I can implement MatBERT
[ ] I can implement MatGPT (including conditional generation)
[ ] I can implement AdapterLayer
[ ] I can implement LoRALayer
[ ] I can implement Masked Atom Prediction
[ ] I can implement Contrastive Learning
[ ] I can build a materials inverse design system
[ ] I can implement a synthesizability check
[ ] I can build a validation pipeline for generated materials
[ ] I can implement beam search

Debugging Skills (5 items)

[ ] I can evaluate the sample quality of a diffusion model
[ ] I can verify the validity of generated SMILES
[ ] I can evaluate the condition-achievement level of conditional generation
[ ] I can compare the performance of LoRA/Adapter with full fine-tuning
[ ] I can visualize and analyze the effects of pretraining

Application Ability (5 items)

[ ] I can apply generative models to novel materials discovery
[ ] I can implement multi-objective optimization (multiple properties)
[ ] I can build a loop combining generative and prediction models
[ ] I can embed domain knowledge (chemical rules, crystallography) into generation
[ ] I can prioritize experimental candidates

Data Processing (5 items)

[ ] I can convert between SMILES and SELFIES
[ ] I can check molecular validity (RDKit)
[ ] I can normalize composition formulas
[ ] I can implement data augmentation
[ ] I can post-process (filter) generated data

Evaluation Skills (5 items)

[ ] I can measure the validity of a generative model
[ ] I can quantitatively evaluate diversity
[ ] I can evaluate novelty
[ ] I can measure condition satisfaction
[ ] I can compute a synthesizability score

Theoretical Background (5 items)

[ ] I have read the diffusion model paper (Ho et al., 2020)
[ ] I have read the ChemBERTa/MatBERT papers
[ ] I have read the LoRA paper (Hu et al., 2021)
[ ] I have read at least one paper on materials inverse design
[ ] I understand the theory of generative models (VAE, GAN, Diffusion)

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%)

🔗 References

Papers

Ho et al. (2020) "Denoising Diffusion Probabilistic Models" arXiv:2006.11239
Chen et al. (2022) "Matformer: Nested Transformer for Elastic Inference"
Xie et al. (2021) "Crystal Diffusion Variational Autoencoder" arXiv:2110.06197
Stokes et al. (2020) "A Deep Learning Approach to Antibiotic Discovery" Nature
Hu et al. (2021) "LoRA: Low-Rank Adaptation of Large Language Models" arXiv:2106.09685
Chithrananda et al. (2020) "ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction" arXiv:2010.09885

Tools

Hugging Face Diffusers
RDKit - molecular processing
Materials Project API
SELFIES - molecular representation
PyMatGen - materials science

Next Series

Introduction to Reinforcement Learning: Applying reinforcement learning to materials exploration
Introduction to GNNs: Representing molecules and materials with graph neural networks
Introduction to Foundation Models: LLaMA, GPT-4, Claude for Materials

Author: Yusuke Hashimoto (Tohoku University) Last Updated: October 19, 2025 Series: Introduction to Transformers and Foundation Models (4 chapters, complete)

License: CC BY 4.0

Chapter 4: Generative Models and Inverse Design

📋 What You Will Learn in This Chapter

4.1 What Are Generative Models?

The Importance of Generative Models in Materials Science

4.2 Principles of Diffusion Models

What Are Diffusion Models?

Visual Understanding

Simple Implementation

4.3 Conditional Generation

Overview

Implementation: Conditional Diffusion

4.4 Molecular Generation: SMILES Generation

Overview

Transformer-based SMILES Generation

4.5 Materials Inverse Design Workflow

Complete Workflow

Implementation Example

4.6 Industrial Applications and Careers

Real-World Success Stories

1. Drug Discovery: Finding a Novel Antibiotic

2. Battery Materials: High Energy Density Electrolytes

3. Catalysts: CO₂ Reduction Catalysts

Career Paths

4.7 Summary

Key Points

Summary of the Series

📝 Exercises

Problem 1: Conceptual Understanding

Problem 2: Implementation

Problem 3: Application

🎓 Congratulations on Completing the Series!

Next Steps

🎯 Details of Materials-Specific Transformers

ChemBERTa: Chemistry BERT

MatBERT: Materials Composition BERT

MatGPT: Materials Generation GPT

🔬 Details of Transfer Learning Strategies

Strategy 1: Full Fine-tuning

Strategy 2: Adapter Tuning

Strategy 3: LoRA (Low-Rank Adaptation)

🎓 Implementing Pretraining for Materials

Pretraining Task 1: Masked Atom Prediction

Pretraining Task 2: Contrastive Learning

✅ Chapter 4 Completion Checklist

Conceptual Understanding (10 items)

Implementation Skills (15 items)

Debugging Skills (5 items)

Application Ability (5 items)

Data Processing (5 items)

Evaluation Skills (5 items)

Theoretical Background (5 items)

Completion Criteria

🔗 References

Papers

Tools

Next Series

Disclaimer