Transformer Architecture

Core Concepts

The Transformer (Vaswani et al., 2017, "Attention Is All You Need") replaced recurrence with self-attention, enabling:

• Parallel processing of all positions（vs RNN's sequential）
• Direct long-range dependencies between any two positions（vs RNN's gradient chain）

Scaled Dot-Product Attention

The fundamental building block:

Three matrices from the same input（self-attention）or different inputs（cross-attention）:

• Query ($Q$): 「我在找什麼？」
• Key ($K$): 「我有什麼？」
• Value ($V$): 「如果 match，我提供什麼？」

直覺：像 database lookup — Query 和每個 Key 做 dot product 得到 relevance score → softmax → weighted sum of Values。

Why Scale by $\sqrt{d_k}$?

When $d_k$ is large, dot products $q \cdot k$ grow in magnitude（variance ≈ $d_k$）。Large dot products push softmax into saturation regions → gradients ≈ 0 → training unstable。

Dividing by $\sqrt{d_k}$ keeps the variance at 1 regardless of dimension → softmax stays in a well-behaved range。

Multi-Head Attention

Run attention $h$ times in parallel with different learned projections:

Why Multiple Heads?

Single attention head 只能學一種 attention pattern。Multi-head 讓 model 同時學多種 relationships：

Total parameters 不增加 — 每個 head 的 dimension = $d_{\text{model}} / h$，concatenate 後恢復到 $d_{\text{model}}$。

Computational Complexity

Positional Encoding

Self-attention is permutation-invariant — it treats input as a set, not a sequence. Positional encoding injects order information:

Sinusoidal encoding (original paper):

Why sin/cos? 兩個理由：

1. 可以 generalize to longer sequences than seen during training（any position has a unique encoding）
2. Relative position 可以被 linear function 表示：$PE_{pos+k}$ is a linear function of $PE_{pos}$

Positional Encoding Variants

Learned 和 sinusoidal 在 practice 差異不大。RoPE 和 ALiBi 是為了 length generalization — 讓 model 能 handle 比 training 更長的 sequences。

Encoder-Decoder Structure

Encoder Layer

Each encoder layer has two sub-layers:

1. Multi-head self-attention: Each position attends to all positions
2. Position-wise feed-forward network: $\text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2$

Both wrapped with residual connection + LayerNorm:

Decoder Layer

Each decoder layer has three sub-layers:

1. Masked self-attention: Each position attends to previous positions only（causal mask）
2. Cross-attention: Decoder queries attend to encoder outputs（Keys + Values from encoder）
3. Position-wise feed-forward network

Causal Mask

Decoder 的 self-attention 不能看到 future tokens（autoregressive constraint）:

加入 mask 後 softmax → future positions 的 attention weight = 0。

Pre-Norm vs Post-Norm

大多數 modern models（GPT, LLaMA）用 Pre-Norm — training 更穩定，不需要 learning rate warmup 也能 converge。

Modern Variants

Encoder-Only: BERT

Decoder-Only: GPT

Encoder-Decoder: T5, BART

Comparison

KV Cache (Inference Optimization)

Autoregressive generation 的效率問題：每生成一個 token，都要重新計算整個 sequence 的 attention — 非常浪費。

KV Cache: Cache previous tokens' Key and Value matrices → new token 只需要計算自己的 Q、K、V，然後 append 到 cache。

這是為什麼 LLM inference 的 memory usage 和 context length 成正比 — KV cache 越長越吃 memory。

Efficient Attention Variants

Standard attention 的 $O(n^2)$ 在 long sequences 上不 scalable。Solutions:

FlashAttention 是目前最 practical 的加速 — 不改 attention 結果（exact），只改 computation order → IO-aware optimization → 實務中幾乎所有 LLM 都用。

Attention Heatmap

Visualize how tokens attend to each other in multi-head attention:

Run the Transformer

This code uses our pure Python implementation — every operation (matmul, softmax, layer norm) written from scratch.

Forward Pass

Greedy Decoding

Inspect Attention Weights

Hands-on: Transformer in PyTorch

Self-Attention

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        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)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, Q, K, V, mask=None):
        B = Q.size(0)
        Q = self.W_q(Q).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)

        out = out.transpose(1, 2).contiguous().view(B, -1, self.num_heads * self.d_k)
        return self.W_o(out)

Encoder Layer

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-Norm variant (modern standard)
        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), mask))
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

Causal Mask

def generate_causal_mask(seq_len):
    """Decoder mask: prevent attending to future tokens."""
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
    return ~mask  # True = allowed, False = masked

Real-World Use Cases

Case 1: 信用卡詐欺 — Transformer for Sequences

BERT-like model on transaction sequences 可以學到 contextual fraud patterns：

# Each transaction = [amount, merchant_category, time_delta, ...]
# Sequence of last 100 transactions → Transformer encoder → [CLS] → P(fraud)
class FraudTransformer(nn.Module):
    def __init__(self, feature_dim, d_model, nhead, num_layers):
        super().__init__()
        self.projection = nn.Linear(feature_dim, d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, batch_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.classifier = nn.Linear(d_model, 1)

    def forward(self, x, mask=None):
        x = self.projection(x)
        x = self.encoder(x, src_key_padding_mask=mask)
        return torch.sigmoid(self.classifier(x[:, -1]))  # last position

面試 follow-up：「Transformer vs LSTM for fraud？」— Transformer 在 longer sequences 上更好（direct attention vs gradient chain），且 training 更快（parallel）。但 real-time scoring 時 RNN 的 incremental inference 更有效率（不需要 KV cache）。

Case 2: 推薦系統 — Self-Attention on User History

SASRec (Self-Attentive Sequential Recommendation): 把 user 的 interaction history 當成 sequence → Transformer encoder → predict next item。

和 RNN-based GRU4Rec 的差異：Transformer 可以直接 attend to 歷史中任何一個 item（不受 recency bias 限制）→ 更好地捕捉 long-term preferences。

Case 3: NLP — The Foundation of LLMs

Interview Signals

What interviewers listen for:

• 你能解釋 Q, K, V 的直覺（database lookup analogy）
• 你知道為什麼 scale by $\sqrt{d_k}$（softmax saturation → gradient issues）
• 你能比較 BERT vs GPT（encoder vs decoder, bidirectional vs autoregressive）
• 你知道 $O(n^2)$ complexity 的意義和 efficient attention 的存在
• 你理解 KV cache 對 inference 的重要性

Practice

Flashcards

Quiz