Từ Transformer nguyên bản đến các LLM hiện đại: Quá trình tiến hóa của kiến trúc Transformer — Phần 3…

Từ Transformer nguyên bản đến các LLM hiện đại: Sự tiến hóa của kiến trúc Transformer — Phần 3: Tối ưu hóa cơ chế Attention (MHA, MQA, GQA, MLA) Tóm tắt nội dung Vấn đề. Khi một LLM tạo văn bản, mỗi token mới đều phải đối chiếu lại toàn bộ các token trước đó. Nhằm tránh việc tính toán lại các giá trị Key (K) và Value (V) của các token trước, mô hình sẽ lưu trữ chúng trong bộ nhớ. Đây là bộ nhớ đệm KV (KV cache). Bộ nhớ đệm này gia tăng theo từng token, từng lớp và từng đầu chú ý (attention head), do đó đối với các chuỗi hội thoại dài, đây trở thành một trong những gánh nặng chi phí bộ nhớ lớn nhất trong quá trình suy luận. Câu hỏi đặt ra. Liệu có thể thu nhỏ KV cache mà không làm suy giảm chất lượng? Bốn giải pháp, theo trình tự thời gian xuất hiện: MHA (Multi-Head Attention) là thiết kế nguyên bản. Mỗi đầu Query sở hữu một cặp Key và Value riêng biệt. Cơ chế này cung cấp cho mô hình các "góc nhìn" độc lập nhất đối với văn bản, song lại tiêu tốn dung lượng lưu trữ lớn nhất. MQA (Multi-Query Attention) cho phép tất cả các đầu Query dùng chung một cặp Key và Value duy nhất. Bộ nhớ đệm giảm xuống mức rất nhỏ, song do mọi đầu đều phải hoạt động dựa trên cùng một dữ liệu K/V, chất lượng đầu ra có thể suy giảm. GQA (Grouped-Query Attention) là giải pháp dung hòa. Các đầu Query được chia thành từng nhóm, và mỗi nhóm dùng chung một cặp Key và Value. Số lượng nhóm có thể tùy chỉnh linh hoạt, giúp cân bằng giữa chất lượng của MHA và khả năng tiết kiệm bộ nhớ của MQA. MLA (Multi-Head Latent Attention) tiếp cận theo hướng khác biệt. Thay vì chia sẻ K/V giữa các đầu, cơ chế này nén chúng thành một vector nhỏ cho mỗi token (trạng thái ẩn - latent) và tái tạo Key cùng Value riêng cho từng đầu từ vector đó khi cần thiết. Mỗi đầu vẫn duy trì góc nhìn độc lập, trong khi dung lượng bộ nhớ đệm giảm đáng kể. Dung lượng bộ nhớ đệm (càng nhỏ càng tốt): MHA >> GQA > MQA; MHA >> MLA (tương đương GQA với 2,25 nhóm trong DeepSeek-V2). Tính đa dạng của các đầu (càng cao càng tốt): MHA = MLA > GQA > MQA. MHA: Cơ sở lý thuyết Trước khi tìm hiểu GQA và MQA, cần nắm rõ cơ chế vận hành chính xác của Multi-Head Attention (MHA). Trong phần trước, có thể nhận thấy KV cache có thể trở thành điểm nghẽn bộ nhớ nghiêm trọng trong quá trình suy luận, đặc biệt khi mỗi đầu chú ý đều duy trì một bộ nhớ đệm Key và Value riêng biệt. Do đó, vấn đề cốt lõi được đặt ra là: Vì sao ngay từ đầu mô hình lại cần nhiều đầu chú ý? Ý tưởng nền tảng của MHA tương đối rõ ràng: Thay vì sử dụng một biểu diễn Query, Key và Value duy nhất cho toàn bộ mô hình, biểu diễn chú ý được phân chia thành nhiều đầu. Mỗi đầu xử lý một phần riêng biệt của biểu diễn và có thể học một dạng quan hệ khác nhau giữa các token. Cơ chế này có thể hiểu tương tự như việc nhiều người cùng đọc một câu văn, nhưng mỗi người lại tập trung vào một khía cạnh riêng. Ví dụ: Đầu 1 → tập trung vào quan hệ cú pháp Đầu 2 → tập trung vào quan hệ ngữ nghĩa Đầu 3 → tập trung vào quan hệ giữa các token lân cận Đầu 4 → tập trung vào quan hệ khoảng cách xa ...Hành vi cụ thể do mô hình tự học; người vận hành không gán mục đích cụ thể cho từng đầu theo phương thức thủ công. Nguồn: Đồ họa của DataMListic. 1. MHA tạo ra các đầu riêng biệt này như thế nào? Giả sử mô hình có các thông số: d_model = 4096 #chiều biểu diễn nhúng (embedding dimension) n_head = 32 #số lượng đầu chú ý Khi đó, mỗi đầu chú ý hoạt động trên một chiều không gian nhỏ hơn: d_head = d_model / n_head #chiều biểu diễn nhúng trên mỗi đầu d_head = 4096 / 32 = 128 Về mặt khái niệm, mô hình trước tiên sẽ chiếu dữ liệu đầu vào thành các biểu diễn Query, Key và Value: Q = XWq K = XWk V = XWv Các biểu diễn thu được sau đó được phân chia thành n_head đầu riêng biệt: Q → Q₁, Q₂, Q₃, ..., Q₃₂ K → K₁, K₂, K₃, ..., K₃₂ V → V₁, V₂, V₃, ..., V₃₂ Mỗi head gồm có: Qᵢ: 128 chiều Kᵢ: 128 chiều Vᵢ: 128 chiều và thực hiện cơ chế attention riêng biệt: Head i: Attentionᵢ = softmax(QᵢKᵢᵀ / √d_head)VᵢCuối cùng, đầu ra từ tất cả các head được kết hợp lại: Head 1 ──┐ Head 2 ──┤ Head 3 ──┤ ... ├──→ Ghép nối (Concatenate) → Phép chiếu đầu ra (Output Projection) Head 32 ─┘Đây là nguyên lý cơ bản của cơ chế Đa đầu chú ý (Multi-Head Attention). 2. Vì sao cần nhiều Head? Ưu điểm nằm ở chỗ mô hình không bắt buộc phải biểu diễn mọi mối quan hệ trong một không gian attention duy nhất. Thay vào đó, mô hình có được nhiều không gian biểu diễn đã học: Đầu vào (Input) ↓ Phép chiếu Q, K, V ↓ ┌─────────┼─────────┐ ↓ ↓ ↓ Head 1 Head 2 Head 3 ... Head N ↓ ↓ ↓ Attention Attention Attention └─────────┼─────────┘ ↓ Ghép nối (Concatenate) ↓ Phép chiếu đầu ra (Output Projection)Điều này mang lại cho mô hình nhiều góc nhìn đã học được trên cùng một chuỗi dữ liệu. Tuy nhiên, điều này cũng đi kèm chi phí đánh đổi. Do mỗi head sở hữu các biểu diễn Key và Value riêng biệt, MHA cũng đòi hỏi một bộ nhớ đệm KV Cache riêng cho từng head trong quá trình suy luận tự hồi quy (autoregressive inference). Ví dụ, với 32 head: Head 1 → K₁, V₁ Head 2 → K₂, V₂ Head 3 → K₃, V₃ ... Head 32 → K₃₂, V₃₂Toàn bộ các biểu diễn K/V này đều phải được lưu trữ cho các token trước đó. Đây chính là vấn đề đã được thảo luận tại Phần 2. Việc có nhiều attention head hơn có thể mang lại tính linh hoạt cao hơn trong biểu diễn, song cũng đồng nghĩa với: Nhiều KV head hơn ↓ KV Cache lớn hơn ↓ Mức tiêu tốn bộ nhớ cao hơn trong quá trình suy luận3. Cần bao nhiêu Attention Head? Một câu hỏi tất yếu được đặt ra: Nếu việc sử dụng nhiều attention head mang lại hiệu quả, một mô hình Transformer cần bao nhiêu head?Số lượng attention head là một siêu tham số kiến trúc (architectural hyperparameter). Thông số này được ấn định khi thiết kế kiến trúc mô hình và duy trì cố định trong suốt quá trình huấn luyện cũng như suy luận. Ví dụ, một mô hình có thể được cấu hình như sau: d_model = 4.096 n_head = 32 Kết quả thu được: d_head = d_model / n_head = 4.096 / 32 = 128 Theo đó, mỗi attention head sẽ xử lý các biểu diễn Query, Key và Value có kích thước 128 chiều. Một mô hình khác có thể sử dụng: d_head = 4.096 / 64 = 64Do đó, việc tăng số lượng head nhìn chung đồng nghĩa với việc mỗi head riêng lẻ sẽ hoạt động trên một không gian con nhỏ hơn. Ngoài ra, không có bất kỳ một công thức chung duy nhất nào quy định: "Một mô hình 7B phải có chính xác 32 head."Thay vào đó, các nhà thiết kế mô hình lựa chọn n_head dựa trên một số yếu tố cân nhắc: 1. Chất lượng mô hình 2. Chi phí huấn luyện 3. Tốc độ suy luận 4. Dung lượng bộ nhớ sử dụng 5. Kích thước KV CacheĐối với nhiều kiến trúc Transformer, giới thiết kế ưu tiên kích thước head đủ lớn để biểu diễn các mối quan hệ hữu ích, đồng thời sử dụng lượng head vừa đủ để cung cấp nhiều mẫu hình chú ý đã học. Ví dụ: d_model = 4.096 Phương án 1: n_head = 32 d_head = 128 Phương án 2: n_head = 64 d_head = 64Cả hai phương án đều khả thi về mặt toán học, song chúng t