Mamba vs Transformer: Khi Nào Nên Dùng State Space Model Thay Vì Attention?

Tóm tắt

• Transformer sử dụng Self-Attention với độ phức tạp O(L^2), mang lại khả năng suy luận mạnh nhưng tốn kém khi xử lý chuỗi dài.
• Mamba ứng dụng Selective State Space Model với thuật toán selective scan giúp giảm độ phức tạp xuống O(L) tuyến tính.
• Mamba lưu trữ trạng thái ẩn kích thước cố định, giải quyết bài toán bùng nổ KV-cache trong quá trình sinh văn bản.
• Tiêu chí lựa chọn phụ thuộc vào độ dài chuỗi, tài nguyên phần cứng và yêu cầu về khả năng suy luận logic phức tạp.

Sự thống trị của kiến trúc Transformer trong lĩnh vực trí tuệ nhân tạo và xử lý ngôn ngữ tự nhiên đã kéo dài nhiều năm qua. Nhờ có cơ chế Self-Attention, các mô hình lớn như GPT hay BERT có thể kết nối bất kỳ từ nào trong câu với mọi từ khác một cách trực tiếp. Tuy nhiên, cái giá phải trả cho khả năng này không hề nhỏ: chi phí tính toán tăng theo bình phương chiều dài chuỗi, khiến việc xử lý văn bản cực dài hoặc video độ phân giải cao trở nên quá tải đối với phần cứng. Trong bối cảnh đó, các mô hình không gian trạng thái (State Space Models - SSM), tiêu biểu là kiến trúc Mamba, nổi lên như một đối trọng đáng gờm [1]. Bài viết này sẽ phân tích bản chất toán học, cơ chế truyền thông tin và đưa ra các tiêu chí thực tế giúp bạn lựa chọn giữa Mamba và Transformer trong các dự án AI.

Bài toán mô hình hóa chuỗi và giới hạn của Attention

Để hiểu vì sao cộng đồng AI lại quan tâm đến Mamba, trước tiên chúng ta phải nhìn thẳng vào điểm yếu cốt lõi của Transformer. Xét một chuỗi đầu vào gồm $L$ token. Trong một lớp Self-Attention chuẩn, mô hình phải tính toán ma trận tương quan giữa tất cả các cặp token với nhau. Số lượng phép tính và bộ nhớ trung gian tiêu tốn cho quá trình này tuân theo độ phức tạp thời gian bậc hai:

$$𝒪(L^2)$$

Điều này có nghĩa là khi độ dài ngữ cảnh tăng gấp đôi, chi phí tính toán sẽ tăng lên gấp bốn lần. Dù các kỹ thuật như FlashAttention đã giúp tối ưu hóa đáng kể tốc độ thực thi và giảm thiểu lượng truy cập bộ nhớ băng thông rộng, giới hạn toán học cơ bản này vẫn tồn tại. Khi ứng dụng chuyển dịch sang các tác vụ đòi hỏi xử lý tài liệu dài hàng triệu token, sách giáo khoa, mã nguồn phần mềm lớn hoặc chuỗi gene sinh học, độ phức tạp $𝒪(L^2)$ nhanh chóng chạm ngưỡng bão hòa phần cứng.

Ngoài ra, cơ chế lưu trữ bộ nhớ đệm (KV-cache) trong quá trình sinh văn bản (autoregressive generation) của Transformer cũng phình to theo số lượng token đã sinh ra, gây áp lực rất lớn lên bộ nhớ VRAM của GPU. Từ góc độ cấu trúc chuỗi, câu hỏi đặt ra là liệu chúng ta có thể truyền tải thông tin từ quá khứ đến hiện tại mà không cần phải so sánh mọi token với nhau hay không?

State Space Model (SSM) và sự tiến hóa đến Mamba

Các mô hình không gian trạng thái bắt nguồn từ lý thuyết điều khiển tự động cổ điển, nơi các hệ thống động lực học liên tục được mô tả thông qua các phương trình vi phân chuyển dịch trạng thái ẩn. Khi được chuyển hóa sang lĩnh vực học sâu rời rạc hóa, SSM đóng vai trò là một hàm ánh xạ chuỗi đầu vào $x(t)\in {ℝ}^D$ thành chuỗi đầu ra $y(t)\in {ℝ}^D$ thông qua một trạng thái ẩn tiềm ẩn $h(t)\in {ℝ}^N$.

Trong các phiên bản SSM truyền thống như S4, hệ thống sử dụng một tập các ma trận tham số cố định $A,B,C$ để cập nhật trạng thái ẩn dọc theo thời gian hoặc vị trí token. Tuy nhiên, điểm hạn chế lớn của các mô hình SSM tuyến tính thuần túy là chúng không có khả năng "chọn lọc" nội dung. Mọi thông tin đi qua đều bị ép vào trạng thái ẩn với trọng số bất biến, khiến mô hình khó lòng tập trung vào các chi tiết quan trọng hoặc bỏ qua các từ nhiễu.

Mamba giải quyết triệt để điểm yếu này bằng cơ chế Selective SSM. Điểm đột phá nằm ở chỗ các tham số kiểm soát quá trình chuyển đổi trạng thái ($B,C$ và bước thời gian $\Delta$) không còn tĩnh nữa, mà trở thành hàm phụ thuộc trực tiếp vào đầu vào hiện tại. Nhờ có tính chọn lọc này, Mamba có thể tự quyết định xem thông tin nào cần được ghi nhớ sâu vào trạng thái ẩn, và thông tin nào nên bị lãng quên ngay lập tức.

Cơ chế hoạt động của Selective Scan trong Mamba

Sự khác biệt lớn nhất giữa Mamba và các kiến trúc trước đó nằm ở thuật toán Selective Scan. Hãy tưởng tượng quá trình duyệt qua một văn bản dài giống như việc đọc sách: thay vì ghi nhớ mọi từ với mức độ quan trọng như nhau, não bộ con người tự động đánh dấu các từ khóa chính và bỏ qua các mạo từ không quan trọng.

Luồng xử lý chuỗi: Mamba (SSM) so với Transformer (Attention)

Về mặt toán học, khi tham số $B$ và $C$ phụ thuộc vào đầu vào $x$, phép toán tích chập toàn cục quen thuộc trong các mô hình SSM tuyến tính không còn áp dụng được nữa vì tính chất bất biến theo thời gian (time-invariance) đã bị phá vỡ. Để khắc phục điều này mà vẫn đảm bảo tốc độ tính toán cao trên phần cứng hiện đại, Mamba sử dụng thuật toán quét chọn lọc tối ưu hóa phần cứng (hardware-aware selective scan algorithm) .

Thuật toán này tính toán trạng thái ẩn tuần tự về mặt logic nhưng được thiết kế khéo léo để tận dụng bộ nhớ SRAM cực nhanh bên trong nhân GPU thay vì liên tục ghi đọc trên băng thông DRAM chậm chạp. Kết quả là Mamba đạt được độ phức tạp tính toán tuyến tính theo chiều dài chuỗi:

$$𝒪(L)$$

Khi kích thước mô hình được giữ cố định, chi phí tính toán tăng trưởng tuyến tính, mở ra khả năng mở rộng cửa sổ ngữ cảnh lên tới hàng triệu token mà không gặp hiện tượng bùng nổ tài nguyên như Self-Attention.

So sánh chi tiết: Mamba vs Transformer

Để đưa ra quyết định kiến trúc chính xác, chúng ta cần đặt Mamba và Transformer lên bàn cân thông qua các tiêu chí kỹ thuật cốt lõi.

Tiêu chí so sánh	Transformer (Self-Attention)	Mamba (Selective SSM)
Độ phức tạp tính toán (Training)	$𝒪(L^2)$	$𝒪(L)$
Tốc độ sinh (Inference Generation)	Chậm hơn do phải truy xuất toàn bộ KV-cache cũ	Cực kỳ nhanh và cố định bộ nhớ trạng thái
Khả năng bắt quan hệ toàn cục	Hoàn hảo, mọi token tương tác trực tiếp với nhau	Tốt, thông tin nén qua trạng thái ẩn liên tục
Khả năng chọn lọc nội dung	Dựa trên trọng số attention score mềm	Dựa trên ma trận trạng thái thay đổi theo đầu vào
Độ sẵn có của hệ sinh thái	Rất lớn, tối ưu hóa sâu trên mọi phần cứng	Đang phát triển mạnh, cần các kernel tùy biến

Một điểm cần lưu ý đặc biệt là khả năng lưu trữ trạng thái khi sinh văn bản. Trong Transformer, khi câu trả lời dài lên, kích thước của KV-cache tăng lên tuyến tính theo số token được sinh, đòi hỏi quản lý bộ nhớ phức tạp (ví dụ như cơ chế PagedAttention). Ngược lại, Mamba duy trì một trạng thái ẩn có kích thước cố định bất kể chuỗi đầu ra dài đến đâu. Điều này làm cho việc triển khai inference của Mamba trở nên vô cùng nhẹ nhàng và tiết kiệm bộ nhớ.

Khi nào nên dùng State Space Model thay vì Attention?

Việc lựa chọn giữa Mamba và Transformer phụ thuộc hoàn toàn vào đặc thù bài toán và bản chất của dữ liệu đầu vào. Không có kiến trúc nào là tối ưu tuyệt đối cho mọi tình huống.

Trường hợp nên ưu tiên Mamba

1. Xử lý chuỗi cực dài (Ultra-long Sequences): Nếu ứng dụng của bạn yêu cầu phân tích các tài liệu văn bản dài hàng chục ngàn từ, mã nguồn phần mềm nguyên vẹn của một dự án lớn, hoặc dữ liệu chuỗi thời gian (time-series) với tần số lấy mẫu cao, độ phức tạp tuyến tính $𝒪(L)$ của Mamba sẽ vượt trội hoàn toàn về tốc độ và chi phí phần cứng.
2. Hệ thống biên (Edge Devices) có tài nguyên hạn chế: Khi triển khai mô hình ngôn ngữ hoặc mô hình chuỗi trên các thiết bị phần cứng có dung lượng VRAM khiêm tốn, việc không phải duy trì KV-cache dung lượng lớn giúp Mamba hoạt động mượt mà hơn.
3. Tác vụ dòng dữ liệu liên tục (Streaming Data): Các bài toán dự báo thời tiết, phân tích tín hiệu âm thanh hoặc cảm biến IoT đòi hỏi mô hình cập nhật trạng thái liên tục theo thời gian thực mà không bị giới hạn bởi một cửa sổ ngữ cảnh cố định.

Trường hợp nên giữ lại Transformer

1. Tác vụ suy luận phức tạp và tổng hợp tri thức cao: Self-Attention thể hiện ưu thế tuyệt vời trong các bài toán đòi hỏi mô hình phải thực hiện các bước suy luận logic nhiều chiều, kết nối các sự kiện rời rạc nằm ở hai đầu xa nhau của văn bản mà không bị suy giảm thông tin qua lớp nén trạng thái ẩn.
2. Hệ sinh thái và độ ổn định: Transformer hiện là tiêu chuẩn vàng với hàng ngàn mô hình pretrained sẵn có, các công cụ tinh chỉnh (fine-tuning) hoàn thiện, và khả năng tương thích cao với mọi framework học sâu hiện hành.
3. Các bài toán RAG và hỏi đáp tài liệu chuẩn: Mặc dù Mamba đang thâm nhập mạnh mẽ vào lĩnh vực xử lý tài liệu dài, các hệ sinh thái Retrieval-Augmented Generation hiện đại vẫn được tối ưu hóa sâu cho embedding và attention dựa trên Transformer.

Tương lai của kiến trúc kết hợp (Hybrid Models)

Sự phân chia rạch ròi giữa Mamba và Transformer đang dần mờ đi trong các nghiên cứu gần đây. Thay vì đối đầu nhau, xu hướng hiện tại của giới nghiên cứu AI là kết hợp ưu điểm của cả hai kiến trúc. Các mô hình lai (Hybrid Architectures) thường đan xen các lớp Selective SSM với các lớp Self-Attention trong cùng một mạng nơ-ron.

Cách tiếp cận này cho phép tầng SSM đảm nhận nhiệm vụ quét qua toàn bộ chuỗi dài với chi phí $𝒪(L)$ để lọc và nén thông tin thô, trong khi các lớp Attention được đặt ở những vị trí chiến lược để thực hiện các phép so sánh ngữ nghĩa phức tạp. Sự kết hợp này hứa hẹn sẽ mang lại các thế hệ mô hình vừa có tốc độ xử lý nhanh của Mamba vừa giữ được khả năng suy luận sắc bén của Transformer.

Kết luận

Mamba và Transformer đại diện cho hai triết lý truyền thông tin khác biệt trong học sâu: một bên dựa trên không gian trạng thái nén tuyến tính theo thời gian, bên còn lại dựa trên mạng lưới tương tác toàn cục bậc hai. Transformer vẫn là lựa chọn an toàn và mạnh mẽ cho phần lớn các bài toán NLP tiêu chuẩn nhờ hệ sinh thái chín muồi. Tuy nhiên, khi đối mặt với các bài toán chuỗi siêu dài và yêu cầu tối ưu hóa tài nguyên phần cứng nghiêm ngặt, State Space Model như Mamba chính là giải pháp thay thế xứng đáng mở ra kỷ nguyên mới cho hiệu năng tính toán trong trí tuệ nhân tạo.

Giới thiệu Mamba: State Space Model cho chuỗi dài

LSTM là gì? Hiểu rõ kiến trúc và ứng dụng trong NLP