Mamba: mô hình không gian trạng thái thách thức Transformer

Mamba là mô hình không gian trạng thái (State Space Model) thách thức Transformer

• AI hiện đang thống trị thế giới, và ở trung tâm của nó là Transformer
• Mamba thuộc về một lớp mô hình thay thế có tên là mô hình không gian trạng thái (SSM)
• Mamba có hiệu năng và khả năng mở rộng tương tự Transformer, đồng thời có thể chạy trên các chuỗi rất dài
• Điểm đáng chú ý là Mamba loại bỏ “nút thắt cổ chai bậc hai” của “cơ chế attention”, nhờ đó hỗ trợ ngữ cảnh dài
• Mamba chạy nhanh hơn Transformer tới 5 lần

Vấn đề của Transformer - chỉ attention thôi có thể là chưa đủ

• Transformer cho phép mọi token tham chiếu tới các token trước đó, nên khi ngữ cảnh tăng lên thì mô hình sẽ chậm hơn
• Việc lưu trữ KV cache này cũng đòi hỏi độ phức tạp không gian O(n)
• Đã có các kỹ thuật giúp giảm nút thắt cổ chai của Transformer hiện tại, nhưng để giải quyết tận gốc thì cần một cách tiếp cận khác

Backbone của mô hình nền tảng

• Hai thành phần quan trọng của một backbone kiến trúc ML tốt
  • Giao tiếp (Communication) giữa các token
  • Tính toán (Computation) bên trong token
• Khối Transformer được cấu thành từ attention và MLPs
• Mamba sử dụng SSM lấy cảm hứng từ lý thuyết điều khiển cho phần giao tiếp, đồng thời giữ lại kiểu chiếu MLP cho phần tính toán

Động lực của Mamba - quay lại với Temple Run

• Trạng thái (state) là các biến cần thiết để xác định hành vi tương lai của một hệ thống
• Trạng thái là phần nén của mọi thứ cần biết về quá khứ, và được chuyển thành một quá trình quyết định Markov

Rời rạc hóa - sống trong một thế giới đã được lượng tử hóa

• Quá trình chuyển phương trình vi phân theo thời gian liên tục thành phương trình sai phân theo thời gian rời rạc được gọi là rời rạc hóa (discretisation)
• Mamba sử dụng phương pháp rời rạc hóa zero-order hold (ZOH)

Hiểu các ma trận SSM

• Các ma trận A, B, C, D lần lượt đảm nhiệm vai trò chuyển trạng thái, ánh xạ đầu vào mới vào trạng thái, ánh xạ trạng thái sang đầu ra SSM, và truyền đầu vào mới sang đầu ra

Hiệu quả và hiệu năng: Attention is Focus, Selectivity is Prioritisation (Attention là tập trung, Selectivity là ưu tiên hóa)

• Transformer rất hiệu quả về mặt kết quả nhưng không thật sự hiệu quả về mặt tính toán
• Kiến trúc Mamba đưa ra một lời giải giúp đẩy đường biên Pareto giữa hiệu quả và hiệu năng

Cơ chế chọn lọc

• Tính chọn lọc (Selectivity) cho phép mỗi token được chuyển thành trạng thái theo đúng nhu cầu của nó
• Mamba biến các ma trận A, B, C thành hàm của x, khiến chúng không còn tĩnh mà phụ thuộc vào ngữ cảnh

Vấn đề của tính chọn lọc

• Khi áp dụng cơ chế chọn lọc, việc tính toán có thể chậm hơn so với SSM không chọn lọc
• Nhờ tối ưu hóa phần cứng, Mamba có thể chạy nhanh hơn Transformer có kích thước tương đương

Học máy và kinh tế chính trị - kích thước trạng thái nên lớn đến mức nào?

• Sự đánh đổi giữa hiệu năng và hiệu quả của mô hình chuỗi được đặc trưng bởi mức độ nén trạng thái tốt đến đâu
• Biểu diễn trạng thái là yếu tố quan trọng, và việc nén trạng thái một cách chọn lọc, động là chìa khóa

Luồng thông tin của Transformer so với Mamba

• Transformer học thông qua dữ liệu huấn luyện và dữ liệu ngữ cảnh
• Với Mamba, dữ liệu huấn luyện và dữ liệu ngữ cảnh được nén/lọc để có thể truy cập

Thay thế trạng thái như một mô hình prompting mới

• Khi dùng các mô hình như Mamba, có thể chia sẻ thư viện trạng thái được tạo ra từ dữ liệu chuyên môn
• Trạng thái cho phép áp dụng việc học ngữ cảnh vô hạn trong thời gian suy luận mà không cần backprop

Mamba và khả năng diễn giải cơ chế

• Khả năng diễn giải của Mamba tập trung vào việc hiểu cách thông tin di chuyển giữa các token

Mamba và SSM sẽ làm gì tiếp theo

• Các mô hình như Mamba có khả năng thể hiện xuất sắc trong các kịch bản cần ngữ cảnh rất dài và bộ nhớ dài hạn

Tác nhân và an toàn AI

• Các mô hình ngôn ngữ về bản chất là an toàn, nhưng khả năng suy luận chuỗi dài hạn làm sống lại tầm quan trọng của an toàn AI dựa trên tác nhân

Sự phối hợp tốt nhất giữa Transformer và Mamba

• Việc kết hợp ngữ cảnh dài của Mamba với độ phân giải cao của Transformer trên các chuỗi ngắn là điều có giá trị

Ý kiến của GN⁺

• Mamba giải quyết nút thắt cổ chai của Transformer và đưa ra một phương án thay thế hiệu quả cho xử lý chuỗi dài
• Công nghệ này có thể đặc biệt hữu ích trong các lĩnh vực mà chuỗi dữ liệu dài là quan trọng, như y tế, di truyền học và xử lý ngôn ngữ tự nhiên
• Cần thêm nghiên cứu để xác minh liệu cơ chế chọn lọc của Mamba có thực sự hiệu quả hay không
• Tính chọn lọc của Mamba có thể giúp tìm ra điểm cân bằng giữa độ chính xác cao và hiệu suất mà Transformer mang lại