Phát triển các Mô hình Ngôn ngữ Lớn (LLM) được xây dựng từ các mô hình transformer chỉ có bộ giải mã đã đóng vai trò quan trọng trong việc chuyển đổi lĩnh vực 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à nhiều hơn nữa. Tuy nhiên, mặc dù chúng có khả năng mở rộng và hiệu suất mạnh mẽ, các LLM được xây dựng từ các mô hình transformer chỉ có bộ giải mã vẫn còn những hạn chế đáng kể. Mặc dù có khả năng diễn đạt, cơ chế chú ý trong các mô hình transformer có nguồn gốc từ transformer đòi hỏi nhiều tài nguyên tính toán trong cả quá trình suy luận và đào tạo, đòi hỏi bộ nhớ tăng theo chiều dài chuỗi và các phép tính (FLOPs) tăng theo cấp số nhân. Những yêu cầu tính toán cao này hạn chế chiều dài ngữ cảnh của các mô hình, làm tăng chi phí của các nhiệm vụ tạo tự động theo cấp số nhân khi mô hình tăng quy mô, và cản trở khả năng của các mô hình học 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ả.

Trong thời gian gần đây, Mô hình Không gian Trạng thái (SSM) đã chứng minh khả năng và hiệu suất đáng kể, cạnh tranh với các mô hình kiến trúc transformer trong các thử nghiệm mô hình hóa lớn, đồng thời đạt được độ phức tạp bộ nhớ theo cấp số nhân với chiều 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 mới được phát hành, đã thể hiện hiệu suất vượt trội trong một loạt các nhiệm 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 đã thể hiện hiệu suất ấn tượng trong khi giảm đáng kể độ trễ và chi phí tính toán của suy luận, mặc dù với chi phí của một dấu chân bộ nhớ lớn hơn. Xây dựng trên 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 lợi thế của cả hai khuôn khổ. Các thí nghiệm trên BlackMamba đã chứng minh khả năng của nó trong việc vượt trội so với khuôn khổ Mamba hiện có và các mô hình baseline transformer trong cả đào tạo FLOPs và suy luận.

Bài viết này nhằm mục đích giới thiệu khuôn khổ BlackMamba một cách sâu sắc. Chúng tôi khám phá cơ chế, phương pháp và kiến trúc của khuôn khổ, cùng với so sánh với các khuôn khổ tạo ảnh và video tiên tiến. Hãy bắt đầu.

BlackMamba: Giới thiệu về MoE cho Mô hình Không gian Trạng thái

Sự tiến bộ của các Mô hình Ngôn ngữ Lớn (LLM), đặc biệt là những mô hình dựa trên kiến trúc transformer chỉ có bộ giải mã, đã ảnh hưởng đáng kể đến lĩnh vực Xử lý Ngôn ngữ Tự nhiên (NLP) và mở rộng vào 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, mặc dù chúng có khả năng mở rộng và hiệu suất mạnh mẽ, các LLM dựa trên kiến trúc transformer chỉ có bộ giải mã vẫn còn những hạn chế đáng kể. Cơ chế chú ý, một tính năng chính của các mô hình transformer, đòi hỏi nhiều tài nguyên tính toán cho cả suy luận và đào tạo.

Các nỗ lực đáng kể đã được thực hiện trong những năm gần đây để vượt qua những hạn chế này, và sự chú ý đã được chuyển hướng đến việc tạo ra các kiến trúc thay thế cho các mô hình transformer truyền thống với SSM và MoE được coi là các kiến trúc ứng viên hứa hẹn nhất. Lợi thế chính của việc ưa chuộng Mô hình Không gian Trạng thái so với các mô hình kiến trúc transformer là độ phức tạp tính toán tuyến tính đối với chiều dài chuỗi đầu vào được cung cấp bởi SSM, so với độ phức tạp cấp số nhân được cung cấp bởi các transformer. Về mặt lý thuyết, độ phức tạp tính toán tuyến tính đối với chiều 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 các mô hình kiến trúc transformer cho một ngân sách FLOPS hoặc Floating-point operations per second nhất định, và cho phép tạo tự động theo cấp số nhân không đổi mà không cần bộ nhớ KV.

Tiến bộ gần đây đã giới thiệu ý tưởng về việc mở rộng chiều dài ngữ cảnh, cho phép các transformer được đào tạo trên một quy mô có thể thực hiện được trước khi được áp dụng cho các ngữ cảnh dài hơn trong quá trình suy luận. Mặc dù 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ớ KV, khiến nó trở thành một nỗ lực tốn kém.

Những nỗ lực nghiên cứu gần đây đã tập trung vào việc tăng cường khả năng diễn đạt của các mô hình không gian trạng thái bằng cách tích hợp các cơ chế cổng đầu vào phụ thuộc vào đầu vào, tương tự như các ma trận QKV được tìm thấy trong các cơ chế chú ý.

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

Mô hình Không gian Trạng thái

Mô hình Không gian Trạng thái thuộc nhóm các mô hình chuỗi có độ phức tạp tuyến tính đối với chiều dài 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 các Mạng nơ-ron hồi quy và Mạng nơ-ron tích chập rather than kiến trúc chú ý, và được lấy cảm hứng từ một hệ thống động liên tục ánh xạ một hàm một chiều thông qua một không gian ẩn ngầm.

Tính chất hồi quy của Mô hình Không gian Trạng thái đã là lý do tại sao nó vẫn chưa được áp dụng rộng rãi trên các phần cứng AI song song như GPU. Tuy nhiên, sự xuất hiện của các SSM như RWKV và Mamba đã sử dụng các nhân quét song song để ánh xạ các hoạt động hồi quy một cách hiệu quả đến GPU, từ đó cho phép đào tạo các kiến trúc mới với hiệu suất tương đương với các mô hình transformer.

Những nỗ lực này nhằm mục đích bảo tồn sự tiến bộ tuyến tính của sự hồi quy không gian trạng thái, cho phép thực hiện hiệu quả thông qua quá trình quét chọn lọc hoặc quét tích hợp.

Mô hình Mixture of Expert

Mô hình Mixture of Expert (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ác tham số một cách chọn lọc trong quá trình truyền tiếp. Thay vì sử dụng tất cả các tham số, các mô hình này chỉ đạo các token đến các chuyên gia MLP cụ thể.

Kiến trúc

BlackMamba sử dụng một mô hình transformer tiêu chuẩn bao gồm các khối MLP và các khối chú ý được xen kẽ dọc theo một dòng dư.

Đào tạo và Dữ liệu

Mô hình BlackMamba được đào tạo trên hơn 300 tỷ token trên một tập dữ liệu tùy chỉnh và sử dụng hàm kích hoạt SwiGLU cho các perceptron đa lớp của chuyên gia.

BlackMamba: Kết quả

Để đảm bảo một so sánh công bằng giữa Mamba và BlackMamba, các nhà phát triển đã đào tạo cả hai mô hình với cùng các tham số đào tạo trên cùng dữ liệu đào tạo.

Lời kết

Trong bài viết này, chúng tôi đã 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 Mixture of Expert để tận dụng lợi thế của cả hai khuôn khổ. Các thí nghiệm trên BlackMamba đã chứng minh khả năng của nó trong việc vượt trội so với khuôn khổ Mamba hiện có và các mô hình baseline transformer trong cả đào tạo FLOPs và suy luận.