Sự trỗi dậy của Mixture-of-Experts cho các mô hình ngôn ngữ lớn hiệu quả

Trong lĩnh vực xử lý ngôn ngữ tự nhiên (NLP), việc xây dựng các mô hình ngôn ngữ lớn hơn và mạnh mẽ hơn đã trở thành một lực đẩy chính đằng sau nhiều tiến bộ gần đây. Tuy nhiên, khi các mô hình này tăng về kích thước, các yêu cầu tính toán cho việc đào tạo và suy luận trở nên ngày càng đòi hỏi, đẩy giới hạn của các tài nguyên phần cứng có sẵn.

Đó là nơi Mixture-of-Experts (MoE) xuất hiện, một kỹ thuật hứa hẹn sẽ giảm bớt gánh nặng tính toán này đồng thời cho phép đào tạo các mô hình ngôn ngữ lớn hơn và mạnh mẽ hơn. Dưới đây, chúng ta sẽ thảo luận về MoE, khám phá nguồn gốc, cơ chế hoạt động và ứng dụng của nó trong các mô hình ngôn ngữ dựa trên transformer.

Nguồn gốc của Mixture-of-Experts

Khái niệm Mixture-of-Experts (MoE) có thể được truy nguồn gốc từ những năm 1990 khi các nhà nghiên cứu khám phá ý tưởng về tính toán có điều kiện, nơi các phần của mạng nơ-ron được kích hoạt có chọn lọc dựa trên dữ liệu đầu vào. Một trong những công trình tiên phong trong lĩnh vực này là bài báo “Adaptive Mixture of Local Experts” của Jacobs et al. vào năm 1991, đã đề xuất một khuôn khổ học có giám sát cho một tập hợp các mạng nơ-ron, mỗi mạng chuyên về một khu vực khác nhau của không gian đầu vào.

Ý tưởng cốt lõi đằng sau MoE là có nhiều mạng “chuyên gia” chuyên về xử lý một tập hợp con của dữ liệu đầu vào. Một cơ chế cổng, thường là một mạng nơ-ron riêng, quyết định chuyên gia nào nên xử lý một đầu vào nhất định. Cách tiếp cận này cho phép mô hình phân bổ tài nguyên tính toán của nó một cách hiệu quả hơn bằng cách kích hoạt chỉ các chuyên gia có liên quan cho mỗi đầu vào, thay vì sử dụng toàn bộ khả năng của mô hình cho mọi đầu vào.

Trong những năm qua, các nhà nghiên cứu đã khám phá và mở rộng ý tưởng về tính toán có điều kiện, dẫn đến các phát triển như MoE phân cấp, xấp xỉ thấp cho tính toán có điều kiện và các kỹ thuật để ước tính gradient thông qua nơ-ron ngẫu nhiên và hàm kích hoạt ngưỡng cứng.

Mixture-of-Experts trong Transformers

Mixture of Experts

Mặc dù ý tưởng về MoE đã tồn tại trong nhiều thập kỷ, việc áp dụng nó vào các mô hình ngôn ngữ dựa trên transformer là tương đối gần đây. Transformers, đã trở thành tiêu chuẩn de facto cho các mô hình ngôn ngữ hiện đại, bao gồm nhiều lớp, mỗi lớp chứa một cơ chế chú ý tự và một mạng nơ-ron truyền thẳng (FFN).

Sự đổi mới chính trong việc áp dụng MoE vào transformers là thay thế các lớp FFN dày bằng các lớp MoE thưa, mỗi lớp bao gồm nhiều chuyên gia FFN và một cơ chế cổng. Cơ chế cổng quyết định chuyên gia nào nên xử lý mỗi token đầu vào, cho phép mô hình kích hoạt có chọn lọc chỉ một tập hợp con các chuyên gia cho một chuỗi đầu vào nhất định.

Một trong những công trình đầu tiên chứng minh tiềm năng của MoE trong transformers là bài báo “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” của Shazeer et al. vào năm 2017. Công trình này giới thiệu khái niệm về một lớp MoE có cổng thưa, sử dụng một cơ chế cổng thêm độ thưa và nhiễu vào quá trình chọn chuyên gia, đảm bảo rằng chỉ một tập hợp con các chuyên gia được kích hoạt cho mỗi đầu vào.

Kể từ đó, nhiều công trình khác đã thúc đẩy việc áp dụng MoE vào transformers, giải quyết các thách thức như sự không ổn định trong đào tạo, cân bằng tải và suy luận hiệu quả. Các ví dụ đáng chú ý bao gồm Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) và GLaM (Du et al., 2022).

Lợi ích của Mixture-of-Experts cho các mô hình ngôn ngữ

Lợi ích chính của việc sử dụng MoE trong các mô hình ngôn ngữ là khả năng tăng kích thước mô hình đồng thời duy trì chi phí tính toán tương đối không đổi trong quá trình suy luận. Bằng cách kích hoạt có chọn lọc chỉ một tập hợp con các chuyên gia cho mỗi token đầu vào, các mô hình MoE có thể đạt được sức biểu đạt của các mô hình dày lớn hơn nhiều trong khi yêu cầu ít tính toán hơn.

Ví dụ, hãy xem xét một mô hình ngôn ngữ có một lớp FFN dày với 7 tỷ tham số. Nếu chúng ta thay thế lớp này bằng một lớp MoE bao gồm tám chuyên gia, mỗi chuyên gia có 7 tỷ tham số, tổng số tham số sẽ tăng lên 56 tỷ. Tuy nhiên, trong quá trình suy luận, nếu chúng ta chỉ kích hoạt hai chuyên gia cho mỗi token, chi phí tính toán sẽ tương đương với một mô hình dày có 14 tỷ tham số, vì nó tính toán hai phép nhân ma trận 7 tỷ tham số.

Hiệu quả tính toán này trong quá trình suy luận đặc biệt có giá trị trong các kịch bản triển khai nơi tài nguyên bị giới hạn, chẳng hạn như thiết bị di động hoặc môi trường tính toán biên. Ngoài ra, yêu cầu tính toán giảm trong quá trình đào tạo có thể dẫn đến tiết kiệm năng lượng đáng kể và giảm lượng khí thải carbon, phù hợp với sự nhấn mạnh ngày càng tăng về các thực hành AI bền vững.

Thách thức và xem xét

Mặc dù các mô hình MoE mang lại những lợi ích hấp dẫn, việc áp dụng và triển khai chúng cũng đi kèm với một số thách thức và xem xét:

1. Sự không ổn định trong đào tạo: Các mô hình MoE được biết đến là dễ bị không ổn định trong đào tạo hơn so với các mô hình dày. Vấn đề này phát sinh từ tính chất thưa và có điều kiện của việc kích hoạt chuyên gia, có thể dẫn đến thách thức trong việc truyền播 gradient và hội tụ. Các kỹ thuật như mất mát z của router (Zoph et al., 2022) đã được đề xuất để giảm thiểu những bất ổn này, nhưng vẫn cần nghiên cứu thêm.
2. Đào tạo lại và quá拟: Các mô hình MoE có xu hướng quá拟 dễ dàng hơn trong quá trình đào tạo lại, đặc biệt khi nhiệm vụ hạ nguồn có một tập dữ liệu tương đối nhỏ. Hành vi này được cho là do khả năng và tính thưa của các mô hình MoE, có thể dẫn đến việc chuyên môn hóa quá mức vào dữ liệu đào tạo. Cần có các chiến lược điều chỉnh và đào tạo lại cẩn thận để giảm thiểu vấn đề này.
3. Yêu cầu bộ nhớ: Mặc dù các mô hình MoE có thể giảm chi phí tính toán trong quá trình suy luận, chúng thường có yêu cầu bộ nhớ cao hơn so với các mô hình dày có kích thước tương tự. Điều này là do tất cả các trọng số chuyên gia cần được tải vào bộ nhớ, ngay cả khi chỉ một tập hợp con được kích hoạt cho mỗi đầu vào. Các hạn chế về bộ nhớ có thể giới hạn khả năng mở rộng của các mô hình MoE trên các thiết bị có tài nguyên hạn chế.
4. Cân bằng tải: Để đạt được hiệu quả tính toán tối ưu, điều quan trọng là phải cân bằng tải trên các chuyên gia, đảm bảo rằng không có chuyên gia nào bị quá tải trong khi các chuyên gia khác không được sử dụng. Việc cân bằng tải này thường được thực hiện thông qua các mất mát phụ trong quá trình đào tạo và điều chỉnh cẩn thận hệ số容量, quyết định số lượng token tối đa có thể được gán cho mỗi chuyên gia.
5. Chi phí truyền thông: Trong các kịch bản đào tạo và suy luận phân tán, các mô hình MoE có thể giới thiệu thêm chi phí truyền thông do nhu cầu trao đổi thông tin kích hoạt và gradient giữa các chuyên gia trên các thiết bị hoặc gia tốc khác nhau. Cần có các chiến lược truyền thông hiệu quả và thiết kế mô hình nhận thức về phần cứng để giảm thiểu chi phí này.

Mặc dù có những thách thức này, tiềm năng của các mô hình MoE trong việc cho phép các mô hình ngôn ngữ lớn hơn và mạnh mẽ hơn đã thúc đẩy các nỗ lực nghiên cứu đáng kể để giải quyết và giảm thiểu những vấn đề này.

Ví dụ: Mixtral 8x7B và GLaM

Để minh họa cho việc áp dụng MoE trong các mô hình ngôn ngữ, hãy xem xét hai ví dụ đáng chú ý: Mixtral 8x7B và GLaM.

Mixtral 8x7B là một biến thể MoE của mô hình ngôn ngữ Mistral, được phát triển bởi Anthropic. Nó bao gồm tám chuyên gia, mỗi chuyên gia có 7 tỷ tham số, dẫn đến tổng cộng 56 tỷ tham số. Tuy nhiên, trong quá trình suy luận, chỉ hai chuyên gia được kích hoạt cho mỗi token, hiệu quả giảm chi phí tính toán xuống mức của một mô hình dày có 14 tỷ tham số.

Mixtral 8x7B đã chứng minh hiệu suất ấn tượng, vượt qua mô hình Llama 70 tỷ tham số trong khi cung cấp thời gian suy luận nhanh hơn nhiều. Một phiên bản được điều chỉnh theo hướng dẫn của Mixtral 8x7B, gọi là Mixtral-8x7B-Instruct-v0.1, cũng đã được phát hành, tăng cường thêm khả năng của nó trong việc tuân theo các hướng dẫn ngôn ngữ tự nhiên.

Một ví dụ khác đáng chú ý là GLaM (Google Language Model), một mô hình MoE lớn được phát triển bởi Google. GLaM sử dụng kiến trúc transformer chỉ có bộ giải mã và được đào tạo trên một tập dữ liệu khổng lồ gồm 1,6 nghìn tỷ token. Mô hình đạt được hiệu suất ấn tượng trên các đánh giá few-shot và one-shot,匹 với chất lượng của GPT-3 trong khi chỉ sử dụng một phần ba năng lượng cần thiết để đào tạo GPT-3.

Sự thành công của GLaM có thể được quy cho kiến trúc MoE hiệu quả của nó, cho phép đào tạo một mô hình với số lượng tham số khổng lồ trong khi duy trì yêu cầu tính toán hợp lý. Mô hình cũng chứng minh tiềm năng của các mô hình MoE để trở nên hiệu quả hơn về năng lượng và bền vững về môi trường so với các mô hình dày tương tự.

Kiến trúc Grok-1

GROK MIXTURE OF EXPERT

Grok-1 là một mô hình MoE dựa trên transformer với kiến trúc độc đáo được thiết kế để tối ưu hóa hiệu quả và hiệu suất. Hãy cùng khám phá các thông số chính:

1. Tham số: Với 314 tỷ tham số, Grok-1 là mô hình LLM mở lớn nhất hiện nay. Tuy nhiên, nhờ kiến trúc MoE, chỉ 25% trọng số (khoảng 86 tỷ tham số) được kích hoạt tại bất kỳ thời điểm nào, tăng cường khả năng xử lý.
2. Kiến trúc: Grok-1 sử dụng kiến trúc Mixture-of-8-Experts, với mỗi token được xử lý bởi hai chuyên gia trong quá trình suy luận.
3. Lớp: Mô hình bao gồm 64 lớp transformer, mỗi lớp bao gồm chú ý tự và khối dày.
4. Token hóa: Grok-1 sử dụng tokenizer SentencePiece với kích thước từ vựng 131.072 token.
5. Nhúng và mã hóa vị trí: Mô hình có nhúng 6.144 chiều và sử dụng mã hóa vị trí quay, cho phép diễn giải dữ liệu động hơn so với mã hóa vị trí cố định truyền thống.
6. Chú ý: Grok-1 sử dụng 48 đầu chú ý cho truy vấn và 8 đầu chú ý cho khóa và giá trị, mỗi đầu có kích thước 128.
7. Độ dài ngữ cảnh: Mô hình có thể xử lý chuỗi lên đến 8.192 token, sử dụng độ chính xác bfloat16 cho tính toán hiệu quả.

Hiệu suất và chi tiết triển khai

Grok-1 đã chứng minh hiệu suất ấn tượng, vượt qua LLaMa 2 70B và Mixtral 8x7B với điểm MMLU là 73%,展示 hiệu quả và chính xác của nó trên các thử nghiệm khác nhau.

Tuy nhiên, cần lưu ý rằng Grok-1 yêu cầu tài nguyên GPU đáng kể do kích thước khổng lồ của nó. Triển khai hiện tại trong bản phát hành mã nguồn mở tập trung vào việc xác thực tính chính xác của mô hình và sử dụng một triển khai lớp MoE không hiệu quả để tránh cần các hạt nhân tùy chỉnh.

Dù vậy, mô hình hỗ trợ phân mảnh kích hoạt và lượng tử hóa 8 bit, có thể tối ưu hóa hiệu suất và giảm yêu cầu bộ nhớ.

Trong một động thái đáng chú ý, xAI đã phát hành Grok-1 theo giấy phép Apache 2.0, làm cho trọng số và kiến trúc của nó có sẵn cho cộng đồng toàn cầu để sử dụng và đóng góp.

Bản phát hành mã nguồn mở bao gồm một kho mã ví dụ JAX展示 cách tải và chạy mô hình Grok-1. Người dùng có thể tải trọng số kiểm tra sử dụng một khách hàng torrent hoặc trực tiếp thông qua HuggingFace Hub, giúp dễ dàng truy cập vào mô hình đột phá này.

Tương lai của Mixture-of-Experts trong các mô hình ngôn ngữ

Khi nhu cầu về các mô hình ngôn ngữ lớn hơn và mạnh mẽ hơn tiếp tục tăng, việc áp dụng các kỹ thuật MoE dự kiến sẽ tăng tốc. Các nỗ lực nghiên cứu đang diễn ra tập trung vào giải quyết các thách thức còn lại, như cải thiện sự ổn định trong đào tạo, giảm thiểu quá拟 trong quá trình đào tạo lại và tối ưu hóa yêu cầu bộ nhớ và truyền thông.

Một hướng đi đầy hứa hẹn là việc khám phá các kiến trúc MoE phân cấp, nơi mỗi chuyên gia chính nó bao gồm nhiều chuyên gia con. Cách tiếp cận này có thể cho phép mở rộng và hiệu quả tính toán thậm chí còn lớn hơn trong khi duy trì sức biểu đạt của các mô hình lớn.

Ngoài ra, sự phát triển của hệ thống phần cứng và phần mềm được tối ưu hóa cho các mô hình MoE là một lĩnh vực nghiên cứu tích cực. Các gia tốc và khuôn khổ đào tạo phân tán được thiết kế để xử lý hiệu quả các mẫu tính toán thưa và có điều kiện của các mô hình MoE có thể tăng cường hiệu suất và khả năng mở rộng của chúng.

Hơn nữa, việc tích hợp các kỹ thuật MoE với các tiến bộ khác trong lĩnh vực mô hình ngôn ngữ, như các cơ chế chú ý thưa, chiến lược token hóa hiệu quả và các biểu diễn đa phương tiện, có thể dẫn đến các mô hình ngôn ngữ thậm chí còn mạnh mẽ và linh hoạt hơn, có khả năng giải quyết một loạt các nhiệm vụ.

Kết luận

Kỹ thuật Mixture-of-Experts đã nổi lên như một công cụ mạnh mẽ trong việc tìm kiếm các mô hình ngôn ngữ lớn hơn và mạnh mẽ hơn. Bằng cách kích hoạt có chọn lọc các chuyên gia dựa trên dữ liệu đầu vào, các mô hình MoE mang lại một giải pháp hứa hẹn cho các thách thức tính toán liên quan đến việc mở rộng các mô hình dày. Mặc dù vẫn còn những thách thức cần vượt qua, như sự không ổn định trong đào tạo, quá拟 và yêu cầu bộ nhớ, tiềm năng của các mô hình MoE về hiệu quả tính toán, khả năng mở rộng và tính bền vững về môi trường làm cho chúng trở thành một lĩnh vực nghiên cứu và phát triển thú vị.

Khi lĩnh vực xử lý ngôn ngữ tự nhiên tiếp tục đẩy ranh giới của những gì có thể, việc áp dụng các kỹ thuật MoE dự kiến sẽ đóng vai trò quan trọng trong việc cho phép thế hệ mô hình ngôn ngữ tiếp theo. Bằng cách kết hợp MoE với các tiến bộ khác trong kiến trúc mô hình, kỹ thuật đào tạo và tối ưu hóa phần cứng, chúng ta có thể mong đợi các mô hình ngôn ngữ thậm chí còn mạnh mẽ và linh hoạt hơn, có thể hiểu và giao tiếp với con người một cách tự nhiên và mượt mà.