Attention và Transformer

Từ trung bình cố định tới trung bình theo câu hỏi

Attention nhận query (câu hỏi), key (thông tin dùng để so khớp) và value (thông tin muốn lấy). Điểm giống nhau giữa query và key quyết định trọng số của value tương ứng. Hình dung truy vấn trong CP: key dùng để tìm, value là dữ liệu cần trả. Khác hash map, attention thường trả tổ hợp mềm của nhiều value thay vì chọn một địa chỉ chính xác.

Đặt $Q$ shape (Lq, dk), $K$ (Lk, dk), $V$ (Lk, dv). Tính

Softmax theo chiều key: mỗi hàng $A$ có tổng 1, nói một query phân phối trọng lượng cho các key ra sao. Query và key cần cùng $d_k$ để nhân vô hướng; value có thể có độ rộng $d_v$ khác. Hệ số $1/\sqrt{d_k}$ giữ thang score khi độ rộng tăng; mục riêng ngay sau ví dụ tính tay giải thích con số này.

Một attention trên hai key

Lấy một query $q=[1]$, hai key $k_1=[1]$, $k_2=[0]$, nên $d_k=1$ và score là $[1,0]$. Softmax cho $a_1=e/(e+1)\approx0.731$, $a_2=1/(e+1)\approx0.269$. Với $v_1=[2,0]$, $v_2=[0,4]$, kết quả là $[2a_1,4a_2]\approx[1.462,1.076]$. Key quyết định hệ số, value cung cấp nội dung.

Khi trọng số đã có, đầu ra chỉ là trung bình có trọng số: $(0.5, 0.3, 0.2)$ trên các value 2, 4, 10 cho $1 + 1.2 + 2 = 4.2$. Nếu mọi key giống hệt nhau thì mọi score bằng nhau, và đầu ra là trung bình cộng các value, query là gì cũng vậy.

import numpy as np
import torch
from torch import nn

torch.set_num_threads(1)
torch.manual_seed(41)
Q = np.array([[1.]])
K = np.array([[1.], [0.]])
V = np.array([[2., 0.], [0., 4.]])
scores = Q @ K.T / np.sqrt(Q.shape[-1])
shifted = scores - scores.max(axis=-1, keepdims=True)
weights = np.exp(shifted)
weights /= weights.sum(axis=-1, keepdims=True)
output = weights @ V
print('weights:', np.round(weights, 6).tolist())
print('output:', np.round(output, 6).tolist())
tq, tk, tv = [torch.tensor(a) for a in [Q, K, V]]
torch_output = (tq @ tk.T / np.sqrt(Q.shape[-1])).softmax(dim=-1) @ tv
print('match:', np.allclose(output, torch_output.numpy()))

Kết quảweights: [[0.731059, 0.268941]]
output: [[1.462117, 1.075766]]
match: True

Trừ max trước exp không đổi softmax vì mọi số mũ cùng chia một hệ số; nó tránh exp của số quá lớn. Đây là cùng tư duy phép số ổn định như loss bài 3.3. Dòng match xác nhận phép tự cài trùng với phép tensor của PyTorch.

Score quá lớn làm softmax bão hòa

Score $q^\top k=\sum_{i=1}^{d_k}q_ik_i$ là tổng của $d_k$ tích. Nếu các thành phần của $q$ và $k$ độc lập, trung bình 0, phương sai 1, mỗi tích có phương sai 1, nên score có phương sai $d_k$ và độ lệch chuẩn $\sqrt{d_k}$. Score càng chênh nhau, softmax càng dồn trọng số vào một key. Khi phân phối gần one-hot, softmax bão hòa: gradient tới các key còn lại gần bằng 0, giống tanh bão hòa ở bài 3.3. Chia cho $\sqrt{d_k}$ đưa độ lệch chuẩn về khoảng 1 ở mọi độ rộng.

Thí nghiệm dưới lấy 2000 query ngẫu nhiên, mỗi query 8 key, đo độ lệch chuẩn của score và trung bình trọng số lớn nhất trong mỗi hàng:

rng = np.random.default_rng(0)

def mean_top_weight(scores):
    shifted = scores - scores.max(axis=-1, keepdims=True)
    weights = np.exp(shifted)
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights.max(axis=-1).mean()

for d_k in [4, 64, 1024]:
    q = rng.standard_normal((2000, 1, d_k))   # 2000 query
    k = rng.standard_normal((2000, 8, d_k))   # 8 key cho mỗi query
    raw = (q * k).sum(axis=-1)                # score chưa chia, shape (2000, 8)
    scaled = raw / np.sqrt(d_k)
    print(f'd_k={d_k:<4}  std {raw.std():5.2f} -> {scaled.std():4.2f}   '
          f'max weight {mean_top_weight(raw):.3f} -> {mean_top_weight(scaled):.3f}')

Kết quảd_k=4     std  2.00 -> 1.00   max weight 0.524 -> 0.342
d_k=64    std  7.96 -> 0.99   max weight 0.878 -> 0.360
d_k=1024  std 31.83 -> 0.99   max weight 0.968 -> 0.358

Không chia, độ lệch chuẩn tăng đúng như $\sqrt{d_k}$ (2, 8, 32) và key mạnh nhất chiếm trung bình 0.524, 0.878 rồi 0.968 trọng số, trong khi chia đều cho 8 key chỉ là 0.125. Có chia, độ lệch chuẩn ở quanh 1 và trọng số lớn nhất giữ quanh 0.35 ở cả ba độ rộng: thang score không còn phụ thuộc kích thước model.

Lập luận dựa trên giả định độc lập lúc khởi tạo; trong quá trình học, model vẫn có thể làm score lớn khi thật sự cần một phân phối nhọn. Đẩy tới cùng, nhân mọi score với $c \to \infty$ thì softmax thành one-hot và đầu ra chỉ còn value của key có score lớn nhất.

Self-attention học cách tạo câu hỏi

Trong self-attention, $Q$, $K$, $V$ đều được tạo từ cùng chuỗi $X$ nhưng qua ba phép chiếu học được: $Q=XW_Q$, $K=XW_K$, $V=XW_V$. Cùng token có ba vai trò khác nhau. Query của một vị trí hỏi nên lấy thông tin gì; key của các vị trí giúp so khớp; value mang dữ liệu được trộn. Ba ma trận được học từ loss, nên từng phần tử của chúng không mang một quy tắc ngữ nghĩa đặt sẵn.

Multi-head attention dùng nhiều bộ chiếu, tạo nhiều cách trộn song song, ghép các kết quả rồi chiếu lại. Mỗi head có softmax riêng, nên một query có $H$ phân phối trọng số cùng lúc: một head có thể dồn vào chủ ngữ, head khác vào từ đứng kề. Nhiều head không bảo đảm mỗi head ứng với một khái niệm con người đặt tên được; chúng chỉ là thêm cách tương tác để model học.

Với embed_dim $d$ và số head $H$ của lớp PyTorch trong bài, cần $d$ chia hết cho $H$, mỗi head rộng $d/H$. Ba phép chiếu $W_Q,W_K,W_V$ và phép chiếu ra $W_O$ đều là $d\times d$, nên lớp có $4d^2+4d$ tham số kể cả bias, không phụ thuộc $H$: với $d=512$ là $4\cdot512^2+4\cdot512=1\,050\,624$ cho cả 8 lẫn 16 head.

Mask bỏ vị trí nào?

Có hai loại mask, giải quyết hai lý do khác nhau và có thể dùng cùng lúc:

• Padding mask bỏ các key PAD.
• Causal mask bỏ key tương lai để vị trí $t$ không nhìn $t+1$ trở đi, cần khi mô hình dự đoán token tiếp theo. Encoder đọc toàn câu thường không dùng causal mask; decoder sinh từng token (tự hồi quy) cần nó.

BERT là encoder, mỗi token nhìn cả hai phía; GPT là decoder, mỗi token chỉ nhìn bên trái. Causal mask viết bằng torch.tril(torch.ones(T, T)).bool() hay ~torch.triu(torch.ones(T, T), diagonal=1).bool() đều như nhau: vị trí $i$ nhìn $j$ khi $j \le i$.

Ta đưa score của vị trí cấm về $-\infty$ trước softmax, nên trọng số đó bằng 0. Nếu một hàng bị cấm mọi key, softmax không còn phân phối hợp lệ; tránh câu rỗng hoặc chọn quy tắc rõ cho trường hợp ấy. Nếu mask sau softmax mà không chuẩn hóa lại, tổng weights không còn 1 và không còn cùng phép attention.

Ví dụ score $[1, 2, 3]$ với key thứ ba là PAD: softmax chỉ còn trên $[1, 2]$, cho $(0.269, 0.731, 0)$. Với mask đúng, thêm bao nhiêu PAD thì biểu diễn của token thật cũng không đổi.

Quy ước True/False của mask khác nhau giữa các hàm:

• Trong nn.MultiheadAttention, key_padding_mask=True nghĩa là bỏ key; attn_mask=True nghĩa là cấm cặp query-key. Causal mask vì vậy thường viết torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1), True ở các ô trên đường chéo, tức key tương lai. Code bài 3.8 có valid_mask True ở token thật, nên truyền key_padding_mask=~valid_mask.
• F.scaled_dot_product_attention dùng quy ước ngược lại: True nghĩa là được tham gia attention.

Đọc tài liệu của đúng hàm đang gọi; không chuyển mask nguyên xi chỉ vì tên tham số giống nhau.

Một encoder block nhỏ

Một encoder block ghép bốn thành phần:

• Attention trộn thông tin giữa các token, cho output cùng shape với input.
• Residual (kết nối tắt) cộng input vào output, giữ một đường thông tin và gradient đi thẳng.
• LayerNorm (bài 3.6) chuẩn hóa theo feature trong từng token.
• Feed-forward network (FFN) là một MLP nhỏ áp riêng cho từng vị trí, biến đổi feature sau khi đã trộn ngữ cảnh.

Ta cài kiểu post-norm: norm(x + attention(x)), rồi norm(x + FFN(x)). Nhiều model dùng pre-norm, x + attention(norm(x)): đường residual đi thẳng không qua norm nào, nên mạng rất sâu học ổn định hơn và ít phụ thuộc warmup. Vị trí norm là một lựa chọn kiến trúc, phải giữ đúng khi nạp checkpoint.

class TinyEncoder(nn.Module):
    def __init__(self, width=4, heads=2):
        super().__init__()
        self.attention = nn.MultiheadAttention(width, heads, batch_first=True, dropout=0.)
        self.norm1 = nn.LayerNorm(width)
        self.ffn = nn.Sequential(nn.Linear(width, 8), nn.ReLU(), nn.Linear(8, width))
        self.norm2 = nn.LayerNorm(width)

    def forward(self, x, padding_mask=None, causal_mask=None):
        attended, weights = self.attention(x, x, x,
            key_padding_mask=padding_mask, attn_mask=causal_mask,
            need_weights=True, average_attn_weights=False)
        x = self.norm1(x + attended)
        x = self.norm2(x + self.ffn(x))
        return x, weights

block = TinyEncoder()
block.eval()
x = torch.randn(2, 3, 4)
padding_mask = torch.tensor([[False, False, True], [False, False, False]])
with torch.no_grad():
    encoded, attn = block(x, padding_mask=padding_mask)
print('output:', tuple(encoded.shape))
print('weights:', tuple(attn.shape))
print('PAD key max weight:', round(attn[0, :, :, 2].abs().max().item(), 6))
print('weight row sums:', torch.allclose(attn.sum(dim=-1), torch.ones(2, 2, 3), atol=1e-6))

Kết quảoutput: (2, 3, 4)
weights: (2, 2, 3, 3)
PAD key max weight: 0.0
weight row sums: True

Đầu ra (2,3,4) giữ batch, độ dài và độ rộng; weights (2,2,3,3) gồm batch, head, query, key vì không lấy mean qua head. Cột PAD của mẫu đầu bằng 0 cho cả hai head. Hàng query tại PAD vẫn có thể sinh vector: key padding mask chỉ bỏ key. Khi gộp thành vector câu hoặc tính token loss, còn phải dùng valid mask để bỏ query PAD ở bước đó.

Thử causal mask bằng perturbation

Một phép kiểm có ý nghĩa hơn chỉ nhìn ma trận mask là thay token tương lai và xem quá khứ có đổi không. Với ba vị trí, cấm phần tam giác trên. Ta sửa token cuối rồi so đầu ra hai token đầu, giữ eval và dropout=0:

causal = torch.triu(torch.ones(3, 3, dtype=torch.bool), diagonal=1)
one_sequence = x[1:2].clone()
changed = one_sequence.clone()
changed[:, 2, :] += torch.tensor([10., -3., 7., 2.])
with torch.no_grad():
    before, _ = block(one_sequence, causal_mask=causal)
    after, _ = block(changed, causal_mask=causal)
print('past unchanged:', torch.allclose(before[:, :2], after[:, :2], atol=1e-6))
print('last position changed:', not torch.allclose(before[:, 2], after[:, 2], atol=1e-6))

Kết quảpast unchanged: True
last position changed: True

Dòng đầu xác nhận hai vị trí đầu không dùng token cuối. Dòng thứ hai là đối chứng: vị trí cuối phải đổi, nếu không thì dòng đầu có thể True chỉ vì block bỏ qua phần input bị sửa. Cách kiểm này có cùng tinh thần tạo input đối kháng nhỏ trong CP: thay phần bị cấm rồi xem output có chịu ảnh hưởng. Với model nhiều layer, chạy cùng test trên toàn model, vì chỉ một khâu trộn cả chuỗi (pooling, chuẩn hóa theo chiều độ dài) cũng đủ làm tương lai lọt vào.

Thứ tự cần thông tin vị trí

Phép kiểm sau dùng block không có thông tin vị trí và không mask, ở chế độ eval với dropout=0. Khi ấy self-attention và FFN có tính equivariant theo hoán vị: đảo thứ tự token đầu vào thì các vector đầu ra đảo theo đúng cách đó. Một mean pooling đặt sau block còn bất biến (invariant): đảo token không đổi vector câu. Block vẫn tạo vector riêng ở mỗi vị trí, nhưng không tự biết token nào đứng thứ hai.

Nếu có mask, phải hoán vị mask cùng hàng query và cột key; một causal mask giữ cố định đã mang thông tin thứ tự nên phá tính chất này.

permutation = torch.tensor([2, 0, 1])
with torch.no_grad():
    original, _ = block(one_sequence)
    permuted, _ = block(one_sequence[:, permutation])
print('equivariant without positions:', torch.allclose(
    permuted, original[:, permutation], atol=1e-6))

Kết quảequivariant without positions: True

True ở đây là tính chất toán học, không phải lỗi thư viện. Positional encoding cộng vào mỗi token một vector cho biết nơi nó xuất hiện. Ba cách thường gặp là bảng vị trí học được, sin/cos cố định và relative position. Bản sin/cos dùng $PE_{p,2i}=\sin(p/10\,000^{2i/d})$, $PE_{p,2i+1}=\cos(p/10\,000^{2i/d})$; với $d=4$, hai tần số là 1 và 0.01. Đoạn sau kiểm nó:

positions = torch.arange(3, dtype=torch.float32).unsqueeze(1)
frequencies = torch.exp(-np.log(10000.) * torch.arange(0, 4, 2) / 4)
pe = torch.zeros(3, 4)
pe[:, 0::2] = torch.sin(positions * frequencies)
pe[:, 1::2] = torch.cos(positions * frequencies)
with torch.no_grad():
    ordered, _ = block(one_sequence + pe)
    swapped, _ = block(one_sequence[:, permutation] + pe)
print('position 0:', pe[0].tolist())
print('equivariant with fixed slots:', torch.allclose(swapped, ordered[:, permutation], atol=1e-6))

Kết quảposition 0: [0.0, 1.0, 0.0, 1.0]
equivariant with fixed slots: False

Khi đảo token nhưng giữ bảng vị trí tại các slot 0, 1, 2, ta đổi quan hệ token–vị trí, nên dòng cuối thành False: block giờ phân biệt được thứ tự. Sin/cos chỉ cung cấp tín hiệu; model dùng tín hiệu ấy đến đâu là do huấn luyện.

Relative position không gắn vị trí tuyệt đối vào token mà đưa khoảng cách có hướng $j-i$ vào score, chẳng hạn $q_i^\top(k_j+r_{j-i})/\sqrt{d_k}$ với $r_{j-i}$ là vector học được. Mọi cặp cách nhau cùng một khoảng dùng chung một vector, hợp với dữ liệu mà quan hệ phụ thuộc khoảng cách hơn vị trí tuyệt đối, như chuỗi thời gian. Bảng vị trí học được thì là một nn.Embedding(max_len, d), nên chuỗi dài hơn max_len cần quy tắc rõ từ trước: bảng 512 vị trí không có hàng cho token thứ 600.

Với $T$ token, khoảng cách $j - i$ nhận $2T - 1$ giá trị, từ $-(T-1)$ tới $T-1$. Sin/cos tính được ở mọi vị trí, và dịch một khoảng cố định là một phép quay trên từng cặp chiều, nên model dễ đọc ra khoảng cách tương đối. Cùng công thức sin/cos còn dùng để đưa một con số như bước khử nhiễu $t$ vào mạng (bài 5.8).

Kiểm graph và đọc attention đúng mức

Để bảo đảm encoder nhỏ có đường gradient, ta nối một loss kiểm với mục tiêu ngẫu nhiên có shape đúng. Không dùng chỉ bình phương đầu ra LayerNorm rồi kết luận gradient lớn hay nhỏ, vì norm đã cố định một số thống kê đầu ra.

block.train()
block.zero_grad(set_to_none=True)
encoded, _ = block(x, padding_mask=padding_mask)
target = torch.randn_like(encoded)
valid = ~padding_mask
loss = ((encoded - target).square().mean(dim=-1) * valid).sum() / valid.sum()
loss.backward()
print('finite loss:', torch.isfinite(loss).item())
print('all parameter gradients finite:', all(p.grad is not None and torch.isfinite(p.grad).all()
                                             for p in block.parameters()))

Kết quảfinite loss: True
all parameter gradients finite: True

Hai True xác nhận loss hữu hạn và mọi tham số đều nhận gradient hữu hạn, kể cả khi có mask PAD: block sẵn sàng để huấn luyện.

Attention weights là hệ số trộn tại một layer/head; chúng không tự là lời giải thích nhân quả cho dự đoán cuối, vì model còn residual, FFN, nhiều head và layer khác. Muốn biết token ảnh hưởng kết quả thế nào, cần perturbation hoặc ablation: xóa hay thay token rồi đo dự đoán đổi bao nhiêu.

Chi phí của attention

Attention mọi cặp tốn khoảng $O(BL^2d)$ tính toán ở phần score và trộn, cùng $O(BHL^2)$ bộ nhớ cho score nếu lưu đủ ma trận; các phép chiếu thêm $O(BLd^2)$. Tăng $L$ gấp đôi thì phần score gấp bốn. Padding cũng trả theo $L^2$: một câu rất dài kéo cả batch lên độ dài của nó, nên xếp các câu dài gần nhau vào cùng batch (bucketing), cắt hoặc chia cửa sổ văn bản dài trước khi nghĩ tới phần cứng lớn hơn.

Bậc hai đến từ việc softmax buộc tạo đủ ma trận $L\times L$: $\exp(q_i^\top k_j)$ không viết được thành tích vô hướng $\phi(q_i)^\top\phi(k_j)$ với $\phi$ hữu hạn chiều.

Linear attention thay hẳn nó bằng $\phi(q_i)^\top\phi(k_j)$ với một feature map $\phi$ hữu hạn chiều; khi đó có thể đổi thứ tự nhân ma trận như bài toán nhân chuỗi ma trận quen thuộc trong CP, và chi phí còn tuyến tính theo $L$, cỡ $O(Ld^2)$. Với mask nhân quả, chỉ cần giữ hai tổng chạy $S_t = \sum_{j \le t}\phi(k_j)v_j^\top$ và $z_t = \sum_{j \le t}\phi(k_j)$, nên sinh mỗi token tốn bộ nhớ không đổi. $\phi$ phải dương, như $1 + \operatorname{ELU}(x)$, để trọng số không âm và mẫu số không bằng 0.

Với dữ liệu ít, model phức tạp chưa chắc thắng baseline masked mean; so trên cùng validation trước khi chuyển sang model pretrained ở chương 4.