Giới thiệu
Sự gia tăng sử dụng các mô hình ngôn ngữ lớn (LLM) và sự xuất hiện của các mô hình suy luận mới đã đặt ra thách thức về hiệu quả suy luận quy mô lớn. Một nút thắt cổ chai đáng kể là bộ nhớ đệm Key-Value (KV) trong các cơ chế attention truyền thống, mở rộng tuyến tính theo kích thước batch và độ dài chuỗi, trở thành một ‘kẻ ngốn bộ nhớ’, cản trở việc mở rộng quy mô của LLM.
Mặc dù các biến thể như MQA, GQA và MLA đã xuất hiện để giải quyết vấn đề này, chúng thường gặp khó khăn trong việc duy trì hiệu suất dưới các ràng buộc bộ nhớ nghiêm ngặt hoặc đưa ra các phức tạp gây ra các thách thức kỹ thuật và vấn đề tương thích.
Multi-matrix Factorization Attention (MFA)
Một bài báo gần đây của các nhà nghiên cứu từ Stepes, Đại học Thanh Hoa và các tổ chức khác đã giới thiệu một kiến trúc cơ chế attention mới: Multi-matrix Factorization Attention (MFA) và biến thể MFA-Key-Reuse (MFA-KR). Cơ chế này giảm đáng kể chi phí suy luận của mô hình ngôn ngữ đồng thời cải thiện hiệu suất.
MFA và MFA-KR không chỉ vượt trội MLA về hiệu suất mà còn phù hợp với hiệu suất MHA truyền thống trong khi giảm mức sử dụng KV Cache lên đến 93.7%.
MFA được thiết kế để đơn giản, dễ tái tạo, ít nhạy cảm với siêu tham số và tương thích với nhiều phương pháp Pos-embedding.
Phương pháp và Phân tích MFA
Nhóm nghiên cứu đã phân tích thiết kế chung và khả năng của các cơ chế attention, xác định hai chiều quan trọng liên quan đến khả năng. Phân tích này dẫn đến sự phát triển của các phương pháp phân tích và nguyên tắc thiết kế mới.
Họ đã giới thiệu khái niệm Generalized Multi-Head Attention (GMHA) như một khuôn khổ thống nhất để hiểu các biến thể MHA khác nhau.
Nhóm cũng khám phá việc tính toán và lưu trữ các giá trị key-value từ góc độ suy luận và kiểm tra khả năng của mô hình từ góc độ phân tách.
Fully Parameterized Bilinear Attention (FPBA) được thiết lập là giới hạn trên lý thuyết về hiệu suất. Họ phát hiện ra rằng MHA và các biến thể của nó là các phân tách hạng thấp của FPBA.
So sánh với MQA và MLA
Phân tích tập trung vào hai lược đồ cải tiến đại diện: Multi-Query Attention (MQA) và Multi-Head Latent Attention (MLA).
- MQA sử dụng chiến lược chia sẻ tham số mạnh mẽ hơn, trong đó tất cả các đầu attention chia sẻ cùng một bộ tham số key-value. Điều này làm giảm việc sử dụng bộ nhớ nhưng có thể ảnh hưởng đến khả năng biểu đạt của mô hình.
- MLA giới thiệu một không gian tiềm ẩn được chia sẻ để nén tham số, nhưng sức mạnh biểu đạt thực tế bị giới hạn bởi chiều nhỏ nhất, có nghĩa là việc tăng kích thước trung gian không cải thiện đáng kể hiệu suất.
Các Đổi Mới Chính của MFA
Sự phát triển của MFA được thúc đẩy bởi mục tiêu tạo ra một cơ chế attention giảm thiểu tiêu thụ tài nguyên đồng thời đạt đến giới hạn hiệu suất lý thuyết. Thiết kế của MFA kết hợp ba cải tiến chính:
- Tăng đáng kể số lượng và kích thước của các đầu attention để tối đa hóa khả năng của mô hình.
- Sử dụng chiến lược phân tách hạng thấp mạnh mẽ để duy trì hiệu quả tham số trong khi mở rộng số lượng và kích thước đầu attention.
- Sử dụng thiết kế một đầu key-value để giữ mức tiêu thụ bộ nhớ ở mức tối thiểu, ngay cả khi độ phức tạp của mô hình tăng lên.
Đo lường và So sánh Khả năng
Để phân tích thêm MFA và các cơ chế attention khác, nhóm nghiên cứu đã giới thiệu hai số liệu chính:
- Total Effective Rank (TER): Tích của số lượng đầu attention và hạng phân tách trên mỗi đầu (FRH).
- Shared Latent Subspace Dimension (SLSD): Kích thước của không gian ẩn được chia sẻ bởi tất cả các đầu attention.
MFA đạt được SLSD và TER cao hơn so với MQA. So với MLA, MFA đạt được kích thước bộ nhớ đệm KV nhỏ hơn và TER cao hơn với ngân sách tham số tương tự, đồng thời duy trì SLSD tương đương. So với MHA truyền thống, MFA có TER cao hơn, mặc dù SLSD của nó nhỏ hơn.
Kết quả Thử nghiệm
Các thử nghiệm mở rộng đã được tiến hành để đánh giá hiệu suất của kiến trúc mới ở quy mô lớn hơn, thử nghiệm các mô hình từ 1 tỷ đến 7 tỷ tham số và dữ liệu huấn luyện từ 10 tỷ đến 1 nghìn tỷ.
- MFA thể hiện khả năng mở rộng tương đương với MHA truyền thống, duy trì hiệu suất tuyệt vời ngay cả ở quy mô lớn hơn.
- Mặc dù MFA-KR thể hiện hiệu suất thấp hơn một chút, xu hướng mở rộng của nó phù hợp với MHA.
- Ưu điểm tiết kiệm bộ nhớ của MFA và MFA-KR tiếp tục mở rộng theo kích thước mô hình, với MFA đạt mức tiết kiệm bộ nhớ 87.5% và MFA-KR giảm mức sử dụng bộ nhớ xuống 6.25% ở quy mô lớn nhất.
Nghiên cứu Ablation
Các nghiên cứu ablation đã xác nhận hiệu quả của MFA và MFA-KR. Ưu điểm về hiệu suất của chúng cũng được xác nhận trên nhiều phương pháp mã hóa vị trí chính thống khác nhau.
Triển vọng
MFA cung cấp những cải tiến đáng kể với thiết kế đơn giản, giải quyết hiệu quả nút thắt bộ nhớ trong suy luận LLM mà không cần thêm sự phức tạp về kỹ thuật. Nó tích hợp liền mạch vào hệ sinh thái Transformer hiện có, đẩy nhanh ứng dụng của LLM trong nhiều tình huống khác nhau.