Gần như mọi mô hình ngôn ngữ bạn nghe tên hiện nay, từ BERT, PhoBERT tới các mô hình sinh văn bản, đều xây trên cùng một kiến trúc: Transformer. Bài này không bàn mô hình nào "thông minh" hơn mà mở nắp kiến trúc đó ra. Ta tự cài cơ chế attention bằng PyTorch, đối chiếu từng con số với thư viện chuẩn, dựng một Transformer nhỏ học sắp xếp chữ số, rồi đọc ma trận attention của PhoBERT thật trên một câu tiếng Việt. Mục tiêu là sau bài, bạn nhìn một mô hình Transformer và nói được từng khối làm gì, tốn bao nhiêu tham số và vì sao chuỗi dài lại đắt.

Bài dùng PyTorch cơ bản ở PyTorch cơ bản: tensor, autograd và vòng lặp huấn luyện và biểu diễn từ ở Word embedding: Word2Vec, FastText và biểu diễn nghĩa của từ. Output dưới đây là kết quả thật, chạy trên CPU của máy Apple Silicon với PyTorch 2.2.2 (phát hành 03/2024); phần PhoBERT dùng transformers 4.38.2 và pyvi 0.1.1. Các API dùng trong bài (nn.MultiheadAttention, F.scaled_dot_product_attention) đã có từ PyTorch 2.0. Do khác phiên bản PyTorch và phần cứng, chạy lại có thể lệch vài phần nghìn ở độ chính xác và vài phần trăm ở các ô attention; kết luận không đổi.

Vì sao cần Transformer: giới hạn của RNN

Trước khi có Transformer vào năm 2017, mô hình chuỗi chủ đạo là mạng hồi quy (RNN, LSTM): đọc từng token một, gom thông tin vào một trạng thái ẩn. Hạn chế lộ rõ: token thứ 100 chỉ tính được sau token thứ 99 nên khó song song hoá, và thông tin từ token đầu phải đi qua cả trăm bước mới tới token cuối nên dễ phai.

Bài báo Attention Is All You Need của Vaswani và cộng sự (2017) bỏ hẳn hồi quy và tích chập, chỉ dùng cơ chế attention. Theo tóm tắt của bài, mô hình đạt 28,4 BLEU trên dịch Anh–Đức WMT 2014, hơn kết quả tốt nhất trước đó (kể cả tổ hợp nhiều mô hình) hơn 2 BLEU, và 41,8 BLEU trên Anh–Pháp sau 3,5 ngày huấn luyện trên tám GPU, với chi phí huấn luyện chỉ bằng một phần nhỏ các mô hình tốt nhất đương thời. Điểm mấu chốt: mọi vị trí được xử lý song song, và hai token bất kỳ nối với nhau trực tiếp trong một bước.

Attention trong Transformer: ba ma trận Q, K, V

Với mỗi token, mô hình sinh ba vector bằng ba phép nhân ma trận học được: query Q ("tôi cần thông tin gì"), key K ("tôi chứa thông tin gì") và value V ("nội dung tôi trao đi"). Độ tương hợp giữa token i và j là tích vô hướng của Qi với Kj; softmax biến một hàng điểm thành trọng số cộng bằng 1; đầu ra của token i là tổng có trọng số các V. Công thức trong bài báo: Attention(Q,K,V) = softmax(QK^T / sqrt(d_k)) V. Mã cài đặt (attention.py):

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    w = torch.softmax(scores, dim=-1)
    return w @ v, w

Chạy khối attention của Transformer trên bốn token với d=3 (dùng chính vector đầu vào làm Q, K, V để số đọc được):

trọng số (mỗi hàng tổng = 1):
 tensor([[0.320, 0.180, 0.320, 0.180],
        [0.180, 0.320, 0.320, 0.180],
        [0.230, 0.230, 0.410, 0.129],
        [0.209, 0.209, 0.209, 0.373]])
tổng hàng: tensor([1.000, 1.000, 1.000, 1.000])

Token thứ ba [1,1,0] giống cả token một lẫn hai nên chú ý đều 0,23 vào cả hai và 0,41 vào chính nó. Token thứ tư [0,0,1] vuông góc với ba token còn lại nên mọi điểm của nó với chúng bằng 0, và trọng số rơi đều 0,209.

Vì sao phải chia cho căn bậc hai của d_k

Nếu các thành phần của q và k độc lập, trung bình 0, phương sai 1, thì tích vô hướng q·k có phương sai bằng d_k, tức độ lệch chuẩn bằng căn của d_k. Với d_k lớn, các điểm to làm softmax bão hoà: một phần tử chiếm gần hết trọng số và gradient tới các phần tử còn lại gần 0. Đo với d_k = 64:

std(q.k)=8.00 (≈ sqrt(64)=8)   std(q.k/sqrt(d))=1.00
softmax khi chưa chia: max trọng số trung bình = 0.687 | sau khi chia: 0.096

Chưa chia, phần tử lớn nhất trong 64 chiếm trung bình 68,7% trọng số; chia xong chỉ còn 9,6%, tức attention vẫn "phân tán" và học được.

Multi-head attention trong Transformer: nhiều góc nhìn

Một đầu attention chỉ cho một kiểu trộn. Multi-head attention chia chiều d_model thành h phần, mỗi đầu có phép chiếu Q, K, V riêng với d_k = d_model / h, chạy attention song song rồi nối kết quả và nhân với ma trận đầu ra WO. Mỗi đầu có thể theo dõi một quan hệ khác nhau. Cài đặt:

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.d_k = n_heads, d_model // n_heads
        self.wq, self.wk, self.wv = (nn.Linear(d_model, d_model) for _ in range(3))
        self.wo = nn.Linear(d_model, d_model)

    def split(self, x):  # (B,T,D) -> (B,h,T,d_k)
        B, T, _ = x.shape
        return x.view(B, T, self.h, self.d_k).transpose(1, 2)

    def forward(self, x, mask=None):
        q, k, v = self.split(self.wq(x)), self.split(self.wk(x)), self.split(self.wv(x))
        out, w = scaled_dot_product_attention(q, k, v, mask)
        B, _, T, _ = out.shape
        return self.wo(out.transpose(1, 2).reshape(B, T, -1)), w

Cài tay Transformer dễ sai, nên ta kiểm chứng. Sao chép cùng trọng số vào torch.nn.MultiheadAttention và vào F.scaled_dot_product_attention rồi so kết quả:

chênh lệch lớn nhất đầu ra: 8.94e-08 | trọng số: 5.96e-08
so với F.scaled_dot_product_attention: 7.45e-08

Chênh lệch cỡ 1e-7 là sai số làm tròn float32, nên mã tự cài trùng thư viện. Trong thực tế hãy gọi hàm chuẩn: theo tài liệu PyTorch, scaled_dot_product_attention tự chọn giữa FlashAttention-2, memory-efficient attention và cài đặt C++ theo đầu vào, và nhận attn_mask hoặc is_causal (dùng cả hai cùng lúc sẽ báo lỗi).

Mặt nạ nhân quả: Transformer sinh văn bản không "nhìn trộm"

Khi sinh văn bản từng token, token thứ i không được thấy các token sau nó, nếu không việc huấn luyện thành gian lận. Cách làm là đặt điểm tại các vị trí j lớn hơn i bằng âm vô cực trước softmax, để trọng số đúng bằng 0:

tensor([[1.000, 0.000, 0.000, 0.000, 0.000],
        [0.017, 0.983, 0.000, 0.000, 0.000],
        [0.033, 0.064, 0.903, 0.000, 0.000],
        [0.210, 0.151, 0.169, 0.471, 0.000],
        [0.012, 0.054, 0.063, 0.005, 0.866]])

Đây là khác biệt giữa hai họ Transformer: BERT và PhoBERT là encoder hai chiều (không mặt nạ, mỗi token thấy cả câu), còn các mô hình sinh văn bản kiểu GPT là decoder dùng mặt nạ nhân quả.

Mã hoá vị trí: Transformer cần biết thứ tự

Attention là phép trộn có trọng số, nên nó không có khái niệm "đứng trước, đứng sau". Thử chứng minh: hoán vị ngẫu nhiên các token đầu vào, rồi hoán vị ngược đầu ra, kết quả phải giống hệt nếu attention đồng biến với hoán vị (permutation equivariant: đổi thứ tự đầu vào thì đầu ra đổi thứ tự tương ứng, không có thông tin nào về vị trí):

hoán vị đầu vào rồi hoán vị lại đầu ra giống hệt: True
sau khi cộng sinusoidal position, còn đồng biến với hoán vị? False

Vì vậy Transformer phải được đưa vị trí vào. Bài báo gốc cộng vào embedding một vector sin/cos với tần số khác nhau theo từng chiều: PE(pos,2i)=sin(pos/10000^(2i/d)), PE(pos,2i+1)=cos(...). Vị trí 0 cho [0, 1, 0, 1, ...], vị trí 1 cho [0.841, 0.54, 0.533, 0.846, 0.311, 0.95] ở sáu chiều đầu: cùng một mẫu nhưng lệch pha. Sau khi cộng vector này, hoán vị đầu vào làm đầu ra đổi (kết quả False ở dòng thứ hai), tức mô hình đã phân biệt được thứ tự. Nhiều mô hình khác dùng vị trí học được, nhưng nguyên tắc là một: phải tiêm vị trí từ ngoài.

Dựng Transformer nhỏ: học sắp xếp 8 chữ số

Một khối Transformer encoder gồm multi-head attention, mạng feed-forward hai lớp (mở rộng 4 lần rồi thu về), hai lớp LayerNorm và hai kết nối tắt. Ta dùng dạng pre-LN: chuẩn hoá trước mỗi nhánh, thường giúp huấn luyện ổn định hơn dạng post-LN của bài báo gốc (phân tích trong On Layer Normalization in the Transformer Architecture, Xiong và cộng sự, 2020). Mã (mini_transformer.py):

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn = MultiHeadAttention(D, H)
        self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D)
        self.ff = nn.Sequential(nn.Linear(D, 4 * D), nn.GELU(), nn.Linear(4 * D, D))

    def forward(self, x):
        a, w = self.attn(self.ln1(x))
        x = x + a
        return x + self.ff(self.ln2(x)), w

Bài toán: nhận 8 chữ số ngẫu nhiên từ 0 đến 9 và trả về chúng đã sắp xếp. Mô hình 2 lớp, d=64, 4 đầu, 101.258 tham số, huấn luyện 1.500 bước với batch 128; việc so sánh các phần tử với nhau đúng là sở trường của attention. Ta chạy hai bản Transformer: có và không có mã hoá vị trí.

  [pos=True] bước    0 loss=9.5350
  [pos=True] bước  500 loss=0.1120
  [pos=True] bước 1000 loss=0.0210
  [pos=True] bước 1499 loss=0.0079
  => pos=True: đúng từng vị trí 0.998 | đúng cả chuỗi 0.988 | 101258 tham số | 15s
  [pos=False] bước    0 loss=9.4835
  [pos=False] bước  500 loss=1.7063
  [pos=False] bước 1000 loss=1.6927
  [pos=False] bước 1499 loss=1.6838
  => pos=False: đúng từng vị trí 0.301 | đúng cả chuỗi 0.000 | 101258 tham số | 15s

Trên 5.000 chuỗi mới, Transformer có vị trí đúng nguyên cả chuỗi 98,8% sau 15 giây huấn luyện CPU; bản không có vị trí kẹt ở loss 1,68, đúng cả chuỗi 0%. Nguyên nhân không phải thiếu tham số (cả hai đều 101.258) mà là tính đồng biến với hoán vị ở trên: mỗi vị trí đầu ra cần biết "tôi là phần tử nhỏ thứ mấy".

Đọc ma trận attention của Transformer đã học

Với đầu vào [7, 2, 9, 2, 5, 0, 8, 3], mô hình trả đúng [0, 2, 2, 3, 5, 7, 8, 9]. Ma trận attention lớp 2, trung bình 4 đầu (hàng là vị trí đầu ra, cột là vị trí đầu vào):

tensor([[0.00, 0.18, 0.00, 0.20, 0.01, 0.32, 0.00, 0.28],
        [0.00, 0.26, 0.02, 0.29, 0.01, 0.10, 0.00, 0.31],
        [0.02, 0.21, 0.01, 0.21, 0.15, 0.03, 0.02, 0.35],
        [0.05, 0.12, 0.00, 0.13, 0.27, 0.00, 0.02, 0.40],
        [0.17, 0.08, 0.01, 0.08, 0.46, 0.01, 0.05, 0.13],
        [0.53, 0.01, 0.01, 0.01, 0.23, 0.00, 0.14, 0.06],
        [0.52, 0.00, 0.07, 0.00, 0.02, 0.00, 0.36, 0.01],
        [0.50, 0.00, 0.24, 0.00, 0.00, 0.00, 0.26, 0.00]])

Có thể đọc được một phần cấu trúc. Hàng 4 (đầu ra phải là 5) dồn 0,46 vào cột 4, nơi đầu vào đúng là 5. Hàng 5 (cần 7) dồn 0,53 vào cột 0, nơi đầu vào là 7; hàng 6 (cần 8) chú ý cột 0 và cột 6 (giá trị 8); hàng 7 (cần 9) trải trên cột 0, 2 (giá trị 9) và 6. Bốn hàng đầu (cần 0, 2, 2, 3) dồn vào các cột 1, 3, 5, 7, nơi chứa đúng các số nhỏ đó. Đừng đọc quá: đây là trung bình bốn đầu, không phải thuật toán sắp xếp gọn. Bài Attention is not Explanation (Jain và Wallace, 2019) cảnh báo trọng số attention không phải lúc nào cũng giải thích được dự đoán, nên coi ma trận này là manh mối để kiểm tra giả thuyết, không phải bằng chứng.

Chi phí bậc hai của Transformer theo độ dài chuỗi

Ma trận điểm có kích thước T x T cho mỗi đầu, nên bộ nhớ và thời gian tăng theo bình phương độ dài. Đo một lớp attention của Transformer (d=64, 4 đầu, batch=1) trên CPU:

T=  256  ma trận điểm      1.0 MB (float32, 4 đầu)  thời gian     0.3 ms
T=  512  ma trận điểm      4.2 MB (float32, 4 đầu)  thời gian     1.1 ms
T= 1024  ma trận điểm     16.8 MB (float32, 4 đầu)  thời gian     3.9 ms
T= 2048  ma trận điểm     67.1 MB (float32, 4 đầu)  thời gian    12.7 ms
T= 4096  ma trận điểm    268.4 MB (float32, 4 đầu)  thời gian    45.5 ms

Mỗi lần gấp đôi T, bộ nhớ ma trận điểm tăng đúng 4 lần, thời gian tăng khoảng 3,5–4 lần (từ 12,7 lên 45,5 ms ở hai dòng cuối). Con số này chỉ cho một lớp và một chuỗi; nhân thêm số lớp và batch là hết RAM rất nhanh. Đó là lý do FlashAttention (Dao và cộng sự, 2022) ra đời: nó tính attention chính xác nhưng chia khối để giảm số lần đọc ghi bộ nhớ GPU, báo cáo tăng tốc 3 lần trên GPT-2 với chuỗi 1.000 token.

PhoBERT là Transformer thật: tham số nằm ở đâu

Cuối cùng, kiểm chứng những gì vừa dựng trên mô hình thật. PhoBERT (Nguyễn Quốc Đạt và Nguyễn Tuấn Anh, 2020) là mô hình ngôn ngữ tiền huấn luyện cho tiếng Việt, cho kết quả tốt nhất ở nhiều tác vụ tiếng Việt khi công bố. Đọc cấu hình bằng transformers (phobert_attention.py):

số lớp=12 | số đầu=12 | d_model=768 | d_k=64 | d_ff=3072 | max_pos=258
tổng tham số: 135.0 triệu
  embedding 49.4 triệu (37%), 12 khối encoder 85.1 triệu
  một khối: attention 2.36 triệu | feed-forward 4.72 triệu

So với Transformer gốc bản base (6 lớp, d_model 512, 8 đầu, d_ff 2048), PhoBERT-base sâu gấp đôi và rộng hơn, nhưng d_k vẫn là 64: 768/12. Hai điều rút ra. Thứ nhất, feed-forward chiếm gấp đôi attention trong mỗi khối (4,72 so với 2,36 triệu tham số), nên không chỉ attention làm nên mô hình. Thứ hai, 37% tham số nằm ở bảng embedding của kho từ vựng 64.001 token.

Attention của PhoBERT trên câu tiếng Việt

Câu thử: "Giảng viên giải thích rất dễ hiểu nhưng phòng học quá nóng". PhoBERT yêu cầu đầu vào đã tách từ, nên ta dùng pyvi nối các âm tiết trong từ ghép bằng dấu gạch dưới:

token: ['<s>', 'Giảng_viên', 'giải_thích', 'rất', 'dễ', 'hiểu', 'nhưng', 'phòng', 'học', 'quá', 'nóng', '</s>']
hình dạng attention (lớp, đầu, T, T): (12, 12, 12, 12) | tổng mỗi hàng = 1.0

Một tensor 12 lớp x 12 đầu, mỗi đầu là ma trận 12 x 12 như phần ta tự cài, mỗi hàng cộng bằng 1. Hỏi token "nóng" chú ý nhiều nhất vào đâu ở vài đầu khác nhau:

  lớp  1 đầu  1: nóng:0.34, <s>:0.23, rất:0.07
  lớp  1 đầu  6: Giảng_viên:0.41, hiểu:0.10, phòng:0.07
  lớp  6 đầu  4: nóng:0.29, <s>:0.29, </s>:0.12
  lớp  9 đầu  8: </s>:0.91, quá:0.04, <s>:0.03
  lớp 12 đầu  1: nóng:0.48, <s>:0.25, </s>:0.13
  lớp 12 đầu 10: </s>:0.51, <s>:0.46, phòng:0.01

Các đầu hành xử rất khác nhau: lớp 1 đầu 1 nhìn chính nó, lớp 1 đầu 6 nhìn từ xa tới "Giảng_viên", còn lớp 9 đầu 8 dồn 91% vào token kết thúc câu. Đúng tinh thần công trình What Does BERT Look At? (Clark và cộng sự, 2019): nhiều đầu BERT tập trung vào token phân cách, một số theo độ lệch vị trí cố định, một số trải rộng. Số liệu tổng hợp trên câu này:

Entropy attention trung bình (nat) theo lớp (tối đa = ln T = 2.48):
   [1.85, 1.02, 1.27, 1.51, 1.41, 1.63, 1.58, 1.33, 0.97, 0.99, 1.01, 1.17]
Phần attention dồn vào token <s>, trung bình theo lớp:
   [0.2, 0.55, 0.31, 0.3, 0.31, 0.35, 0.17, 0.04, 0.05, 0.05, 0.05, 0.32]

Entropy đo độ phân tán: tối đa 2,48 nat nếu attention đều tuyệt đối trên 12 token. Lớp 1 là lớp phân tán nhất (1,85), các lớp 9 và 10 tập trung nhất (0,97 và 0,99). Ở lớp 2, trung bình 55% attention đổ vào token đầu câu, cho thấy mô hình dùng token này như chỗ "đỗ" khi một đầu không có gì cần lấy. Đây là kết quả trên một câu: chỉ minh hoạ, muốn kết luận phải đo trên hàng nghìn câu.

Khi nào không nên dùng Transformer

  • Dữ liệu bảng, vài nghìn dòng: gradient boosting như ở Gradient boosting với XGBoost và LightGBM thường thắng và rẻ hơn nhiều.
  • Chuỗi rất dài trên phần cứng nhỏ: Transformer tốn bộ nhớ bậc hai (bảng đo ở trên), nên cần FlashAttention, cửa sổ trượt hoặc kiến trúc khác.
  • Phân loại văn bản đơn giản: TF-IDF kèm mô hình tuyến tính thường đã đủ làm baseline; Transformer chỉ đáng giá khi baseline thật sự thiếu.

Lỗi thường gặp khi tự cài Transformer

  1. Quên chia cho căn d_k: softmax bão hoà, loss khó giảm; thí nghiệm ở trên cho thấy trọng số lớn nhất 0,687 so với 0,096.
  2. Đặt mặt nạ sai chiều: dùng masked_fill với mặt nạ ngược, hoặc mặt nạ trên chiều query thay vì key, làm mô hình nhìn trộm tương lai mà loss vẫn đẹp.
  3. Nhầm kích thước khi tách đầu: quên transpose hoặc reshape sai thứ tự làm các đầu trộn lẫn token; luôn đối chiếu với nn.MultiheadAttention như bài.
  4. Không có thông tin vị trí: mô hình chạy bình thường nhưng bị mù thứ tự, đúng như bản pos=False ở trên.

Checklist hiểu một mô hình Transformer

  • Với mỗi Transformer, biết d_model, số lớp, số đầu, d_k = d_model / số đầu, d_ff; đọc được từ file cấu hình.
  • Phân biệt Transformer encoder (hai chiều, như BERT/PhoBERT) và decoder (mặt nạ nhân quả, sinh văn bản).
  • Biết vị trí được tiêm từ ngoài; kiểm tra bằng thí nghiệm hoán vị.
  • Dùng F.scaled_dot_product_attention trong mã thật; tự cài chỉ để học Transformer.

Mã nằm trong code/s7-27: attention.py, mini_transformer.py, phobert_attention.py cùng các file *_output.txt và requirements.txt. phobert_attention.py cần mô hình vinai/phobert-base và pyvi. Số liệu là output thật nhưng chỉ cho một câu và một hạt giống, dùng để minh hoạ cơ chế.

Bài tiếp theo

Fine-tune PhoBERT với Hugging Face cho phân tích cảm xúc tiếng Việt. Sau khi hiểu kiến trúc Transformer bên trong, bài sau đặt thêm một lớp phân loại lên PhoBERT, tinh chỉnh trên dữ liệu đánh giá tiếng Việt và đo kết quả so với baseline TF-IDF.