Grouped Query Attention (GQA)

Grouped-query attention (GQA) là một dạng attention đa đầu cho transformer, trong đó các nhóm query head dùng chung một key head và một value head. Ví dụ, Llama 3.1 8B ghép 32 query head với 8 key-value head ở mỗi layer. GQA nằm giữa multi-head attention (MHA), trong đó mỗi query head có key và value riêng, và multi-query attention (MQA), trong đó mọi query head dùng chung một key head và một value head. Joshua Ainslie cùng các đồng nghiệp tại Google Research giới thiệu phương pháp này trong bài viết GQA: Huấn luyện mô hình Transformer đa truy vấn tổng quát từ checkpoint đa đầu tại EMNLP 2023, cho thấy model GQA đạt chất lượng gần với MHA và tốc độ giải mã gần với MQA. Các mô hình ngôn ngữ lớn (LLM) mở như Llama 2 70B, dòng Llama 3, Mistral 7B và Gemma 2 sử dụng GQA; các mô hình thị giác-ngôn ngữ (VLM) được xây dựng dựa trên chúng, chẳng hạn PaliGemma 2 của Google sử dụng Gemma 2, cũng kế thừa GQA trong decoder ngôn ngữ.

Cách GQA hoạt động#

Trong self-attention, mỗi token được chiếu thành query, key và value cho mỗi head; mỗi head so sánh query với key của mình. GQA giữ nguyên toàn bộ query head nhưng chiếu ít key head và value head hơn; mỗi nhóm query head thực hiện attention bằng key và value dùng chung của nhóm đó. Với 8 query head và 2 key-value head, mỗi key-value head phục vụ 4 query head:

Phép tính attention bên trong mỗi head vẫn là attention tích vô hướng có scale như cũ. Với một nhóm, layer trở thành MQA; khi số nhóm bằng số query head, layer trở thành MHA. PyTorch hỗ trợ trực tiếp bố cục này thông qua đối số enable_gqa của scaled_dot_product_attention:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

Lệnh in ra Output (1, 8, 16, 64), same as copied heads: True và 256 with GQA, 1024 with MHA: cả 8 query head đều tạo đầu ra từ lượng key và value chỉ bằng một phần tư lượng mà MHA phải lưu trữ.

Vì sao GQA thu nhỏ KV cache#

Trong quá trình tạo văn bản, model giữ key và value của mọi token trước đó trong KV cache rồi đọc toàn bộ cache cho mỗi token mới; vì vậy, tốc độ giải mã bị giới hạn bởi băng thông bộ nhớ nhiều hơn là năng lực tính toán. Kích thước cache tăng theo số key-value head. Llama 3.1 8B có 32 layer, mỗi layer có 8 key-value head với 128 chiều, nên lưu 128 KiB cho mỗi token ở độ chính xác 16 bit, tương đương 16 GiB tại cửa sổ ngữ cảnh tối đa 131.072 token. Với MHA, cùng model đó cần 512 KiB mỗi token và 64 GiB khi đạt độ dài tối đa. Bộ nhớ được giải phóng có thể dùng cho ngữ cảnh dài hơn, batch lớn hơn khi batch liên tục và độ trễ suy luận mỗi token thấp hơn. Giao diện LLM của Ultralytics chấp nhận base_url cho mọi server tương thích OpenAI, chẳng hạn Llama 3.1 8B tự lưu trữ; với 8 key-value head, server có thể giữ số token hội thoại gấp bốn lần trong cùng dung lượng cache so với bố cục MHA 32 head.

Chuyển đổi model sang GQA#

GQA không đòi hỏi phải huấn luyện từ đầu. Các tác giả chuyển đổi checkpoint MHA hiện có bằng cách lấy trung bình các key head và value head trong từng nhóm, sau đó tiếp tục pretraining với khoảng 5% lượng compute ban đầu; họ gọi bước này là uptraining. MQA, biến thể cực hạn với một head dùng chung, được Noam Shazeer giới thiệu vào năm 2019 và có bố cục nhanh nhất, nhưng làm giảm chất lượng và có thể khiến quá trình huấn luyện mất ổn định; số nhóm trung gian giúp khôi phục phần lớn chất lượng đó.

GQA và các kỹ thuật liên quan#

• Attention tiềm ẩn đa đầu (MLA): Được giới thiệu cùng DeepSeek-V2, phương pháp này nén key và value của mỗi token thành một vector tiềm ẩn nhỏ dùng chung.
• PagedAttention: Cấp phát KV cache theo các khối kích thước cố định để tránh phân mảnh, không phụ thuộc vào số head sử dụng cache.
• Attention cửa sổ trượt: Giới hạn số token được lưu trong cache cho mỗi layer; Mistral 7B kết hợp phương pháp này với GQA.
• Flash attention: Kernel GPU tính attention chính xác nhanh hơn và hỗ trợ trực tiếp bố cục GQA.