[AI Blog #30] Kiến Trúc Transformer & Cơ Chế Self-Attention Chuyên Sâu: Giải Mã Toán Học Scaled Dot-Product Attention, Grouped-Query Attention (GQA), RoPE & Bí Quyết Tối Ưu Hóa FlashAttention

1. Bối cảnh thực tế (Context & Problem Statement)

Trước năm 2017, các kiến trúc xử lý ngôn ngữ tự nhiên (NLP) thống trị là RNN (Recurrent Neural Networks) và LSTM (Long Short-Term Memory). Tuy nhiên, RNN/LSTM gặp phải hai "nút thắt cổ chai" chí mạng:

Xử lý tuần tự (Sequential Bottleneck): Để tính toán từ thứ $t$ , mô hình bắt buộc phải chờ tính xong từ thứ $t-1$ . Điều này khiến RNN hoàn toàn không thể tận dụng khả năng tính toán song song khổng lồ của GPU.
Quên ngữ cảnh dài (Long-term Dependency Loss): Thông tin bị nén qua một vector ẩn duy nhất qua nhiều bước, khiến mô hình bị mất dần ngữ cảnh khi độ dài văn bản vượt quá vài trăm từ.

Năm 2017, nhóm nghiên cứu tại Google công bố bài báo lịch sử "Attention Is All You Need", khai sinh ra kiến trúc Transformer. Bằng cách loại bỏ hoàn toàn cấu trúc hồi quy và thay thế bằng cơ chế Self-Attention (Tự chú ý), Transformer cho phép:

Tính toán song song toàn bộ văn bản trong 1 chu kỳ trên GPU.
Mọi từ trong câu có thể tương tác trực tiếp với nhau bất kể khoảng cách địa lý ( $O(1)$ path length).

2. Các Khái Niệm & Cơ Chế Cốt Lõi

2.1. Bản Chất Toán Học Của Scaled Dot-Product Attention

Mỗi từ (Token) sau khi đi qua lớp Embedding sẽ được nhân với 3 ma trận trọng số học được ( $W_Q, W_K, W_V$ ) để tạo thành 3 vector:

Query ( $Q$ - Câu hỏi): "Tôi đang tìm kiếm thông tin gì?"
Key ( $K$ - Từ khóa): "Tôi chứa đựng thông tin gì để người khác tìm?"
Value ( $V$ - Giá trị thực): "Nội dung thông tin thực sự của tôi là gì?"

$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) V$

Query (Q)   Key (K)
    \       /
     \     /
   [ Q * K^T ]   (Tính độ tương đồng Dot-Product giữa mọi cặp từ)
        |
        v
   [ / sqrt(d_k) ] (Chia cho căn bậc hai chiều không gian để ổn định phương sai)
        |
        v
   [ Softmax ]    (Chuyển đổi thành phân phối xác suất trọng số chú ý 0 - 1)
        |
        v
   [ * Value (V) ] (Tính tổng có trọng số của các vector Value)
        |
        v
   [ Attention Output ]

Tại sao bắt buộc phải chia cho $\sqrt{d_k}$ ?
- Giả sử các phần tử của $Q$ và $K$ là các biến ngẫu nhiên độc lập có kỳ vọng bằng $0$ và phương sai bằng $1$ .
- Tích vô hướng $Q \cdot K = \sum_{i=1}^{d_k} Q_i K_i$ sẽ có kỳ vọng bằng $0$ nhưng phương sai tăng vọt lên bằng $d_k$ !
- Khi số chiều $d_k$ lớn (ví dụ $d_k = 128$ ), giá trị tích vô hướng sẽ rất lớn $\rightarrow$ Đẩy hàm Softmax vào vùng bão hòa cực đoan (Saturation Zone), nơi đạo hàm tiến sát về $0$ $\rightarrow$ Gây tê liệt hoàn toàn quá trình Backpropagation (Vanishing Gradient)!
- Phép chia cho $\sqrt{d_k}$ đưa phương sai trở lại đúng bằng $1$ , giúp dòng chảy gradient luôn khỏe mạnh.

2.2. Tiến Hóa Của Các Cấu Trúc Attention (MHA $\rightarrow$ MQA $\rightarrow$ GQA)

Trong quá trình tạo văn bản (Autoregressive Generation), mô hình phải lưu lại toàn bộ các vector Key và Value của các từ trước đó vào bộ nhớ VRAM (KV Cache) để không phải tính lại.

Multi-Head Attention (MHA - Chuẩn ban đầu): Mỗi Head có bộ $Q, K, V$ riêng biệt (ví dụ 32 Query Heads đi kèm 32 Key Heads và 32 Value Heads).
$\rightarrow$ Dung lượng KV Cache khổng lồ, làm tràn VRAM GPU khi tăng Context Length ( $32\text{k} - 128\text{k}$ tokens).
Multi-Query Attention (MQA): Tất cả các Query Heads dùng chung duy nhất 1 Key Head và 1 Value Head.
$\rightarrow$ Tiết kiệm $90\%+$ dung lượng KV Cache nhưng làm suy giảm chất lượng mô hình.
Grouped-Query Attention (GQA - Tiêu chuẩn vàng của LLaMA 3, Mistral): Chia 32 Query Heads thành các nhóm (ví dụ 8 nhóm), mỗi nhóm gồm 4 Query Heads dùng chung 1 cặp Key-Value Head.
$\rightarrow$ Đạt tốc độ suy luận siêu nhanh và tiết kiệm bộ nhớ như MQA nhưng giữ nguyên độ chính xác cao tương đương MHA!

MHA (32 Q, 32 K, 32 V)          GQA (32 Q, 8 K, 8 V)            MQA (32 Q, 1 K, 1 V)
[Q Q Q Q] ... [Q Q Q Q]        [Q Q Q Q] ... [Q Q Q Q]        [Q Q Q Q ... Q Q Q Q]
 | | | |       | | | |          \ | / /       \ | / /           \ \ \ \ / / / /
[K K K K] ... [K K K K]            [K]           [K]                  [K]
[V V V V] ... [V V V V]            [V]           [V]                  [V]
 (Tốn VRAM nhất)                (Cân bằng hoàn hảo)            (Chất lượng giảm nhẹ)

2.3. Mã Hóa Vị Trí Hiện Đại: Rotary Position Embedding (RoPE)

Khác với mạng CNN hay RNN, phép tính Self-Attention hoàn toàn bất biến theo vị trí (Permutation Invariant - hoán đổi vị trí các từ thì kết quả vẫn như nhau). Do đó, ta bắt buộc phải đưa thông tin thứ tự vào mô hình.

Absolute Sinusoidal (Transformer gốc): Cộng trực tiếp vector vị trí $P_i$ vào vector từ: $x_i + P_i$ .
Rotary Position Embedding (RoPE - Chuẩn mực hiện đại):
- Thay vì phép cộng, RoPE xoay vector $Q$ và $K$ trong không gian 2D bằng một ma trận trực giao (Rotation Matrix) theo góc tỷ lệ thuận với vị trí của từ ( $m\theta$ ): $\tilde{q}_m = R_{\Theta, m}^d q_m, \quad \tilde{k}_n = R_{\Theta, n}^d k_n$
- Tích vô hướng sau khi xoay tự động thỏa mãn tính chất tương đối: $\langle \tilde{q}_m, \tilde{k}_n \rangle = q_m^T R_{\Theta, n-m}^d k_n = g(q, k, m - n)$
- $\rightarrow$ Mô hình chỉ quan tâm đến khoảng cách tương đối $(m - n)$ giữa 2 từ, cho phép mở rộng cửa sổ ngữ cảnh (Context Window Extension) lên $128\text{k} - 1\text{M}$ tokens cực kỳ dễ dàng thông qua RoPE Scaling (YaRN, NTK-aware).

2.4. Bí Quyết Tăng Tốc Tầng Phần Cứng: FlashAttention (Triệt Tiêu HBM Bottleneck)

Vấn đề của Attention chuẩn: Việc tính toán $S = Q K^T$ ( $N \times N$ ) và $\text{softmax}(S)$ đòi hỏi liên tục ghi và đọc các ma trận khổng lồ ra bộ nhớ chính của GPU (High Bandwidth Memory - HBM). Băng thông truyền dữ liệu giữa HBM và chip tính toán (SRAM) trở thành điểm nghẽn chính (Memory-Bound).
FlashAttention (Dao et al.):
- Áp dụng kỹ thuật Tiling (Chia nhỏ ma trận) để nạp vừa vặn các khối dữ liệu vào bộ nhớ đệm siêu tốc SRAM bên trong GPU.
- Tận dụng thuật toán Online Softmax để tính toán kết quả Attention từng phần mà không bao giờ phải lưu trữ ma trận $N \times N$ khổng lồ ra HBM đắt đỏ.
- $\rightarrow$ Tăng tốc độ huấn luyện $2\times - 4\times$ và giảm độ phức tạp bộ nhớ từ $O(N^2)$ xuống $O(N)$ !

3. Các Bảng Markdown So Sánh Chi Tiết

📊 Bảng 1: So Sánh 3 Nhánh Kiến Trúc Transformer: Encoder-Only vs. Decoder-Only vs. Encoder-Decoder

Tiêu chí	Encoder-Only (BERT, RoBERTa)	Decoder-Only (GPT-4, LLaMA 3, Mistral)	Encoder-Decoder (T5, BART)
Cơ chế Attention	Bidirectional Self-Attention (Nhìn thấy cả từ trước và sau)	Causal Masked Self-Attention (Chỉ nhìn thấy các từ quá khứ)	Encoder: 2 chiều; Decoder: 1 chiều + Cross-Attention
Mục tiêu huấn luyện	Masked Language Modeling (Điền từ vào chỗ trống `[MASK]`)	Next-Token Prediction (Dự đoán từ tiếp theo)	Sequence-to-Sequence Reconstruction (Phục hồi câu bị xáo trộn)
Thế mạnh vượt trội	Phân loại văn bản, Trích xuất thực thể (NER), Semantic Search	Tự do sinh văn bản, Chatbot, Lập luận logic, Code Generation	Dịch thuật (Translation), Tóm tắt văn bản (Summarization)
Vị thế hiện nay	Làm Embedding Models, RAG Bi-encoders	Thống trị $95\%+$ các mô hình Generative AI & Foundation LLMs	Dùng trong các bài toán Seq2Seq chuyên biệt

📊 Bảng 2: So Sánh Các Biến Thể Attention: MHA vs. MQA vs. GQA

Tiêu chí	Multi-Head Attention (MHA)	Multi-Query Attention (MQA)	Grouped-Query Attention (GQA)
Số lượng Key/Value Heads	Bằng số Query Heads ( $H_K = H_Q$ )	Duy nhất 1 Head ( $H_K = 1$ )	Nhóm nhỏ ( $1 < H_K < H_Q$ )
Dung lượng KV Cache	$100\%$ (Chuẩn gốc)	Giảm tới $95\%$ (Siêu nhẹ)	Giảm $75\% - 87.5\%$
Tốc độ sinh Token (Inference)	Chậm khi Context dài	Cực nhanh	Rất nhanh (Tương đương MQA)
Khả năng duy trì độ chính xác	Tối đa	Có thể bị sụt giảm trên bài toán khó	Gần như tương đương $100\%$ so với MHA
Mô hình tiêu biểu	GPT-3, Transformer 2017, BERT	Falcon-7B, StarCoder	LLaMA 2/3 (70B), Mistral-7B, Mixtral 8x7B, Gemma

📊 Bảng 3: So Sánh Các Phương Pháp Mã Hóa Vị Trí (Positional Encodings)

Phương pháp	Cơ chế thực hiện	Khả năng mở rộng Context (Extrapolation)	Chi phí tính toán	Mô hình áp dụng
Absolute Sinusoidal	Cộng hàm $\sin/\cos$ cố định vào Input Embedding	Rất kém (Không hiểu được vị trí vượt quá lúc Train)	$O(1)$	Transformer gốc (2017)
Learned Positional	Học vector vị trí như tham số mô hình	Kém (Cố định độ dài tối đa)	$O(1)$	BERT, GPT-2
ALiBi (Linear Biases)	Trừ trực tiếp một giá trị phạt tuyến tính vào ma trận Attention	Xuất sắc (Huấn luyện $2\text{k}$ , suy luận $8\text{k}$ )	$O(1)$	MPT-7B, BLOOM
RoPE (Rotary)	Xoay vector $Q, K$ trong không gian 2D theo góc $m\theta$	Tối ưu nhất khi kết hợp Scaling	Rất nhẹ (Chỉ tốn phép nhân vector)	LLaMA 1/2/3, Mistral, Qwen, DeepSeek

4. Best Practices & Kỹ Thuật Tối Ưu Hóa LLM Serving (Pro-Tips)

💡 Khắc phục phân mảnh bộ nhớ VRAM với PagedAttention (vLLM Engine)

Vấn đề: Trong quá trình sinh từ, độ dài câu trả lời của mỗi user là không biết trước. Cơ chế cấp phát bộ nhớ truyền thống buộc phải dành sẵn các khối VRAM liền kề liên tục cho độ dài tối đa (ví dụ $4096$ tokens) $\rightarrow$ Gây lãng phí và phân mảnh tới $60\% - 80\%$ VRAM GPU!
Giải pháp PagedAttention (Ý tưởng từ Virtual Memory của OS):
- Chia nhỏ KV Cache thành các Trang bộ nhớ (Memory Pages) có kích thước cố định (ví dụ 16 tokens/page).
- Các trang này được lưu rải rác không cần liền kề trong VRAM và được quản lý qua một Bảng ánh xạ trang (Page Table).
- $\rightarrow$ Tăng số lượng request xử lý đồng thời (Throughput) lên gấp $2\times - 4\times$ trên cùng một GPU!

💡 Chiến lược Kích hoạt Mạng nơ-ron: Thay thế ReLU bằng SwiGLU

Các mô hình LLM hiện đại đều thay thế Feed-Forward Network (FFN) truyền thống bằng SwiGLU (Swish Gated Linear Unit): $\text{FFN}_{\text{SwiGLU}}(x) = \Big( \text{Swish}(x W_1) \otimes (x W_3) \Big) W_2$

Phép nhân có cổng (Gated mechanism) cho phép mô hình kiểm soát lượng thông tin được truyền qua một cách mềm dẻo, giúp tăng đáng kể năng lực suy luận logic mà không làm tăng quá nhiều tham số.

5. Góc Ôn Luyện Phỏng Vấn (Interview Corner)

❓ Câu hỏi 1: Tại sao trong công thức Scaled Dot-Product Attention, chúng ta bắt buộc phải chia cho $\sqrt{d_k}$ ? Nếu bỏ qua hệ số này, điều gì sẽ xảy ra với quá trình huấn luyện mô hình khi kích thước hidden dimension ( $d_k$ ) tăng lớn?

Gợi ý trả lời:
- Bản chất xác suất:
  - Giả sử các thành phần của vector Query $Q$ và Key $K$ là các biến ngẫu nhiên độc lập có trung bình bằng $0$ và phương sai bằng $1$ .
  - Tích vô hướng $Q \cdot K = \sum_{i=1}^{d_k} Q_i K_i$ là tổng của $d_k$ biến ngẫu nhiên độc lập. Theo tính chất phương sai: $\text{Var}(Q \cdot K) = \sum_{i=1}^{d_k} \text{Var}(Q_i K_i) = d_k$ .
  - Khi số chiều $d_k$ lớn (ví dụ $d_k = 64$ hoặc $128$ ), độ lớn của tích vô hướng sẽ tăng rất mạnh, khiến độ lệch chuẩn của nó là $\sqrt{d_k}$ .
- Hậu quả nếu không chia cho $\sqrt{d_k}$ :
  - Các giá trị đầu vào của hàm Softmax sẽ có sự chênh lệch cực lớn (ví dụ một giá trị đạt $+50$ trong khi các giá trị khác chỉ quanh $0$ ).
  - Hàm Softmax khi nhận các giá trị quá lớn sẽ rơi vào vùng bão hòa cực đoan (Saturation): Giá trị xác suất lớn nhất tiến sát về $1.0$ , trong khi tất cả các giá trị khác tiến về $0$ .
  - Tại vùng bão hòa, đạo hàm của Softmax xấp xỉ bằng $0$ ( $\text{Softmax}'(z) \approx 0$ ) $\rightarrow$ Gây ra hiện tượng Biến mất đạo hàm (Vanishing Gradient), khiến các tầng nơ-ron phía trước ngừng học hoàn toàn.
- Kết luận: Chia cho $\sqrt{d_k}$ giúp kéo phương sai của tích vô hướng trở về $1.0$ , giữ cho hàm Softmax luôn nằm ở vùng có độ dốc gradient tốt nhất.

❓ Câu hỏi 2: So sánh Multi-Head Attention (MHA) và Grouped-Query Attention (GQA). Tại sao các kiến trúc LLM mã nguồn mở hàng đầu hiện nay (như LLaMA 3, Mistral) đều chuyển sang sử dụng GQA thay vì MHA truyền thống?

Gợi ý trả lời:
- Sự khác biệt cấu trúc:
  - Trong MHA, nếu mô hình có $32$ Query Heads thì cũng bắt buộc phải có $32$ Key Heads và $32$ Value Heads độc lập.
  - Trong GQA, $32$ Query Heads được chia thành $8$ nhóm, mỗi nhóm gồm $4$ Query Heads chia sẻ dùng chung duy nhất $1$ Key Head và $1$ Value Head $\rightarrow$ Tổng cộng chỉ cần $8$ Key Heads và $8$ Value Heads.
- Lý do chuyển đổi sang GQA:
  1. Tiết kiệm $75\%$ dung lượng KV Cache: Trong giai đoạn tạo văn bản (Inference), toàn bộ các vector $K, V$ của các từ trước đó phải được giữ lại trong VRAM GPU. Giảm từ 32 xuống 8 Heads giúp giải phóng phần lớn VRAM.
  2. Mở rộng Context Window khổng lồ: Nhờ dung lượng KV Cache nhỏ hơn nhiều, mô hình có thể mở rộng độ dài ngữ cảnh từ $4\text{k}$ lên $32\text{k} - 128\text{k}$ tokens mà không bị tràn bộ nhớ GPU.
  3. Tăng thông lượng (Throughput) suy luận: Giảm lượng dữ liệu cần nạp từ HBM vào chip tính toán, giúp GPU xử lý được nhiều luồng người dùng đồng thời hơn gấp $2\times - 3\times$ trong khi chất lượng mô hình gần như không suy giảm so với MHA.

❓ Câu hỏi 3: Rotary Position Embedding (RoPE) hoạt động như thế nào? Tại sao RoPE lại vượt trội hơn hẳn Absolute Positional Embedding truyền thống trong việc mở rộng độ dài ngữ cảnh (Context Window Extension)?

Gợi ý trả lời:
- Cơ chế hoạt động:
  - RoPE mã hóa vị trí bằng cách xoay từng cặp phần tử $(x_{2i}, x_{2i+1})$ của vector Query và Key trong mặt phẳng tọa độ 2D một góc $\theta_i = m \cdot \omega_i$ (trong đó $m$ là chỉ số vị trí của từ, $\omega_i = 10000^{-2i/d}$ ).
  - Khi tính tích vô hướng $\langle \text{RoPE}(Q, m), \text{RoPE}(K, n) \rangle$ , nhờ tính chất lượng giác, kết quả trả về chỉ phụ thuộc vào khoảng cách tương đối giữa hai vị trí: $(m - n)$ .
- Lý do vượt trội khi mở rộng Context:
  - Absolute Embedding: Chỉ học các vector vị trí tĩnh từ $0$ đến $L_{\text{max}}$ . Khi gặp văn bản dài hơn $L_{\text{max}}$ , mô hình hoàn toàn không có vector vị trí cho các token mới $\rightarrow$ Thất bại hoàn toàn.
  - RoPE: Biểu diễn vị trí dưới dạng góc quay liên tục. Để mở rộng cửa sổ ngữ cảnh (ví dụ từ $4\text{k}$ lên $32\text{k}$ tokens), ta chỉ cần áp dụng kỹ thuật RoPE Interpolation / YaRN (chia nhỏ góc quay theo hệ số $s = 8$ để nén dải vị trí $32\text{k}$ vừa vặn vào dải góc quay $4\text{k}$ mà mô hình đã học). Mô hình chỉ cần Fine-tune nhẹ vài trăm bước là có thể hiểu được ngữ cảnh dài mượt mà.

6. Tóm Tắt Ghi Nhớ Nhanh (Key Takeaways)

Self-Attention cho phép xử lý song song toàn bộ chuỗi trên GPU với công thức cốt lõi $\text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$ .
Phép chia $\sqrt{d_k}$ là bắt buộc để ngăn hàm Softmax rơi vào vùng bão hòa gradient.
Grouped-Query Attention (GQA) là chuẩn mực của LLMs hiện đại (LLaMA 3, Mistral) giúp tiết kiệm $75\%+$ bộ nhớ KV Cache.
Rotary Position Embedding (RoPE) mã hóa vị trí tương đối thông qua phép xoay vector 2D, cho phép mở rộng Context Window lên đến hàng trăm nghìn tokens.
FlashAttention là bước đột phá phần cứng GPU, tối ưu hóa bộ nhớ SRAM để tăng tốc độ huấn luyện và suy luận Attention lên nhiều lần.