Transformer & Attention Mechanism

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaledDotProductAttention(nn.Module):
    """Scaled Dot-Product Attention."""
    def __init__(self, d_k):
        super().__init__()
        self.d_k = d_k
    
    def forward(self, Q, K, V, mask=None):
        # Q, K, V shape: (batch, heads, seq_len, d_k)
        
        # Step 1: Hitung attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        # scores shape: (batch, heads, seq_len, seq_len)
        
        # Step 2: Apply mask (untuk decoder - cegah melihat masa depan)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Step 3: Softmax
        attn_weights = F.softmax(scores, dim=-1)
        
        # Step 4: Weighted sum
        output = torch.matmul(attn_weights, V)
        # output shape: (batch, heads, seq_len, d_k)
        
        return output, attn_weights


class MultiHeadAttention(nn.Module):
    """Multi-Head Attention."""
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # dimensi per head
        
        # Linear projections
        self.W_q = nn.Linear(d_model, d_model)  # Query projection
        self.W_k = nn.Linear(d_model, d_model)  # Key projection
        self.W_v = nn.Linear(d_model, d_model)  # Value projection
        self.W_o = nn.Linear(d_model, d_model)  # Output projection
        
        self.attention = ScaledDotProductAttention(self.d_k)
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 1. Linear projection
        Q = self.W_q(Q)  # (batch, seq_len, d_model)
        K = self.W_k(K)
        V = self.W_v(V)
        
        # 2. Split into multiple heads
        # (batch, seq_len, d_model) → (batch, n_heads, seq_len, d_k)
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # 3. Scaled dot-product attention
        output, attn_weights = self.attention(Q, K, V, mask)
        
        # 4. Concatenate heads
        # (batch, n_heads, seq_len, d_k) → (batch, seq_len, d_model)
        output = output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )
        
        # 5. Final linear projection
        output = self.W_o(output)
        
        return output, attn_weights


class TransformerBlock(nn.Module):
    """Single Transformer Encoder Block."""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        
        # Multi-Head Self-Attention
        self.attention = MultiHeadAttention(d_model, n_heads)
        
        # Feed-Forward Network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        
        # Layer Normalization
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        
        # Dropout
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # 1. Self-Attention + Residual + LayerNorm
        attn_out, attn_weights = self.attention(x, x, x, mask)
        x = self.ln1(x + self.dropout(attn_out))
        
        # 2. FFN + Residual + LayerNorm
        ffn_out = self.ffn(x)
        x = self.ln2(x + self.dropout(ffn_out))
        
        return x, attn_weights


# =============================================
# TEST: Jalankan Multi-Head Attention
# =============================================
print("=" * 60)
print("TESTING MULTI-HEAD SELF-ATTENTION")
print("=" * 60)

# Hyperparameters
d_model = 512
n_heads = 8
d_ff = 2048
seq_len = 10
batch_size = 2

# Create model
block = TransformerBlock(d_model, n_heads, d_ff)

# Random input (simulating token embeddings)
x = torch.randn(batch_size, seq_len, d_model)

# Forward pass
output, attn_weights = block(x)

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

# Attention weight visualization
print(f"\nAttention weights (Head 0, Sample 0):")
print(f"Shape: {attn_weights[0, 0].shape} (seq_len × seq_len)")
print(f"Row sum: {attn_weights[0, 0].sum(dim=-1).tolist()}")  # ≈ 1.0 per row

# Parameter count
params = sum(p.numel() for p in block.parameters())
print(f"\nParameters in Transformer block: {params:,}")

# Multi-head attention shapes summary
print("\n=== Shape Summary ===")
print(f"d_model = {d_model}")
print(f"n_heads = {n_heads}")
print(f"d_k = d_v = {d_model // n_heads}")
print(f"Q, K, V per head: ({seq_len}, {d_model // n_heads})")
print(f"Attention matrix: ({seq_len}, {seq_len})")
print(f"Output per head: ({seq_len}, {d_model // n_heads})")
print(f"Concat output: ({seq_len}, {d_model})")