Transformer block: attention + MLP + layer norm + residual

Bài 11 xong multi-head attention — bạn đã hiểu cách query, key, value hoạt động, cách nhiều head song song học nhiều dạng quan hệ khác nhau giữa các token. Đó là thứ làm Transformer khác biệt hoàn toàn so với RNN.

Nhưng attention chỉ là một nửa.

Transformer block thực sự gồm bốn thành phần: multi-head attention, MLP (feed-forward network), layer normalization, và residual connection. Bỏ đi bất kỳ thành phần nào, model hoặc không train nổi, hoặc không đạt độ sâu đủ để học được gì có ý nghĩa.

Bài này ghép cả bốn thành phần đó lại thành một block hoàn chỉnh, giải thích thứ tự quan trọng, rồi stack 12 block lại để ra GPT-2. Cuối bài: code skeleton forward pass từ đầu tới cuối.

Mental model tổng quát

Một Transformer block là vòng lặp xử lý hai lần:

x (input, shape: B x N x D)
|
+----> LayerNorm --> MultiHead Attention --> +
|                                            |
+<------------------------------------------+   (residual 1)
|
+----> LayerNorm --> MLP (Feed-Forward) ------> +
|                                                |
+<----------------------------------------------+   (residual 2)
|
out (cùng shape: B x N x D)

Input vào và output ra cùng shape. Block không thay đổi kích thước, chỉ thay đổi nội dung — mỗi vector token được “enrich” thêm thông tin từ context và được xử lý thêm bởi MLP. Đây là lý do có thể stack N block tuỳ ý: output của block này là input của block kế tiếp, không cần adapter hay reshape.

Phần 1: MLP — thinking per position

Attention giải quyết câu hỏi: token này nên “nhìn” vào token nào khác? Sau khi attention trộn context xong, mỗi token có một vector đã được enrich. Nhưng enrich theo nghĩa “pha trộn thông tin”, chưa phải “suy nghĩ” gì về thông tin đó.

MLP làm phần đó — xử lý từng position độc lập, biến thông tin đã được mix thành dạng hữu ích hơn.

Architecture của MLP trong Transformer:

Linear (D → 4D) → GeLU → Linear (4D → D)

Hai phép biến đổi tuyến tính với một activation phi tuyến ở giữa. Điểm đáng chú ý là expansion ratio 4x: nếu d_model = 768 (GPT-2 Small), MLP expand lên 3072, xử lý trong không gian đó, rồi contract về 768. Không gian lớn hơn cho model nhiều “room” hơn để biểu diễn các phép biến đổi phức tạp.

import numpy as np

def gelu(x):
    # Approximation dùng trong GPT-2
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))

def mlp_forward(x, W1, b1, W2, b2):
    # x:  (B, N, D)
    # W1: (D, 4D),  b1: (4D,)
    # W2: (4D, D),  b2: (D,)
    h = gelu(x @ W1 + b1)   # expand: (B, N, 4D)
    return h @ W2 + b2       # contract: (B, N, D)

Vì sao GeLU thay vì ReLU? ReLU đặt hard zero cho mọi giá trị âm — thông tin bị mất hoàn toàn. GeLU làm mềm hơn: giá trị âm nhỏ vẫn có thể pass qua với hệ số nhỏ. Trong thực tế training, GeLU cho kết quả tốt hơn ReLU ở scale lớn.

Các modern variants: Llama dùng SwiGLU thay vì GeLU — MLP trở thành 3 ma trận thay vì 2, với một nhánh gate nhân vào. PaLM dùng GeGLU. Tất cả đều là biến thể của ý tưởng gốc: expand → activate → contract. Sẽ nói chi tiết ở bài riêng về Llama architecture.

Phần 2: Residual connection — highway cho gradient

Bạn có một hàm f(x) nào đó — có thể là attention, có thể là MLP. Thay vì chỉ trả ra f(x), residual connection cộng thêm input gốc:

output = f(x) + x

Một dòng cộng, nhưng ảnh hưởng lớn đến khả năng train được của model sâu.

Vấn đề không có residual: khi stack nhiều layer, gradient phải truyền ngược từ output về đến input qua chuỗi nhân ma trận dài. Mỗi lần nhân, gradient hoặc shrink (vanish) hoặc explode. Với 12 layer như GPT-2, gradient của layer đầu tiên nhận được gần như bằng không — weights không cập nhật, model không học được gì.

Residual giải quyết như thế nào: trong backward pass, gradient của residual d(f(x) + x)/dx = f'(x) + 1. Cái +1 đó là “highway” — gradient luôn có một con đường thẳng về input mà không qua bất kỳ phép biến đổi nào. Dù f'(x) có vanish đến đâu, gradient vẫn có thể truyền ngược về các layer trước.

Ý tưởng đến từ ResNet (He et al., 2015) — paper đầu tiên train được mạng 100+ layer cho image recognition. Transformer paper (2017) kế thừa ngay và nó trở thành chuẩn mực cho mọi deep network hiện đại.

Kiểm chứng thực tế: remove residual khỏi GPT-2, train cùng dataset — model bắt đầu diverge sau khoảng 5-6 layer. Với residual, 12 layer train ổn định, 96 layer (GPT-3 architecture) cũng train được.

Phần 3: Layer normalization — ổn định activations

Qua nhiều layer biến đổi, các giá trị trong vector có thể grow hoặc shrink không kiểm soát. Một layer trả ra giá trị magnitude 1e-3, layer tiếp theo thấy đó nhân với weights và ra magnitude 1e-6, layer sau nữa lại nhân và ra 1e-9 — gradient vanish, hoặc ngược lại, activation explode lên 1e+6 và training diverge.

Layer normalization giữ cho activation ở magnitude ổn định bằng cách normalize theo chiều feature (chiều D) cho từng token:

LN(x) = gamma * (x - mean(x)) / (std(x) + eps) + beta

Trong đó:

• mean(x) và std(x) tính trên D chiều của token đó, không tính across batch hay sequence
• gamma và beta là hai vector học được, cùng shape (D,) — cho phép model “re-scale” sau normalize
• eps = 1e-5 tránh chia cho zero

Code:

def layer_norm(x, gamma, beta, eps=1e-5):
    # x: (B, N, D)
    mean = x.mean(axis=-1, keepdims=True)        # (B, N, 1)
    var  = x.var(axis=-1, keepdims=True)          # (B, N, 1)
    x_hat = (x - mean) / np.sqrt(var + eps)      # (B, N, D), normalized
    return gamma * x_hat + beta                   # (B, N, D), rescaled

LayerNorm vs BatchNorm: BatchNorm normalize across batch dimension — hoạt động tốt cho CNN với batch size lớn, nhưng có vấn đề khi sequence length thay đổi và khi batch size nhỏ. LayerNorm normalize per sample, per token — không phụ thuộc batch size, không phụ thuộc độ dài sequence. Transformer dùng LayerNorm vì lý do đó.

RMSNorm (Llama, Mistral): bỏ luôn bước trừ mean, chỉ chia cho RMS (root mean square). Ít tính toán hơn ~15%, trong thực nghiệm cho kết quả tương đương. Ngày nay đa số LLM mới đều dùng RMSNorm thay vì LayerNorm đầy đủ.

Phần 4: Pre-norm vs post-norm — thứ tự quan trọng

Transformer paper gốc (Vaswani et al., 2017) dùng post-norm:

x = LayerNorm(x + Attention(x))   # post-norm attention
x = LayerNorm(x + MLP(x))         # post-norm MLP

LayerNorm đặt sau khi cộng residual. Trông hợp lý nhưng trong thực tế, post-norm cần warmup learning rate rất cẩn thận và dễ diverge khi scale lên sâu hơn.

GPT-2 và hầu hết mọi LLM hiện đại dùng pre-norm:

x = x + Attention(LayerNorm(x))   # pre-norm attention
x = x + MLP(LayerNorm(x))         # pre-norm MLP

LayerNorm đặt bên trong trước mỗi sublayer, trước khi input đi vào attention hoặc MLP. Residual cộng trực tiếp với output thô, không qua norm.

Lý do pre-norm tốt hơn: ở đầu training, weights còn ngẫu nhiên, output của attention/MLP rất noisy. Post-norm normalize sau khi cộng residual — nếu signal noisy quá, norm cũng không cứu được. Pre-norm normalize trước, đảm bảo input vào attention/MLP luôn ở scale ổn định ngay từ step đầu tiên. Training khởi động sớm hơn, ổn định hơn, không cần warmup quá dài.

Hiện nay: pre-norm là mặc định trong mọi LLM production. Post-norm chỉ còn trong các paper so sánh baseline.

Phần 5: Ghép thành Transformer block hoàn chỉnh

Bốn thành phần, thứ tự pre-norm, hai residual:

def transformer_block(x, params, num_heads):
    """
    x: (B, N, D) — batch, sequence length, model dim
    params: dict chứa weights của block này
    num_heads: số attention heads
    """

    # --- Attention sublayer ---
    # 1. Normalize input trước attention
    y = layer_norm(x, params["ln1_gamma"], params["ln1_beta"])

    # 2. Multi-head attention (từ bài 11)
    y = multi_head_attention(
        y,
        params["Wq"], params["Wk"], params["Wv"], params["Wo"],
        num_heads
    )

    # 3. Cộng residual
    x = x + y

    # --- MLP sublayer ---
    # 4. Normalize input trước MLP
    y = layer_norm(x, params["ln2_gamma"], params["ln2_beta"])

    # 5. Feed-forward: expand → GeLU → contract
    y = gelu(y @ params["W1"] + params["b1"]) @ params["W2"] + params["b2"]

    # 6. Cộng residual
    x = x + y

    return x

Chú ý: x ở step 3 và step 6 là hai x khác nhau. Step 3 tạo x mới bằng cách cộng residual cho attention output. Step 6 cộng residual MLP vào x đã có từ step 3. Đây là chain của hai residual connection nối tiếp nhau.

Params của một block:

• ln1_gamma, ln1_beta: (D,) — LayerNorm 1
• ln2_gamma, ln2_beta: (D,) — LayerNorm 2
• Wq, Wk, Wv: (D, D) — Query, Key, Value projections
• Wo: (D, D) — Output projection của attention
• W1: (D, 4D), b1: (4D,) — MLP expand
• W2: (4D, D), b2: (D,) — MLP contract

Tổng params một block với D = 768: 4 * D * D (attention) + 8 * D * D (MLP) + 4 * D (norms) ≈ 12 * 768² ≈ 7 triệu params.

Phần 6: Stack block — từ block đến GPT

GPT chỉ là:

token IDs
    |
    v
[ Token Embedding ] + [ Position Embedding ]   (B, N, D)
    |
    v
[ Transformer Block 1 ]                        (B, N, D)
    |
    v
[ Transformer Block 2 ]                        (B, N, D)
    |
    ...
    |
    v
[ Transformer Block L ]                        (B, N, D)
    |
    v
[ Final LayerNorm ]                            (B, N, D)
    |
    v
[ Output Projection ]   (D -> vocab_size)      (B, N, V)
    |
    v
logits → softmax → sample → next token

Bốn model GPT-2, khác nhau ở số layer và d_model:

Model	Layers (L)	d_model (D)	Heads	Approx Params
GPT-2 Small	12	768	12	124M
GPT-2 Medium	24	1024	16	355M
GPT-2 Large	36	1280	20	774M
GPT-2 XL	48	1600	25	1.5B

Llama-3-8B để so sánh: 32 layers, d_model = 4096, 32 heads, 8B params — dùng RMSNorm, SwiGLU, Grouped Query Attention thay vì multi-head thuần, và không có position embedding tuyến tính mà dùng RoPE. Các khác biệt đó là chủ đề của bài riêng về Llama architecture.

Phần 7: Hands-on — build GPT skeleton

Đây là forward pass hoàn chỉnh, từ token IDs đến logits. Chưa training, chưa backprop — chỉ kiểm chứng shapes đúng và đếm params:

import numpy as np

# --------------- Activation ---------------
def gelu(x):
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))

# --------------- Layer Norm ---------------
def layer_norm(x, gamma, beta, eps=1e-5):
    mean = x.mean(axis=-1, keepdims=True)
    std  = np.sqrt(x.var(axis=-1, keepdims=True) + eps)
    return gamma * (x - mean) / std + beta

# --------------- Attention ----------------
def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def multi_head_attention(x, Wq, Wk, Wv, Wo, num_heads):
    B, N, D = x.shape
    head_dim = D // num_heads

    Q = x @ Wq   # (B, N, D)
    K = x @ Wk
    V = x @ Wv

    # split heads
    def split(t):
        return t.reshape(B, N, num_heads, head_dim).transpose(0, 2, 1, 3)

    Q, K, V = split(Q), split(K), split(V)  # (B, H, N, head_dim)

    scale = head_dim ** -0.5
    scores = Q @ K.transpose(0, 1, 3, 2) * scale  # (B, H, N, N)

    # causal mask
    mask = np.triu(np.full((N, N), -1e9), k=1)
    scores = scores + mask

    attn = softmax(scores)           # (B, H, N, N)
    out  = attn @ V                  # (B, H, N, head_dim)

    # merge heads
    out = out.transpose(0, 2, 1, 3).reshape(B, N, D)
    return out @ Wo

# --------------- Transformer Block --------
def transformer_block(x, p, num_heads):
    y = layer_norm(x, p["ln1_g"], p["ln1_b"])
    y = multi_head_attention(y, p["Wq"], p["Wk"], p["Wv"], p["Wo"], num_heads)
    x = x + y

    y = layer_norm(x, p["ln2_g"], p["ln2_b"])
    y = gelu(y @ p["W1"] + p["b1"]) @ p["W2"] + p["b2"]
    x = x + y
    return x

# --------------- GPT forward pass ---------
def gpt_forward(token_ids, params, config):
    B, N     = token_ids.shape
    D        = config["d_model"]
    V        = config["vocab_size"]
    L        = config["num_layers"]
    H        = config["num_heads"]

    # Embedding
    tok_emb = params["wte"][token_ids]          # (B, N, D)
    pos_ids = np.arange(N)
    pos_emb = params["wpe"][pos_ids]            # (N, D)
    x = tok_emb + pos_emb                       # (B, N, D)

    # Transformer blocks
    for i in range(L):
        x = transformer_block(x, params["blocks"][i], H)

    # Final norm + projection
    x = layer_norm(x, params["ln_f_g"], params["ln_f_b"])
    logits = x @ params["wte"].T               # (B, N, V) — weight tying

    return logits

# --------------- Init random weights ------
def init_gpt2_small():
    D, V, L = 768, 50257, 12
    rng = np.random.default_rng(42)

    def r(*shape): return rng.standard_normal(shape).astype(np.float32) * 0.02

    blocks = []
    for _ in range(L):
        blocks.append({
            "ln1_g": np.ones(D, dtype=np.float32),
            "ln1_b": np.zeros(D, dtype=np.float32),
            "Wq": r(D, D), "Wk": r(D, D), "Wv": r(D, D), "Wo": r(D, D),
            "ln2_g": np.ones(D, dtype=np.float32),
            "ln2_b": np.zeros(D, dtype=np.float32),
            "W1": r(D, 4*D), "b1": np.zeros(4*D, dtype=np.float32),
            "W2": r(4*D, D), "b2": np.zeros(D, dtype=np.float32),
        })

    params = {
        "wte": r(V, D),   # token embedding
        "wpe": r(1024, D), # position embedding (max 1024 tokens)
        "blocks": blocks,
        "ln_f_g": np.ones(D, dtype=np.float32),
        "ln_f_b": np.zeros(D, dtype=np.float32),
    }
    return params

# --------------- Count params -------------
def count_params(params):
    total = 0
    total += params["wte"].size    # token embedding (V x D)
    total += params["wpe"].size    # position embedding
    for b in params["blocks"]:
        for v in b.values():
            total += v.size
    total += params["ln_f_g"].size + params["ln_f_b"].size
    return total

# --------------- Test ---------------------
config = {"d_model": 768, "vocab_size": 50257, "num_layers": 12, "num_heads": 12}
params = init_gpt2_small()

# Forward pass: batch=2, seq_len=8
token_ids = np.array([[1, 2, 3, 4, 5, 6, 7, 8],
                       [9, 10, 11, 12, 13, 14, 15, 16]])
logits = gpt_forward(token_ids, params, config)

print(f"logits shape: {logits.shape}")   # (2, 8, 50257)
print(f"total params: {count_params(params):,}")  # ~124M

Chạy đoạn code trên: logits shape: (2, 8, 50257) — batch 2, sequence 8, vocab 50,257. Số params tính ra khoảng 124 triệu — đúng GPT-2 Small. Forward pass của GPT-2 không phải magic: chỉ là embedding lookup, cộng với position embedding, qua 12 Transformer block, final norm, và matmul với embedding matrix ngược lại (weight tying).

Cheatsheet

Thành phần	Vị trí trong block	Chức năng chính	Shape IO
LayerNorm	Trước mỗi sublayer (pre-norm)	Ổn định activation magnitude	(B,N,D) → (B,N,D)
Multi-head Attention	Sublayer 1	Mix context giữa các token	(B,N,D) → (B,N,D)
Residual 1	Sau attention	Highway cho gradient, bảo toàn info	cộng x vào y
MLP (FFN)	Sublayer 2	Xử lý per-position, D → 4D → D	(B,N,D) → (B,N,D)
Residual 2	Sau MLP	Highway thứ hai	cộng x vào y

Sáu điều phải nhớ:

• Transformer block luôn trả ra cùng shape với input — có thể stack N lần tuỳ ý
• Pre-norm (LN trước sublayer) là standard hiện đại — train ổn định hơn post-norm
• Residual cộng input thẳng vào output — gradient có highway về, model sâu mới train được
• MLP expand 4x: D → 4D → D — activation function phi tuyến ở giữa
• GPT-2 Small = 12 block × ~7M params/block + embedding ≈ 124M tổng
• Weight tying: embedding matrix đầu vào và output projection dùng chung — tiết kiệm V×D params

Lời kết

Bây giờ bạn đã có đủ thành phần để đọc kiến trúc của bất kỳ LLM nào trên paper: Transformer block = LayerNorm + Attention + residual + LayerNorm + MLP + residual, stack L lần. Phần lớn sự khác biệt giữa GPT-2, Llama, Mistral, Gemma chỉ là biến thể nhỏ — RMSNorm thay LayerNorm, SwiGLU thay GeLU, GQA thay MHA, RoPE thay learned position embedding. Core vẫn là những thứ bài này vừa build.

Bài 13 sẽ là nanoGPT — Andrej Karpathy viết GPT đầy đủ trong khoảng 300 dòng PyTorch, train trên Shakespeare. Đây là lúc mọi thứ lý thuyết trong series gặp nhau trong một file code thực sự chạy được.

Trước khi sang bài 13: thử đếm tay số params của GPT-2 Small. Công thức:

Token embedding:       V × D    = 50257 × 768  = 38.6M
Position embedding:    1024 × D = 1024  × 768  =  0.8M
Mỗi block:
  Attention (Wq+Wk+Wv+Wo): 4 × D² = 4 × 768² = 2.36M
  MLP (W1+W2):              8 × D² = 8 × 768² = 4.72M
  Layer norms:              4 × D  ≈ 0 (negligible)
  Block tổng:               ≈ 7.08M
12 blocks:                12 × 7.08M = 85M
Final LN:                  ≈ 0
Total (trước weight tying): 38.6 + 0.8 + 85 ≈ 124.4M

Khớp với con số official 124M của OpenAI. Weight tying (dùng lại token embedding cho output projection) là lý do params không tăng thêm V×D nữa — không có Unembedding riêng, chỉ dùng wte.T.