Ba Ma trận Cho Mỗi Head, Hoặc Một Ma trận Lớn

Quan điểm Theo Từng Head

Mỗi head cần có phép chiếu query, phép chiếu key, & phép chiếu value riêng của nó. Đối với head h:

Q_h = X · W_Q^h    trong đó W_Q^h có hình dạng [d_model, head_dim]
K_h = X · W_K^h    trong đó W_K^h có kích thước [d_model, head_dim]
V_h = X · W_V^h    trong đó W_V^h có kích thước [d_model, head_dim]

X mang kích thước đầu vào [batch, seq_len, d_model]. Sau khi chiếu, Q_h, K_h, V_h mỗi cái mang kích thước [batch, seq_len, head_dim].

Fused View

Các ma trận per-head nằm cạnh nhau trong bộ nhớ. Một ma trận fused duy nhất W_Q có kích thước [d_model, d_model] tạo ra tất cả các head cùng lúc:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

Phép nhân ma trận hợp nhất chỉ gửi một lệnh gọi BLAS thay vì 12. Các lõi tensor CUDA đạt thông lượng đỉnh trên các phép nhân ma trận kích thước này; các phép nhân ma trận từng đầu sẽ không tận dụng hết phần cứng.

Số lượng Tham số

Ba ma trận hợp nhất W_Q, W_K, W_V, mỗi ma trận d_model × d_model. Cộng với ma trận chiếu đầu ra W_O, cũng d_model × d_model. Đối với ANDREA-120M:

tham số mỗi lớp attention = 4 × 768² = 2,359,296 ≈ 2.36M
tham số qua 12 lớp      = 12 × 2.36M ≈ 28.3M

Khoảng một phần tư tổng số tham số của ANDREA-120M nằm trong các phép chiếu attention. Ba phần tư còn lại nằm trong sublayers MLP & embeddings.

Đặt tên cho các phép chiếu

Mười hai Vector Trở thành Một

Sau Khi Mỗi Head Tính Toán

Mỗi head tạo ra một tensor đầu ra có hình dạng [batch, seq_len, head_dim]. Mười hai head tạo ra mười hai tensor như vậy. Nối dọc theo chiều đặc trưng xếp chúng lại với nhau:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # for ANDREA-120M

Concat đảo ngược quá trình split. Tổng kích thước đặc trưng trở lại d_model. Không mất thông tin ở các chiều; sự khác biệt nằm ở nội dung mà mỗi chunk chứa: chunk của head 1 phản ánh mẫu attention đã học của head 1.

Ma trận Chiếu Đầu Ra W_O

Chỉ ghép nối thôi thì các head vẫn bị cô lập: đầu ra của head 4 nằm cạnh đầu ra của head 7, không head nào biết về head kia. Ma trận chiếu đầu ra W_O có kích thước [d_model, d_model] trộn lẫn chúng:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

Sau W_O, mỗi chiều đầu ra mang một tổ hợp tuyến tính đã học của tất cả mười hai head. Thông tin chảy tự do giữa các head thông qua phép nhân ma trận đơn này.

Tại sao các Head Chuyên biệt

Không có gì trong kiến trúc buộc head 4 học thì động từ hoặc head 9 học dấu câu ghép đôi. Sự chuyên biệt hóa nổi lên từ áp lực gradient: trong quá trình huấn luyện, các head đóng góp dư thừa nhận gradient nhỏ hơn so với các head đóng góp độc đáo. Qua hàng nghìn bước, mỗi head ổn định vào một niche giảm tổng loss hiệu quả nhất.

Thực nghiệm cho thấy các transformer đã huấn luyện có các head xử lý: mẫu vị trí (head nhìn vào token trước đó), mẫu cú pháp (head nhìn vào dấu ngoặc đóng khớp), mẫu ngữ nghĩa (head nhìn vào thực thể được đặt tên gần nhất). Không có nhãn nào huấn luyện sự chuyên biệt hóa này. Tín hiệu huấn luyện đơn thuần, lan truyền qua W_O, sắp xếp các head.

Tại sao Mười hai Head, Không Phải Một Head Rộng Hơn

Cách CUDA Lưu Trữ Các Heads

Một Tensor Duy Nhất, Được Reshape

Engine huấn luyện của ANDREA microgpt_cuda.cu không cấp phát mười hai buffer riêng biệt cho mười hai heads. Nó cấp phát một tensor hợp nhất & coi chiều head như một mẫu stride:

// sau Q = X · W_Q (một matmul, hợp nhất qua các heads)
// Q có shape [batch, seq_len, d_model]
// reshape thành    [batch, seq_len, n_head, head_dim]
// transpose thành  [batch, n_head, seq_len, head_dim]
// mỗi head giờ liền kề trong hai chiều bên trong

The transpose moves n_head ahead of seq_len. Why? Because the next operation (Q_h · K_h^T) needs each head's seq_len × head_dim slice contiguous in memory. CUDA matmuls run faster on contiguous tensors.

Một Kernel, Nhiều Heads

Một CUDA kernel attention đơn lẻ chạy song song trên tất cả các head. Mỗi thread block xử lý một cặp (batch, head); các thread bên trong block hợp tác trên tile seq_len × head_dim. Kernel không bao giờ biết nó xử lý nhiều head; launch grid xử lý sự song song.

Cấu hình Phản ánh Phần cứng

Lựa chọn n_head=12, head_dim=64 của ANDREA-120M phù hợp với RTX 4090 tensor cores, vốn ưu tiên các matmul tile theo bội số của 16. head_dim=64 = 4 × 16 khớp chính xác với hình dạng tile. head_dim=32 (ANDREA-12M) cũng khớp nhưng sử dụng dưới công suất tile. head_dim=72 sẽ không khớp & buộc phải dùng fallback kernel.

Hình ảnh Tổng quan

Năm bước. Ba phép nhân ma trận chạm vào đầu vào (Q, K, V projections). Một phép nhân ma trận chạm vào các head đã nối (W_O). Một attention kernel xử lý tất cả các head song song. ANDREA-120M chạy tất cả năm bước một lần mỗi layer, mười hai layer sâu, mỗi forward pass.

Đọc Pipeline