Sự phát triển của Mô hình ngôn ngữ lớn (LLM) được xây dựng từ các mô hình biến áp chỉ có bộ giải mã đã đóng một vai trò quan trọng trong việc chuyển đổi miền Xử lý ngôn ngữ tự nhiên (NLP), cũng như thúc đẩy các ứng dụng học sâu đa dạng bao gồm học tăng cường, phân tích chuỗi thời gian, xử lý hình ảnh và hơn thế nữa. Tuy nhiên, bất chấp khả năng mở rộng và hiệu suất mạnh mẽ, LLM được xây dựng từ mô hình máy biến áp chỉ có bộ giải mã vẫn gặp phải những thiếu sót đáng kể. Mặc dù mang tính biểu cảm nhưng cơ chế chú ý trong LLM có nguồn gốc từ máy biến áp đòi hỏi tài nguyên tính toán cao trong cả quá trình suy luận và huấn luyện, đòi hỏi bộ nhớ đáng kể cho độ dài chuỗi và FLOP bậc hai. Yêu cầu tính toán cao này giới hạn độ dài ngữ cảnh của các mô hình máy biến áp, khiến cho các tác vụ tạo tự hồi quy trở nên đắt đỏ tương ứng với quy mô và cản trở việc học từ các luồng dữ liệu liên tục cũng như khả năng xử lý chuỗi thực sự không giới hạn.

Gần đây, Mô hình Không gian Nhà nước (SSM) đã chứng tỏ khả năng và hiệu suất vượt trội, cạnh tranh với các mô hình kiến ​​trúc máy biến áp ở các tiêu chuẩn mô hình hóa quy mô lớn trong khi vẫn đạt được độ phức tạp của bộ nhớ như một hàm của độ dài chuỗi và thời gian tuyến tính. Hơn nữa, Mamba, một Mô hình Không gian Trạng thái được phát hành gần đây, đã cho thấy hiệu suất vượt trội trong một loạt các tác vụ mô hình hóa ngôn ngữ và xử lý chuỗi dài. Đồng thời, các mô hình Mixture of Expert (MoE) cũng cho thấy hiệu suất ấn tượng đồng thời giảm đáng kể độ trễ và chi phí tính toán của suy luận, mặc dù phải trả giá bằng dung lượng bộ nhớ lớn hơn. Dựa trên các mô hình Mamba và MoE, bài viết này sẽ thảo luận về BlackMamba, một kiến ​​trúc mới kết hợp Mô hình Không gian Trạng thái Mamba với các mô hình MoE để tận dụng các lợi ích do cả hai khung mang lại. Các thử nghiệm trên BlackMamba đã chứng minh khả năng của nó vượt trội hơn khung Mamba hiện có và các đường cơ sở biến đổi trong cả FLOP huấn luyện và suy luận. Hiệu suất vượt trội của khung BlackMamba cho thấy nó có thể kết hợp hiệu quả các khả năng của khung Mamba và MoE, mang lại khả năng suy luận nhanh chóng và tiết kiệm chi phí từ MoE với khả năng tạo độ phức tạp tuyến tính từ Mamba.

Bài viết này nhằm mục đích trình bày sâu về khung BlackMamba. Chúng tôi khám phá cơ chế, phương pháp và kiến ​​trúc của khung, cùng với sự so sánh của nó với các khung tạo hình ảnh và video hiện đại. Bắt đầu nào.

BlackMamba: Giới thiệu về MoE cho các mô hình không gian trạng thái

Sự phát triển của các Mô hình ngôn ngữ lớn (LLM), đặc biệt là các mô hình dựa trên kiến ​​trúc biến áp chỉ có bộ giải mã, đã ảnh hưởng đáng kể đến Xử lý ngôn ngữ tự nhiên (NLP) lĩnh vực này và mở rộng sang các ứng dụng học sâu khác nhau, bao gồm học tăng cường, phân tích chuỗi thời gian, xử lý hình ảnh, v.v. Tuy nhiên, bất chấp khả năng mở rộng và hiệu suất mạnh mẽ, các LLM dựa trên máy biến áp chỉ có bộ giải mã này gặp phải những thách thức đáng chú ý. Cơ chế chú ý, một tính năng chính của máy biến áp LLMs, đòi hỏi nguồn lực tính toán rộng rãi cho cả suy luận và đào tạo. Điều này liên quan đến nhu cầu về bộ nhớ tăng theo độ dài chuỗi và các phép toán tính toán (FLOP) tăng theo bậc hai. Các nhu cầu tính toán chuyên sâu như vậy hạn chế độ dài ngữ cảnh của mô hình, tăng chi phí cho các nhiệm vụ tạo tự hồi quy khi quy mô mô hình và cản trở khả năng học hỏi của mô hình từ các luồng dữ liệu liên tục hoặc xử lý các chuỗi có độ dài không giới hạn một cách hiệu quả. 

Những nỗ lực đáng kể đã được thực hiện trong vài năm qua nhằm khắc phục những hạn chế này và sự chú ý đã chuyển sang việc đưa ra các giải pháp thay thế kiến ​​trúc cho các mô hình biến áp chú ý dày đặc chuẩn mực với các mô hình SSM và MoE là những kiến ​​trúc ứng cử viên hứa hẹn nhất. Lợi ích chính thu được bằng cách ưu tiên Mô hình không gian trạng thái hơn các mô hình kiến ​​trúc máy biến áp là độ phức tạp tính toán tuyến tính đối với độ dài chuỗi đầu vào do SSM cung cấp, trái ngược với độ phức tạp bậc hai do máy biến áp cung cấp. Về mặt lý thuyết, độ phức tạp tính toán tuyến tính liên quan đến độ dài chuỗi đầu vào cho phép Mô hình không gian trạng thái xử lý các chuỗi lớn hơn so với mô hình kiến ​​trúc máy biến áp cho một FLOPS hoặc hoạt động dấu phẩy động nhất định trên mỗi giây và hiển thị hằng số tạo tự hồi quy trong điện toán mà không cần bộ nhớ đệm KV. Các Mô hình Không gian Trạng thái được phát triển gần đây bao gồm Mamba, RetNet và một số mô hình khác đã chứng minh khả năng suy luận và đào tạo chuỗi dài hiệu quả, cùng với hiệu suất thực hiện nhiệm vụ mô hình hóa ngôn ngữ cạnh tranh đối với các máy biến áp có đặc tính chia tỷ lệ tương tự. Mặt khác, kiến ​​trúc Hỗn hợp các mô hình Chuyên gia đang trở nên phổ biến như một giải pháp thay thế cho các máy biến áp dày đặc vì nó tạo điều kiện giảm đáng kể FLOP suy luận và huấn luyện cần thiết để đạt được chất lượng tương đương với một mô hình dày đặc. Các mô hình MoE (Hỗn hợp các chuyên gia) hoạt động bằng cách chỉ kích hoạt một lựa chọn thưa thớt của tổng số tham số trong một lần chuyển tiếp. Họ sử dụng chức năng định tuyến để xác định 'chuyên gia' nào được gọi hành động dựa trên bối cảnh nhất định. Cách tiếp cận này tạo ra sự tách biệt giữa chi phí tính toán của suy luận và tổng số tham số, cho phép nâng cao hiệu suất trong phạm vi ngân sách suy luận cố định, mặc dù số lượng tham số tăng lên và yêu cầu bộ nhớ lớn hơn.

Sự tiến bộ trong kiến ​​trúc này mang lại những lợi ích đáng chú ý so với máy biến áp truyền thống và thể hiện một hướng phát triển thú vị hơn nữa. Chúng tôi thừa nhận rằng việc tích hợp những cải tiến này vào mô hình Mamba-MoE kết hợp có thể tăng tốc đáng kể khả năng và hiệu quả của mô hình hóa ngôn ngữ vượt xa các mô hình biến áp tiêu chuẩn. Những ưu điểm dự đoán của kiến ​​trúc Mamba-MoE so với mô hình máy biến áp dày đặc truyền thống bao gồm:

Ba: Đạt được độ phức tạp tính toán tuyến tính tương ứng với độ dài chuỗi đầu vào cho cả giai đoạn huấn luyện và suy luận. Nó cho phép quá trình tạo tự hồi quy diễn ra trong một khung thời gian không đổi và với mức sử dụng bộ nhớ không đổi.

MoE: Cung cấp tốc độ suy luận và hiệu quả tính toán đào tạo tương đương với mô hình cơ sở dày đặc, nhỏ hơn trong khi vẫn duy trì mức chất lượng mô hình cạnh tranh với mô hình có số lượng tham số tương đương như phiên bản dày đặc hơn.

Như đã nói, điều cần thiết là phải nói rằng các mô hình kiến ​​trúc máy biến áp vẫn ở trạng thái hiện đại và đã chứng tỏ hiệu suất mạnh mẽ và nhất quán đáng chú ý đối với các tác vụ mô hình hóa ngôn ngữ và các tác vụ xử lý chuỗi. Về cốt lõi, kiến ​​trúc máy biến áp sử dụng khả năng tự chú ý để thực hiện so sánh bậc hai về tất cả các điểm tương đồng của tích số chấm giữa các phần nhúng của các mã thông báo khác nhau trong một chuỗi và thực hiện ánh xạ tuyến tính tới vectơ đầu ra. Mô hình máy biến áp bao gồm các khối tự chú ý được xếp chồng lên nhau giữa các khối MLP hoặc Perceptron nhiều lớp, bao gồm thêm MLP hai lớp với chức năng kích hoạt nhất định. 

BlackMamba: Kiến trúc và phương pháp

Mô hình Không gian Nhà nước

Mô hình không gian trạng thái thuộc nhóm mô hình chuỗi có độ phức tạp tuyến tính xét theo độ dài của chuỗi đầu vào. Kiến trúc của Mô hình không gian trạng thái phù hợp hơn với Mạng thần kinh tái phát và Mạng thần kinh tích chập hơn là kiến ​​trúc dựa trên sự chú ý và được lấy cảm hứng từ một hệ thống động lực liên tục ánh xạ hàm 1 chiều thông qua một không gian tiềm ẩn. Một hệ thống động lực tuyến tính giúp cho việc tính toán song song trở nên hiệu quả hơn bằng cách sử dụng quét kết hợp hoặc quét tích chập. Trong các tình huống thực tế, tính chất lặp đi lặp lại của Mô hình không gian trạng thái là lý do tại sao nó vẫn được áp dụng trên phần cứng AI có tính song song cao như GPU. Tuy nhiên, sự xuất hiện của SSM như RWKV và Mamba đã sử dụng nhân quét song song để ánh xạ các hoạt động định kỳ một cách hiệu quả tới GPU, từ đó tạo điều kiện thuận lợi cho việc đào tạo các kiến ​​trúc mới với hiệu quả tương đương với hiệu quả mà các mô hình máy biến áp đạt được. 

Độ phức tạp bậc hai vốn có liên quan đến độ dài chuỗi trong máy biến áp là một hạn chế nổi tiếng cản trở khả năng lý luận và hiểu biết trong các bối cảnh rất dài. Những đổi mới gần đây đã đưa ra ý tưởng mở rộng độ dài ngữ cảnh, cho phép các máy biến áp được huấn luyện ở quy mô khả thi trước khi áp dụng vào các ngữ cảnh dài hơn nhiều trong quá trình suy luận. Bất chấp những tiến bộ này, quá trình suy luận vẫn đòi hỏi một lượng tài nguyên tính toán và bộ nhớ đáng kể, đặc biệt là để duy trì bộ nhớ đệm Khóa-Giá trị (KV), khiến nó trở thành một nỗ lực tiêu tốn nhiều tài nguyên. Những nỗ lực nghiên cứu gần đây đã tập trung vào việc nâng cao khả năng biểu đạt của các mô hình không gian trạng thái bằng cách kết hợp các cơ chế cổng phụ thuộc vào đầu vào, giống như ma trận Truy vấn, Khóa, Giá trị (QKV) được tìm thấy trong các cơ chế chú ý. 

Những nỗ lực này nhằm mục đích duy trì tiến trình tuyến tính vốn có của đệ quy không gian trạng thái, cho phép thực hiện hiệu quả thông qua quá trình tích chập hoặc quá trình quét chọn lọc. Cách tiếp cận này thu hẹp đáng kể sự chênh lệch hiệu suất với máy biến áp trong các ứng dụng thực tế. Trong số những tiến bộ này, Mamba nổi bật như một mô hình không gian trạng thái phản ánh các mục tiêu của nghiên cứu trước đó, cho thấy mức hiệu suất ấn tượng có thể so sánh với máy biến áp ở quy mô lên tới 2.8 tỷ thông số. Nó đạt được điều này bằng cách áp dụng việc chọn cổng phụ thuộc vào đầu vào cho các đầu vào của đệ quy mô hình không gian trạng thái (SSM), đồng thời đảm bảo tính toán hiệu quả thông qua việc sử dụng các hạt nhân quét chọn lọc riêng biệt.

Sự kết hợp của các mô hình chuyên gia

Các mô hình Hỗn hợp Chuyên gia (MoE) đạt được sự tách biệt giữa chi phí suy luận và tổng số tham số bằng cách kích hoạt có chọn lọc các tham số trong quá trình chuyển tiếp. Thay vì sử dụng tất cả các tham số, các mô hình này hướng mã thông báo đến các chuyên gia Perceptron đa lớp (MLP) cụ thể. Lý tưởng nhất là mỗi chuyên gia được điều chỉnh để xử lý một loại đầu vào cụ thể, với cơ chế định tuyến, về cơ bản là một mạng thần kinh nhỏ gọn, xác định chuyên gia phù hợp nhất cho mỗi mã thông báo. Cách tiếp cận này nhằm mục đích duy trì sức mạnh biểu đạt toàn diện của một mô hình có số lượng tham số tương đương trong cấu hình dày đặc hơn nhưng với nhu cầu tính toán giảm đáng kể. Thông thường, bộ định tuyến là ánh xạ của các lớp tuyến tính từ mã thông báo đến chỉ số chuyên gia với mỗi chuyên gia chỉ đơn giản là một Perceptron đa lớp biến áp tiêu chuẩn. Tuy nhiên, các nhà phát triển vẫn chưa tìm ra phương pháp đào tạo tối ưu cho bộ định tuyến vì vấn đề phân công chuyên gia là không thể phân biệt được và các mô hình Mixture of Expert thường gặp khó khăn với việc cân bằng tải và đào tạo sự ổn định giữa các chuyên gia khác nhau để đạt được hiệu quả phần cứng. 

Kiến trúc

Về cốt lõi, BlackMamba sử dụng mô hình máy biến áp tiêu chuẩn bao gồm các khối MLP xen kẽ và các khối chú ý được thêm vào theo trình tự dọc theo luồng dư. Giờ đây, phần lớn các mô hình Mixture of Expert chỉ cần thay thế các khối perceptron nhiều lớp bằng một lớp chuyên gia được định tuyến. Mặt khác, khung BlackMamba không chỉ thay thế khối tri giác nhiều lớp trong máy biến áp bằng lớp chuyên gia được định tuyến mà còn thay thế lớp chú ý bằng lớp Mô hình không gian trạng thái Mamba. Kiến trúc của khung BlackMamba được thể hiện trong hình sau. 

Đào tạo và tập dữ liệu

Mô hình BlackMamba được đào tạo trên hơn 300 tỷ mã thông báo trên tập dữ liệu tùy chỉnh và sử dụng chức năng kích hoạt SwiGLU cho các tri giác đa lớp chuyên gia. Khung đào tạo với 8 chuyên gia, một con số mà các nhà phát triển nhận thấy là sự cân bằng phù hợp và cân bằng giữa dung lượng bộ nhớ và chi phí suy luận của mô hình. Tập dữ liệu tùy chỉnh được sử dụng để huấn luyện khung BlackMamba bao gồm hỗn hợp các tập dữ liệu nguồn mở hiện có bao gồm Starcode, SlimPajama, Pile, v.v. Bảng sau đây thể hiện trọng số của từng tập dữ liệu được sử dụng để đào tạo khung BlackMamba. Nhìn chung, có 1.8 nghìn tỷ token trong tập dữ liệu. 

BlackMamba : Kết quả

Để đảm bảo so sánh công bằng giữa Mamba và BlackMamba, các nhà phát triển đã huấn luyện cả hai mô hình với cùng thông số huấn luyện trên cùng một dữ liệu huấn luyện. Khung BlackMamba có thể hoạt động tốt hơn cả mô hình Mamba và máy biến áp với kích thước mô hình chuyển tiếp giống hệt nhau tại thời điểm suy luận cũng như huấn luyện các phép toán Dấu phẩy động mỗi giây. Hình dưới đây minh họa thời gian cần thiết để tự động tạo một chuỗi có độ dài nhất định từ dấu nhắc một mã thông báo ban đầu dưới dạng hàm của độ dài chuỗi. 

Hơn nữa, lợi ích về độ trễ của cả hai mô hình Mixture of Expert và Mamba được kết hợp trong khung BlackMamba dẫn đến thời gian suy luận nhanh hơn đáng kể khi so sánh với các mô hình máy biến áp, mô hình Mamba thuần túy và mô hình MoE. Hơn nữa, lợi thế suy luận của khung BlackMamba tỷ lệ thuận với độ dài chuỗi, khiến BlackMamba cực kỳ hiệu quả khi tạo chuỗi dài. Tiếp theo, hình dưới đây minh họa số lượng mã thông báo được gán cho các mô hình BlackMamba với thông số lần lượt là 340 triệu và 640 triệu. Như có thể thấy, phần lớn các lớp thể hiện mức độ cân bằng chuyên môn cao nhờ thuật toán Sinkhorn cải tiến được triển khai bởi các mô hình BlackMamba. 

Bảng sau đây trình bày điểm đánh giá của khung BlackMamba so với nhiều mô hình ngôn ngữ được đào tạo trước nguồn mở. Như có thể thấy, khung BlackMamba có thể cạnh tranh và hoạt động tốt hơn với phần lớn các khung trên tất cả các đường cơ sở. Hơn nữa, điều đáng chú ý là các mô hình hoạt động tốt hơn BlackMamba có số lượng tham số cao hơn đáng kể và khoảng cách về hiệu suất là rất nhỏ, cho thấy khả năng của khung BlackMamba với ít tham số hơn. 

Kết luận:

Trong bài viết này, chúng ta đã nói về BlackMamba, một kiến ​​trúc mới kết hợp Mô hình không gian trạng thái Mamba với các mô hình Hỗn hợp chuyên gia để tận dụng những lợi ích mà cả hai khung này mang lại. Các thử nghiệm trên BlackMamba đã chứng minh rằng nó hoạt động tốt hơn khung Mamba hiện có và các đường cơ sở biến áp trong cả FLOP huấn luyện và suy luận. Hiệu suất vượt trội của khung BlackMamba chứng tỏ rằng nó có thể kế thừa và kết hợp các khả năng của khung Mamba và MoE một cách đặc biệt tốt vì nó kết hợp khả năng suy luận nhanh và rẻ từ MoE với khả năng tạo độ phức tạp tuyến tính từ Mamba. Chúng tôi đã nói về cách kiến ​​trúc của khung BlackMamba có thể hoạt động tốt hơn các Mô hình Ngôn ngữ Lớn được đào tạo mạnh mẽ, khung Mamba hiện có và các mô hình Hỗn hợp Chuyên gia về FLOP đào tạo và chi phí suy luận. Hơn nữa, khung BlackMamba cũng kế thừa FLOP thế hệ và giảm đào tạo từ cả hai mô hình Mixture of Expert và khung Mamba cùng một lúc.