Mamba: Định nghĩa lại Mô hình Chuỗi và Outforming Kiến trúc Transformer

Bài viết này tìm hiểu Mamba, mô hình không gian trạng thái (SSM) do Albert Gu và Tri Dao phát triển, và cách mô hình này cải tiến việc mô hình hóa chuỗi. Mamba xử lý hiệu quả các chuỗi phức tạp trong ngôn ngữ, hệ gen và âm thanh. Không gian trạng thái có chọn lọc cho phép chi phí tính toán tăng tuyến tính theo độ dài chuỗi.

Mamba giải quyết những khó khăn tính toán của Transformer truyền thống khi chuỗi dài hơn. Cách tiếp cận chọn lọc giúp suy luận nhanh và tăng thông lượng. Điểm nổi bật là tốc độ xử lý, lớp SSM có chọn lọc và thiết kế phù hợp phần cứng, lấy cảm hứng từ FlashAttention. Trong những kết quả được bài gốc trình bày, Mamba vượt nhiều mô hình hiện có, kể cả một số mô hình Transformer.

Transformer và Mamba

Transformer như GPT-4 đã thiết lập những chuẩn mực trong xử lý ngôn ngữ tự nhiên, nhưng chi phí tính toán tăng khi chuỗi dài hơn. Mamba sử dụng kiến trúc khác để xử lý chuỗi dài hiệu quả hơn.

Transformer xử lý dữ liệu dạng chuỗi, chẳng hạn văn bản. Khác với các mô hình trước xử lý tuần tự từng phần tử, Transformer có thể xử lý song song các vị trí trong quá trình huấn luyện, từ đó nắm bắt quan hệ phức tạp. Cơ chế chú ý cho phép mô hình tập trung vào những phần khác nhau của đầu vào khi dự đoán.

Các biểu diễn truy vấn, khóa và giá trị được tạo từ đầu vào. Việc so sánh các phần tử tạo ra trọng số chú ý, thể hiện mức độ liên quan của từng phần tử đối với dự đoán. Trong kiến trúc bộ mã hóa–bộ giải mã được mô tả ở đây, bộ mã hóa xử lý đầu vào và bộ giải mã tạo đầu ra. Đây là một cấu hình Transformer cụ thể, không phải cấu trúc của mọi mô hình trong họ này.

Mỗi lớp mã hóa có cơ chế tự chú ý đa đầu và mạng truyền thẳng kết nối đầy đủ áp dụng riêng tại từng vị trí. Chuẩn hóa và kết nối phần dư hỗ trợ huấn luyện mạng sâu. Bộ giải mã bổ sung chú ý lên đầu ra của bộ mã hóa. Mặt nạ nhân quả bảo đảm dự đoán tại một vị trí chỉ dùng các vị trí trước đó, duy trì tính tự hồi quy.

Mamba sử dụng không gian trạng thái có chọn lọc thay cho cách tiếp cận chú ý này. Lưu ý về nguồn: câu so sánh tương ứng trong bài tiếng Anh bị cắt giữa một từ; bản dịch không tự bổ sung kết luận còn thiếu.

Các bước chính của Transformer gồm:

• Xử lý đầu vào: chuyển dữ liệu thành các biểu diễn nhúng, kết hợp thông tin về vị trí của phần tử trong chuỗi.
• Chú ý: tính điểm thể hiện mức độ cần tập trung vào những phần khác của đầu vào.
• Mã hóa và giải mã: trong cấu hình đang xét, nhiều lớp xử lý đầu vào và tạo đầu ra.
• Chú ý đa đầu: xem xét đồng thời nhiều quan hệ từ các không gian biểu diễn khác nhau để học nhiều ngữ cảnh.
• Mạng truyền thẳng theo vị trí: áp dụng cùng một phép biến đổi độc lập tại từng vị trí, kết hợp đầu vào qua kết nối phần dư rồi chuẩn hóa lớp.
• Tạo đầu ra: bộ giải mã dự đoán chuỗi đầu ra dựa trên ngữ cảnh của bộ mã hóa và phần đã tạo trước đó.

Xử lý song song và cơ chế chú ý giúp Transformer thực hiện các nhiệm vụ như dịch và sinh văn bản. Mamba giải quyết chi phí của chuỗi dài bằng không gian trạng thái có chọn lọc, suy luận nhanh hơn và mức tăng chi phí tuyến tính theo độ dài chuỗi.

Mamba

Mamba rời khỏi thiết kế dùng các khối chú ý và perceptron nhiều lớp (MLP) truyền thống, hướng tới mô hình nhẹ hơn, nhanh hơn và đồng nhất hơn.

• SSM có chọn lọc: lọc thông tin không liên quan và tập trung vào dữ liệu hữu ích, hỗ trợ suy luận phụ thuộc nội dung.
• Thuật toán tối ưu cho phần cứng: thuật toán song song tận dụng phần cứng hiện đại, nhất là GPU, để tính toán nhanh hơn và giảm nhu cầu bộ nhớ.
• Kiến trúc đơn giản: tích hợp SSM có chọn lọc và loại bỏ các khối chú ý, MLP truyền thống giúp cải thiện khả năng mở rộng và hiệu năng.

Các kết quả được nêu bao gồm ngôn ngữ, âm thanh và hệ gen, cả tiền huấn luyện lẫn nhiệm vụ chuyên biệt. Trong các so sánh mô hình hóa ngôn ngữ được báo cáo, Mamba đạt hoặc vượt hiệu năng của Transformer lớn hơn. Mã nguồn và mô hình tiền huấn luyện được công khai trên GitHub.

Sao chép thông thường là nhiệm vụ đơn giản đối với mô hình tuyến tính. Sao chép có chọn lọc và các đầu quy nạp cần bộ nhớ động có nhận biết nội dung.

Mô hình không gian trạng thái có cấu trúc S4 kết hợp đặc điểm của mạng hồi quy (RNN), mạng tích chập (CNN) và mô hình không gian trạng thái cổ điển. S4 lấy cảm hứng từ hệ liên tục ánh xạ hàm một chiều hoặc chuỗi thông qua trạng thái ẩn. Trong học sâu, đây là một phương pháp xây dựng mô hình chuỗi hiệu quả và linh hoạt.

Động lực học của mô hình S4

SSM có cấu trúc nhận chuỗi x và tạo đầu ra y, dùng các tham số học được A, B, C cùng bước thời gian Δ. Các tham số liên tục được rời rạc hóa trước khi áp dụng phép toán SSM. Ở dạng bất biến theo thời gian cơ bản, tham số không thay đổi giữa các bước.

Vai trò của rời rạc hóa

Rời rạc hóa chuyển tham số liên tục thành tham số rời rạc bằng những công thức xác định, giữ liên hệ với hệ thời gian liên tục. Quá trình này mang lại các tính chất như bất biến đối với độ phân giải và chuẩn hóa phù hợp, giúp tăng độ ổn định và hiệu năng. Nó cũng có liên hệ với cơ chế cổng trong RNN, vốn điều khiển luồng thông tin qua mạng.

Tuyến tính và bất biến theo thời gian (LTI)

Một đặc điểm cốt lõi của S4 là động lực học nhất quán theo thời gian, với tham số cố định ở mọi bước. Tính chất LTI tạo cơ sở cho phép truy hồi và tích chập, mang lại khung mô hình hóa chuỗi đơn giản nhưng mạnh mẽ.

Vượt qua những hạn chế cơ bản

Tính bất biến theo thời gian hạn chế khả năng mô hình hóa dữ liệu cần động lực học thích ứng. Công trình nghiên cứu đưa vào tham số biến thiên theo thời gian, loại bỏ ràng buộc LTI để xử lý nhiều loại chuỗi và nhiệm vụ hơn.

“Mô hình không gian trạng thái” là thuật ngữ rộng dùng cho các quá trình hồi quy có trạng thái ẩn trong nhiều lĩnh vực. Trong học sâu, S4 hay SSM có cấu trúc chỉ một lớp mô hình cụ thể được tối ưu tính toán nhưng vẫn mô hình hóa được chuỗi phức tạp. Chúng có thể tích hợp vào mạng huấn luyện đầu cuối như các phép biến đổi chuỗi độc lập, với vai trò tương tự lớp tích chập trong CNN.

SSM và SSM kết hợp lựa chọn

Vì sao mô hình chuỗi cần tính chọn lọc?

SSM có cấu trúc

Một khía cạnh quan trọng là nén ngữ cảnh thành trạng thái có thể quản lý. Mô hình biết tập trung hoặc lọc đầu vào có thể giữ lại thông tin hữu ích hiệu quả hơn trong biểu diễn nén này. Tính chọn lọc giúp điều khiển thích ứng luồng thông tin dọc theo chuỗi, cần thiết cho các nhiệm vụ ngôn ngữ phức tạp và nhiều lĩnh vực khác.

SSM có chọn lọc cho phép tham số phụ thuộc đầu vào, tạo ra động lực học biến thiên theo thời gian. Vì vậy, chúng không còn dùng được cách tính bằng tích chập của SSM bất biến theo thời gian mà chuyển sang cơ chế truy hồi có chi phí tuyến tính.

Trong biến thể SSM kết hợp lựa chọn (S6), các tham số B, C và Δ phụ thuộc chuỗi đầu vào x. Quá trình rời rạc hóa tính đến lựa chọn này; phép quét xử lý chuỗi và điều chỉnh sự tập trung linh hoạt qua từng bước thời gian.

Những kết quả nổi bật của Mamba

Biểu đồ của nguồn trình bày Mamba dẫn đầu trong từng phép đánh giá được hiển thị.

Mamba đạt tốc độ suy luận và độ chính xác cao trong những thử nghiệm được báo cáo. Thiết kế tận dụng tốt ngữ cảnh dài hơn trong mô hình hóa DNA và âm thanh, vượt các mô hình trước trên nhiệm vụ phức tạp cần quan hệ phụ thuộc xa. Những đánh giá zero-shot trên nhiều nhiệm vụ cũng cho thấy tính linh hoạt, hiệu quả và khả năng mở rộng.

Bắt đầu với Mamba

Bài viết năm 2023 nêu các yêu cầu Linux, GPU NVIDIA, PyTorch 1.12+ và CUDA 11.6+. Việc cài đặt dùng pip với các gói từ kho mã Mamba. Khi gặp vấn đề tương thích phiên bản PyTorch, nguồn đề xuất cờ --no-build-isolation. Các mô hình được huấn luyện trên những tập dữ liệu lớn như Pile và SlimPajama nhằm đáp ứng nhiều nhu cầu tính toán.

Các cấp giao diện gồm lớp SSM có chọn lọc, khối Mamba và cấu trúc mô hình ngôn ngữ hoàn chỉnh. Khối Mamba sử dụng lớp Conv1d nhân quả và có thể tích hợp vào thiết kế mạng nơ-ron. Lưu ý về nguồn: bài tiếng Anh nhắc đến ví dụ Python khởi tạo mô hình và xử lý dữ liệu, nhưng không chứa đoạn mã đó.

Các mô hình tiền huấn luyện trên Hugging Face có từ 130 triệu đến 2,8 tỷ tham số, được huấn luyện trên Pile và SlimPajama. Chúng tuân theo các chuẩn kích thước của GPT-3, phục vụ nhiều yêu cầu tính toán và hiệu năng. Nguồn báo cáo thông lượng và độ chính xác cao, giúp Mamba trở thành lựa chọn cạnh tranh cho mô hình hóa ngôn ngữ và những ứng dụng khác.

Tác động của Mamba

Mamba cung cấp một lựa chọn mạnh mẽ bên cạnh Transformer để xử lý dữ liệu giàu thông tin. Thiết kế phù hợp phần cứng hiện đại, tối ưu bộ nhớ và tính toán song song. Mã nguồn mở cùng mô hình tiền huấn luyện giúp các nhà nghiên cứu và lập trình viên AI, học sâu dễ dàng tiếp cận.