Đọc paper Attention Is All You Need, chín năm sau

Paper "Attention Is All You Need" (Vaswani và cộng sự, NIPS 2017, hội nghị nay tên là NeurIPS) đề xuất Transformer: mô hình dịch máy bỏ hẳn RNN và convolution, chỉ dùng attention. GPT-2, PaLM và Llama 3 đều ghi trong paper của mình rằng kiến trúc dựa trên Transformer này. Bài này tính attention bằng numpy với 3 token, đo vì sao phải chia cho √d_k, rồi đối chiếu từng lựa chọn năm 2017 với các LLM hôm nay. Số liệu về paper lấy từ bản arXiv v7 (2023-08-02), có ghi mục.

1. Paper giải quyết vấn đề gì #

RNN tính trạng thái ẩn h_t từ h_(t−1) và token thứ t (mục 1). Token thứ 100 phải đợi 99 bước trước nó, nên trong một câu không chạy song song được. Paper gọi đây là ràng buộc cơ bản của tính toán tuần tự.

Vấn đề thứ hai là đường đi giữa hai token xa nhau. Trong RNN, thông tin từ token đầu tới token thứ n qua n − 1 bước biến đổi. Đường càng dài, phụ thuộc xa càng khó học (mục 4). ConvS2S và ByteNet song song được, nhưng số phép nối hai vị trí vẫn tăng theo khoảng cách: tuyến tính với ConvS2S, logarit với ByteNet (mục 2).

Bảng 1 của paper: n là độ dài chuỗi, d là số chiều biểu diễn, k là kích thước kernel, r là kích thước vùng lân cận của self-attention giới hạn.

Self-attention đưa hai cột cuối về O(1) và trả giá ở cột đầu. Tỉ số chi phí giữa self-attention và recurrent là (n²·d) / (n·d²) = n/d. Self-attention rẻ hơn khi n < d. Tính theo bậc, bỏ hằng số, với d = 512:

n = 50:    self-attention 50² × 512    = 1.280.000       recurrent 50 × 512²    = 13.107.200      → rẻ hơn 10,24 lần
n = 8.192: self-attention 8.192² × 512 = 34.359.738.368  recurrent 8.192 × 512² = 2.147.483.648   → đắt hơn 16 lần

Năm 2017, paper nhận xét n < d là trường hợp thường gặp với câu trong dịch máy (mục 4). Ngữ cảnh dài của LLM hôm nay đảo ngược điều kiện đó. Mục 5 quay lại chỗ này.

2. Scaled dot-product attention, tính tay với 3 token #

Attention ánh xạ một query và một tập cặp key–value thành output. Output là tổng có trọng số của các value, trọng số tính từ độ khớp giữa query và từng key (mục 3.2). Gom query, key, value thành ma trận Q, K, V, ta có công thức (1) ở mục 3.2.1:

Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V
Q: n × d_k     K: n × d_k     V: n × d_v
Q·Kᵀ: n × n, mỗi query nhân với mỗi key; softmax lấy theo từng hàng, mỗi hàng cộng lại bằng 1

Ví dụ dùng 3 token, d_k = d_v = 4. Số được chọn để tính nhẩm được. Ba cột đầu của V là ma trận đơn vị, nên ba cột đầu của output chính là trọng số attention. Cột cuối của V là 1, 2, 3 để thấy phép lấy trung bình.

import numpy as np

np.set_printoptions(precision=3, suppress=True)

# 3 token, d_k = d_v = 4. Mỗi hàng là vector của một token.
Q = np.array([[1, 0, 1, 0],
              [0, 2, 0, 2],
              [1, 1, 1, 1]], dtype=float)
K = np.array([[1, 0, 1, 0],
              [0, 1, 0, 1],
              [1, 1, 0, 0]], dtype=float)
V = np.array([[1, 0, 0, 1],
              [0, 1, 0, 2],
              [0, 0, 1, 3]], dtype=float)

def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)  # trừ max từng hàng để exp không tràn số
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    weights = softmax(scores)
    return scores, weights, weights @ V

scores, weights, out = attention(Q, K, V)
print("scores = QK^T / sqrt(4):\n", scores)
print("weights = softmax(scores):\n", weights)
print("tong moi hang:", weights.sum(axis=1))
print("output = weights @ V:\n", out)

Kết quả, Python 3.14.3 và numpy 2.4.4:

scores = QK^T / sqrt(4):
 [[1.  0.  0.5]
 [0.  2.  1. ]
 [1.  1.  1. ]]
weights = softmax(scores):
 [[0.506 0.186 0.307]
 [0.09  0.665 0.245]
 [0.333 0.333 0.333]]
tong moi hang: [1. 1. 1.]
output = weights @ V:
 [[0.506 0.186 0.307 1.801]
 [0.09  0.665 0.245 2.155]
 [0.333 0.333 0.333 2.   ]]

• Hàng 1, q = [1, 0, 1, 0]. Tích vô hướng với ba key là 2, 0, 1. Chia √4 = 2 được 1; 0; 0,5. Softmax cho 0,506; 0,186; 0,307. Output là 0,506·v₁ + 0,186·v₂ + 0,307·v₃. Cột cuối 1,801 là trung bình có trọng số của 1, 2, 3, nghiêng về 1 vì token 1 được chú ý nhiều nhất.
• Hàng 2, q = [0, 2, 0, 2]. Tích vô hướng là 0, 4, 2, chia 2 được 0; 2; 1. Key 2 khớp nhất và nhận trọng số 0,665. Cột cuối 2,155 nghiêng về 2.
• Hàng 3, q = [1, 1, 1, 1]. Tích với mọi key đều bằng 2, nên điểm đều bằng 1 và trọng số đều bằng 1/3. Output là trung bình cộng ba hàng của V. Softmax chỉ nhìn chênh lệch giữa các điểm, không nhìn độ lớn tuyệt đối.

Ba hàng tính độc lập, gói trong hai phép nhân ma trận, nên chạy song song được trong một câu. Trong decoder, điểm ở vị trí j > i được đặt thành −∞ trước softmax (mục 3.2.3). Trọng số ở đó thành 0, nên token i không nhìn thấy token sau nó.

3. Vì sao chia cho √d_k #

Paper không chứng minh, chỉ đưa một giả thuyết và một phép tính ở chú thích 4 của mục 3.2.1. Giả sử các thành phần của q và k độc lập, trung bình 0, phương sai 1. Khi đó:

q·k = q₁k₁ + q₂k₂ + ... + q_(d_k)·k_(d_k)
E[q·k] = 0       Var(q·k) = d_k
d_k = 64  → độ lệch chuẩn √64  = 8
d_k = 512 → độ lệch chuẩn √512 ≈ 22,6

Điểm chênh nhau vài chục đơn vị đẩy softmax về gần one-hot, vùng mà paper gọi là có "extremely small gradients". Paper dẫn Britz và cộng sự (2017): khi d_k lớn, additive attention thắng dot-product attention không chia. Chia cho √d_k kéo phương sai về 1.

Đoạn code dưới đo phương sai của q·k, rồi đo softmax của một query trên 8 key: trọng số lớn nhất, và độ lớn gradient tính bằng chuẩn Frobenius của Jacobian J = diag(p) − p·pᵀ.

import numpy as np

rng = np.random.default_rng(1706)

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

for d_k in (4, 512):
    # 1) Phương sai của q·k khi mỗi thành phần ~ N(0, 1), 100.000 cặp
    q = rng.standard_normal((100_000, d_k))
    k = rng.standard_normal((100_000, d_k))
    dot = (q * k).sum(axis=1)
    print(f"d_k={d_k:3d}  var(q.k)={dot.var():7.2f}  var(q.k/sqrt(d_k))={(dot / np.sqrt(d_k)).var():.3f}")

    # 2) Softmax trên 8 key, 2.000 query: trọng số lớn nhất và độ lớn gradient
    Qm = rng.standard_normal((2_000, d_k))
    Km = rng.standard_normal((8, d_k))
    for name, s in (("khong chia", Qm @ Km.T), ("chia sqrt ", Qm @ Km.T / np.sqrt(d_k))):
        p = softmax(s)
        # Jacobian của softmax: J = diag(p) - p p^T, đo bằng chuẩn Frobenius
        J = p[:, :, None] * np.eye(8) - p[:, :, None] * p[:, None, :]
        grad = np.linalg.norm(J, axis=(1, 2))
        print(f"   {name}: max weight TB={p.max(axis=1).mean():.3f}"
              f"  ||J|| trung vi={np.median(grad):.2e}")

Kết quả, Python 3.14.3, numpy 2.4.4, seed 1706:

d_k=  4  var(q.k)=   4.02  var(q.k/sqrt(d_k))=1.006
   khong chia: max weight TB=0.551  ||J|| trung vi=3.79e-01
   chia sqrt : max weight TB=0.360  ||J|| trung vi=3.67e-01
d_k=512  var(q.k)= 507.38  var(q.k/sqrt(d_k))=0.991
   khong chia: max weight TB=0.958  ||J|| trung vi=7.12e-05
   chia sqrt : max weight TB=0.365  ||J|| trung vi=3.69e-01

• Phương sai đo được 4,02 và 507,38, khớp với d_k. Sau khi chia, cả hai về quanh 1.
• Với d_k = 4, chia hay không gần như như nhau (‖J‖ 0,379 và 0,367). Độ lệch chuẩn 2 chưa đủ làm softmax bão hòa.
• Với d_k = 512 không chia, trọng số lớn nhất trung bình 0,958, ‖J‖ trung vị 7,12·10⁻⁵. Có chia: 0,365 và 0,369. Gradient qua softmax nhỏ đi khoảng 5.000 lần.

Giới hạn của phép đo: d_k của paper là 64 mỗi head, 512 chọn để thấy rõ hiệu ứng. Giả định phương sai 1 chỉ gần đúng lúc khởi tạo.

4. Phần còn lại của kiến trúc #

Multi-head attention (mục 3.2.2). Thay vì một attention trên vector 512 chiều, paper chiếu Q, K, V xuống h không gian con nhỏ hơn và chạy h attention song song:

MultiHead(Q, K, V) = Concat(head₁, ..., head_h) · Wᴼ
headᵢ = Attention(Q·Wᵢ^Q, K·Wᵢ^K, V·Wᵢ^V)
h = 8,  d_k = d_v = d_model / h = 512 / 8 = 64
Wᵢ^Q, Wᵢ^K, Wᵢ^V: 512 × 64       Wᴼ: (8 × 64) × 512 = 512 × 512

Mỗi head là một lần tính như mục 2, trên phép chiếu riêng 64 chiều. Tổng chi phí gần bằng một head 512 chiều. Bảng 3 hàng (A), đo BLEU trên tập dev newstest2013: 1 head được 24,9, kém 0,9 so với 25,8 của cấu hình 8 head. 32 head cũng giảm, còn 25,4.

Positional encoding (mục 3.5). Attention không biết thứ tự. Hoán vị các token thì output hoán vị theo. Paper cộng vào embedding một vector vị trí cố định:

PE(pos, 2i)   = sin(pos / 10.000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10.000^(2i/d_model))
pos = 3, i = 0:    sin(3) = 0,141     cos(3) = −0,990
pos = 3, i = 255:  3 / 10.000^(510/512) = 3 / 9.646,6 ≈ 0,000311
                   sin ≈ 0,000311     cos ≈ 1,000

Chiều thấp đổi nhanh theo vị trí, chiều cao đổi rất chậm: bước sóng tăng theo cấp số nhân từ 2π đến 10.000·2π. Bảng 3 hàng (E): embedding vị trí học được cho 25,7, gần bằng 25,8 của sinusoid. Paper chọn sinusoid vì nó "may allow" ngoại suy ra chuỗi dài hơn lúc train, nhưng không có thí nghiệm kiểm chứng.

Encoder–decoder (mục 3.1, 3.3). Encoder gồm N = 6 lớp, mỗi lớp hai sublayer: multi-head self-attention, rồi mạng feed-forward áp riêng cho từng vị trí, FFN(x) = max(0, x·W₁ + b₁)·W₂ + b₂, với d_ff = 2.048. Decoder cũng 6 lớp, thêm sublayer thứ ba là encoder–decoder attention: query từ decoder, key và value từ output của encoder. Self-attention trong decoder có mask như cuối mục 2.

Residual và LayerNorm. Paper đặt LayerNorm sau phép cộng residual, gọi là post-LN:

Post-LN, paper mục 3.1:   x ← LayerNorm(x + Sublayer(x))
Pre-LN, GPT-2 và Llama:   x ← x + Sublayer(LayerNorm(x))

Huấn luyện và kết quả (mục 5, 6.1, Bảng 2, Bảng 3). Hai cấu hình, train trên một máy 8 GPU NVIDIA P100:

Optimizer là Adam với β₁ = 0,9, β₂ = 0,98, ε = 10⁻⁹. Learning rate theo công thức lrate = d_model^(−0,5) × min(step^(−0,5), step × 4.000^(−1,5)): tăng tuyến tính trong 4.000 bước warmup rồi giảm theo nghịch đảo căn bậc hai của số bước (mục 5.3). Đỉnh ở bước 4.000 là 512^(−0,5) × 4.000^(−0,5) ≈ 0,0442 × 0,0158 ≈ 7,0·10⁻⁴. Label smoothing ε_ls = 0,1 (mục 5.4).

5. Chín năm sau: cái gì còn, cái gì đổi #

Công thức (1) giữ nguyên. Paper Llama 3 (2024) gọi mô hình là "a standard, dense Transformer architecture (Vaswani et al., 2017)". Q, K, V, multi-head, residual và FFN theo từng vị trí vẫn ở đó. Phần đã đổi nằm quanh công thức:

Decoder-only. Bỏ encoder và encoder–decoder attention, giữ masked self-attention và FFN, train bằng dự đoán token kế tiếp. PaLM (2022) ghi rõ dùng Transformer "in a decoder-only setup" và xếp GPT-3 vào cùng lớp mô hình đó.

Pre-LN. GPT-2 (2019) chuyển LayerNorm lên đầu mỗi sublayer. Xiong và cộng sự (ICML 2020) chỉ ra rằng với post-LN, gradient gần lớp output lớn lúc khởi tạo nên cần warmup. Pre-LN có gradient ổn định và train được không cần warmup. Theo phân tích đó, 4.000 bước warmup ở mục 5.3 là phần cần thiết của cấu hình post-LN gốc. LLaMA (2023) dùng pre-normalization với RMSNorm "to improve the training stability".

RoPE. Su và cộng sự (2021) mã hóa vị trí bằng phép xoay áp lên q và k thay vì cộng vào embedding. Tích q·k sau khi xoay chỉ phụ thuộc vào nội dung của q, k và khoảng cách giữa hai vị trí. LLaMA thay embedding vị trí tuyệt đối bằng RoPE ở mỗi lớp. PaLM chọn RoPE vì hiệu quả tốt hơn trên chuỗi dài.

MQA và GQA. Lúc sinh văn bản, mỗi bước chỉ có một query mới nhưng phải đọc lại K và V của mọi token trước đó, gọi là KV cache. Shazeer (2019) chỉ ra nút thắt là băng thông bộ nhớ khi nạp K, V, và đề xuất multi-query attention (MQA): mọi head dùng chung một cặp K, V. Grouped-query attention (GQA, Ainslie và cộng sự 2023) dùng số KV head nằm giữa 1 và số query head, chất lượng gần multi-head, tốc độ gần MQA. PaLM dùng MQA. Llama 2 dùng GQA cho bản 34B và 70B, Llama 3 dùng GQA với 8 KV head ở cả ba cỡ. Tính KV cache cho Llama 3 8B theo Bảng 3 của paper Llama 3:

KV cache mỗi token = 2 (K và V) × số lớp × số KV head × d_head × 2 byte (float16)
Llama 3 8B: 32 lớp, 32 query head, 8 KV head, d_model 4.096 → d_head = 4.096 / 32 = 128
GQA, 8 KV head:   2 × 32 × 8  × 128 × 2 = 131.072 byte = 128 KiB mỗi token
MHA, 32 KV head:  2 × 32 × 32 × 128 × 2 = 524.288 byte = 512 KiB mỗi token
Chuỗi 128.000 token: 15,6 GiB với GQA, 62,5 GiB nếu là MHA
(ước lượng lý thuyết cho một chuỗi, chưa tính phần quản lý bộ nhớ của engine suy luận)

O(n²) và FlashAttention. Đoạn numpy ở mục 2 tạo nguyên ma trận điểm: 3 token cho 9 số. Với chuỗi dài, mỗi head một ma trận n × n, float16 chiếm 2 byte mỗi phần tử:

n = 8.192:    8.192²   = 67.108.864 phần tử      × 2 = 134.217.728 byte     = 128 MiB
n = 128.000:  128.000² = 16.384.000.000 phần tử  × 2 = 32.768.000.000 byte  ≈ 30,5 GiB
Bản base, h = 8, n = 8.192:          8 × 128 MiB  = 1 GiB cho một lớp
Llama 3 8B, 32 head, n = 128.000:    32 × 30,5 GiB ≈ 976,6 GiB cho một lớp

n = 128.000 là cỡ cửa sổ ngữ cảnh 128K mà paper Llama 3 công bố. Cách "tạo nguyên ma trận rồi softmax" không đi tới độ dài đó. FlashAttention (Dao và cộng sự, 2022) là thuật toán attention chính xác và IO-aware: chia Q, K, V thành khối vừa SRAM trên chip GPU, tính softmax theo từng khối, không ghi ma trận n × n ra HBM, và tính lại ở lượt backward thay vì lưu. Output giống công thức (1), chỉ khác sai số làm tròn. Bộ nhớ thêm vào tuyến tính theo n, số phép tính vẫn O(n²·d). FlashAttention đổi cách tính, không đổi công thức.

6. Đọc paper thế nào: phần nào kỹ, phần nào lướt #

Paper dài 15 trang tính cả phụ lục. Thứ tự đọc cho kỹ sư cần hiểu mô hình để dùng hoặc tối ưu:

• Đọc kỹ mục 3.2.1 cùng chú thích 4, mục 3.2.2, rồi tự tính lại ví dụ 3 token với số của mình. Đây là phần còn nguyên giá trị.
• Đọc kỹ Bảng 1 và mục 4. Cột độ phức tạp giải thích vì sao ngữ cảnh dài vẫn tốn kém.
• Đọc kỹ mục 3.1 và 3.2.3: ba kiểu attention, mask, vị trí LayerNorm. Sau đó so với bảng ở mục 5 bài này.
• Đọc có chọn lọc Bảng 3, hàng (A), (B), (E): bao nhiêu head, d_k bao nhiêu, sinusoid hay học. Số đo trên tập dev, không so trực tiếp với Bảng 2.
• Đọc nếu cần tái lập mục 5 và 6.1: batch khoảng 25.000 token nguồn và 25.000 token đích, lịch learning rate, trung bình 5 checkpoint cuối (base) hoặc 20 (big), beam search kích thước 4, α = 0,6.
• Lướt mục 2, mục 3.4, mục 6.3 và Hình 3 đến 5 ở phụ lục. Cột FLOPs của Bảng 2 là ước lượng từ thời gian × số GPU × một mức TFLOPS giả định (chú thích 5).

Những chỗ hay hiểu sai #

• "Chia √d_k là để chuẩn hóa thành xác suất." Softmax làm việc đó. √d_k giữ phương sai của điểm quanh 1 để softmax không bão hòa (chú thích 4).
• "Multi-head tốn gấp h lần một head." Mỗi head có d_model/h chiều. Tổng chi phí xấp xỉ một head đủ chiều (mục 3.2.2).
• "Transformer gốc là mô hình kiểu GPT." Bản gốc là encoder–decoder cho dịch máy. Decoder-only là biến thể các LLM dùng về sau.
• "FlashAttention là attention xấp xỉ." FlashAttention tính chính xác công thức (1). Nó giảm truy cập HBM, không giảm số phép tính bậc n².

Đọc tiếp #

• Embedding và tìm kiếm ngữ nghĩa tiếng Việt: biến câu thành vector và tìm theo nghĩa.
• Đánh giá một hệ RAG: đo chất lượng một hệ ghép tìm kiếm với LLM.

Nguồn #

• Vaswani và cộng sự, Attention Is All You Need, NIPS 2017, bản arXiv v7 (2023-08-02). Mục 1–6, Bảng 1–3, chú thích 4 và 5.
• Radford và cộng sự, Language Models are Unsupervised Multitask Learners, 2019. Mục 2.3: LayerNorm chuyển lên đầu sublayer.
• Brown và cộng sự, Language Models are Few-Shot Learners, 2020. Mục 2.1: kiến trúc GPT-3, pre-normalization.
• Xiong và cộng sự, On Layer Normalization in the Transformer Architecture, ICML 2020.
• Su và cộng sự, RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021.
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019.
• Ainslie và cộng sự, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023.
• Chowdhery và cộng sự, PaLM: Scaling Language Modeling with Pathways, 2022. Mục 1 và 2.
• Touvron và cộng sự, LLaMA: Open and Efficient Foundation Language Models, 2023. Mục 2.2.
• Touvron và cộng sự, Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023. Bảng 1.
• Llama Team, AI @ Meta, The Llama 3 Herd of Models, 2024. Mục 3.2, Bảng 3.
• Dao và cộng sự, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022.