Kiến trúc Transformer: Giải phẫu cỗ máy đằng sau mọi LLM

Khi đọc tin tức công nghệ, bạn thấy hàng loạt tên gọi mới xuất hiện mỗi tháng: GPT-4, Llama-3, Claude, Qwen, Gemma, DeepSeek… Có vẻ như thế giới AI đang thay đổi chóng mặt. Nhưng nếu mở mã nguồn của tất cả các mô hình đó ra, bạn sẽ bắt gặp một sự thật đáng kinh ngạc: về mặt kiến trúc cốt lõi, hầu như tất cả chúng đều là cùng một cỗ máy — một chồng các khối Transformer Decoder lặp đi lặp lại.

Vậy chính xác chuyện gì diễn ra bên trong một khối đó? Trong chương này, chúng ta sẽ đi tìm lời giải cho hai câu hỏi cốt lõi:

• Bên trong một khối Transformer Decoder hiện đại thực sự gồm những phép toán ma trận nào, và mỗi phép toán đóng vai trò gì trong việc xử lý ngôn ngữ?
• Tại sao kiến trúc Transformer gốc năm 2017 của Vaswani và cộng sự lại được tinh chỉnh thành “công thức chuẩn hiện đại” (Pre-RMSNorm, RoPE, SwiGLU) dùng trong Llama-3 và hầu hết các LLM ngày nay?

1. Cơ chế Chú ý Tích vô hướng có Tỷ lệ và Mặt nạ Nhân quả

Hãy nhớ lại điểm dừng ở cuối Chương 0: trong mạng RNN/LSTM, thông tin phải đi tuần tự qua từng từ $t-1 \to t$, khiến GPU không thể tính toán song song trên toàn câu và đạo hàm bị suy giảm khi khoảng cách giữa hai từ quá xa. Tháng 6 năm 2017, nhóm nghiên cứu của Vaswani và cộng sự tại Google đã đặt một câu hỏi táo bạo: Điều gì xảy ra nếu ta vứt bỏ hoàn toàn cấu trúc truy hồi, và chỉ giữ lại cơ chế chú ý? (Vaswani et al., 2017).

1.1 Trực giác truy xuất mềm: Truy vấn (Query), Khóa (Key) và Giá trị (Value)

Hãy xét câu tiếng Việt: “Con mèo ngồi trên chiếu vì nó buồn ngủ”. Khi đọc đến đại từ “nó”, làm thế nào một mạng nơ-ron biết “nó” đang chỉ “con mèo” chứ không phải “chiếu”?

Trong cơ chế Tự chú ý (Self-Attention), mỗi token ở vị trí $i$ có vector biểu diễn đầu vào $x_i \in \mathbb{R}^{d_{\text{model}}}$. Để các từ giao tiếp với nhau, mỗi từ đóng đồng thời ba vai trò khác nhau thông qua ba phép chiếu tuyến tính học được:

1. Truy vấn (Query — $q_i = x_i W_Q$): Vector thể hiện “thông tin mà vị trí $i$ đang đi tìm” (ví dụ: đại từ “nó” phát ra một truy vấn tìm danh từ chỉ động vật đứng trước).
2. Khóa (Key — $k_j = x_j W_K$): Vector thể hiện “nhãn nhận diện của vị trí $j$” để trả lời các truy vấn (từ “mèo” giơ cao tấm biển: “tôi là danh từ chỉ con vật”).
3. Giá trị (Value — $v_j = x_j W_V$): Vector chứa “nội dung ngữ nghĩa mà vị trí $j$ sẽ truyền đi” nếu được chú ý tới.

Khi xếp toàn bộ $L$ token của một chuỗi thành ma trận $X \in \mathbb{R}^{L \times d_{\text{model}}}$, ta tính được cả ba ma trận cho mọi vị trí cùng một lúc bằng ba phép nhân ma trận:

trong đó $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ và $W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$.

1.2 Tại sao phải chia cho $\sqrt{d_k}$? Bài toán phương sai và bão hòa Softmax

Để đo mức độ khớp nhau giữa truy vấn $q_i$ và khóa $k_j$, cách tự nhiên nhất là lấy tích vô hướng $q_i \cdot k_j = \sum_{r=1}^{d_k} q_{i,r} k_{j,r}$. Nhưng tại sao công thức của Vaswani et al. (Vaswani et al., 2017) lại bắt buộc phải chia tích vô hướng này cho $\sqrt{d_k}$?

Hãy làm một phép tính xác suất nhỏ. Giả sử các thành phần của $q_i$ và $k_j$ là các biến ngẫu nhiên độc lập có trung bình bằng $0$ và phương sai bằng $1$. Khi đó, mỗi tích $q_{i,r} k_{j,r}$ có trung bình bằng $0$ và phương sai bằng $1$. Tổng của $d_k$ số hạng độc lập như vậy sẽ có:

Nếu $d_k = 64$ hoặc $128$, độ lệch chuẩn của tích vô hướng sẽ là $\sqrt{d_k} = 8$ hoặc $11.3$. Điều này có nghĩa là một số phần tử trong tích $QK^\top$ sẽ vọt lên $+20$ hoặc tụt xuống $-20$. Khi đưa các con số chênh lệch lớn như vậy vào hàm $\text{softmax}(z)_j = \frac{e^{z_j}}{\sum_{n} e^{z_n}}$, phần tử lớn nhất sẽ nuốt trọn xác suất $\approx 1.0$, còn mọi phần tử khác nhận xác suất $\approx 0.0$. Trong vùng bão hòa (saturation) này, đạo hàm của hàm Softmax gần như bằng $0$, làm đóng băng quá trình học! Chia cho $\sqrt{d_k}$ đưa phương sai của điểm số trở lại bằng $1$ bất kể số chiều $d_k$ lớn đến đâu.

1.3 Mặt nạ Nhân quả (Causal Mask) cho mô hình sinh văn bản tự hồi quy

Trong một mô hình ngôn ngữ tự hồi quy (Decoder-only như GPT hay Llama), nhiệm vụ tại vị trí $i$ là dự đoán token $i+1$. Do đó, khi tính biểu diễn tại vị trí $i$, mô hình tuyệt đối không được phép nhìn thấy các token tương lai ở vị trí $j > i$.

Làm thế nào để vừa tính toàn bộ ma trận $QK^\top \in \mathbb{R}^{L \times L}$ song song trong một bước duy nhất trên GPU, vừa ngăn thông tin rò rỉ từ tương lai về quá khứ? Ta cộng thêm một ma trận mặt nạ nhân quả (Causal Mask) $M \in \mathbb{R}^{L \times L}$ trước khi lấy hàm Softmax:

Vì $e^{-\infty} = 0$, sau khi đi qua hàm Softmax theo từng hàng, mọi trọng số chú ý $\alpha_{i,j}$ ở vùng tam giác trên ($j > i$) đều bằng $0$ tuyệt đối, còn các trọng số hợp lệ ở quá khứ và hiện tại ($j \le i$) tự động được chuẩn hóa sao cho tổng bằng $1$.

Hãy kiểm chứng trực tiếp cơ chế này bằng PyTorch trong Bài tập 2.1.

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

torch.manual_seed(42)
B, L, d_k = 2, 6, 16
Q = torch.randn(B, L, d_k)
K = torch.randn(B, L, d_k)
V = torch.randn(B, L, d_k)
print("Kích thước Q, K, V:", tuple(Q.shape))

Kích thước Q, K, V: (2, 6, 16)

Kích thước đầu ra out: (2, 6, 16)
Tổng trọng số mỗi hàng: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
Tổng trọng số tam giác trên (j > i): 0.000000
Sai khác tại vị trí k <= 2 khi sửa token 4, 5: 0.00e+00
Sai khác tại vị trí k > 2 khi sửa token 4, 5:  6.41e-01

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

torch.manual_seed(42)
B, L, d_k = 2, 6, 16
Q = torch.randn(B, L, d_k)
K = torch.randn(B, L, d_k)
V = torch.randn(B, L, d_k)
print("Kích thước Q, K, V:", tuple(Q.shape))

Kích thước Q, K, V: (2, 6, 16)

# (1) Hàm tính Scaled Dot-Product Causal Self-Attention
def causal_attention(q, k, v):
    L_seq = q.size(-2)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.size(-1))
    mask = torch.triu(
        torch.ones(L_seq, L_seq, dtype=torch.bool, device=q.device),
        diagonal=1,
    )
    scores = scores.masked_fill(mask, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ v, weights

out1, w1 = causal_attention(Q, K, V)
row_sums = [round(x, 4) for x in w1[0].sum(dim=-1).tolist()]
upper_sum = torch.triu(w1, diagonal=1).sum().item()

print("Kích thước đầu ra out:", tuple(out1.shape))
print("Tổng trọng số mỗi hàng:", row_sums)
print(f"Tổng trọng số tam giác trên (j > i): {upper_sum:.6f}")

# (2) Kiểm chứng tính nhân quả: thay đổi các token ở vị trí t >= 4
Q2, K2, V2 = Q.clone(), K.clone(), V.clone()
Q2[:, 4:] = torch.randn_like(Q2[:, 4:])
K2[:, 4:] = torch.randn_like(K2[:, 4:])
V2[:, 4:] = torch.randn_like(V2[:, 4:])

out2, _ = causal_attention(Q2, K2, V2)
diff_past = (out1[:, :3] - out2[:, :3]).abs().max().item()
diff_future = (out1[:, 4:] - out2[:, 4:]).abs().max().item()
print(f"Sai khác tại vị trí k <= 2 khi sửa token 4, 5: {diff_past:.2e}")
print(f"Sai khác tại vị trí k > 2 khi sửa token 4, 5:  {diff_future:.2e}")

Kích thước đầu ra out: (2, 6, 16)
Tổng trọng số mỗi hàng: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
Tổng trọng số tam giác trên (j > i): 0.000000
Sai khác tại vị trí k <= 2 khi sửa token 4, 5: 0.00e+00
Sai khác tại vị trí k > 2 khi sửa token 4, 5:  6.41e-01

2. Chú ý Đa đầu (Multi-Head Attention) và Luồng Thặng dư

Một phép tự chú ý đơn lẻ (Single-Head Attention) buộc mỗi token phải gộp mọi mối quan hệ ngữ pháp, ngữ nghĩa, quy chiếu vào chung một bộ trọng số trung bình $\alpha_{i,1} v_1 + \dots + \alpha_{i,i} v_i$. Nhưng trong ngôn ngữ tự nhiên, tại cùng một thời điểm, một từ cần chú ý đến nhiều khía cạnh khác nhau: đầu chú ý này tìm chủ ngữ của động từ, đầu chú ý kia nối tính từ với danh từ, đầu thứ ba theo dõi dấu câu.

2.1 Tách không gian thành $H$ đầu chú ý song song

Thay vì thực hiện một phép chú ý duy nhất trên không gian $d_{\text{model}}$ chiều, Vaswani et al. (Vaswani et al., 2017) chia $d_{\text{model}}$ thành $H$ không gian con (subspaces) độc lập, mỗi không gian con có số chiều:

Với mỗi đầu $h \in \{1, \dots, H\}$, ta có bộ ma trận chiếu riêng $W_Q^{(h)}, W_K^{(h)}, W_V^{(h)} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{head}}}$ để tính:

Sau đó, đầu ra của cả $H$ đầu được ghép nối (concatenate) lại dọc theo chiều đặc trưng và nhân với ma trận chiếu đầu ra $W_O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$:

2.2 Góc nhìn Mạch Cơ học (Mechanistic Interpretability): Luồng Thặng dư là một Kênh Giao tiếp Chung

Năm 2021, nhóm nghiên cứu của Nelson Elhage và Chris Olah tại Anthropic đã đề xuất một cách viết lại phương trình Multi-Head Attention giúp làm sáng tỏ bản chất hoạt động của Transformer (Elhage & others, 2021).

Nếu ta tách ma trận đầu ra $W_O \in \mathbb{R}^{(H \cdot d_{\text{head}}) \times d_{\text{model}}}$ thành $H$ khối hàng $W_O^{(h)} \in \mathbb{R}^{d_{\text{head}} \times d_{\text{model}}}$, thì phép ghép nối rồi nhân $W_O$ hoàn toàn tương đương với tổng của $H$ đầu chú ý độc lập:

Kết hợp với nhánh cộng thặng dư $x_\ell = x_{\ell-1} + \mathrm{MHA}(\mathrm{Norm}(x_{\ell-1}))$, ta nhận ra:

• Luồng Thặng dư (Residual Stream) là một dải băng chuyền có kích thước cố định $(B, L, d_{\text{model}})$ chạy xuyên suốt từ tầng đầu tiên đến tầng cuối cùng.
• Mỗi đầu chú ý $h$ đọc một không gian con từ luồng thặng dư, di chuyển thông tin giữa các vị trí (nhờ mạch QK $W_Q^{(h)} (W_K^{(h)})^\top$), biến đổi nội dung (nhờ mạch OV $W_V^{(h)} W_O^{(h)}$), rồi cộng đè kết quả của nó trở lại vào luồng thặng dư!

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

torch.manual_seed(7)
B, L, d_model, n_heads = 2, 8, 64, 4
x = torch.randn(B, L, d_model)
print(f"Đầu vào x: {tuple(x.shape)}, d_head = {d_model // n_heads}")

Đầu vào x: (2, 8, 64), d_head = 16

Kích thước Q sau khi tách head: (2, 4, 8, 16)
Kích thước bản đồ chú ý attn:   (2, 4, 8, 8)
Kích thước đầu ra MHA:          (2, 8, 64)
Tổng số tham số (4 * d_model^2): 16384
Khoảng cách trung bình giữa các head: 0.2841

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

torch.manual_seed(7)
B, L, d_model, n_heads = 2, 8, 64, 4
x = torch.randn(B, L, d_model)
print(f"Đầu vào x: {tuple(x.shape)}, d_head = {d_model // n_heads}")

Đầu vào x: (2, 8, 64), d_head = 16

class CausalMultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, L, D = x.shape
        q = self.w_q(x).view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        k = self.w_k(x).view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        v = self.w_v(x).view(B, L, self.n_heads, self.d_head).transpose(1, 2)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        mask = torch.triu(
            torch.ones(L, L, dtype=torch.bool, device=x.device), diagonal=1
        )
        attn = F.softmax(scores.masked_fill(mask, float("-inf")), dim=-1)

        out = (attn @ v).transpose(1, 2).contiguous().view(B, L, D)
        return self.w_o(out), attn, q

mha = CausalMultiHeadAttention(d_model, n_heads)
out, attn, q = mha(x)
n_params = sum(p.numel() for p in mha.parameters())
head_diff = (attn[0, 0] - attn[0, 1]).abs().mean().item()

print("Kích thước Q sau khi tách head:", tuple(q.shape))
print("Kích thước bản đồ chú ý attn:  ", tuple(attn.shape))
print("Kích thước đầu ra MHA:         ", tuple(out.shape))
print("Tổng số tham số (4 * d_model^2):", n_params)
print(f"Khoảng cách trung bình giữa các head: {head_diff:.4f}")

Kích thước Q sau khi tách head: (2, 4, 8, 16)
Kích thước bản đồ chú ý attn:   (2, 4, 8, 8)
Kích thước đầu ra MHA:          (2, 8, 64)
Tổng số tham số (4 * d_model^2): 16384
Khoảng cách trung bình giữa các head: 0.2841

3. Mã hóa Vị trí: Từ Sinusoidal đến Rotary Position Embedding (RoPE)

Hãy nhìn lại phương trình Self-Attention: nếu không có mặt nạ nhân quả và ta tráo đổi vị trí của hai từ trong câu đầu vào, các vector đầu ra cũng chỉ đổi chỗ tương ứng mà giá trị không hề thay đổi (tính chất đồng biến với phép hoán vị — permutation equivariance). Nói cách khác, bản thân phép chú ý không hề có khái niệm về thứ tự từ!

3.1 Hạn chế của Mã hóa Vị trí Tuyệt đối (Absolute Positional Encoding)

Trong bài báo gốc năm 2017 (Vaswani et al., 2017) và trong GPT-2, thông tin vị trí được đưa vào bằng cách cộng một vector vị trí tuyệt đối $p_m \in \mathbb{R}^{d_{\text{model}}}$ vào vector nhúng của token ở vị trí $m$ ngay tại tầng đầu vào: $x_m = e_m + p_m$.

Cách tiếp cận cộng trực tiếp này có hai nhược điểm lớn:

1. Nhiễm bẩn ngữ nghĩa trên luồng thặng dư: Vector vị trí $p_m$ bị trộn lẫn với vector ngữ nghĩa $e_m$ và bị biến đổi qua nhiều tầng phi tuyến.
2. Kém linh hoạt với khoảng cách tương đối: Trong ngôn ngữ, điều quan trọng nhất giữa hai từ không phải là chúng đứng ở từ thứ $5$ và từ thứ $2$ kể từ đầu văn bản, mà là chúng cách nhau $m - n = 3$ từ!

Làm thế nào để thiết kế một phép biến đổi vị trí $f(q, m)$ và $f(k, n)$ sao cho tích vô hướng chú ý $\langle f(q, m), f(k, n) \rangle$ chỉ phụ thuộc vào vector $q, k$ và hiệu vị trí tương đối $m - n$?

3.2 Phép quay vị trí RoPE trên mặt phẳng phức

Năm 2021, Jianlin Su và cộng sự đã giải bài toán này một cách tuyệt đẹp bằng hình học thông qua Rotary Position Embedding (RoPE) (Su & others, 2024).

Xét trước hết trường hợp vector $q, k \in \mathbb{R}^2$ có 2 chiều. Nếu tại vị trí $m$, ta quay vector $q$ một góc $m\theta$ trên mặt phẳng 2D bằng ma trận quay trực giao $R_{m\theta}$:

thì khi tính tích vô hướng giữa $R_{m\theta} q$ (ở vị trí $m$) và $R_{n\theta} k$ (ở vị trí $n$), nhờ tính chất của ma trận quay $R_{m\theta}^\top R_{n\theta} = R_{(n-m)\theta}$, ta có ngay:

Tích vô hướng này chỉ phụ thuộc vào góc lệch tương đối $(m - n)\theta$! Dù hai từ nằm ở vị trí $(5, 2)$ hay bị dịch đi $1000$ token tới vị trí $(1005, 1002)$, góc giữa hai vector sau khi quay vẫn giữ nguyên không đổi!

Với vector nhiều chiều $d_{\text{head}}$ (là một số chẵn), RoPE đơn giản chia $d_{\text{head}}$ tọa độ thành $d_{\text{head}}/2$ cặp 2 chiều độc lập $(x_{2i}, x_{2i+1})$, và quay mỗi cặp $i \in \{0, 1, \dots, d_{\text{head}}/2 - 1\}$ với tốc độ góc (tần số) giảm dần theo cấp số nhân:

• Các cặp chiều đầu tiên ($i$ nhỏ, $\theta_i \approx 1$) quay rất nhanh — đóng vai trò như kim giây đồng hồ, giúp phân biệt sắc nét các từ đứng sát cạnh nhau.
• Các cặp chiều cuối cùng ($i$ lớn, $\theta_i \ll 1$) quay cực chậm — đóng vai trò như kim giờ, giúp mô hình nắm bắt cấu trúc ngữ cảnh dài hàng nghìn token.

Hãy cài đặt trực tiếp phép quay RoPE và kiểm nghiệm tính bất biến dịch chuyển bằng thực nghiệm trong Bài tập 2.3.

import torch

torch.manual_seed(123)
d_head = 16
q = torch.randn(d_head)
k = torch.randn(d_head)
print("Vector q và k (16 chiều), chuẩn L2:",
      round(q.norm().item(), 4), round(k.norm().item(), 4))

Vector q và k (16 chiều), chuẩn L2: 3.7483 3.9112

Chuẩn L2 của q sau khi quay ở m=5: 3.7483 (sai khác 0.00e+00)
Điểm chú ý <R_5 q, R_2 k>   (m-n = 3):  2.184937
Điểm chú ý <R_17 q, R_14 k> (m-n = 3):  2.184937
Sai số cực đại khi dịch chuyển Delta in [1..20]: 9.54e-07

import torch

torch.manual_seed(123)
d_head = 16
q = torch.randn(d_head)
k = torch.randn(d_head)
print("Vector q và k (16 chiều), chuẩn L2:",
      round(q.norm().item(), 4), round(k.norm().item(), 4))

Vector q và k (16 chiều), chuẩn L2: 3.7483 3.9112

# (1) Cài đặt phép quay vị trí phức RoPE (Rotary Position Embedding)
def apply_rope_1d(vec, pos, base=10000.0):
    d = vec.shape[-1]
    i = torch.arange(0, d, 2, dtype=vec.dtype, device=vec.device)
    theta = base ** (-i / d)
    angles = pos * theta
    cos_a, sin_a = torch.cos(angles), torch.sin(angles)

    x_even = vec[..., 0::2]
    x_odd = vec[..., 1::2]
    rot_even = x_even * cos_a - x_odd * sin_a
    rot_odd = x_even * sin_a + x_odd * cos_a
    return torch.stack([rot_even, rot_odd], dim=-1).flatten(-2)

q_rot5 = apply_rope_1d(q, pos=5)
norm_diff = (q.norm() - q_rot5.norm()).abs().item()
print(f"Chuẩn L2 của q sau khi quay ở m=5: {q_rot5.norm():.4f} "
      f"(sai khác {norm_diff:.2e})")

# (2) Kiểm chứng bất biến dịch chuyển <R_{m+d} q, R_{n+d} k> = <R_m q, R_n k>
m, n = 5, 2
base_score = torch.dot(apply_rope_1d(q, m), apply_rope_1d(k, n)).item()
shifted_score = torch.dot(
    apply_rope_1d(q, m + 12), apply_rope_1d(k, n + 12)
).item()

max_err = max(
    abs(torch.dot(apply_rope_1d(q, m + d),
                  apply_rope_1d(k, n + d)).item() - base_score)
    for d in range(1, 21)
)
print(f"Điểm chú ý <R_5 q, R_2 k>   (m-n = 3):  {base_score:.6f}")
print(f"Điểm chú ý <R_17 q, R_14 k> (m-n = 3):  {shifted_score:.6f}")
print(f"Sai số cực đại khi dịch chuyển Delta in [1..20]: {max_err:.2e}")

Chuẩn L2 của q sau khi quay ở m=5: 3.7483 (sai khác 0.00e+00)
Điểm chú ý <R_5 q, R_2 k>   (m-n = 3):  2.184937
Điểm chú ý <R_17 q, R_14 k> (m-n = 3):  2.184937
Sai số cực đại khi dịch chuyển Delta in [1..20]: 9.54e-07

4. Mạng Feed-Forward, Chuẩn hóa và Công thức Llama Hiện đại

Nếu Self-Attention là nơi các token trao đổi thông tin với nhau across positions, thì nửa còn lại của mỗi khối Transformer — mạng Feed-Forward (FFN / MLP) — lại tác động độc lập lên từng token một để xử lý thông tin vừa thu thập được và tra cứu tri thức lưu trong trọng số.

Từ kiến trúc Transformer gốc năm 2017 (Vaswani et al., 2017) đến kiến trúc chuẩn của dòng mô hình LLaMA (Touvron & others, 2023) và Llama-3 (Dubey & others, 2024), cộng đồng nghiên cứu đã thực hiện ba cải tiến quyết định về chuẩn hóa và mạng FFN.

4.1 Từ Post-LayerNorm sang Pre-RMSNorm: Đường cao tốc cho Đạo hàm

Trong bài báo gốc năm 2017, lớp chuẩn hóa LayerNorm (Ba & Hinton, 2016) được đặt sau phép cộng thặng dư (Post-LN):

Năm 2020, Xiong và cộng sự (Xiong & others, 2020) đã chứng minh bằng giải tích rằng thiết kế Post-LN khiến đạo hàm truyền ngược qua mỗi lớp LayerNorm bị chia nhỏ theo hệ số $\frac{1}{\sqrt{1 + \sigma^2}}$, làm đạo hàm ở các tầng đầu tiên suy giảm theo hàm mũ theo độ sâu $L$, còn đạo hàm ở tầng cuối lại quá lớn. Đó là lý do Transformer gốc bắt buộc phải có giai đoạn khởi động tốc độ học (learning-rate warmup) rất nhạy cảm, nếu không mô hình sẽ phân kỳ ngay từ những bước đầu.

Giải pháp hiện đại là chuyển lớp chuẩn hóa lên trước nhánh con (Pre-Norm), giữ cho luồng thặng dư hoàn toàn thông suốt từ đầu đến cuối:

Khi lấy đạo hàm từ tầng cuối $x_L = x_0 + \sum_{\ell=1}^L F_\ell(\mathrm{Norm}(x_{\ell-1}))$ về bất kỳ tầng $\ell$ nào, luôn có số hạng đạo hàm đơn vị $I$ truyền thẳng không bị cản trở!

Hơn nữa, năm 2019, Biao Zhang và Rico Sennrich (Zhang & Sennrich, 2019) nhận thấy rằng phép trừ trung bình $\mu$ trong LayerNorm không đóng góp gì cho hiệu năng của Transformer mà chỉ tốn thêm phép tính. Họ đề xuất RMSNorm (Root Mean Square Layer Normalization) — chỉ chuẩn hóa theo căn bậc hai của trung bình bình phương:

trong đó $\gamma \in \mathbb{R}^d$ là vector hệ số tỷ lệ học được (không có hệ số dịch $\beta$).

4.2 Từ ReLU và GELU sang Cổng Tuyến tính SwiGLU

Trong Transformer năm 2017, mạng FFN gồm hai phép chiếu tuyến tính kẹp giữa một hàm kích hoạt ReLU (hoặc GELU (Hendrycks & Gimpel, 2016) trong GPT-2):

Năm 2020, Noam Shazeer (một trong các đồng tác giả của Transformer gốc) đã thử nghiệm thay thế hàm kích hoạt đơn lẻ bằng cơ chế Đơn vị Tuyến tính có Cổng (Gated Linear Units — GLU), cụ thể là biến thể SwiGLU sử dụng hàm $\mathrm{SiLU}(z) = z \cdot \sigma(z)$ (còn gọi là hàm Swish) (Shazeer, 2020):

trong đó $\odot$ là phép nhân từng phần tử (element-wise product), $W_{\text{gate}}, W_{\text{up}} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$ và $W_{\text{down}} \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$.

Tại sao SwiGLU cần tới 3 ma trận ($W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}$) thay vì 2 ma trận? Vì nhánh $x W_{\text{gate}}$ đi qua hàm kích hoạt mượt $\mathrm{SiLU}$ đóng vai trò như một cánh cổng mềm quyết định cho bao nhiêu thông tin từ nhánh tuyến tính $x W_{\text{up}}$ đi qua, tạo thành tích bậc hai giữa các đặc trưng đầu vào. Để tổng số tham số $3 \cdot d_{\text{model}} \cdot d_{\text{ff}}$ xấp xỉ bằng $2 \cdot d_{\text{model}} \cdot (4 d_{\text{model}}) = 8 d_{\text{model}}^2$ của FFN cũ, chiều trung gian $d_{\text{ff}}$ trong SwiGLU thường được chọn bằng:

4.3 Bảng đối chiếu: Kiến trúc Transformer 2017 vs GPT-2 vs Llama-3

Bảng dưới đây tóm tắt sự tiến hóa từ kiến trúc Transformer gốc đến công thức chuẩn hiện đại (Dubey & others, 2024; Touvron & others, 2023):

5. Giải phẫu Tham số và Độ phức tạp Tính toán (FLOPs)

Trước khi bắt tay lắp ráp một mô hình hoàn chỉnh, một kỹ năng không thể thiếu của kỹ sư LLM là nhẩm tính chính xác số lượng tham số ($N$) và số phép tính dấu phẩy động (FLOPs) chỉ từ vài con số cấu hình.

5.1 Công thức đếm tham số của một mô hình ngôn ngữ

Giả sử mô hình có $L_{\text{layers}}$ tầng, chiều mô hình $d_{\text{model}} = d$, chiều ẩn FFN là $d_{\text{ff}}$, và từ vựng có $|V|$ token:

1. Nhúng đầu vào ($W_E$) và Đầu dự đoán ($W_{\text{LM}}$): Mỗi ma trận có kích thước $|V| \times d$. Nếu không chia sẻ trọng số (untied embeddings như Llama-3), chúng chiếm $2 |V| d$ tham số.
2. Mỗi khối Multi-Head Attention: Gồm 4 ma trận $W_Q, W_K, W_V, W_O$. Với MHA tiêu chuẩn, mỗi ma trận có kích thước $d \times d$, tổng cộng là $4 d^2$ tham số mỗi tầng (hoặc $d(d + 2 d_{\text{kv}}) + d^2$ nếu dùng Grouped-Query Attention như Llama-3, chúng ta sẽ học kỹ ở Chương 5).
3. Mỗi khối FFN:
   • Với FFN truyền thống ($d_{\text{ff}} = 4d$, 2 ma trận): $2 \cdot d \cdot 4d = 8 d^2$ tham số.
   • Với SwiGLU ($3$ ma trận $W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}$): $3 \cdot d \cdot d_{\text{ff}}$ tham số (nếu $d_{\text{ff}} \approx \frac{8}{3}d$, con số này cũng bằng $8 d^2$).
4. Hai lớp RMSNorm mỗi tầng: Chỉ chiếm $2 d$ tham số mỗi tầng (rất nhỏ so với $12 d^2$).

Như vậy, tổng số tham số truyền qua các tầng tuyến tính (không tính bảng tra cứu Embedding) của toàn mô hình là:

5.2 Tại sao quy tắc $C \approx 6 N D$ lại đúng?

Khi nhân một vector $1 \times d_{\text{in}}$ với một ma trận trọng số $W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$ (chứa $P = d_{\text{in}} d_{\text{out}}$ tham số), mỗi phần tử đầu ra cần $d_{\text{in}}$ phép nhân và $d_{\text{in}}$ phép cộng, tức là mỗi tham số trong ma trận đóng góp đúng $2$ FLOPs cho mỗi token đi qua!

• Lượt truyền tiến (Forward pass): Đi qua toàn bộ $N$ tham số của mô hình một lần cho mỗi token tốn $\approx 2N$ FLOPs/token.
• Lượt truyền ngược (Backward pass): Để tính đạo hàm theo kích hoạt đầu vào ($\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} W^\top$, tốn $2N$ FLOPs) và đạo hàm theo trọng số ($\frac{\partial \mathcal{L}}{\partial W} = x^\top \frac{\partial \mathcal{L}}{\partial y}$, tốn thêm $2N$ FLOPs), lượt truyền ngược tốn gấp đôi lượt truyền tiến, tức là $\approx 4N$ FLOPs/token.

Cộng cả hai lượt lại, chi phí huấn luyện một mô hình có $N$ tham số trên một tập dữ liệu gồm $D$ token tuân theo công thức vàng mà chúng ta sẽ dùng xuyên suốt Chương 3:

Bây giờ, hãy kết hợp toàn bộ kiến thức trong chương này để lắp ráp hoàn chỉnh một khối LlamaBlock từ đầu trong Bài tập 2.5!

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

torch.manual_seed(0)
cfg = dict(d_model=64, n_heads=4, d_ff=172)
x = torch.randn(2, 12, cfg["d_model"])
print("Cấu hình khối Mini-Llama:", cfg)

Cấu hình khối Mini-Llama: {'d_model': 64, 'n_heads': 4, 'd_ff': 172}

Kích thước đầu ra TransformerBlock: (2, 12, 64)
Tham số RMSNorm (2 * d_model):       128
Tham số Attention (4 * d_model^2):   16384
Tham số SwiGLU FFN (3 * d * d_ff):   33024
Tổng tham số lý thuyết vs thực tế:   49536 == 49536 (Khớp 100%!)

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

torch.manual_seed(0)
cfg = dict(d_model=64, n_heads=4, d_ff=172)
x = torch.randn(2, 12, cfg["d_model"])
print("Cấu hình khối Mini-Llama:", cfg)

Cấu hình khối Mini-Llama: {'d_model': 64, 'n_heads': 4, 'd_ff': 172}

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x * rms * self.weight

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

class LlamaBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn_qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.attn_out = nn.Linear(d_model, d_model, bias=False)
        self.norm2 = RMSNorm(d_model)
        self.ffn = SwiGLU(d_model, d_ff)
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

    def forward(self, x):
        B, L, D = x.shape
        q, k, v = self.attn_qkv(self.norm1(x)).chunk(3, dim=-1)
        q = q.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        mask = torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1)
        attn = F.softmax(scores.masked_fill(mask, float("-inf")), dim=-1)
        h = x + self.attn_out((attn @ v).transpose(1, 2).reshape(B, L, D))
        return h + self.ffn(self.norm2(h))

block = LlamaBlock(**cfg)
out = block(x)
d, d_ff = cfg["d_model"], cfg["d_ff"]
theo = 2 * d + 4 * d * d + 3 * d * d_ff
actual = sum(p.numel() for p in block.parameters())

print("Kích thước đầu ra TransformerBlock:", tuple(out.shape))
print(f"Tham số RMSNorm (2 * d_model):       {2 * d}")
print(f"Tham số Attention (4 * d_model^2):   {4 * d * d}")
print(f"Tham số SwiGLU FFN (3 * d * d_ff):   {3 * d * d_ff}")
print(f"Tổng tham số lý thuyết vs thực tế:   {theo} == {actual} (Khớp 100%!)")

Kích thước đầu ra TransformerBlock: (2, 12, 64)
Tham số RMSNorm (2 * d_model):       128
Tham số Attention (4 * d_model^2):   16384
Tham số SwiGLU FFN (3 * d * d_ff):   33024
Tổng tham số lý thuyết vs thực tế:   49536 == 49536 (Khớp 100%!)

Nếu muốn đào sâu hơn nữa vào cơ học bên trong khối Transformer, bạn có thể thử sức với các bài tập mở rộng dưới đây trước khi bước sang Chương 3: Tiền huấn luyện (Pre-training).

Tài liệu tham khảo

1. 2016 Ba, Kiros & Hinton. Layer Normalization. arXiv. 4
2. 2016 Hendrycks & Gimpel. Gaussian Error Linear Units (GELUs). arXiv. 4
3. 2017 Vaswani et al. Attention Is All You Need. NeurIPS. 1
4. 2019 Zhang & Sennrich. Root Mean Square Layer Normalization. NeurIPS. 4
5. 2020 Shazeer. GLU Variants Improve Transformer. arXiv. 4
6. 2020 Xiong et al. On Layer Normalization in the Transformer Architecture. ICML. 4
7. 2021 Elhage et al. A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread. 2
8. 2021 Press, Smith & Lewis. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022. 3
9. 2023 Touvron et al. LLaMA: Open and Efficient Foundation Language Models. arXiv. 4, 5
10. 2024 Dubey et al. The Llama 3 Herd of Models. arXiv. 5
11. 2024 Su et al. RoFormer: Enhanced Transformer with Rotary Position Embedding. Neurocomputing (arXiv 2021). 3