Implementing Attention from Scratch
Single-Head Attention
``python
import torch
import torch.nn as nn
import math
class Attention(nn.Module): def __init__(self, d_model): super().__init__() self.d_model = d_model 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, mask=None): Q = self.W_q(x) K = self.W_k(x) V = self.W_v(x) # Compute attention scores scores = torch.matmul(Q, K.transpose(-2, -1)) scores = scores / math.sqrt(self.d_model) # Apply mask if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # Softmax and weighted sum attn = torch.softmax(scores, dim=-1) return torch.matmul(attn, V) ``
Multi-Head Attention
Split into h heads, attend, concatenate, project.
Complete Transformer Block
Combine attention with FFN, residuals, and LayerNorm.