Trong kiến trúc Transformer, cơ chế self-attention cho phép mô hình "nhìn" toàn bộ chuỗi đầu vào để xác định từ nào quan trọng nhất với từ đang xét. Nhưng một head attention đơn lẻ có hạn chế lớn: nó chỉ có thể tập trung vào một kiểu quan hệ tại một thời điểm. Hãy hình dung bạn đọc câu "Con mèo ngồi trên thảm vì nó mệt". Để hiểu từ "nó", bạn cần đồng thời xử lý quan hệ cú pháp ("nó" là đại từ thay thế cho "con mèo"), quan hệ ngữ nghĩa ("mệt" giải thích lý do), và quan hệ vị trí (khoảng cách giữa các từ). Một head đơn lẻ buộc phải dồn mọi tín hiệu vào một bộ trọng số, khiến thông tin dễ bị thất thoát.

Chia nhỏ để xử lý: tách head và chiếu vào không gian con

Multi-head attention giải quyết vấn đề này bằng cách tách không gian vector thành nhiều không gian con (subspace) nhỏ hơn để xử lý song song. Cụ thể, với vector có chiều d = 512 và 8 head, mỗi head sẽ làm việc trong một không gian 64 chiều riêng biệt. Mỗi head có bộ ma trận chiếu (projection matrices) W_Q, W_K, W_V riêng, biến đổi Query, Key, Value gốc thành các phiên bản thu nhỏ trong không gian con của mình. Quá trình này hoàn toàn song song, nghĩa là tất cả 8 head tính toán cùng lúc, không chờ đợi nhau. Điểm mấu chốt là mô hình học các ma trận chiếu này trong quá trình huấn luyện, nên mỗi head tự tìm ra cách biểu diễn hữu ích nhất.

Chuyên môn hóa tự nhiên: Mỗi head một vai trò

Nghiên cứu thực nghiệm trên các mô hình AI ngôn ngữ lớn cho thấy các head tự chuyên môn hóa trong quá trình huấn luyện. Head thứ nhất có thể học cách theo dõi quan hệ chủ-vị (ai làm gì), head thứ hai chuyên xử lý đồng tham chiếu ("nó" trỏ về đâu), head thứ ba nhạy cảm với khoảng cách vị trí giữa các token. Sự đa dạng trong cách biểu diễn này giống như một nhóm chuyên gia: mỗi người phân tích cùng một dữ liệu từ một góc khác nhau, rồi ghép lại thành bức tranh toàn diện hơn bất kỳ cá nhân nào.

Ghép nối và chiếu tuyến tính: Bước cuối hoàn thiện

Sau khi tất cả head hoàn tất tính toán, kết quả của chúng được ghép lại thành một vector duy nhất có số chiều bằng ban đầu (8 head × 64 chiều = 512 chiều). Vector này đi qua một lớp chiếu tuyến tính cuối cùng (output linear projection) W_O, cho phép mô hình học cách kết hợp thông tin từ các head một cách tối ưu. Bước chiếu này không chỉ đơn thuần ghép nối, nó tạo điều kiện cho mô hình trộn lẫn thông tin giữa các head, phát hiện những mối quan hệ mà không head đơn lẻ nào nắm bắt được.

Multi-head attention là một trong những đột phá thiết kế quan trọng nhất của bài báo "Attention Is All You Need" (2017), và nguyên lý chia nhỏ để xử lý song song này vẫn là nền tảng cho mọi mô hình ngôn ngữ lớn hiện đại, từ GPT đến Gemini.