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

Tóm tắt nội dung

• Mamba là kiến trúc mô hình hóa chuỗi dựa trên mô hình không gian trạng thái chọn lọc, không dựa vào self-attention.
• Điểm quan trọng nhất là các tham số kiểm soát trạng thái có thể thay đổi theo đầu vào, giúp mô hình học cách giữ hoặc loại bỏ thông tin.
• Selective scan cho phép Mamba xử lý chuỗi với chi phí tăng tuyến tính theo chiều dài chuỗi khi kích thước mô hình được giữ cố định.
• Mamba không đơn giản là một Transformer nhanh hơn; hai kiến trúc lưu trữ và truyền thông tin quá khứ theo những cách fundamentally khác nhau.

Mamba là một kiến trúc mô hình hóa chuỗi dựa trên mô hình không gian trạng thái (State Space Model, SSM), được Albert Gu và Tri Dao giới thiệu trong công trình Mamba: Linear-Time Sequence Modeling with Selective State Spaces. Khác với Transformer, Mamba không sử dụng self-attention làm cơ chế cốt lõi mà truyền thông tin qua một trạng thái ẩn được cập nhật dọc theo chuỗi [1].

Điểm đột phá không chỉ nằm ở việc dùng SSM. Các mô hình SSM có cấu trúc đã tồn tại trước Mamba, nổi bật là S4 [2]. Đóng góp quan trọng của Mamba là làm cho một số tham số của SSM phụ thuộc vào dữ liệu đầu vào, từ đó mô hình có thể chọn thông tin cần ghi nhớ và thông tin nên bỏ qua.

Trong paper gốc, Mamba là tên của kiến trúc, không được định nghĩa như một từ viết tắt. Vì vậy, cách hiểu đúng về Mamba nên bắt đầu từ SSM, cơ chế chọn lọc và selective scan thay vì cố gắng ánh xạ nó sang một pipeline retrieval hoặc một biến thể của RAG.

Mamba giải quyết vấn đề gì?

Transformer [4] có khả năng mô hình hóa quan hệ giữa các token rất mạnh nhờ self-attention. Với một lớp attention chuẩn, mỗi vị trí có thể tính mức liên quan với các vị trí khác trong cửa sổ ngữ cảnh. Đổi lại, số phép tính của phần attention tăng theo bậc hai của chiều dài chuỗi L nếu các kích thước còn lại được giữ cố định.

Attention chuẩn: O(L²)

Selective SSM: O(L) khi kích thước trạng thái cố định

Các kỹ thuật như FlashAttention có thể giảm mạnh lượng bộ nhớ trung gian và cải thiện hiệu suất thực tế của attention, nhưng không thay đổi bản chất rằng số tương tác query-key của full attention tăng theo L². Trong suy luận tự hồi quy, Transformer còn thường phải duy trì bộ nhớ đệm khóa-giá trị (KV cache) cho các token trước đó.

SSM tiếp cận bài toán theo hướng khác. Thay vì giữ biểu diễn của từng token quá khứ để truy cập trực tiếp, nó liên tục nén lịch sử vào một trạng thái ẩn. Mỗi token mới cập nhật trạng thái này và đầu ra được tính từ trạng thái hiện tại.

Ý tưởng này cho phép chi phí tính toán tăng tuyến tính theo chiều dài chuỗi. Tuy nhiên, các SSM tuyến tính bất biến theo thời gian trước Mamba gặp khó khăn với dữ liệu rời rạc giàu thông tin như ngôn ngữ vì cùng một quy luật động lực được áp dụng cho mọi token. Mamba giải quyết chính điểm yếu đó bằng selection mechanism [1]. 

Từ State Space Model đến Selective State Space Model

Một SSM cổ điển mô tả cách trạng thái của một hệ thống thay đổi theo thời gian. Với đầu vào x(t), trạng thái ẩn h(t) và đầu ra y(t), mô hình liên tục cơ bản có dạng:

h′(t) = Ah(t) + Bx(t)

y(t) = Ch(t)

Ma trận A quyết định trạng thái cũ tiến hóa như thế nào. B điều khiển cách đầu vào mới đi vào trạng thái, còn C xác định phần nào của trạng thái được đọc ra thành đầu ra. Đây là mô hình thời gian liên tục, trong khi NLP làm việc với một chuỗi token rời rạc.

Vì vậy, SSM phải được rời rạc hóa (discretization). Với bước thời gian Δ và quy tắc giữ bậc không (zero-order hold), paper Mamba trình bày:

Ā = exp(ΔA)

B̄ = (ΔA)⁻¹[exp(ΔA) − I]ΔB

Sau khi rời rạc hóa, quá trình xử lý chuỗi có thể viết thành một recurrence:

h_t = Āh_t−1 + B̄x_t

y_t = Ch_t

Trong các SSM tuyến tính bất biến theo thời gian (Linear Time-Invariant, LTI), A, B, C và Δ về cơ bản không thay đổi theo từng vị trí trong chuỗi. Do đó cùng một hệ động lực được sử dụng dù token hiện tại là một từ khóa quan trọng, một dấu câu hay một token ít thông tin. Đây là một hạn chế đáng kể khi mô hình hóa ngôn ngữ [1].

Selection: để dữ liệu quyết định trạng thái nên thay đổi ra sao

Mamba làm cho các tham số B, C và Δ phụ thuộc vào đầu vào. Có thể biểu diễn trực giác của cơ chế này như sau:

B_t = s_B(x_t)

C_t = s_C(x_t)

Δ_t = softplus(θ_Δ + s_Δ(x_t))

Ma trận A vẫn là tham số học được của mô hình, trong khi B_t, C_t và Δ_t thay đổi theo token. Điều này biến hệ thống thành một SSM biến thiên theo dữ liệu thay vì một bộ lọc cố định.

Trực giác của Δ đặc biệt hữu ích. Trong phân tích của Gu và Dao, Δ có quan hệ gần với cơ chế cổng của RNN: giá trị nhỏ có thể khiến trạng thái cũ tồn tại lâu hơn, còn giá trị lớn có thể làm hệ thống tập trung mạnh hơn vào đầu vào hiện tại và làm suy giảm lịch sử trước đó [1].

Hình 1. Luồng dữ liệu của Selective SSM. Sơ đồ được tái dựng và giản lược dựa trên cơ chế trong Fig. 1 của Gu và Dao [1].

Một Mamba block xử lý dữ liệu như thế nào?

Selective SSM chỉ là phép biến đổi chuỗi cốt lõi. Một Mamba block hoàn chỉnh còn kết hợp các phép chiếu tuyến tính, tích chập một chiều cục bộ, hàm kích hoạt và cơ chế cổng. Paper gốc xây dựng block này bằng cách hợp nhất ý tưởng từ kiến trúc SSM trước đó với cấu trúc gated MLP, sau đó xếp chồng cùng normalization và residual connection [1].

Ở mức trực giác, biểu diễn đầu vào được chiếu thành hai nhánh. Nhánh chính đi qua causal Conv1D để trộn thông tin cục bộ, sau đó qua Selective SSM để truyền thông tin dọc chuỗi. Nhánh còn lại tạo tín hiệu cổng. Hai nhánh được kết hợp theo từng phần tử trước phép chiếu đầu ra.

Hình 2. Kiến trúc Mamba block ở mức khái niệm, tái dựng và giản lược dựa trên Fig. 3 của paper Mamba [1].

Causal Conv1D không thay thế Selective SSM. Nó cung cấp một phép trộn cục bộ ngắn hạn trước khi nhánh SSM xử lý quan hệ qua khoảng cách dài hơn. Selective SSM sau đó quyết định cách nội dung hiện tại tác động lên trạng thái tích lũy.

Cũng cần phân biệt giữa SSM layer và mô hình Mamba hoàn chỉnh. Một SSM chỉ định nghĩa phép biến đổi trạng thái; Mamba là mạng sâu gồm nhiều block, phép chiếu, normalization và residual connection. Việc nói “Mamba là một SSM” hữu ích ở mức khái quát, nhưng không có nghĩa toàn bộ mạng chỉ gồm phương trình recurrence.

Selective scan giúp Mamba xử lý chuỗi dài ra sao?

Với SSM bất biến theo thời gian, recurrence có thể được biến đổi thành convolution, tạo ra cách huấn luyện song song hiệu quả. Khi B, C và Δ trở thành hàm của đầu vào trong Mamba, kernel convolution không còn cố định. Đây là lý do Mamba cần một thuật toán tính toán khác.

Giải pháp là selective scan: recurrence được tổ chức thành một phép scan song song phù hợp với phần cứng GPU. Thiết kế của Gu và Dao tránh tạo toàn bộ tensor trạng thái mở rộng trong bộ nhớ băng thông cao của GPU; các bước rời rạc hóa và recurrence được hợp nhất, còn trạng thái trung gian cần cho lan truyền ngược có thể được tính lại thay vì lưu tất cả [4].

Về độ phức tạp, với kích thước mô hình D và kích thước trạng thái N được giữ cố định, mỗi vị trí trong chuỗi chỉ cần một lượng công việc hữu hạn để cập nhật trạng thái. Vì vậy tổng chi phí của phép scan tăng tuyến tính theo chiều dài L.

T_Mamba(L) ∝ L

Trong suy luận tự hồi quy, trạng thái từ bước trước có thể được tái sử dụng để xử lý token tiếp theo. Mô hình không cần lưu KV cache của toàn bộ các token đã đi qua. Vì kích thước trạng thái cố định khi kiến trúc được cố định, lượng trạng thái cần duy trì cho mỗi layer không tăng theo chiều dài lịch sử [1].

Không nên hiểu “O(L)” là “luôn nhanh hơn Transformer”. Tốc độ thực tế còn phụ thuộc kích thước mô hình, batch size, chiều dài chuỗi, precision, GPU, kernel và implementation. Mamba có cấu trúc thuận lợi cho chuỗi dài, nhưng kết luận về latency hoặc throughput phải được đo trên chính workload cần triển khai.

Ví dụ PyTorch: một Selective SSM tối giản

Để thấy selection hoạt động cụ thể, ta có thể xây một lớp nhỏ bằng PyTorch. Ví dụ dưới đây không phải implementation Mamba chính thức; nó chỉ mô phỏng phần lõi của Selective SSM với A dạng đường chéo, còn B, C và Δ được tạo từ từng input.

import torch
import torch.nn as nn
import torch.nn.functional as F

class TinySelectiveSSM(nn.Module): def **init**(self, d_model, d_state): super().**init**()

```
    # A có dạng [D, N] và được giữ âm để minh họa động lực ổn định
    self.A_log = nn.Parameter(torch.zeros(d_model, d_state))

    # B_t, C_t và delta_t được sinh từ input hiện tại
    self.to_B = nn.Linear(d_model, d_state, bias=False)
    self.to_C = nn.Linear(d_model, d_state, bias=False)
    self.to_delta = nn.Linear(d_model, d_model)

def forward(self, x):
    batch, length, d_model = x.shape

    A = -torch.exp(self.A_log)
    h = x.new_zeros(batch, d_model, A.shape[1])
    outputs = []

    for t in range(length):
        x_t = x[:, t]

        B_t = self.to_B(x_t)
        C_t = self.to_C(x_t)
        delta_t = F.softplus(self.to_delta(x_t))

        # Zero-order hold với A dạng đường chéo
        delta_A = delta_t.unsqueeze(-1) * A.unsqueeze(0)
        A_bar = torch.exp(delta_A)

        B_bar = torch.expm1(delta_A) / A.unsqueeze(0)
        B_bar = B_bar * B_t.unsqueeze(1)

        # Cập nhật trạng thái rồi đọc output
        h = A_bar * h + B_bar * x_t.unsqueeze(-1)
        y_t = (h * C_t.unsqueeze(1)).sum(dim=-1)
        outputs.append(y_t)

    return torch.stack(outputs, dim=1)
```

torch.manual_seed(0)

x = torch.randn(2, 6, 8)
layer = TinySelectiveSSM(d_model=8, d_state=4)
y = layer(x)

print(y.shape) 

Kết quả là torch.Size([2, 6, 8]): hai chuỗi, mỗi chuỗi sáu bước và mỗi bước trả lại tám đặc trưng. Quan trọng hơn kích thước output là việc B_t, C_t và delta_t được tính lại tại từng bước dựa trên x_t.

Vòng lặp Python trong ví dụ cố tình viết tuần tự để dễ đọc và không đại diện cho hiệu suất của Mamba thực tế. Implementation Mamba sử dụng selective scan song song và các kernel hợp nhất nhằm khai thác GPU hiệu quả. Đây là một ví dụ điển hình cho sự khác biệt giữa phương trình toán học và implementation hiệu năng cao.

Mamba khác Transformer ở điểm nào?

Khác biệt quan trọng nhất không phải “có attention hay không” mà là cách quá khứ được lưu và truy cập. Transformer giữ các biểu diễn khóa và giá trị của những token trước để token hiện tại có thể truy cập chúng thông qua attention. Mamba liên tục nén lịch sử vào trạng thái hồi quy.

Khía cạnh	Transformer với full attention	Mamba
Cơ chế chính	Self-attention giữa các token.	Selective State Space Model.
Thông tin quá khứ	Có thể truy cập trực tiếp các biểu diễn token trong cửa sổ ngữ cảnh.	Được nén dần vào trạng thái ẩn.
Chi phí theo L	Full attention có số phép tính bậc hai theo L.	Selective scan tăng tuyến tính theo L khi D và N cố định.
Suy luận tự hồi quy	Thường sử dụng KV cache tăng theo số token đã lưu.	Duy trì trạng thái có kích thước cố định theo kiến trúc.
Truy xuất token quá khứ	Attention tạo đường truyền trực tiếp giữa các vị trí.	Thông tin phải được bảo toàn qua chuỗi cập nhật trạng thái.

Khi nào Mamba đáng cân nhắc?

Mamba đặc biệt đáng nghiên cứu khi chiều dài chuỗi lớn và chi phí của full attention trở thành nút thắt, hoặc khi hệ thống sinh tự hồi quy cần tránh KV cache tăng dần theo chiều dài ngữ cảnh. Thiết kế SSM cũng không phụ thuộc riêng vào văn bản: paper gốc đánh giá Mamba trên language modeling, audio và genomics [1].

Tuy nhiên, không nên chọn Mamba chỉ dựa trên ký hiệu O(L). Attention có một ưu điểm khái niệm rất mạnh: một token có thể truy cập trực tiếp một token khác trong context. Mamba phải truyền thông tin đó qua trạng thái nén; selection giúp quá trình nén linh hoạt hơn nhưng không làm biến mất giới hạn dung lượng của trạng thái.

Mamba không phải điểm kết thúc của hướng nghiên cứu SSM

Mamba-2 tiếp tục phát triển ý tưởng này thông qua đối ngẫu không gian trạng thái (State Space Duality, SSD). Dao và Gu chỉ ra các mối liên hệ toán học sâu hơn giữa structured SSM và một số dạng attention, đồng thời thiết kế một lớp Mamba-2 đơn giản hơn; paper báo cáo phần lõi nhanh hơn Mamba đời đầu trong các cấu hình tác giả đánh giá [3].

Điểm này cũng cho thấy không nên nhìn Mamba và Transformer như hai ý tưởng hoàn toàn không liên quan. Chúng sử dụng những phép tính khác nhau trong kiến trúc thực tế, nhưng nghiên cứu về SSD cho thấy có thể mô tả nhiều cơ chế sequence modeling trong các khung toán học có quan hệ chặt chẽ với nhau.

Kết luận

Mamba là một hướng tiếp cận sequence modeling dựa trên Selective State Space Model. Thay vì lưu toàn bộ quá khứ để attention truy cập trực tiếp, mô hình học cách cập nhật một trạng thái nén; B, C và Δ phụ thuộc vào đầu vào giúp quá trình này có khả năng chọn lọc thông tin.

Ba ý cần nắm vững trước khi đi sâu vào Mamba là: SSM biểu diễn lịch sử bằng trạng thái, selection làm động lực trạng thái phụ thuộc dữ liệu, và selective scan biến recurrence này thành phép tính hiệu quả trên phần cứng hiện đại. Khi đã hiểu ba thành phần đó, kiến trúc Mamba block, ưu thế trên chuỗi dài và sự khác biệt với Transformer trở nên trực quan hơn nhiều.

Tài liệu tham khảo

1. A. Gu and T. Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces,” in Proc. Conference on Language Modeling (COLM), 2024. 
2. A. Gu, K. Goel, and C. Ré, “Efficiently Modeling Long Sequences with Structured State Spaces,” in Proc. International Conference on Learning Representations (ICLR), 2022. 
3. T. Dao and A. Gu, “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality,” in Proc. 41st International Conference on Machine Learning, PMLR, vol. 235, pp. 10041–10071, 2024. 
4. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention Is All You Need,” in Advances in Neural Information Processing Systems, vol. 30, pp. 5998–6008, 2017.