Transformer Architecture

1# RNN processes sequences step by step
2# Cannot parallelize over sequence length
3for t in range(sequence_length):
4    hidden_state = rnn_cell(hidden_state, input[t])

1import torch
2import torch.nn as nn
3import math
4
5class SelfAttention(nn.Module):
6    def __init__(self, embed_dim):
7        super().__init__()
8        self.embed_dim = embed_dim
9
10        # Three learned projection matrices
11        self.query = nn.Linear(embed_dim, embed_dim)
12        self.key = nn.Linear(embed_dim, embed_dim)
13        self.value = nn.Linear(embed_dim, embed_dim)
14
15    def forward(self, x):
16        # x shape: (batch_size, seq_len, embed_dim)
17
18        # Project inputs to queries, keys, values
19        Q = self.query(x)  # (batch, seq_len, embed_dim)
20        K = self.key(x)    # (batch, seq_len, embed_dim)
21        V = self.value(x)  # (batch, seq_len, embed_dim)
22
23        # Compute attention scores
24        # QK^T gives similarity between all pairs of tokens
25        scores = torch.matmul(Q, K.transpose(-2, -1))  # (batch, seq_len, seq_len)
26
27        # Scale by sqrt(d_k) to prevent vanishing gradients
28        scores = scores / math.sqrt(self.embed_dim)
29
30        # Apply softmax to get attention weights
31        attention_weights = torch.softmax(scores, dim=-1)
32
33        # Weighted sum of values
34        output = torch.matmul(attention_weights, V)  # (batch, seq_len, embed_dim)
35
36        return output, attention_weights

1score(q, k) = q · k / √d_k

1attention_weights = softmax(scores)

1output = attention_weights × V

1# Example: Attention in action
2sentence = ["The", "cat", "sat", "on", "the", "mat"]
3
4# For the token "cat", attention might look like:
5attention_for_cat = {
6    "The": 0.15,    # Some attention to determiner
7    "cat": 0.40,    # Strong self-attention
8    "sat": 0.25,    # Attention to action
9    "on": 0.05,     # Weak attention to preposition
10    "the": 0.05,    # Weak attention to determiner
11    "mat": 0.10     # Some attention to object
12}
13# Total = 1.0 (probability distribution)

1class MultiHeadAttention(nn.Module):
2    def __init__(self, embed_dim, num_heads):
3        super().__init__()
4        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
5
6        self.embed_dim = embed_dim
7        self.num_heads = num_heads
8        self.head_dim = embed_dim // num_heads
9
10        # Projections for all heads combined
11        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
12        self.output_proj = nn.Linear(embed_dim, embed_dim)
13
14    def forward(self, x, mask=None):
15        batch_size, seq_len, embed_dim = x.shape
16
17        # Project and split into Q, K, V
18        qkv = self.qkv_proj(x)  # (batch, seq_len, 3 * embed_dim)
19        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
20        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch, num_heads, seq_len, head_dim)
21        Q, K, V = qkv[0], qkv[1], qkv[2]
22
23        # Scaled dot-product attention for each head
24        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
25
26        # Apply mask if provided (for padding or causal attention)
27        if mask is not None:
28            scores = scores.masked_fill(mask == 0, float('-inf'))
29
30        attention_weights = torch.softmax(scores, dim=-1)
31        attention_output = torch.matmul(attention_weights, V)
32
33        # Concatenate heads and project
34        attention_output = attention_output.permute(0, 2, 1, 3).contiguous()
35        attention_output = attention_output.reshape(batch_size, seq_len, embed_dim)
36        output = self.output_proj(attention_output)
37
38        return output, attention_weights

1class PositionalEncoding(nn.Module):
2    def __init__(self, embed_dim, max_len=5000):
3        super().__init__()
4
5        # Create positional encoding matrix
6        pe = torch.zeros(max_len, embed_dim)
7        position = torch.arange(0, max_len).unsqueeze(1).float()
8
9        # Compute the positional encodings
10        div_term = torch.exp(torch.arange(0, embed_dim, 2).float() *
11                             -(math.log(10000.0) / embed_dim))
12
13        # Apply sine to even indices
14        pe[:, 0::2] = torch.sin(position * div_term)
15        # Apply cosine to odd indices
16        pe[:, 1::2] = torch.cos(position * div_term)
17
18        pe = pe.unsqueeze(0)  # (1, max_len, embed_dim)
19        self.register_buffer('pe', pe)
20
21    def forward(self, x):
22        # x shape: (batch_size, seq_len, embed_dim)
23        seq_len = x.size(1)
24        # Add positional encoding to input embeddings
25        x = x + self.pe[:, :seq_len, :]
26        return x

1PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
2PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

1class LearnedPositionalEncoding(nn.Module):
2    def __init__(self, embed_dim, max_len=5000):
3        super().__init__()
4        self.pos_embedding = nn.Embedding(max_len, embed_dim)
5
6    def forward(self, x):
7        batch_size, seq_len, _ = x.shape
8        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)
9        pos_encodings = self.pos_embedding(positions)  # (1, seq_len, embed_dim)
10        return x + pos_encodings

1class FeedForward(nn.Module):
2    def __init__(self, embed_dim, ff_dim, dropout=0.1):
3        super().__init__()
4        self.net = nn.Sequential(
5            nn.Linear(embed_dim, ff_dim),
6            nn.GELU(),  # Gaussian Error Linear Unit
7            nn.Dropout(dropout),
8            nn.Linear(ff_dim, embed_dim),
9            nn.Dropout(dropout)
10        )
11
12    def forward(self, x):
13        return self.net(x)

1class TransformerEncoderBlock(nn.Module):
2    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
3        super().__init__()
4
5        # Multi-head self-attention
6        self.attention = MultiHeadAttention(embed_dim, num_heads)
7
8        # Feed-forward network
9        self.feed_forward = FeedForward(embed_dim, ff_dim, dropout)
10
11        # Layer normalization
12        self.norm1 = nn.LayerNorm(embed_dim)
13        self.norm2 = nn.LayerNorm(embed_dim)
14
15        # Dropout
16        self.dropout = nn.Dropout(dropout)
17
18    def forward(self, x, mask=None):
19        # Self-attention with residual connection and layer norm
20        attention_output, _ = self.attention(x, mask)
21        x = self.norm1(x + self.dropout(attention_output))
22
23        # Feed-forward with residual connection and layer norm
24        ff_output = self.feed_forward(x)
25        x = self.norm2(x + self.dropout(ff_output))
26
27        return x

1class TransformerEncoder(nn.Module):
2    def __init__(self, vocab_size, embed_dim, num_heads, ff_dim,
3                 num_layers, max_len=5000, dropout=0.1):
4        super().__init__()
5
6        # Embedding layers
7        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
8        self.positional_encoding = PositionalEncoding(embed_dim, max_len)
9
10        # Stack of encoder blocks
11        self.layers = nn.ModuleList([
12            TransformerEncoderBlock(embed_dim, num_heads, ff_dim, dropout)
13            for _ in range(num_layers)
14        ])
15
16        self.dropout = nn.Dropout(dropout)
17
18    def forward(self, x, mask=None):
19        # x shape: (batch_size, seq_len) - token indices
20
21        # Embed tokens and add positional encoding
22        x = self.token_embedding(x) * math.sqrt(self.token_embedding.embedding_dim)
23        x = self.positional_encoding(x)
24        x = self.dropout(x)
25
26        # Pass through encoder blocks
27        for layer in self.layers:
28            x = layer(x, mask)
29
30        return x  # (batch_size, seq_len, embed_dim)

1class TransformerDecoderBlock(nn.Module):
2    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
3        super().__init__()
4
5        # Masked self-attention (for target sequence)
6        self.self_attention = MultiHeadAttention(embed_dim, num_heads)
7
8        # Cross-attention (attending to encoder output)
9        self.cross_attention = MultiHeadAttention(embed_dim, num_heads)
10
11        # Feed-forward network
12        self.feed_forward = FeedForward(embed_dim, ff_dim, dropout)
13
14        # Layer normalization
15        self.norm1 = nn.LayerNorm(embed_dim)
16        self.norm2 = nn.LayerNorm(embed_dim)
17        self.norm3 = nn.LayerNorm(embed_dim)
18
19        self.dropout = nn.Dropout(dropout)
20
21    def forward(self, x, encoder_output, src_mask=None, tgt_mask=None):
22        # Masked self-attention
23        self_attn_output, _ = self.self_attention(x, tgt_mask)
24        x = self.norm1(x + self.dropout(self_attn_output))
25
26        # Cross-attention (query from decoder, key/value from encoder)
27        cross_attn_output, _ = self.cross_attention(
28            x, encoder_output, encoder_output, src_mask
29        )
30        x = self.norm2(x + self.dropout(cross_attn_output))
31
32        # Feed-forward
33        ff_output = self.feed_forward(x)
34        x = self.norm3(x + self.dropout(ff_output))
35
36        return x

1def create_causal_mask(seq_len, device):
2    """Create mask that prevents attending to future positions"""
3    mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1)
4    mask = mask.masked_fill(mask == 1, float('-inf'))
5    return mask
6
7# During generation: token i can only attend to tokens 0..i
8# Example for seq_len=4:
9# [[0, -inf, -inf, -inf],
10#  [0,    0, -inf, -inf],
11#  [0,    0,    0, -inf],
12#  [0,    0,    0,    0]]

1import torch
2import torch.nn as nn
3from torch.utils.data import Dataset, DataLoader
4
5class TextClassificationTransformer(nn.Module):
6    def __init__(self, vocab_size, num_classes, embed_dim=256,
7                 num_heads=8, ff_dim=1024, num_layers=4,
8                 max_len=512, dropout=0.1):
9        super().__init__()
10
11        # Transformer encoder
12        self.encoder = TransformerEncoder(
13            vocab_size, embed_dim, num_heads, ff_dim,
14            num_layers, max_len, dropout
15        )
16
17        # Classification head
18        self.classifier = nn.Sequential(
19            nn.Linear(embed_dim, embed_dim),
20            nn.ReLU(),
21            nn.Dropout(dropout),
22            nn.Linear(embed_dim, num_classes)
23        )
24
25    def forward(self, x, mask=None):
26        # x: (batch_size, seq_len)
27        encoder_output = self.encoder(x, mask)
28
29        # Pool encoder output (use [CLS] token or mean pooling)
30        # Here we use mean pooling
31        pooled = encoder_output.mean(dim=1)  # (batch_size, embed_dim)
32
33        # Classify
34        logits = self.classifier(pooled)  # (batch_size, num_classes)
35
36        return logits
37
38# Training loop
39def train_transformer_classifier(model, train_loader, val_loader,
40                                 num_epochs=10, lr=0.001):
41    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
42    model = model.to(device)
43
44    criterion = nn.CrossEntropyLoss()
45    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
46    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)
47
48    for epoch in range(num_epochs):
49        # Training
50        model.train()
51        train_loss = 0.0
52        correct = 0
53        total = 0
54
55        for batch in train_loader:
56            inputs, labels = batch['input_ids'].to(device), batch['labels'].to(device)
57            mask = batch.get('attention_mask', None)
58
59            optimizer.zero_grad()
60            outputs = model(inputs, mask)
61            loss = criterion(outputs, labels)
62            loss.backward()
63
64            # Gradient clipping
65            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
66
67            optimizer.step()
68
69            train_loss += loss.item()
70            _, predicted = outputs.max(1)
71            total += labels.size(0)
72            correct += predicted.eq(labels).sum().item()
73
74        train_acc = 100. * correct / total
75
76        # Validation
77        model.eval()
78        val_loss = 0.0
79        correct = 0
80        total = 0
81
82        with torch.no_grad():
83            for batch in val_loader:
84                inputs, labels = batch['input_ids'].to(device), batch['labels'].to(device)
85                mask = batch.get('attention_mask', None)
86
87                outputs = model(inputs, mask)
88                loss = criterion(outputs, labels)
89
90                val_loss += loss.item()
91                _, predicted = outputs.max(1)
92                total += labels.size(0)
93                correct += predicted.eq(labels).sum().item()
94
95        val_acc = 100. * correct / total
96
97        scheduler.step()
98
99        print(f"Epoch {epoch+1}/{num_epochs}")
100        print(f"Train Loss: {train_loss/len(train_loader):.4f}, Train Acc: {train_acc:.2f}%")
101        print(f"Val Loss: {val_loss/len(val_loader):.4f}, Val Acc: {val_acc:.2f}%")
102        print()
103
104# Example usage
105vocab_size = 10000
106num_classes = 5
107model = TextClassificationTransformer(vocab_size, num_classes)
108
109# Assuming you have train_loader and val_loader prepared
110# train_transformer_classifier(model, train_loader, val_loader)

1def get_lr(step, d_model, warmup_steps=4000):
2    return (d_model ** -0.5) * min(step ** -0.5, step * warmup_steps ** -1.5)

1criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

1torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

1from torch.cuda.amp import autocast, GradScaler
2
3scaler = GradScaler()
4
5with autocast():
6    outputs = model(inputs)
7    loss = criterion(outputs, labels)
8
9scaler.scale(loss).backward()
10scaler.step(optimizer)
11scaler.update()