Query, key và value kết hợp context như thế nào

1Định nghĩa

Self-Attention là khối bên trong mỗi lớp Transformer cho phép mỗi vị trí token lấy thông tin có chọn lọc từ các vị trí đứng trước nó. Với vector ẩn đã được chuẩn hóa, khối này tạo ba vector từ mỗi token: Query (câu hỏi của vị trí hiện tại), Key (nhãn tra cứu của từng vị trí) và Value (nội dung sẽ lấy ra). Key và Value thường được chia sẻ theo nhóm trong GQA.

Tiếp theo, mô hình xoay Query và Key theo mã vị trí (RoPE), nhân vô hướng Query với Key để ra điểm tương quan, dùng mặt nạ nhân quả gán −∞ cho các vị trí tương lai rồi lấy softmax để có trọng số chú ý. Cuối cùng trọng số nhân với Value và chiếu qua ma trận đầu ra thành vector ngữ cảnh; trong suy luận, các cặp Key/Value này là dữ liệu được lưu lại trong KV cache.

Ví dụ đời thường: Một bước attention giống như lúc bạn tìm sách trong thư viện: cầm tờ phiếu ghi chủ đề cần tìm (Query), đi dọc kệ so với nhãn dán trên gáy từng cuốn (Key) để chấm điểm cuốn nào hợp nhất, rồi mở các cuốn điểm cao nhất ra chép nội dung bên trong (Value) về tổng hợp. Khác thư viện, ở đây thủ thư không cho bạn nhìn các kệ đặt phía sau vị trí của bạn, đó là mặt nạ nhân quả.

2Luồng

Đây là ô 2A trên sơ đồ. Mũi tên vào mang bảng vector (T × d) đã chuẩn hóa; mũi tên ra mang bảng cùng shape (T × d) đã được trộn thông tin giữa các vị trí, rồi sang MLP ở ô 2B. Trong suy luận, cặp K, V của các token cũ được lưu lại làm KV cache.

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Tất cả (22)1. Dữ liệu & nhãn2. Forward3. Học từ sai số4. Quy mô & phương pháp5. Suy luận

Ẩn khung overlay

1A1B1C2A2B2C2D3A3B3C3D3E4A4B4C4D4E4F4G5A5B5C

2A2 · Forward through the Transformer

embedding [T, d] → Q K V → softmax(QKᵀ/√d)·V · causal mask

2A · Embedding + Attention

Tra ID ra vector rồi cho mỗi vị trí lấy thông tin từ các vị trí trước bằng attention.

Nhận gì

Token ID (input) từ ô 1B, qua mũi tên sang phải. Đây là ô đầu của cột 2.

Làm gì, vì sao

Embedding tra mỗi ID ra một vector d số, tạo ma trận [T, d] (T vị trí, d chiều). Rồi attention cho mỗi vị trí tạo Query (đang tìm gì), Key (có gì) và Value (nội dung), so Query với các Key bằng softmax(QKᵀ/√d) và trộn Value theo trọng số đó. Mặt nạ nhân quả (hình tam giác) chặn việc nhìn về sau.

Ra gì

Ma trận [T, d] đã trộn thông tin từ các vị trí trước, đi xuống ô 2B.

Ví dụ, cách đọc ô

Đọc nhãn: embedding [T, d] → Q K V → softmax(QKᵀ/√d)·V. Mỗi hàng trọng số tổng bằng 1 và chỉ có giá trị ở các vị trí được phép nhìn (phần tam giác).

Lưu ý

Trọng số attention là phân phối trên vị trí, không phải xác suất của token. Thứ tự token được mã hóa riêng và attention thật thường chia nhiều head; ô này vẽ gọn một đường.

Đọc bài 2A · Embedding + Attention↗Giai đoạn 3 trong Bài Lab: Forward: embedding, attention, MLP, norm → logits↗

flowchart TD
  step1["Vector ẩn đã chuẩn hóa"]
  step2["Chiếu thành Q, K, V và gắn vị trí"]
  step3["Điểm số, mask và softmax"]
  step4["Trộn Value và chiếu ra"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: ô 2A "Embedding + Attention" trong cột 2 · Forward qua Transformer. Nó nhận vector của các vị trí từ bảng embedding (hoặc từ lớp trước) sau bước chuẩn hóa ở ô 2C, và đưa kết quả sang MLP ở ô 2B qua đường cộng residual. Tổng quan cột nằm ở bài forward.

Vấn đề cần giải. Vector của một token ngay sau embedding chỉ biết mỗi mình nó: chữ “cá” trong “mèo ăn cá” và trong “mua cá” có vector giống hệt nhau. Để dự đoán chữ kế tiếp, mỗi vị trí phải nhìn được những vị trí đứng trước và lấy đúng những gì liên quan. Attention là cơ chế làm việc đó.

Khái niệm tối thiểu và cách đọc trục. Bảng X có T hàng (mỗi hàng là một vị trí) và d cột (d số mô tả vị trí đó). Ba ma trận trọng số W_Q, W_K, W_V biến X thành Q, K, V; với một head, mỗi bảng có T hàng và d_head cột. Bảng điểm S = Q nhân K chuyển vị có T hàng, T cột: hàng i là câu hỏi của vị trí i, cột j là vị trí j bị hỏi, ô (i, j) là mức hợp nhau. Softmax chạy trên từng hàng của S.

Từng thao tác và lý do. (1) Chiếu Q, K, V: dùng ba ma trận khác nhau để 'câu hỏi', 'nhãn' và 'nội dung' có thể khác nhau. (2) Gắn vị trí bằng RoPE: nếu không có, attention coi “mèo ăn” và “ăn mèo” như nhau vì chỉ so khớp nội dung. (3) Chia cho căn bậc hai của d_head: tích vô hướng càng nhiều chiều càng lớn, chia lại giữ cho softmax khỏi bão hòa (một ô chiếm gần hết trọng số). (4) Che tương lai: ô (i, j) với j > i đặt −∞ để sau softmax bằng 0. (5) Softmax từng hàng thành trọng số cộng lại bằng 1. (6) Nhân trọng số với V để ra vector ngữ cảnh, ghép các head rồi nhân với ma trận đầu ra W_O.

Ví dụ có số (minh họa, bỏ RoPE cho gọn). Ba vị trí “mèo, ăn, cá”, d_head = 2. Q = K = [[1, 0], [0, 1], [1, 1]] và V = [[1, 0], [0, 2], [3, 1]], mỗi hàng ứng với một vị trí. Hàng 0 (mèo) chỉ nhìn được chính nó: điểm [0,707], trọng số [1], đầu ra bằng V0 = [1, 0]. Hàng 1 (ăn): điểm với K0, K1 là [0, 0,707], trọng số [0,330; 0,670], đầu ra = 0,330 × [1, 0] + 0,670 × [0, 2] = [0,330; 1,340]. Hàng 2 (cá): q = [1, 1], điểm với K0, K1, K2 là [1, 1, 2] chia 1,414 = [0,707; 0,707; 1,414], trọng số [0,248; 0,248; 0,504], đầu ra = [1,759; 1,000].

Cách đọc. Hàng 2 dành 50,4% chú ý cho chính vị trí 2 (vì điểm cao nhất) và mỗi vị trí còn lại 24,8%; đầu ra là trung bình có trọng số của các Value nên nằm 'giữa' V0, V1, V2 và nghiêng về V2. Nếu bỏ mặt nạ, hàng 0 sẽ nhìn cả ba vị trí với trọng số xấp xỉ [0,401; 0,198; 0,401] và 'thấy' tương lai, đó chính là lỗi rò rỉ mà ô 1B cùng mặt nạ nhân quả ngăn chặn. Trong ma trận trọng số 3 × 3, nửa trên bên phải đường chéo phải đúng bằng 0.

Điều có thể và không thể kết luận. Bạn có thể khẳng định: trọng số mỗi hàng không âm, cộng lại bằng 1, và các vị trí bị che có trọng số 0. Bạn không thể khẳng định: bản đồ nhiệt (heatmap) của trọng số giải thích vì sao mô hình trả lời như vậy, vì còn nhiều lớp, nhiều head và các khối khác cùng ảnh hưởng đầu ra. Cũng đừng nhầm trọng số attention (chia sẻ chú ý giữa các vị trí) với xác suất của token trong bảng từ vựng (chỉ có sau LM head và softmax).

Thực thi. Attention đầy đủ tạo ma trận điểm T × T cho mỗi head; với T = 4.096 là hơn 16,7 triệu phần tử cho mỗi head ở mỗi lớp. FlashAttention chia ma trận thành các ô và tính softmax tăng dần để không phải ghi cả ma trận ra bộ nhớ lớn, vẫn cho kết quả attention tương đương (xem kernel). Trong suy luận, Key và Value của token cũ được giữ trong KV cache nên mỗi token mới chỉ tính Query, Key, Value của riêng nó. Trong huấn luyện, backward cần các giá trị này (hoặc phải tính lại), xem gradient qua một lớp.

Để tự quan sát attention weights thật của một forward và kiểm tra các vị trí bị che bằng 0, mở giai đoạn 3 · Forward (khối attention) của Bài Lab tổng hợp. Hãy tự đoán phần nào của ma trận trọng số phải bằng 0 trước khi nhìn kết quả.

4Profiling

• Tính đúng và rò rỉ nhân quả

  So đầu ra với một cài đặt tham chiếu đơn giản trên cùng Q, K, V. Sau đó thử đổi một token ở vị trí tương lai và kiểm tra logits của các vị trí đứng trước không đổi trong sai số cho phép. Cả hai phép kiểm tra chỉ ra phép tính đúng, không đảm bảo mô hình học tốt.

• Bộ nhớ và thời gian theo độ dài

  Đo thời gian và bộ nhớ khi tăng T với số head và d_head cố định; ma trận điểm tăng theo T bình phương nếu không dùng kernel tiết kiệm đọc ghi. Phân biệt attention trên cả prompt (T hàng truy vấn) với attention khi decode có cache (một hàng truy vấn so với T khóa), vì hai trường hợp bị giới hạn bởi những thứ khác nhau.

• Đọc trọng số một cách cẩn trọng

  Nếu in trọng số, kiểm tra tổng mỗi hàng bằng 1 và phần bị che bằng 0 trước khi diễn giải. Ghi rõ lớp và head đang xem; một head của một lớp không nói lên toàn bộ mô hình. Không dùng một heatmap để khẳng định lý do của câu trả lời.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Query, Key, Value (Q, K, V)

Ba vector tạo ra từ cùng một token bằng ba ma trận khác nhau: Query là 'câu hỏi' của vị trí hiện tại, Key là 'nhãn tra cứu' của mỗi vị trí, Value là nội dung được lấy ra.

Attention scores và weights

Scores là tích vô hướng Q·K (đã chia căn bậc hai của d_head); weights là kết quả softmax của scores trên mỗi hàng, không âm và cộng lại bằng 1.

Rotary Position Embedding (RoPE)

Cách đưa thông tin vị trí vào bằng phép quay áp lên Q và K, để điểm số phụ thuộc vào khoảng cách tương đối giữa hai vị trí.

Causal mask

Gán −∞ cho điểm của các vị trí tương lai trước softmax để trọng số của chúng đúng bằng 0.

Grouped-Query Attention (GQA)

Nhiều Query head dùng chung một nhóm Key/Value head để giảm bộ nhớ; mô hình phải được thiết kế và huấn luyện đúng như vậy.

Head

Một bộ Q, K, V độc lập. Nhiều head chạy song song rồi ghép kết quả, mỗi head có thể chú ý theo kiểu khác nhau.

7Bài liên quan

8Nguồn