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 động lực 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 đến 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 vùng khác nhau của không gian đầu vào.

Ý tưởng cốt lõi đằng sau MoE là có nhiều “chuyên gia” mạng, mỗi chuyên gia chịu trách nhiệm 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 mình một cách hiệu quả hơn bằng cách kích hoạt chỉ các chuyên gia 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 đã có từ hàng 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 feed-forward (FFN).

Sự đổi mới chính trong việc áp dụng MoE vào transformer 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ủa 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 transformer 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ủa 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 transformer, 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 trong khi 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ủa 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 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ó thực hiệ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 phương pháp 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 nhiều lợi ích, 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 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-router (Zoph et al., 2022) đã được đề xuất để giảm thiểu những không ổn định này, nhưng vẫn cần nghiên cứu thêm.
2. Đào tạo tinh chỉnh 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 tinh chỉnh, đặ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 quá chuyên biệt hóa vào dữ liệu đào tạo. Cần có các chiến lược tinh chỉnh và quy chuẩn hóa 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ố của chuyên gia cần được tải vào bộ nhớ, ngay cả khi chỉ một tập hợp con của chuyên gia đượ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, việc cân bằng tải trên các chuyên gia là rất quan trọng, đả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 tận dụng. Cân bằng tải này thường được đạt được 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 của yếu tố khả năng, 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 trên các chuyên gia cư trú 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 ứng dụng thực tế của 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 14 tỷ tham số.

Mixtral 8x7B đã thể hiện 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 điều chỉnh của Mixtral 8x7B, gọi là Mixtral-8x7B-Instruct-v0.1, cũng đã được phát hành, nâng cao hơn nữa 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 (GOOGL ) Language Model), một mô hình ngôn ngữ 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ồ 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, sánh ngang 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ì các 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 trong việc trở nên tiết kiệm năng lượng và bền vững môi trường hơn so với các mô hình dày tương đương.

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 suất và hiệu quả. Hãy cùng khám phá các thông số kỹ thuật 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 kết hợp chú ý tự và các 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ó cá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. Chiều 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 thực hiện

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

Tuy nhiên, cần lưu ý rằng Grok-1 đòi hỏi tài nguyên GPU đáng kể do kích thước khổng lồ của nó. Việc 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 nhận 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 nhu cầu về kernel tùy chỉnh.

Dù vậy, mô hình hỗ trợ phân mảnh kích hoạt và định lượng 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 dưới 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 cho thấy 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 bằng cách 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ẽ thu được động lực hơn nữa. 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 tinh chỉnh và tối ưu hóa yêu cầu bộ nhớ và truyền thông.

Một hướng nghiên cứu 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 bản thân cũng bao gồm nhiều chuyên gia con. Cách tiếp cận này có thể cho phép khả năng 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 các 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à khung đà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ể nâng cao hiệu suất và khả năng mở rộng của chúng hơn nữa.

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ơ chế chú ý thưa, chiến lược token hóa hiệu quả, và 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à đa năng 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 cung cấp 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 tăng kích thước của 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ớ, nhưng lợi ích 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 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ệ tiếp theo của các mô hình ngôn ngữ. 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à đa năng hơn, có khả năng hiểu và giao tiếp với con người một cách tự nhiên và liền mạch.