Bài 11: Multi-Head Attention — nhiều góc nhìn song song

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

• "Head" trong multi-head attention là gì — mỗi head chia sẻ và tách biệt cái gì với head khác.
• Vì sao Vaswani et al. 2017 chọn \( h \) head với \( d_k = d_{\text{model}} / h \) thay vì \( h \) head full-dim.
• Tổng số tham số của 1 layer MHA với \( d_{\text{model}} = 512, h = 8 \) (đáp số khoảng \( 4 \cdot d_{\text{model}}^2 \)).
• Tại sao implement MHA thực tế chỉ cần 4 ma trận \( W^Q, W^K, W^V, W^O \) cỡ \( d_{\text{model}} \times d_{\text{model}} \) — không phải \( 3h + 1 \) ma trận nhỏ.
• Multi-Query Attention (MQA), Grouped-Query Attention (GQA), Multi-Latent Attention (MLA) khác multi-head gốc thế nào, dùng ở model nào.
• Viết một nn.Module MHA ~30 dòng PyTorch, đồng thời gọi được nn.MultiheadAttention built-in.

Bài 10 đã có self-attention 1 head; bài này nhân lên \( h \) head song song và trộn lại bằng \( W^O \). Sau bài này bạn đã có hai sub-layer cốt lõi của Transformer block — Bài 13 sẽ ghép chúng với FFN, residual, LayerNorm thành block hoàn chỉnh.

Self-attention 1 head với input \( X \in \mathbb{R}^{n \times d_{\text{model}}} \):

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

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

Output \( (n, d_v) \) — mỗi hàng là contextualized representation của token tương ứng. Một head học một pattern attention duy nhất. Pattern này được quyết định bởi \( W^Q, W^K \) (cách định nghĩa similarity) và \( W^V \) (cách extract value).

Câu hỏi tự nhiên: nếu một head chỉ học được một pattern, vậy chạy nhiều head song song để cover nhiều pattern khác nhau cùng lúc thì sao?

Một câu chứa đồng thời nhiều loại quan hệ:

• Quan hệ cú pháp — subject ↔ verb, modifier ↔ head noun.
• Quan hệ ngữ nghĩa — coreference, anaphora ("it" → "animal" ở Bài 10).
• Quan hệ vị trí — token liền trước, n-gram cục bộ.
• Quan hệ với special token — [CLS], [SEP], BOS / EOS làm "kênh" tổng hợp toàn câu.

Một head với một bộ \( W^Q, W^K \) chỉ định nghĩa được một metric similarity duy nhất. Cố ép nó cover mọi quan hệ ở trên dẫn tới attention "trung bình hóa" — chất lượng kém.

Giải pháp Vaswani et al. 2017: chạy \( h \) head song song, mỗi head có \( W^Q_i, W^K_i, W^V_i \) riêng. Mỗi head tự do học một metric similarity của riêng nó. Sau đó concat output và project bằng \( W^O \) để gộp lại.

Loại bỏ khỏi đầu cách hiểu "head = attention heat map". Đúng hơn: head = một subspace của \( d_{\text{model}} \) trong đó model định nghĩa một quan hệ token-token riêng.

Với input \( Q, K, V \) (trong self-attention thì cả ba đều bằng \( X \)) và số head \( h \):

\[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O \]

Trong đó head thứ \( i \) là một self-attention/cross-attention riêng:

\[ \text{head}_i = \text{Attention}\!\left(Q W^Q_i, \ K W^K_i, \ V W^V_i\right) \]

Các ma trận học được:

• \( W^Q_i, W^K_i \in \mathbb{R}^{d_{\text{model}} \times d_k} \), \( W^V_i \in \mathbb{R}^{d_{\text{model}} \times d_v} \) — riêng cho mỗi head.
• \( W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}} \) — projection trộn output của \( h \) head về lại không gian \( d_{\text{model}} \).

Concat ở đây nghĩa là xếp theo trục feature: nếu mỗi head xuất \( (n, d_v) \), thì concat \( h \) head được \( (n, h d_v) \). \( W^O \) đưa về \( (n, d_{\text{model}}) \) — đúng shape input để chồng tiếp layer kế tiếp.

\( W^O \) không chỉ là phép biến đổi tuyến tính phụ. Nó là chỗ trộn thông tin giữa các head. Bỏ \( W^O \), \( h \) head sẽ "sống độc lập" — model mất khả năng kết hợp pattern từ head khác nhau.

Lựa chọn dimension trong Vaswani et al. 2017 là chia đều \( d_{\text{model}} \) cho \( h \) head:

\[ d_k = d_v = \frac{d_{\text{model}}}{h} \]

Ví dụ Transformer base: \( d_{\text{model}} = 512, h = 8 \Rightarrow d_k = 64 \).

Lý do của cách chia này:

• Concat \( h \) head ra đúng \( h \cdot d_k = d_{\text{model}} \) — match input dim, không cần đổi shape ngoài.
• Tổng FLOPs của attention trên \( h \) head \( \approx \) FLOPs của 1 head full-dim \( d_{\text{model}} \) (xem Bước 6, 18). MHA "miễn phí" về compute so với single-head cùng \( d_{\text{model}} \).
• Param đếm cho MHA cũng tương đương: \( \sum_i 3 \cdot d_{\text{model}} \cdot d_k = 3 \cdot d_{\text{model}}^2 \), không phụ thuộc \( h \).

Có biến thể khác — ví dụ T5 chọn \( d_k \) độc lập (\( d_k = 64 \) bất kể \( d_{\text{model}} \)). Khi đó \( h d_k \ne d_{\text{model}} \) và \( W^O \) chỉnh shape về \( d_{\text{model}} \) tương ứng. Trong khuôn khổ bài này, giữ default \( d_k = d_{\text{model}} / h \).

Đếm tham số 1 layer MHA, bỏ qua bias (đa số implementation hiện đại — Llama, GPT-NeoX — bỏ bias trong projection attention):

\[ |W^Q| + |W^K| + |W^V| = h \cdot (3 \cdot d_{\text{model}} \cdot d_k) = 3 \cdot d_{\text{model}}^2 \]

(vì \( h \cdot d_k = d_{\text{model}} \))

\[ |W^O| = (h d_v) \cdot d_{\text{model}} = d_{\text{model}}^2 \]

Tổng:

\[ \text{Params}_{\text{MHA}} = 4 \cdot d_{\text{model}}^2 \]

Ví dụ:

• \( d_{\text{model}} = 512 \): \( 4 \times 262{,}144 \approx 1.05 \) triệu tham số / layer MHA.
• \( d_{\text{model}} = 768 \) (BERT base, GPT-2 small): \( 4 \times 589{,}824 \approx 2.36 \)M.
• \( d_{\text{model}} = 4096 \) (Llama 3 8B): \( 4 \times 16{,}777{,}216 \approx 67 \)M (cho 1 layer; toàn model 32 layer → ~2.1B param chỉ riêng attention).

Param không phụ thuộc \( h \) — chia 8 head hay 16 head cùng tổng param. Cái thay đổi là cấu trúc tính toán bên trong, không phải kích thước.

(Với GQA/MQA, \( W^K, W^V \) nhỏ hơn — Bước 10, 11. Lúc đó tổng param < \( 4 d^2 \).)

Nhìn công thức Bước 4, có cảm giác cần \( 3h \) ma trận riêng. Thực tế gộp lại làm 3 ma trận full-dim:

1. 1 ma trận \( W^Q \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}} \) đóng vai trò "concat" của \( h \) ma trận con \( W^Q_i \). Tính \( Q = X W^Q \) một lần ra \( (B, n, d_{\text{model}}) \).
2. Reshape \( Q \) thành \( (B, n, h, d_k) \) rồi transpose thành \( (B, h, n, d_k) \) — tách \( h \) head qua trục thứ 2.
3. Áp attention batched trên trục head: \( Q K^\top / \sqrt{d_k} \) cho ra \( (B, h, n, n) \), softmax, nhân \( V \) ra \( (B, h, n, d_k) \).
4. Transpose ngược lại \( (B, n, h, d_k) \) → reshape \( (B, n, d_{\text{model}}) \) — đây chính là concat của \( h \) head.
5. Project bằng \( W^O \): \( (B, n, d_{\text{model}}) \cdot (d_{\text{model}}, d_{\text{model}}) \to (B, n, d_{\text{model}}) \).

GPU thực thi cực nhanh nhờ batched matmul trên trục head. Việc tách head chỉ là reshape — không có copy bộ nhớ thật sự trong PyTorch khi tensor contiguous.

Clark et al. 2019 ("What Does BERT Look At?", arXiv:1906.04341) phân tích từng head của BERT base (12 layer × 12 head) và quan sát các nhóm pattern:

• Syntax head: head attend từ động từ về chủ ngữ, từ giới từ về object, từ tính từ về danh từ được bổ nghĩa. Gần như phục hồi cây dependency parse.
• Semantic head: head xử lý coreference (pronoun → antecedent), liên kết tên riêng được nhắc lại nhiều lần.
• Local head: head attend token liền trước hoặc liền sau — bắt n-gram, cụm từ ngắn.
• Special token head: head dồn attention vào [CLS], [SEP]. Đóng vai trò "kênh broadcast" tổng hợp toàn câu.
• Trivial / redundant head: head attend đều khắp hoặc luôn attend chính token đó — không mang thông tin có ích.

Quan trọng: phân chia này không được giám sát trực tiếp. Pre-training (MLM cho BERT, next-token cho GPT) đủ để các pattern phát sinh. Multi-head cho model khả năng chứa đa pattern; head nào không có ích sẽ bị optimizer "bỏ rơi" với weight nhỏ.

Voita et al. 2019 ("Analyzing Multi-Head Self-Attention", arXiv:1905.09418) thực nghiệm prune từng head trong mô hình WMT translation Transformer:

• Đo độ "quan trọng" mỗi head qua gradient và tác động khi gate head về 0.
• Có thể prune ~70-80% head mà BLEU score chỉ giảm dưới 1 point.
• Một số head essential — bỏ là chất lượng sụp đổ. Đa số head dư thừa (redundant) — chia sẻ pattern với head khác.

Hệ quả thực tế:

• Multi-head có vẻ "overparameterized" — không phải mọi head đều cần.
• Nhưng training với nhiều head dễ hơn vì gradient flow tốt hơn (mỗi pattern có "đường thoát" riêng). Prune sau khi train rẻ hơn nhiều so với train ngay model nhỏ.
• Đây là một động lực để các biến thể chia sẻ head ra đời (MQA, GQA) — vẫn giữ đa pattern trên \( Q \) nhưng giảm chi phí \( K, V \).

Michel et al. 2019 ("Are Sixteen Heads Really Better than One?", arXiv:1905.10650) độc lập confirm cùng kết luận: phần lớn head có thể prune.

Shazeer 2019 ("Fast Transformer Decoding: One Write-Head is All You Need", arXiv:1911.02150) đề xuất MQA: giữ \( h \) head riêng cho \( Q \), nhưng \( K \) và \( V \) chia sẻ 1 head duy nhất cho mọi query head.

Công thức head thứ \( i \) trong MQA:

\[ \text{head}_i^{\text{MQA}} = \text{Attention}\!\left(Q W^Q_i, \ K W^K, \ V W^V\right) \]

Trong đó \( W^K, W^V \) không có chỉ số \( i \) — dùng chung cho mọi head.

Hệ quả:

• KV cache giảm \( h \) lần — quan trọng cực kỳ khi inference autoregressive với context dài.
• Memory bandwidth khi đọc KV cache giảm tỉ lệ thuận — tốc độ token/s tăng đáng kể.
• Trade-off: chất lượng giảm nhẹ vs MHA gốc (Shazeer 2019 báo BLEU drop ~0.5 trên WMT EnDe).

Model dùng MQA: PaLM (Google 2022), Falcon, ban đầu là Gopher / Chinchilla nội bộ DeepMind.

Ainslie et al. 2023 ("GQA: Training Generalized Multi-Query Transformer Models", arXiv:2305.13245) đưa ra trung gian giữa MHA và MQA:

• Chia \( h \) query head thành \( g \) nhóm (\( 1 < g < h \)).
• Mỗi nhóm chia sẻ 1 \( K, V \) — tức có \( g \) bộ \( W^K, W^V \).
• \( g = h \) → MHA gốc; \( g = 1 \) → MQA.

Kết quả thực nghiệm Ainslie 2023:

• Quality gần MHA (drop < 0.2 point trên multi-task).
• Speed inference gần MQA.
• \( g = 8 \) là cấu hình "ngọt nhất" trên đa số benchmark.

Model dùng GQA năm 2024-2025:

Model            d_model    Q heads    KV heads (g)
Llama 2 70B      8192       64         8
Llama 3 8B       4096       32         8
Llama 3 70B      8192       64         8
Llama 3 405B     16384      128        16
Mistral 7B       4096       32         8
Mixtral 8x7B     4096       32         8
Qwen 2.5 7B      3584       28         4
DeepSeek-V2      5120       128        128 (dùng MLA, Bước 12)

So sánh KV cache cho context 4K, \( d_{\text{model}} = 4096, h = 32 \), batch 1, fp16:

• MHA (\( h_{kv} = 32 \)): \( 2 \times 32 \times 4096 \times 128 \times 2 \) byte \( \approx 67 \) MB.
• GQA \( g = 8 \): \( \approx 67 / 4 = 16.8 \) MB.
• MQA (\( g = 1 \)): \( \approx 67 / 32 = 2.1 \) MB.

Khi batch 256 và context 32K, sự khác biệt thành hàng chục GB — quyết định model có chạy được trên 1 GPU hay không.

DeepSeek-V2 (DeepSeek-AI 2024, arXiv:2405.04434) đề xuất Multi-Latent Attention: nén \( K, V \) xuống một latent chiều thấp \( d_c \ll d_{\text{model}} \) trước khi cache, rồi up-project khi tính attention.

Sơ đồ:

X → W^DKV → c (latent, d_c chiều, được cache)
c → W^UK → K
c → W^UV → V

Điểm chính:

• Cache \( c \) chiều \( d_c \) thay vì cache \( K, V \) chiều \( h \cdot d_k = d_{\text{model}} \). Với DeepSeek-V2, \( d_c \approx d_{\text{model}} / 8 \) — KV cache giảm ~8 lần.
• Tích hợp trick "absorb" \( W^{UK}, W^{UV} \) vào \( W^Q, W^O \) lúc inference — không cần up-project thực tế ở mỗi token.
• Chất lượng theo paper báo cáo tương đương MHA, vượt GQA trên nhiều benchmark.

MLA được DeepSeek-V2, V3, R1 sử dụng. Đây là hướng nghiên cứu mới — cộng đồng đang đánh giá xem có thay thế GQA làm default cho LLM lớn không.

Module MHA tự cài, theo cách reshape Bước 7. Code PyTorch 2.x:

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


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_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: torch.Tensor, mask: torch.Tensor | None = None):
        B, N, _ = X.shape
        # (B, N, d_model) → (B, num_heads, N, d_k)
        Q = self.W_Q(X).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_K(X).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_V(X).view(B, N, self.num_heads, self.d_k).transpose(1, 2)

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

        # Concat: (B, h, N, d_k) → (B, N, h, d_k) → (B, N, d_model)
        out = out.transpose(1, 2).contiguous().view(B, N, -1)
        return self.W_O(out)

Test trên input giả:

torch.manual_seed(0)
B, N, d_model, h = 1, 10, 512, 8
X = torch.randn(B, N, d_model)

mha = MultiHeadAttention(d_model, h)
out = mha(X)
print(out.shape)              # torch.Size([1, 10, 512])

n_params = sum(p.numel() for p in mha.parameters())
print(f"Params: {n_params:,}")  # 4 × 512² = 1,048,576

Kết quả param đếm đúng \( 4 d^2 \) như Bước 6 (không có bias).

PyTorch cung cấp nn.MultiheadAttention sẵn — dùng trong production thay vì cài tay:

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True, bias=False)
X = torch.randn(1, 10, 512)
out, attn_weights = mha(X, X, X)  # self-attention: Q = K = V = X
print(out.shape)          # torch.Size([1, 10, 512])
print(attn_weights.shape) # torch.Size([1, 10, 10])  — trung bình h head

Chú ý:

• batch_first=True (PyTorch ≥ 1.9) — input shape \( (B, N, d) \) thay vì \( (N, B, d) \) mặc định cũ.
• Trả về 2 thứ: output và attention weights. Khi need_weights=False, chỉ trả output (nhanh hơn vì có thể dùng SDPA backend).
• attn_weights mặc định là trung bình \( h \) head. Muốn lấy từng head, set average_attn_weights=False → shape \( (B, h, N, N) \).

Trong layer Transformer của PyTorch (nn.TransformerEncoderLayer), MultiheadAttention được gọi từ ngầm — bạn không cần khởi tạo trực tiếp khi xếp block.

Model           d_model    Layer    h (Q heads)    KV heads    Variant
BERT base       768        12       12             12          MHA
BERT large      1024       24       16             16          MHA
GPT-2 small     768        12       12             12          MHA
GPT-2 XL        1600       48       25             25          MHA
GPT-3 175B      12288      96       96             96          MHA
Llama 2 7B      4096       32       32             32          MHA
Llama 2 70B     8192       80       64             8           GQA
Llama 3 8B      4096       32       32             8           GQA
Llama 3 70B     8192       80       64             8           GQA
Mistral 7B      4096       32       32             8           GQA
Mixtral 8x7B    4096       32       32             8           GQA + MoE
Gemma 7B        3072       28       16             16          MHA
Qwen 2.5 7B     3584       28       28             4           GQA
PaLM 540B       18432      118      48             1           MQA
DeepSeek-V2     5120       60       128            -           MLA
DeepSeek-V3     7168       61       128            -           MLA

Quan sát:

• Model 2017-2022 phổ biến dùng MHA "đơn giản".
• Từ 2023, model open-source >= 7B chuyển sang GQA \( g = 8 \) — tiết kiệm KV cache mà gần như không mất chất lượng.
• Model proprietary lớn (PaLM) chọn MQA cực đoan để tối ưu serving.
• 2024-2025: DeepSeek dẫn đầu hướng MLA — vẫn còn mới so với GQA.
• \( d_k \) thường giữ ở 64-128 — không tăng theo \( d_{\text{model}} \). Tăng \( d_{\text{model}} \) thường đi cùng tăng số head.

Với MHA \( h \) head, ma trận attention không còn 1 ma trận \( n \times n \) mà là tensor \( (h, n, n) \) — \( h \) heatmap riêng. Cách quan sát:

• Plot \( h \) heatmap trong lưới (subplot grid). Nhìn head nào có diagonal, head nào có cột dọc, head nào "phân bố đều".
• Quan sát theo layer: layer thấp thường nhiều head local, layer cao nhiều head semantic / coreference.

Tool:

• BertViz (Vig 2019, github.com/jessevig/bertviz): view "head", "model", "neuron" cho BERT / GPT-2 trong Jupyter notebook.
• Transformers Interpret: tích hợp Captum cho HF Transformers.
• AttentionViz (Yeh et al. 2023, arXiv:2305.03210): visualize query-key embedding 2D thay vì heatmap.

Khi visualize lần đầu, kỳ vọng phân biệt rõ giữa head — đừng kỳ vọng head 1 luôn là syntax head, head 2 luôn là coreference. Vị trí cụ thể của từng pattern thay đổi theo seed train và checkpoint.

Trong kiến trúc encoder-decoder (T5, BART, mBART, NLLB), decoder có sub-layer cross-attention: \( Q \) lấy từ decoder hidden state, \( K, V \) lấy từ encoder output. Công thức MHA giữ nguyên — chỉ thay nguồn:

\[ \text{CrossMHA}(Q_{\text{dec}}, K_{\text{enc}}, V_{\text{enc}}) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O \]

\[ \text{head}_i = \text{Attention}\!\left(Q_{\text{dec}} W^Q_i, \ K_{\text{enc}} W^K_i, \ V_{\text{enc}} W^V_i\right) \]

Chú ý:

• \( Q \) có chiều dài \( n_{\text{dec}} \), \( K, V \) có chiều dài \( n_{\text{enc}} \) — ma trận attention \( (n_{\text{dec}}, n_{\text{enc}}) \), không vuông.
• Không có causal mask trong cross-attention — decoder có quyền nhìn toàn bộ encoder.
• \( K, V \) thường chỉ được tính 1 lần (cho mọi step decode) và cache lại — tối ưu inference.

Decoder-only LLM (GPT, Llama) không có cross-attention — chỉ self-attention causal.

So sánh với 1 head \( d_{\text{model}} \) (Bài 10):

\[ \text{Forward time} = O(n^2 d + n d^2) \]

• \( n^2 d \) — từ \( QK^\top \) và \( \alpha V \), tổng trên \( h \) head: \( h \cdot n^2 d_k = n^2 d \). Không phụ thuộc \( h \) (vì \( h d_k = d \)).
• \( n d^2 \) — 4 projection \( W^Q, W^K, W^V, W^O \).

Memory cho attention matrix:

\[ \text{Memory}_{\text{attn}} = O(h \cdot n^2) \]

Đây là chỗ MHA "đắt" hơn single-head: với \( h \) head, mỗi head có 1 ma trận attention \( n \times n \) riêng. Memory tăng tuyến tính theo \( h \). FlashAttention giải quyết bằng cách không vật chất hóa ma trận này trong HBM (Bước 19).

FLOPs forward 1 layer MHA (bỏ FFN):

\[ \text{FLOPs}_{\text{MHA}} \approx 4 n d^2 + 2 n^2 d \]

Với \( n = 2048, d = 4096 \) (Llama 3 8B style): khoảng \( 1.4 \times 10^{11} \) FLOPs / layer / forward / batch=1. Toàn model 32 layer → \( 4.5 \times 10^{12} \) FLOPs. RTX 4090 cung cấp ~83 TFLOPs fp16 sustained — forward 1 chuỗi ~55 ms (thực tế chậm hơn do memory bound).

Một số kỹ thuật MHA chạy nhanh và đỡ tốn memory trên GPU hiện đại (A100, H100, B200, RTX 4090):

• FlashAttention 2 / 3 (Dao 2023, Shah 2024): fused kernel IO-aware, tile attention vào SRAM, không vật chất hóa ma trận \( (h, n, n) \) trong HBM. Memory \( O(n) \) thay vì \( O(n^2) \), tốc độ tăng 2-4x. Mặc định trong PyTorch SDPA khi điều kiện thỏa mãn.
• PagedAttention (Kwon et al. 2023, vLLM, arXiv:2309.06180): chia KV cache thành block giống virtual memory paging, cho phép serving nhiều request cùng lúc với ít fragment.
• xFormers (Meta): collection các kernel attention tối ưu, hỗ trợ nhiều biến thể (memory-efficient, block-sparse).
• PyTorch SDPA (torch.nn.functional.scaled_dot_product_attention): auto chọn backend FlashAttention / memory-efficient / math tùy điều kiện đầu vào.
• Tensor Parallelism (Megatron-LM, Shoeybi 2019): chia head ngang qua nhiều GPU — head 1-16 trên GPU 0, head 17-32 trên GPU 1, v.v. Allreduce sau projection \( W^O \).

Code production thực sự dùng các thư viện này — không cài MHA tay bằng matmul thô. Đoạn code ở Bước 13 chỉ để hiểu cơ chế.

MHA là sub-layer đầu tiên trong mỗi Transformer block. Cấu trúc pre-LN phổ biến hiện nay:

input X (B, n, d_model)
  │
  ├─→ LayerNorm → MultiHeadAttention(causal mask?) → +X   (residual)
  │
  ├─→ LayerNorm → FFN (2-layer MLP, hidden 4d)     → +    (residual)
  │
output (B, n, d_model)

Output giữ shape \( (n, d_{\text{model}}) \) — đầu vào của block kế tiếp. Bài 13 deep dive cấu trúc đầy đủ: pre-LN vs post-LN, FFN, residual, dropout, các biến thể như RMSNorm / SwiGLU / Pre-Norm trong Llama.

Trong decoder-only LLM, mỗi block có duy nhất 1 sub-layer MHA (causal). Trong encoder-decoder (T5), block decoder có 2 sub-layer attention: 1 self-attention causal + 1 cross-attention với encoder output.

1. Implement MultiHeadAttention ở Bước 13 cho \( d_{\text{model}} = 512, h = 8 \). Test forward với input shape \( (1, 10, 512) \), in shape output (kỳ vọng \( (1, 10, 512) \)) và đếm param tổng (kỳ vọng \( \approx 1{,}048{,}576 \) = \( 4 \times 512^2 \), không bias).
2. Sửa class trên cho phép nhận causal mask. Sinh mask tam giác dưới shape \( (1, 1, N, N) \), test trên cùng input. Verify rằng attention pattern ở trên đường chéo bằng 0.
3. Tính tay param count cho 1 layer MHA của Llama 3 8B: \( d_{\text{model}} = 4096, h = 32 \), \( h_{kv} = 8 \) (GQA). So sánh với MHA đầy đủ \( h_{kv} = 32 \) — tỉ lệ tiết kiệm bao nhiêu?
4. Tính KV cache MHA vs GQA \( g = 8 \) vs MQA cho cấu hình: context 4K token, \( d_{\text{model}} = 4096 \), \( h = 32 \), \( d_k = 128 \), batch 1, fp16. Đơn vị MB. Mở rộng cho batch 32 và context 32K — viết bảng tổng KV cache (MB) cho 9 ô (3 cấu hình × 3 cặp batch/context).
5. Lấy nn.MultiheadAttention, set need_weights=True, average_attn_weights=False để lấy attention từng head. So sánh output với MultiHeadAttention tự cài ở Bước 13 (set cùng seed, copy weight tương ứng). Sai số tuyệt đối có lớn hơn 1e-5 không?
6. Đọc Section 3.2.2 của Vaswani et al. 2017 ("Multi-Head Attention"). Ghi câu hỏi nếu chưa rõ phần nào — đặc biệt phần "dot products grow large in magnitude" và lý do chọn 8 head.
7. (Đọc thêm) Skim Ainslie et al. 2023 (GQA, arXiv:2305.13245) Section 3.1 + Table 1. Tại sao \( g = 8 \) lại "ngọt" hơn các giá trị khác trên benchmark của họ?

• Multi-Head Attention chạy \( h \) self-attention song song, mỗi head có riêng \( W^Q_i, W^K_i, W^V_i \). Concat output và project bằng \( W^O \) để trộn về \( d_{\text{model}} \).
• Chia chiều \( d_k = d_{\text{model}} / h \) — concat ra đúng \( d_{\text{model}} \), không tăng FLOPs vs single-head full-dim.
• Param count 1 layer MHA: \( 4 \cdot d_{\text{model}}^2 \) (3 từ Q, K, V + 1 từ \( W^O \)). Không phụ thuộc \( h \).
• Implement efficient: 4 ma trận full-dim, reshape tách head qua trục thứ 2, batched matmul.
• Multi-head làm việc vì mỗi head học một pattern khác nhau: syntax, semantic, coreference, local, special token (Clark 2019). Pattern phát sinh từ pre-training, không giám sát trực tiếp.
• Đa số head có thể prune (Voita 2019, Michel 2019) — một số "essential", hầu hết redundant. Động lực cho MQA / GQA chia sẻ \( K, V \).
• MQA (Shazeer 2019): \( h \) Q head, 1 KV head. KV cache giảm \( h \) lần, chất lượng giảm nhẹ. Dùng trong PaLM, Falcon.
• GQA (Ainslie 2023): \( g \) KV group (\( 1 < g < h \)). Quality gần MHA, speed gần MQA. \( g = 8 \) phổ biến — Llama 2/3, Mistral, Mixtral, Qwen.
• MLA (DeepSeek-V2 2024): nén KV xuống latent \( d_c \ll d_{\text{model}} \). Cache nhỏ hơn cả GQA, chất lượng theo paper tương đương MHA.
• FlashAttention, PagedAttention, PyTorch SDPA là các optimization production thực tế. Code tự cài chỉ để học.

Bài 12 (Positional Encoding) trả lời câu hỏi: self-attention và MHA bất biến với hoán vị token — model làm sao biết "tôi yêu bạn" khác "bạn yêu tôi"? Câu trả lời nằm ở positional encoding gắn vào embedding trước khi đưa qua block đầu tiên.