Chapter

πŸ“– Reading Time: 20-25 min πŸ“Š Difficulty: Beginner πŸ’» Code Examples: 0 πŸ“ Exercises: 0

Chapter 1: The Transformer Revolution and Materials Science

Gain a high-level understanding of the key ideas behind Self-Attention and why it can be applied to materials. Learn the patterns for transferring techniques from existing fields.

πŸ’‘ Note: Attention is a mechanism that "shines a spotlight where it is needed." It can also be applied to the sequential and graph information of materials.

Study time: 20-30 min | Difficulty: Intermediate

πŸ“‹ What You Will Learn in This Chapter


1.1 Why the Transformer Sparked a Revolution

Limitations of Conventional RNNs/CNNs

Problems with RNNs (Recurrent Neural Networks): - Vanishing/exploding gradients over long sequences - Difficult to parallelize (sequential processing required) - Difficulty capturing long-term dependencies

Problems with CNNs (Convolutional Neural Networks): - Can only capture local features - Deep layers are needed to capture long-range relationships - Not well suited to irregular structures such as molecules and materials

The Innovation of the Transformer

Introduced in 2017 in the paper "Attention Is All You Need": - βœ… Directly models the relationships between all elements (Attention mechanism) - βœ… Fully parallelizable (maximizes GPU utilization) - βœ… Efficiently captures long-range dependencies - βœ… Interpretability (visualize important parts via Attention weights)

flowchart LR A[Input Sequence] --> B[Self-Attention] B --> C[Feed Forward] C --> D[Output] B -.-> E[Compute relationships between all elements] E -.-> B style B fill:#e1f5ff

1.2 Principles of the Attention Mechanism

What Is the Attention Mechanism?

Basic concept: A mechanism that learns "where to focus" within the input

Formula: $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

Intuitive Understanding

Library analogy: - Query: "I'm looking for a book on machine learning" - Key: The table of contents / title of each book - Value: The actual content of the book - Attention: "Focus" on and read the most relevant books

Python Implementation: Basic Attention

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Scaled Dot-Product Attention

    Args:
        Q: Query (batch_size, seq_len, d_k)
        K: Key (batch_size, seq_len, d_k)
        V: Value (batch_size, seq_len, d_v)
        mask: Mask (optional)
    """
    d_k = Q.size(-1)

    # 1. Compute the dot product of Q and K (similarity)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
    # scores shape: (batch_size, seq_len_q, seq_len_k)

    # 2. Apply mask (if needed)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    # 3. Normalize with Softmax (Attention weights)
    attention_weights = F.softmax(scores, dim=-1)

    # 4. Weighted sum of Value using the Attention weights
    output = torch.matmul(attention_weights, V)

    return output, attention_weights

# Usage example
batch_size, seq_len, d_model = 2, 5, 64
Q = torch.randn(batch_size, seq_len, d_model)
K = torch.randn(batch_size, seq_len, d_model)
V = torch.randn(batch_size, seq_len, d_model)

output, attn_weights = scaled_dot_product_attention(Q, K, V)
print(f"Output shape: {output.shape}")  # (2, 5, 64)
print(f"Attention weights shape: {attn_weights.shape}")  # (2, 5, 5)

Visualizing Attention Weights

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(attention_weights, tokens=None):
    """
    Visualize Attention weights as a heatmap

    Args:
        attention_weights: Attention weights of shape (seq_len, seq_len)
        tokens: List of tokens (optional)
    """
    plt.figure(figsize=(8, 6))

    # Get the Attention weights of the first head of the first sample
    attn = attention_weights[0].detach().numpy()

    sns.heatmap(attn, cmap='YlOrRd', cbar=True, square=True,
                xticklabels=tokens if tokens else range(attn.shape[0]),
                yticklabels=tokens if tokens else range(attn.shape[0]))

    plt.xlabel('Key (target)')
    plt.ylabel('Query (source)')
    plt.title('Attention Weights')
    plt.tight_layout()
    plt.show()

# Usage example
tokens = ['H', 'C', 'C', 'O', 'H']
visualize_attention(attn_weights, tokens)

1.3 Self-Attention: The Self-Attention Mechanism

What Is Self-Attention?

Definition: A mechanism that applies Attention to the input sequence itself

Characteristics: - Query, Key, and Value are all generated from the same input - Directly models the relationship between any two elements in the sequence - Focuses on highly relevant elements regardless of position

Example of Self-Attention in Molecules

Example of methanol (CH₃OH):

# Atoms: C, H, H, H, O, H
# Self-Attention learns the relationship of each atom with the other atoms
# Example: the O atom has a strong relationship with the C atom

Self-Attention Implementation

import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, d_model):
        super(SelfAttention, self).__init__()
        self.d_model = d_model

        # Linear transformations to Q, K, V
        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)

    def forward(self, x):
        """
        Args:
            x: (batch_size, seq_len, d_model)
        """
        # Generate Q, K, V
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        # Scaled Dot-Product Attention
        output, attn_weights = scaled_dot_product_attention(Q, K, V)

        return output, attn_weights

# Usage example
d_model = 128
seq_len = 10
batch_size = 4

self_attn = SelfAttention(d_model)
x = torch.randn(batch_size, seq_len, d_model)
output, attn_weights = self_attn(x)

print(f"Input shape: {x.shape}")          # (4, 10, 128)
print(f"Output shape: {output.shape}")    # (4, 10, 128)
print(f"Attention shape: {attn_weights.shape}")  # (4, 10, 10)

1.4 Multi-Head Attention: The Multi-Head Attention Mechanism

Why Multi-Head Is Needed

Limitations of a single Attention head: - Can only view relationships from a single perspective - Cannot fully capture complex relationships (chemical bonds, conformations, etc.)

Advantages of Multi-Head Attention: - Learns relationships from multiple different perspectives - Each head captures different features (bonds, distances, angles, etc.) - Enables richer representations

Formula

$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O $$

where $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

Implementation

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Q, K, V transformations
        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)

        # Output transformation
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size = x.size(0)

        # 1. Generate Q, K, V and split by head
        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)
        # Shape: (batch_size, num_heads, seq_len, d_k)

        # 2. Scaled Dot-Product Attention for each head
        output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        # output: (batch_size, num_heads, seq_len, d_k)

        # 3. Concatenate the heads
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        # Shape: (batch_size, seq_len, d_model)

        # 4. Output transformation
        output = self.W_o(output)

        return output, attn_weights

# Usage example
d_model = 512
num_heads = 8
seq_len = 20
batch_size = 2

mha = MultiHeadAttention(d_model, num_heads)
x = torch.randn(batch_size, seq_len, d_model)
output, attn_weights = mha(x)

print(f"Input shape: {x.shape}")          # (2, 20, 512)
print(f"Output shape: {output.shape}")    # (2, 20, 512)
print(f"Attention shape: {attn_weights.shape}")  # (2, 8, 20, 20)

1.5 Positional Encoding: Embedding Positional Information

Why It Is Needed

Problem: Self-Attention has no notion of order - Cannot distinguish "H-C-O" from "O-C-H" - In molecules and materials, the arrangement order of atoms is important

Solution: Add positional information via Positional Encoding

Formula

$$ PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$

$$ PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$

Implementation

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()

        # Create the positional encoding matrix
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Args:
            x: (batch_size, seq_len, d_model)
        """
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len, :]
        return x

# Usage example and visualization
d_model = 128
max_len = 100

pos_enc = PositionalEncoding(d_model, max_len)

# Dummy input
x = torch.zeros(1, 50, d_model)
output = pos_enc(x)

# Visualization
plt.figure(figsize=(12, 4))
plt.plot(pos_enc.pe[0, :50, :8].numpy())
plt.xlabel('Position')
plt.ylabel('Encoding Value')
plt.title('Positional Encoding (first 8 dimensions)')
plt.legend([f'dim {i}' for i in range(8)])
plt.tight_layout()
plt.show()

1.6 The Transformer and BERT/GPT

Overall Transformer Architecture

flowchart TB subgraph Encoder E1[Input Embedding] --> E2[Positional Encoding] E2 --> E3[Multi-Head Attention] E3 --> E4[Add and Norm] E4 --> E5[Feed Forward] E5 --> E6[Add and Norm] end subgraph Decoder D1[Output Embedding] --> D2[Positional Encoding] D2 --> D3[Masked Multi-Head Attention] D3 --> D4[Add and Norm] D4 --> D5[Multi-Head Attention] D5 --> D6[Add and Norm] D6 --> D7[Feed Forward] D7 --> D8[Add and Norm] end E6 -.-> D5 D8 --> O[Output] style E3 fill:#e1f5ff style D3 fill:#ffe1e1 style D5 fill:#e1ffe1

BERT (Bidirectional Encoder Representations from Transformers)

Characteristics: - Uses the Encoder only - Understands context bidirectionally - Pre-training tasks: Masked Language Model (MLM) + Next Sentence Prediction (NSP) - Use cases: Classification, feature extraction, question answering

Applications in materials science: - MatBERT: Property prediction from material composition formulas - ChemBERTa: Learning molecular SMILES representations

GPT (Generative Pre-trained Transformer)

Characteristics: - Uses the Decoder only - Generates text unidirectionally (left to right) - Pre-training task: Next-word prediction - Use cases: Text generation, dialogue, creative tasks

Applications in materials science: - Molecule generation (SMILES string generation) - Automatic generation of material descriptions - Proposing synthesis routes


1.7 Success Cases in Materials Science

1. ChemBERTa: Molecular Representation Learning

Overview: Learning SMILES with BERT

# Molecular SMILES: CC(C)Cc1ccc(cc1)C(C)C(=O)O (ibuprofen)
# Convert to an embedding vector with ChemBERTa -> property prediction

Results: - High-accuracy prediction on small datasets - Shorter development time through transfer learning - Interpretability (visualize important parts via Attention)

2. Matformer: Material Property Prediction

Overview: Processing crystal structures with a Transformer

# Input: atomic coordinates, atomic numbers, lattice constants
# Output: band gap, formation energy

Results: - High accuracy on Materials Project data - Performance equal to or better than GNNs - Good computational efficiency

3. Molecule Generation with Diffusion Models

Overview: Generating novel molecules with conditional diffusion models

# Conditions: solubility > 5 mg/mL, LogP < 3
# Generated: molecular SMILES that satisfy the conditions

Results: - Discovery of promising candidate molecules in drug discovery - Higher diversity than conventional methods - Also considers synthetic accessibility


1.8 Summary

Key Points

  1. Attention mechanism: Directly models the relationships between any elements in a sequence
  2. Self-Attention: Attention applied to the input sequence itself
  3. Multi-Head Attention: Learns relationships from multiple perspectives
  4. Positional Encoding: Embeds positional information
  5. BERT/GPT: Representative Transformer-based pre-trained models
  6. Materials science applications: Molecular/material representation learning, property prediction, generative models

Preparation for the Next Chapter

In Chapter 2, we will study in detail the Transformer architectures specialized for materials science (Matformer, CrystalFormer, ChemBERTa).


πŸ“ Exercises

Exercise 1: Fundamental Understanding (Concept)

Explain the roles of Query, Key, and Value in the Attention mechanism using an analogy other than the library one.

Sample Answer **Search engine analogy**: - **Query**: The search keywords entered by the user - **Key**: The metadata of each web page (title, summary) - **Value**: The actual content of the web page - **Attention**: Rank pages that are highly relevant to the search keywords higher **Molecular analogy**: - **Query**: Which atoms a given atom "wants to interact with" - **Key**: The features of each atom (atomic number, charge, position) - **Value**: Detailed information about each atom - **Attention**: Represents the strength of chemical bonds and interactions

Exercise 2: Implementation (Coding)

Fill in the blanks in the following code to implement Simple Attention (no scaling, no mask).

def simple_attention(Q, K, V):
    """
    A simple Attention mechanism

    Args:
        Q: Query (batch_size, seq_len, d_k)
        K: Key (batch_size, seq_len, d_k)
        V: Value (batch_size, seq_len, d_v)

    Returns:
        output: (batch_size, seq_len, d_v)
        attention_weights: (batch_size, seq_len, seq_len)
    """
    # 1. Compute the dot product of Q and K
    scores = torch.matmul(______, ______.transpose(-2, -1))

    # 2. Normalize with Softmax
    attention_weights = F.softmax(______, dim=-1)

    # 3. Weighted sum of Value using the Attention weights
    output = torch.matmul(______, ______)

    return output, attention_weights
Sample Answer
def simple_attention(Q, K, V):
    # 1. Compute the dot product of Q and K
    scores = torch.matmul(Q, K.transpose(-2, -1))

    # 2. Normalize with Softmax
    attention_weights = F.softmax(scores, dim=-1)

    # 3. Weighted sum of Value using the Attention weights
    output = torch.matmul(attention_weights, V)

    return output, attention_weights

Exercise 3: Application (Discussion)

Consider Self-Attention in the molecule "CCO" (ethanol). Answer the following questions:

  1. Between which atoms do you expect the Attention weights to be highest?
  2. Explain the reason from a chemical perspective.
  3. In Multi-Head Attention, what different kinds of information might each head capture?
Sample Answer 1. **Highest Attention weights**: C-C bond, C-O bond 2. **Chemical reasons**: - There is a strong interaction due to covalent bonding - Electron density is high due to shared electrons - The O atom forms a polar bond with the C atom 3. **Examples of information captured by each head**: - **Head 1**: Chemical bonds (primary bonds) - **Head 2**: Secondary bonds (C-C-O angle) - **Head 3**: Electron density distribution - **Head 4**: Atom type (C vs O vs H) - **Head 5**: Conformational information - **Head 6**: Polar interactions Because each head understands the molecule from a different perspective, a richer representation becomes possible.

πŸ“Š Data Licenses and Terms of Use

Language Datasets

Materials Science Datasets

Best Practices for License Compliance

# Example citation when using a dataset
"""
This work uses data from Materials Project (materialsproject.org),
which is released under CC BY 4.0 license.

Citation:
Jain, A., Ong, S. P., Hautier, G., Chen, W., Richards, W. D.,
Dacek, S., ... & Persson, K. A. (2013).
Commentary: The Materials Project: A materials genome approach
to accelerating materials innovation. APL materials, 1(1).
"""

πŸ”§ Code Reproducibility Guidelines

Environment Setup

# requirements.txt
torch==2.0.1
transformers==4.30.2
numpy==1.24.3
matplotlib==3.7.1
seaborn==0.12.2

# Recommended: full environment reproduction
# conda env export > environment.yml

Settings for Reproducibility

import torch
import numpy as np
import random

def set_seed(seed=42):
    """
    Guarantee full reproducibility

    Args:
        seed: Random seed
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # Make CuDNN behavior deterministic (reduces speed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# Usage example
set_seed(42)

# Record version information
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")

Making Transformer Parameters Explicit

# Manage experiment settings with a dictionary
config = {
    'model': {
        'd_model': 512,
        'num_heads': 8,
        'num_layers': 6,
        'd_ff': 2048,
        'dropout': 0.1,
        'max_seq_len': 512
    },
    'training': {
        'batch_size': 32,
        'learning_rate': 1e-4,
        'num_epochs': 100,
        'warmup_steps': 4000,
        'optimizer': 'Adam',
        'weight_decay': 0.01
    },
    'data': {
        'train_split': 0.8,
        'val_split': 0.1,
        'test_split': 0.1,
        'tokenizer': 'BPE',
        'vocab_size': 50000
    },
    'seed': 42
}

# Save the settings
import json
with open('experiment_config.json', 'w') as f:
    json.dump(config, f, indent=2)

Detailed Configuration of Attention Parameters

class ReproducibleMultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1, bias=True):
        """
        Multi-Head Attention with an emphasis on reproducibility

        Args:
            d_model: Model dimension (512 recommended)
            num_heads: Number of heads (8 recommended, must divide d_model)
            dropout: Dropout rate (0.1 recommended)
            bias: Whether to use bias in the linear layers
        """
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Make the initialization method explicit
        self.W_q = nn.Linear(d_model, d_model, bias=bias)
        self.W_k = nn.Linear(d_model, d_model, bias=bias)
        self.W_v = nn.Linear(d_model, d_model, bias=bias)
        self.W_o = nn.Linear(d_model, d_model, bias=bias)

        # Xavier initialization
        for module in [self.W_q, self.W_k, self.W_v, self.W_o]:
            nn.init.xavier_uniform_(module.weight)
            if bias:
                nn.init.zeros_(module.bias)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # The implementation is the same as the code above
        pass

⚠️ Practical Pitfalls and Countermeasures

1. Attention Mask Errors

Problem: Future information leaks due to incorrect mask application

# ❌ Wrong: the mask is not applied correctly
def wrong_attention(Q, K, V, mask):
    scores = torch.matmul(Q, K.transpose(-2, -1))
    # The mask values are inverted
    scores = scores.masked_fill(mask == 1, -1e9)  # Wrong!
    return F.softmax(scores, dim=-1)

# βœ… Correct: the mask ignores positions where it is 0
def correct_attention(Q, K, V, mask):
    scores = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(Q.size(-1))
    if mask is not None:
        # Set positions where mask==0 to -inf
        scores = scores.masked_fill(mask == 0, float('-inf'))
    return F.softmax(scores, dim=-1)

# Debugging method
print("Attention scores before mask:", scores)
print("Mask:", mask)
print("Attention scores after mask:", scores.masked_fill(mask == 0, float('-inf')))

2. Positional Encoding Implementation Errors

Problem: The dimensions of sin and cos are swapped

# ❌ Wrong: the dimension assignment is reversed
class WrongPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

        pe[:, 1::2] = torch.sin(position * div_term)  # Wrong!
        pe[:, 0::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

# βœ… Correct: sin for even dimensions, cos for odd dimensions
class CorrectPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)  # Correct
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

3. Memory Overflow

Problem: OOM (Out of Memory) with long sequences

# ❌ Problem: processing the entire sequence at once
def memory_intensive_attention(x):
    # x: (batch=64, seq_len=10000, d_model=512)
    # Attention matrix: (64, 10000, 10000) = about 24GB!
    return multi_head_attention(x)

# βœ… Solution 1: Gradient checkpointing
from torch.utils.checkpoint import checkpoint

def memory_efficient_attention(x):
    return checkpoint(multi_head_attention, x)

# βœ… Solution 2: Split the sequence
def chunked_attention(x, chunk_size=512):
    batch, seq_len, d_model = x.shape
    outputs = []

    for i in range(0, seq_len, chunk_size):
        chunk = x[:, i:i+chunk_size, :]
        output = multi_head_attention(chunk)
        outputs.append(output)

    return torch.cat(outputs, dim=1)

# βœ… Solution 3: Sparse Attention (for long-range tasks)
# Use libraries such as Longformer or BigBird

4. Tokenization Problems (Specific to Materials Science)

Problem: SMILES parentheses and branches are not processed correctly

# ❌ Wrong: simple character splitting
def wrong_smiles_tokenize(smiles):
    return list(smiles)  # "C(C)O" -> ['C', '(', 'C', ')', 'O']

# βœ… Correct: use a SMILES tokenizer
from transformers import RobertaTokenizer

tokenizer = RobertaTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")

# Or a regular-expression-based approach
import re
def correct_smiles_tokenize(smiles):
    pattern = r'(\[[^\]]+\]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])'
    return re.findall(pattern, smiles)

# Test
smiles = "CC(C)Cc1ccc(cc1)C(C)C(=O)O"  # ibuprofen
tokens = correct_smiles_tokenize(smiles)
print(f"Tokens: {tokens}")

5. Numerical Instability

Problem: Overflow/underflow in Softmax

# ❌ Problem: exp() overflows for large scores
def unstable_softmax(x):
    return torch.exp(x) / torch.sum(torch.exp(x), dim=-1, keepdim=True)

# βœ… Solution: numerically stable softmax (PyTorch implements this internally)
def stable_softmax(x):
    # Subtract the maximum value to ensure numerical stability
    x_max = torch.max(x, dim=-1, keepdim=True)[0]
    exp_x = torch.exp(x - x_max)
    return exp_x / torch.sum(exp_x, dim=-1, keepdim=True)

# Using PyTorch's F.softmax() is the safest
import torch.nn.functional as F
safe_output = F.softmax(x, dim=-1)

βœ… Chapter 1 Completion Checklist

Conceptual Understanding (10 items)

Mathematical Understanding (5 items)

Implementation Skills (15 items)

Debugging Skills (5 items)

Application Ability (5 items)

Theoretical Background (5 items)

Reproducibility (5 items)

Completion Criteria


πŸ”— References

Papers

Tutorials

Next Chapter

Chapter 2: Transformer Architectures for Materials where we study materials-science-specialized models such as Matformer and ChemBERTa.


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

Disclaimer