Mixtral of Experts

Hướng dẫn đọc

Vấn đề và động lực

Các mô hình ngôn ngữ tự hồi quy dày đặc (dense) đòi hỏi chi phí tính toán tỷ lệ thuận trực tiếp với tổng số lượng tham số cho mỗi token được xử lý. Khi các mô hình mở rộng lên quy mô 70B tham số và lớn hơn, việc triển khai phục vụ (serving) với độ trễ thấp và thông lượng cao trở nên vô cùng tốn kém. Các tác giả đặt mục tiêu mở rộng dung lượng mô hình nhưng vẫn giữ chi phí suy luận trên mỗi token tương đương với các mô hình nhỏ hơn nhiều, tạo ra một mô hình trọng số mở có khả năng tương đương hoặc vượt qua Llama 2 70B cùng các mô hình thương mại đóng vai trò cơ sở so sánh như GPT-3.5.

Ý tưởng cốt lõi

• Sparse Mixture of Experts (SMoE): Thay thế các khối mạng truyền thẳng (FFN) dày đặc trên toàn bộ 32 tầng transformer bằng 8 mạng chuyên gia riêng biệt, chỉ định tuyến mỗi token đến top-2 chuyên gia hàng đầu.
• Tách rời số lượng tham số tổng và kích hoạt: Kiến trúc chứa tổng cộng 47B tham số trên các tầng, nhưng chỉ thực thi 13B tham số kích hoạt cho mỗi lượt truyền xuôi (forward pass) của token, tương đương với mức sử dụng tính toán của một mô hình dense kích thước trung bình.
• Cửa sổ ngữ cảnh đầy đủ 32k: Hỗ trợ tự nhiên độ dài chuỗi 32.768 token với độ chính xác truy xuất 100% trên các benchmark passkey.
• Định tuyến theo cú pháp thay vì chủ đề: Phân tích thực nghiệm về việc định tuyến token cho thấy việc lựa chọn chuyên gia chủ yếu dựa vào chức năng cú pháp và tính cục bộ theo thời gian hơn là các miền chủ đề cấp cao.

Phương pháp (kèm các phương trình chính)

Mixtral sửa đổi kiến trúc Mistral 7B ($32$ tầng, số chiều ẩn $4096$, $32$ attention head, $8$ key-value head) bằng cách thay thế mỗi khối con FFN bằng một mô-đun MoE bao gồm $n=8$ chuyên gia ($E_0,\dots ,E_{n-1}$), mỗi chuyên gia được triển khai dưới dạng một mạng SwiGLU với số chiều trung gian $14336$.

Đối với biểu diễn token đầu vào $x$, một mạng định tuyến sẽ tính toán phân phối cổng thưa (sparse gating) trên các chuyên gia bằng cách sử dụng các giá trị logit top-$K$ ($K=2$): $$G(x):=\text{Softmax}(\text{TopK}(x\cdot W_g))$$ trong đó toán tử Top-K gán giá trị $-\infty$ cho tất cả các tọa độ nằm ngoài $K$ giá trị lớn nhất của logits $\ell =x\cdot W_g\in {\mathbb{R}}^n$: $$(\text{TopK}(\ell ){)}_i:=\{\begin{array}{ll}{\ell }_i& \text{neˆˊu }{\ell }_i\text{ na˘ˋm trong top-}K\text{ tọa độ của }\ell ,\\ -\infty & \text{ngược lại.}\end{array}$$

Đầu ra tổng thể $y$ của tầng được tính bằng tổng có trọng số cổng của các đầu ra từ các chuyên gia: $$y=\sum _{i=0}^{n-1}G(x{)}_i\cdot {\text{SwiGLU}}_i(x)$$ Việc thực thi hiệu quả đạt được thông qua các kernel như MegaBlocks, biểu diễn tính toán MoE dưới dạng phép nhân ma trận thưa, kết hợp với kỹ thuật Song song hóa Chuyên gia (Expert Parallelism) trên nhiều GPU.

Kết quả (kèm các số liệu quan trọng nhất)

• Benchmark tổng quát & suy luận: Trên MMLU (5-shot), Mixtral 8x7B đạt 70,6% so với 69,9% của Llama 2 70B và 70,0% của GPT-3.5. Trên GSM8K (8-shot, maj@8), mô hình đạt 74,4% so với 69,6% của Llama 2 70B. Trên benchmark MATH (4-shot, maj@4), Mixtral đạt 28,4% so với 13,8% của Llama 2 70B.
• Sinh mã lệnh: Mixtral vượt trội đáng kể so với Llama 2 70B trên HumanEval (40,2% so với 29,3%) và MBPP (60,7% so với 49,8%).
• Tác vụ đa ngôn ngữ: Việc tăng tần suất lấy mẫu dữ liệu đa ngôn ngữ giúp Mixtral vượt qua Llama 2 70B trên tiếng Pháp, tiếng Đức, tiếng Tây Ban Nha và tiếng Ý trong ARC-Challenge, HellaSwag và MMLU (ví dụ: MMLU tiếng Pháp: 70,9% so với 64,3%).
• Tinh chỉnh chỉ dẫn & Điểm Arena Elo: Mixtral 8x7B – Instruct (được huấn luyện với SFT và Direct Preference Optimization) đạt 8,30 điểm trên MT-Bench và đạt mức Elo 1121 trên LMSys Chatbot Arena, vượt qua Claude-2.1 (1117), GPT-3.5-Turbo (1117) và Llama-2-70b-chat (1077).
• Truy xuất ngữ cảnh: Đạt độ chính xác truy xuất 100% trên tác vụ truy xuất passkey tổng hợp trên toàn bộ cửa sổ ngữ cảnh 32k.

Hạn chế và các câu hỏi mở

• Dung lượng bộ nhớ so với Chi phí tính toán: Mặc dù chi phí tính toán tương đương với 13B tham số kích hoạt, việc phục vụ mô hình đòi hỏi đủ VRAM để chứa toàn bộ 47B tham số, yêu cầu thiết lập đa GPU hoặc các kỹ thuật offloading.
• Nghẽn cổ chai trong Song song hóa Chuyên gia: Nghiên cứu phát hiện tính cục bộ thời gian cao—các token liên tiếp chia sẻ các phép gán chuyên gia thường xuyên hơn nhiều so với ngẫu nhiên (lên tới 28,4% tỷ lệ lặp lại lựa chọn đầu tiên ở tầng 15 so với 12,5% ngẫu nhiên). Điều này có thể gây mất cân bằng tải và quá tải một số GPU cụ thể khi dùng Expert Parallelism.
• Thiếu sự chuyên môn hóa theo chủ đề: Việc lựa chọn của router không tự động phân vùng theo các miền khái niệm (ví dụ: sinh học so với triết học), để lại câu hỏi mở về cách cấu trúc cơ chế định tuyến nếu muốn các chuyên gia chuyên môn hóa rõ ràng theo từng miền cụ thể.