Liệu Transformer có thực sự cần ba phép chiếu QKV? Nghiên cứu mới mở đường cho AI hiệu năng cao trên thiết bị biên

Một nghiên cứu mới từ arXiv đã kiểm chứng khả năng giảm số lượng phép chiếu trong cơ chế Attention của Transformer. Kết quả cho thấy việc chia sẻ phép chiếu giữa Key và Value có thể giảm tới 96,9% bộ nhớ đệm KV khi kết hợp với MQA, giúp các mô hình AI lớn chạy hiệu quả hơn trên thiết bị di động.

Kiến trúc Transformer đã trở thành xương sống của các mô hình AI hiện đại, từ xử lý ngôn ngữ tự nhiên đến thị giác máy tính. Tuy nhiên, một nghiên cứu mới đăng trên arXiv đã đặt ra câu hỏi đầy thú vị: Liệu cơ chế Attention phức tạp với ba phép chiếu riêng biệt (Query, Key, Value) có thực sự cần thiết, hay chúng ta có thể đơn giản hóa nó để tối ưu hóa hiệu năng mà không làm giảm chất lượng?

Thách thức cấu trúc QKV truyền thống

Trong kiến trúc Transformer chuẩn, cơ chế Attention sử dụng ba phép chiếu tuyến tính độc lập: Query (Q), Key (K) và Value (V). Mặc dù thiết kế này cực kỳ hiệu quả, nó tạo ra gánh nặng lớn về bộ nhớ và tính toán, đặc biệt là trong quá trình suy luận (inference) nơi bộ nhớ đệm KV (KV cache) chiếm dụng phần lớn tài nguyên.

Nghiên cứu thực hiện bởi Ali Kayyam và các cộng sự đã tiến hành đánh giá hệ thống các biến thể chia sẻ phép chiếu để xem liệu có thể cắt giảm chi phí này hay không. Họ tập trung vào ba ràng buộc chính:

Q-K=V: Chia sẻ phép chiếu cho Key và Value.
Q=K-V: Chia sẻ phép chiếu cho Query và Key.
Q=K=V: Chỉ sử dụng một phép chiếu duy nhất cho cả ba thành phần.

Phát hiện quan trọng: Chia sẻ Key-Value mang lại hiệu quả cao nhất

Qua các thí nghiệm trên nhiều nhiệm vụ khác nhau—từ dữ liệu tổng hợp, thị giác máy tính (MNIST, CIFAR, TinyImageNet) đến mô hình hóa ngôn ngữ với quy mô lên đến 1,2 tỷ tham số—nhóm nghiên cứu đã đưa ra kết luận đáng ngạc nhiên.

Biến thể Q-K=V (chia sẻ Key và Value) cho thấy hiệu suất hoạt động ngang bằng hoặc thậm chí vượt trội so với Transformer QKV chuẩn trong nhiều trường hợp. Đặc biệt, trong bài toán mô hình hóa ngôn ngữ, phương pháp này giúp giảm 50% dung lượng bộ nhớ đệm KV với chỉ mức suy giảm nhẹ 3,1% về chỉ số perplexity (độ phức tạp của ngôn ngữ).

Ngược lại, biến thể Q=K-V gặp vấn đề về tính đối xứng của bản đồ attention, làm mất đi tính định hướng cần thiết, mặc dù có thể khắc phục một phần bằng các mã hóa vị trí 2D.

Tác động lớn đến AI trên thiết bị biên (Edge AI)

Điểm sáng lớn nhất của nghiên cứu này nằm ở khả năng kết hợp phương pháp chia sẻ phép chiếu với các kỹ thuật tối ưu hóa hiện có như GQA (Grouped Query Attention) hoặc MQA (Multi-Query Attention).

Khi kết hợp Q-K=V với GQA-4, nghiên cứu ghi nhận mức giảm bộ nhớ đệm lên tới 87,5%. Con số ấn tượng hơn nữa là khi kết hợp với MQA, mức giảm đạt tới 96,9%.

"Chúng tôi chứng minh rằng chia sẻ phép chiếu là bổ sung cho việc chia sẻ đầu (GQA/MQA)... cho phép suy luận thực tế trên thiết bị."

Điều này có ý nghĩa cực kỳ quan trọng đối với xu hướng triển khai AI trên thiết bị biên (on-device AI). Việc giảm thiểu bộ nhớ KV cache giúp các mô hình ngôn ngữ lớn (LLM) có thể chạy mượt mà hơn trên điện thoại thông minh, laptop và các thiết bị IoT mà không cần phụ thuộc quá nhiều vào đám mây, đồng thời tiết kiệm năng lượng và chi phí phần cứng.

Kết luận

Nghiên cứu này đã hệ thống hóa việc chia sẻ phép chiếu như một trường hợp chưa được khám phá nhiều của việc ràng buộc trọng số (weight tying). Nó mở ra hướng đi mới để thiết kế các mô hình Transformer "nhẹ hơn" nhưng vẫn giữ được sức mạnh trí tuệ, đặc biệt hữu ích trong bối cảnh nhu cầu chạy AI cục bộ đang tăng cao.