Bài 10: Self-Attention — token "nhìn" các token khác

Sau bài này, bạn cần trả lời được:

• "Self" trong self-attention nghĩa là gì — khác với attention encoder-decoder của Bài 9 ở điểm nào.
• Vì sao cần ba ma trận projection \( W^Q, W^K, W^V \) khác nhau dù đều áp lên cùng \( X \).
• Số tham số của 1 head self-attention với \( d_{\text{model}} = 512, d_k = 64 \).
• Causal mask làm gì — tại sao decoder GPT cần, còn encoder BERT thì không.
• Self-attention nhanh hơn RNN ở điểm nào, đắt hơn ở điểm nào.
• Cách viết một nn.Module self-attention ~20 dòng PyTorch.

Bài 9 đã tách rời nguồn sinh Q (decoder hidden) và nguồn sinh K, V (encoder hidden). Bài này gộp tất cả về một nguồn duy nhất — bước đầu để hiểu Transformer block ở các bài sau.

Scaled dot-product attention với ba ma trận \( Q, K, V \):

\[ \text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V \]

Trong đó:

• \( Q \in \mathbb{R}^{n_q \times d_k} \): query — "đang hỏi cái gì".
• \( K \in \mathbb{R}^{n_k \times d_k} \): key — "có cái gì để match".
• \( V \in \mathbb{R}^{n_k \times d_v} \): value — "nội dung mang theo".
• \( \sqrt{d_k} \) là scale factor giữ phương sai của tích vô hướng ổn định khi \( d_k \) lớn (Vaswani et al. 2017).

Trong attention encoder-decoder (Bahdanau 2014, Luong 2015), \( Q \) lấy từ decoder hidden state, còn \( K, V \) lấy từ encoder hidden state — hai nguồn khác nhau. Bài này thay đổi đúng một thứ: cả ba ma trận sinh ra từ cùng một chuỗi đầu vào.

Gọi \( X \in \mathbb{R}^{n \times d_{\text{model}}} \) là ma trận đầu vào (sau token embedding + positional encoding, xem Bài 6 và Bài 12), \( n \) là độ dài chuỗi, \( d_{\text{model}} \) là chiều ẩn của model.

Self-attention tính \( Q, K, V \) bằng ba ma trận projection learnable:

\[ Q = X W^Q, \qquad K = X W^K, \qquad V = X W^V \]

với

\[ W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}, \qquad W^V \in \mathbb{R}^{d_{\text{model}} \times d_v} \]

Sau đó áp công thức attention quen thuộc:

\[ \text{SelfAttention}(X) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V \]

Output có shape \( (n, d_v) \). Mỗi hàng output là một contextualized representation của token tương ứng — nó không chỉ chứa thông tin riêng của token đó mà còn pha trộn thông tin từ các token khác trong chuỗi (theo trọng số attention).

Tên "self" đến từ chỗ Q, K, V cùng nguồn (X). Mỗi token vừa là người hỏi, vừa là kho key, vừa là kho value cho phần còn lại của chuỗi.

Cùng một token, nhưng giữ ba "vai trò" khác nhau trong phép tính:

• \( W^Q \) — biến \( x_i \) thành vector "tôi đang hỏi gì". Project \( x_i \) sang không gian query.
• \( W^K \) — biến \( x_j \) thành vector "tôi có gì để match". Project sang không gian key.
• \( W^V \) — biến \( x_j \) thành vector "thông tin tôi mang theo". Project sang không gian value.

Nếu bỏ ba projection và lấy \( Q = K = V = X \) (như một số simplified illustration), score luôn lớn nhất khi \( i = j \) — mỗi token chỉ attend chính nó, model không học được pattern gì có ích.

Ba ma trận \( W^Q, W^K, W^V \) là tham số học trong quá trình training. Chúng cho phép model học cách biến cùng một input vector thành ba biểu diễn riêng biệt cho ba vai trò. Đây là điểm khiến self-attention linh hoạt hơn việc lấy similarity thô.

Xét token vị trí \( i \). Quá trình tính output cho riêng nó:

1. Tạo query \( q_i = x_i W^Q \).
2. Lấy key của tất cả token \( k_j = x_j W^K \) với \( j = 1, \dots, n \).
3. Tính score \( s_{ij} = \dfrac{q_i \cdot k_j}{\sqrt{d_k}} \).
4. Softmax trên cả hàng để được trọng số \( \alpha_{ij} = \dfrac{e^{s_{ij}}}{\sum_{j'} e^{s_{ij'}}} \), \( \sum_j \alpha_{ij} = 1 \).
5. Lấy weighted sum value: \( z_i = \sum_{j=1}^n \alpha_{ij} v_j \), với \( v_j = x_j W^V \).

Kết quả \( z_i \in \mathbb{R}^{d_v} \) là tổ hợp tuyến tính của thông tin từ tất cả token trong chuỗi, mỗi token đóng góp tỉ lệ \( \alpha_{ij} \). Token \( i \) cũng tự attend chính nó (\( j = i \)) — không loại trừ.

Câu ví dụ kinh điển trong bài blog Illustrated Transformer (Alammar 2018):

The animal didn't cross the street because it was tired.

Đại từ "it" có thể tham chiếu tới animal hoặc street. Để giải nghĩa đúng câu, model phải "biết" rằng "it" liên kết với "animal" (vì chính animal mới có thể "tired"). Khi visualize trọng số attention cho token "it" ở một head của BERT / GPT, đa số khối lượng attention dồn về token "animal".

Trong khi đó nếu thay câu thành "...because it was wide", attention pattern thay đổi — "it" attend sang "street". Cùng một pronoun, cùng vị trí trong câu, nhưng head self-attention học được nguyên tắc chọn antecedent dựa trên ngữ cảnh.

Hiện tượng này không hardcode — không có quy tắc nào trong code nói "pronoun phải attend tới noun phù hợp". Pattern phát sinh từ training trên corpus đủ lớn (xem Bài 12 — Clark et al. 2019, observed patterns).

Pseudocode cho self-attention 1 head, batch \( B \), độ dài \( n \):

# X: (B, n, d_model)
Q = X @ W_Q                                  # (B, n, d_k)
K = X @ W_K                                  # (B, n, d_k)
V = X @ W_V                                  # (B, n, d_v)

scores = Q @ K.transpose(-2, -1) / sqrt(d_k) # (B, n, n)
attn   = softmax(scores, dim=-1)             # (B, n, n) — mỗi hàng sum = 1
out    = attn @ V                            # (B, n, d_v)

Diễn giải shape:

• \( Q K^\top \) có shape \( (B, n, n) \) — phần tử \( (i, j) \) là score giữa token \( i \) và token \( j \).
• Softmax áp trên trục cuối (dim=-1) — chuẩn hóa theo hàng để tổng trọng số attention từ một token bằng 1.
• \( \text{attn} \cdot V \) trộn value của mọi token theo trọng số, ra output cùng độ dài.

Tất cả token được tính song song trong cùng vài lệnh matmul — không có vòng lặp tuần tự theo \( i \).

Trong self-attention 1 head, tham số học được chỉ nằm trong ba ma trận projection:

\[ |W^Q| + |W^K| + |W^V| = d_{\text{model}} \cdot d_k + d_{\text{model}} \cdot d_k + d_{\text{model}} \cdot d_v \]

Trường hợp thường gặp \( d_k = d_v \):

\[ \text{params} = 3 \cdot d_{\text{model}} \cdot d_k \]

Ví dụ Vaswani et al. 2017 — base config \( d_{\text{model}} = 512, d_k = d_v = 64 \) (\( h = 8 \) head, mỗi head 64-dim):

\[ 3 \times 512 \times 64 = 98{,}304 \ \text{tham số} \]

Đây là param cho một head. Multi-head 8 head cùng config sẽ là \( 8 \times 98{,}304 \approx 786\text{k} \), cộng thêm projection output \( W^O \) (Bài 11). Bias term thường được bỏ qua hoặc thêm \( 3 d_k \) tham số — không đáng kể.

LLM kiểu GPT sinh token tuần tự từ trái sang phải. Khi predict token \( i+1 \), model chỉ được nhìn các token \( 1, 2, \dots, i \) — không được "rò rỉ" thông tin tương lai về quá khứ trong lúc train.

Cơ chế ép buộc điều này là causal mask (đã đề cập Bài 9): trước softmax, thay tất cả vị trí \( j > i \) trong \( s_{ij} \) bằng \( -\infty \):

\[ s_{ij} \leftarrow \begin{cases} s_{ij} & \text{nếu } j \le i \\ -\infty & \text{nếu } j > i \end{cases} \]

Sau softmax, \( \alpha_{ij} = 0 \) với mọi \( j > i \). Token \( i \) chỉ attend chính nó và các token đứng trước.

Đây là foundation cho generation autoregressive. Model train một lần trên toàn câu (mọi vị trí được tính output song song), nhưng nhờ mask, mỗi vị trí chỉ học từ phần "đã thấy" — y hệt trạng thái khi inference từng token một. Train song song, inference tuần tự.

Các model decoder-only dùng causal mask: GPT-2 / 3 / 4, Llama, Mistral, Gemma, Qwen, DeepSeek...

BERT (Devlin et al. 2018) và họ encoder không sinh token mà tạo representation cho từng vị trí phục vụ task downstream (classification, NER, retrieval, MLM…). Khi đó việc nhìn cả past và future đều có lợi — token "it" trong câu hiểu đúng nghĩa khi attend cả "animal" phía trước và "tired" phía sau.

Với encoder-only, KHÔNG áp mask:

• Mỗi token attend tất cả \( n \) token, bao gồm chính nó.
• Ma trận attention \( n \times n \) đầy, không có vùng tam giác trống.
• Pre-training task: Masked Language Modeling (MLM) — che 15% token và bắt model dự đoán; bidirectional context giúp dự đoán chính xác hơn so với chỉ nhìn một chiều.

Tóm tắt khác biệt code:

decoder (GPT):     scores.masked_fill_(causal_mask, -inf) → softmax
encoder (BERT):    scores → softmax   (chỉ mask padding nếu có)

Cùng một module self-attention, khác đúng dòng mask — đó là sự phân chia encoder vs decoder (Bài 14 sẽ so sánh đầy đủ ba kiến trúc: encoder-only BERT, decoder-only GPT, encoder-decoder T5).

Trọng số attention \( \alpha_{ij} \) không có rule cứng. Pattern nào "có ích" cho việc giảm loss thì \( W^Q, W^K \) sẽ hội tụ về đó qua backpropagation.

Hệ quả thực nghiệm:

• Sau training, \( W^Q W^{K\top} \) là một ma trận tương tác \( d_{\text{model}} \times d_{\text{model}} \) cho biết "feature query nào hợp với feature key nào".
• Nhiều head trong cùng layer học các pattern khác nhau — một head chuyên attend token liền kề, head khác chuyên attend special token, head khác lại theo dependency cú pháp. Đây là động lực của multi-head (Bài 11).
• Pattern không cố định — model trên domain code (StarCoder) có pattern khác model trên text general.

Clark et al. 2019 ("What Does BERT Look At? An Analysis of BERT's Attention", arXiv:1906.04341) phân tích toàn bộ head của BERT base và phát hiện một số pattern phổ biến:

• Local — attend token liền trước hoặc liền sau (head "diagonal-shift"). Bắt n-gram, cụm từ ngắn.
• Global — attend mạnh vào special token ([CLS], [SEP]) như một "kênh broadcast". Trong GPT thường là token đầu chuỗi.
• Syntactic — pattern phụ thuộc cú pháp: head attend từ động từ về chủ ngữ, từ tính từ về danh từ được bổ nghĩa, từ giới từ về object.
• Semantic — coreference / anaphora: pronoun attend antecedent, danh từ định danh attend tên riêng đã xuất hiện.
• Trivial — head attend đều mọi token, hoặc head luôn attend chính nó. Một số head được prune đi mà không ảnh hưởng performance (Voita et al. 2019).

Quan trọng: các pattern này được model học mà không được giám sát trực tiếp. Pre-training trên MLM (BERT) hoặc next-token prediction (GPT) đủ để pattern syntactic / semantic phát sinh.

RNN / LSTM: token \( i \) phụ thuộc \( h_{i-1} \), \( h_{i-1} \) phụ thuộc \( h_{i-2} \), v.v. — chuỗi tính toán dài.

Self-attention: token \( i \) trực tiếp tính score với token \( j \) bất kỳ trong một lần tính. Đường truyền tín hiệu ngắn hơn nhiều.

Vaswani et al. 2017 (Table 1) so sánh ba dimension:

                    Complexity / layer    Sequential ops    Maximum path length
Self-Attention      O(n² · d)             O(1)              O(1)
Recurrent (RNN)     O(n · d²)             O(n)              O(n)
Convolutional       O(k · n · d²)         O(1)              O(log_k n)

"Maximum path length" là độ dài tín hiệu phải đi giữa hai token bất kỳ — quyết định khả năng học dependency xa. RNN có path length \( O(n) \), tín hiệu đi qua \( n \) bước nonlinearity, dễ vanishing gradient. Self-attention path length \( O(1) \) — mọi token kết nối trực tiếp.

Đổi lại self-attention đắt hơn theo \( n^2 \) (xem Bước 16). Với \( n \) nhỏ (< vài nghìn), self-attention thắng RNN cả về tốc độ lẫn chất lượng.

Self-attention chỉ gồm matmul và softmax — đều là phép toán batched matrix mà GPU / TPU thực thi rất nhanh.

Bốn matmul chính:

1. \( Q = X W^Q \) — \( (B, n, d_{\text{model}}) \times (d_{\text{model}}, d_k) \to (B, n, d_k) \).
2. \( K = X W^K \), \( V = X W^V \) — tương tự.
3. \( S = Q K^\top \) — batched matmul, \( (B, n, d_k) \times (B, d_k, n) \to (B, n, n) \).
4. \( Z = \alpha V \) — \( (B, n, n) \times (B, n, d_v) \to (B, n, d_v) \).

Không có vòng lặp theo \( i \) — mọi token tính cùng lúc trong một forward. Đây là ưu thế lớn so với RNN: RNN bắt buộc tuần tự theo time step trong khi train. Self-attention cho phép full parallelism trên seq length, làm GPU không bị "đói" công việc.

Một forward pass self-attention trên batch 32 câu × 512 token chạy vài chục mili-giây trên GPU consumer (RTX 4090). RNN cùng kích thước phải hàng trăm mili-giây vì 512 bước tuần tự.

Module 1 head self-attention bằng PyTorch (v2.x):

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


class SelfAttention(nn.Module):
    def __init__(self, d_model: int, d_k: int):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k

    def forward(self, X: torch.Tensor, mask: torch.Tensor | None = None):
        # X: (B, n, d_model)
        Q = self.W_Q(X)                                   # (B, n, d_k)
        K = self.W_K(X)
        V = self.W_V(X)

        scores = Q @ K.transpose(-2, -1) / self.d_k ** 0.5  # (B, n, n)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)                    # (B, n, n)
        return attn @ V, attn                               # output, weights

Demo forward trên 1 chuỗi giả 10 token:

torch.manual_seed(0)
B, n, d_model, d_k = 1, 10, 64, 16
X = torch.randn(B, n, d_model)

sa = SelfAttention(d_model, d_k)
out, attn = sa(X)
print(out.shape)    # torch.Size([1, 10, 16])
print(attn.shape)   # torch.Size([1, 10, 10])
print(attn.sum(dim=-1))  # mỗi hàng sum ≈ 1.0

Áp causal mask (tam giác dưới = 1, trên = 0):

causal = torch.tril(torch.ones(n, n)).unsqueeze(0)   # (1, n, n)
out_causal, attn_causal = sa(X, mask=causal)
print(attn_causal[0])
# Tam giác trên = 0, tổng mỗi hàng vẫn = 1

Trong production code (BERT, GPT), người ta dùng torch.nn.functional.scaled_dot_product_attention (PyTorch 2.0+) hoặc Flash Attention — fused kernel chạy nhanh hơn và tiết kiệm bộ nhớ. Bản tự cài ở trên chỉ để hiểu cơ chế.

Chi phí self-attention dominate bởi hai matmul liên quan tới attention matrix:

\[ Q K^\top: \ O(n^2 \cdot d_k) \quad \text{và} \quad \alpha V: \ O(n^2 \cdot d_v) \]

Tổng:

\[ \text{Time} = O(n^2 \cdot d), \qquad \text{Memory} = O(n^2) \]

(memory cho attention matrix \( \alpha \) hoặc score \( S \), kích thước \( n \times n \) trên mỗi head, mỗi batch).

Hệ quả: khi \( n \) tăng từ 2k lên 32k, chi phí tăng \( 256 \times \). Đây là lý do context window dài rất "đắt" — model 128k context không phải chỉ giải bài toán bộ nhớ mà còn FLOPs.

Giải pháp 2022-2025:

• FlashAttention (Dao et al. 2022, 2023): tile attention vào SRAM, giảm I/O memory, không thay đổi thuật toán toán học.
• Sparse / linear attention: Longformer, BigBird, Linformer, Performer — giảm complexity về \( O(n \log n) \) hoặc \( O(n) \) bằng cách giới hạn pattern attention.
• State Space Model (Mamba): bỏ attention, dùng recurrence dạng đặc biệt — out of scope bài này.

Self-attention không đứng một mình. Trong mỗi Transformer block, nó là sub-layer đầu tiên, bọc trong residual + LayerNorm:

input X
  │
  ├─→ LayerNorm → SelfAttention(MultiHead) → +X  (residual)
  │
  ├─→ LayerNorm → FFN (2 layer MLP)        → +   (residual)
  │
output

Output \( (n, d_{\text{model}}) \) shape giữ nguyên qua block — cho phép xếp chồng nhiều block (12 với BERT base, 96 với GPT-3, 80 với Llama 3 70B).

Sự sắp xếp "LayerNorm trước attention" gọi là pre-LN, ngược với cách trong paper gốc (post-LN). Pre-LN ổn định train cho model lớn hơn, được chuẩn hóa trong GPT-2 trở đi (Bài 13 sẽ deep dive Transformer block đầy đủ).

Ma trận \( \alpha \in \mathbb{R}^{n \times n} \) là tài liệu trực quan tốt nhất để hiểu một head đang làm gì:

• Hàng \( i \) = attention từ token \( i \) tới mọi token. Tổng hàng = 1.
• Cột \( j \) = mức độ token \( j \) "được nhìn" bởi mọi token.
• Đường chéo nổi = head attend token chính nó / liền kề.
• Cột riêng lẻ nổi (token [CLS], [SEP], BOS) = head global broadcast.
• Vệt tam giác dưới rõ rệt = causal head trong decoder.

Vẽ heatmap đơn giản với matplotlib:

import matplotlib.pyplot as plt
# attn: (n, n)  — lấy ra từ output ở Bước 15
plt.imshow(attn[0].detach().numpy(), cmap="viridis")
plt.colorbar()
plt.xlabel("key position j")
plt.ylabel("query position i")
plt.show()

Tool chuyên dụng:

• BertViz (Vig 2019): visualize multi-head attention của BERT / GPT-2 trực tiếp trong Jupyter, có view neuron / head / model.
• AttentionViz (Yeh et al. 2023): visualize attention dạng query-key embedding trong không gian 2D, scale tới model lớn.
• exBERT: explore attention pattern cho BERT trên web UI.

Vaswani et al. 2017 chia attention thành \( h \) head, mỗi head có riêng \( W^Q, W^K, W^V \). Khi inference token-by-token (autoregressive), \( K \) và \( V \) của các token đã sinh được cache (KV cache) — tốn rất nhiều bộ nhớ với \( h \) head riêng.

• Multi-Query Attention (MQA, Shazeer 2019): \( h \) head riêng \( Q \), nhưng chia sẻ 1 \( K \) và 1 \( V \) chung. KV cache giảm \( h \) lần.
• Grouped-Query Attention (GQA, Ainslie et al. 2023): tách head Q thành \( g \) nhóm, mỗi nhóm share 1 \( K, V \). \( g = h \) thành multi-head gốc, \( g = 1 \) thành MQA. GQA trung gian, giữ chất lượng gần multi-head nhưng cache nhỏ hơn.

Sử dụng trong model 2024-2025:

Model            Num Q heads    Num KV heads (groups)
Llama 2 70B      64             8           (GQA, g=8)
Llama 3 8B       32             8           (GQA, g=8)
Llama 3 70B      64             8           (GQA, g=8)
Mistral 7B       32             8           (GQA, g=8)
Mixtral 8x7B     32             8           (GQA)
Gemma 7B         16             16          (multi-head gốc)
PaLM 2           ...            1           (MQA)

Tiết kiệm KV cache đặc biệt quan trọng khi serve LLM ở batch lớn / context dài. Bài 11 đi vào chi tiết multi-head, các biến thể MQA / GQA sẽ được giới thiệu lại cùng cấu trúc head.

1. Tính tay self-attention cho chuỗi 3 token, \( d_{\text{model}} = 4, d_k = 2 \). Tự chọn \( X, W^Q, W^K, W^V \) bằng số nhỏ (integer 0-2), tính \( Q, K, V \), \( QK^\top / \sqrt{d_k} \), softmax từng hàng (dùng máy tính), và output \( \alpha V \). Kiểm tra shape kết quả là \( (3, 2) \).
2. Mở rộng class SelfAttention ở Bước 15 thêm tham số causal: bool. Khi causal=True, tự sinh tam giác mask thay vì nhận từ ngoài. Test forward với cùng input một lần causal, một lần bidirectional, so sánh hai output có khác nhau không.
3. Tính param count cho self-attention 1 head với \( d_{\text{model}} = 768, d_k = 64 \) (config BERT base) và \( d_{\text{model}} = 4096, d_k = 128 \) (Llama 3 8B). Quy đổi sang MB ở float16.
4. Lấy model BERT base (bert-base-uncased) qua Hugging Face Transformers, encode câu "The animal didn't cross the street because it was tired." và extract attention weight ở layer 8, head 5 (chỉ ví dụ — bạn có thể chọn layer / head khác). Plot heatmap. Token "it" attend mạnh vào token nào?
5. Đo thời gian forward self-attention cho \( n \in \{128, 512, 2048, 8192\} \), giữ \( d_{\text{model}} = 256 \). Kiểm tra thực nghiệm complexity có gần \( O(n^2) \) không (vẽ log-log plot).
6. (Đọc thêm) Đọc lướt Section 3.2 của Vaswani et al. 2017 (arXiv:1706.03762) — đối chiếu công thức bài viết với paper gốc. Tìm 1 điểm trong paper bạn chưa rõ và ghi câu hỏi lại.

• Self-attention: \( Q = X W^Q, K = X W^K, V = X W^V \) — ba projection của cùng input X. Áp công thức scaled dot-product để ra contextualized representation.
• Ba projection có ba vai trò: query "tôi hỏi gì", key "tôi có gì", value "tôi mang gì". Cùng X nhưng ba góc nhìn riêng.
• Mỗi token attend mọi token trong chuỗi, bao gồm chính nó. Output là weighted sum của value theo trọng số softmax.
• Param count 1 head: \( 3 \cdot d_{\text{model}} \cdot d_k \). Vd \( d=512, d_k=64 \) → 98.304 tham số.
• Causal mask phân biệt decoder (GPT) — nhìn quá khứ, vs encoder (BERT) — nhìn cả hai chiều.
• Self-attention path length \( O(1) \), so với RNN \( O(n) \). Đổi lại complexity \( O(n^2 d) \) — đắt với context dài.
• Implement ~20 dòng PyTorch. Production dùng scaled_dot_product_attention hoặc FlashAttention.
• Pattern observed: local, global, syntactic, semantic — học không giám sát qua MLM / next-token prediction.
• Biến thể hiện đại MQA / GQA giảm KV cache cho inference, dùng phổ biến trong Llama 3, Mistral, Mixtral.

Bài 11 (Multi-Head Attention) trả lời câu hỏi: tại sao chạy nhiều head song song lại tốt hơn một head to. Bài 12 (Positional Encoding) bổ sung phần còn thiếu — self-attention bản gốc không biết thứ tự token, cần signal vị trí gắn vào X trước khi đưa vào layer.