Selective Scan trong Mamba tối ưu chuỗi dài thế nào?

Tóm tắt nội dung

• Selection và Selective Scan là hai khái niệm khác nhau: selection làm các tham số SSM phụ thuộc đầu vào, còn scan là thuật toán tính recurrence đó hiệu quả.
• Khi B, C và Δ thay đổi theo từng token, Mamba không còn sử dụng được một convolution kernel cố định như SSM bất biến theo thời gian.
• Selective Scan khai thác tính kết hợp của phép cập nhật trạng thái để song song hóa recurrence, đồng thời dùng kernel fusion và phân cấp bộ nhớ GPU để tránh materialize tensor trạng thái rất lớn trong HBM.
• Chi phí của Selective SSM tăng tuyến tính theo chiều dài L khi kích thước mô hình và state được giữ cố định, nhưng O(L) không đồng nghĩa Mamba luôn nhanh hơn mọi Transformer trong mọi workload.

Selective Scan là mắt xích kỹ thuật giúp kiến trúc Mamba giải quyết một nghịch lý quan trọng: làm thế nào để State Space Model có thể thay đổi động lực theo nội dung của từng token nhưng vẫn xử lý chuỗi dài hiệu quả trên GPU. Nếu chỉ thêm cơ chế chọn lọc mà tính recurrence bằng một vòng lặp tuần tự thông thường, Mamba sẽ mất phần lớn lợi thế về throughput khi huấn luyện.

Giải pháp của Gu và Dao là kết hợp recurrence với một thuật toán scan nhận biết phần cứng (hardware-aware scan), trong đó phép tính được song song hóa, các bước trung gian được hợp nhất và trạng thái mở rộng không bị ghi liên tục ra bộ nhớ HBM [1]. Để hiểu vì sao thiết kế này quan trọng, cần tách rõ ba vấn đề: selection làm thay đổi toán học của SSM ra sao, scan song song hóa recurrence như thế nào và cách triển khai giảm chi phí truyền dữ liệu trên GPU.

Nếu chưa quen với State Space Model và kiến trúc tổng thể của Mamba, có thể đọc trước Giới thiệu Mamba: State Space Model cho chuỗi dài. Bài hiện tại chỉ tập trung vào Selective Scan và cơ chế tối ưu chuỗi dài.

Tại sao Mamba cần Selective Scan?

Một State Space Model rời rạc cơ bản cập nhật trạng thái theo recurrence:

h_t = Āh_t−1 + B̄x_t

y_t = Ch_t

Ở đây h_t là trạng thái ẩn, x_t là input tại vị trí t, còn Ā, B̄ và C kiểm soát quá trình duy trì, ghi và đọc thông tin từ state. Trong các Structured State Space Model như S4, các tham số động lực này không thay đổi theo từng token, tức hệ thống có tính bất biến theo thời gian (Linear Time-Invariant, LTI) [2].

Tính LTI cho phép mở recurrence thành một convolution. Giả sử trạng thái ban đầu bằng 0:

y_t = Σ_j=0^t CĀ^t−jB̄x_j

Hệ số tác động của x_j lên y_t chỉ phụ thuộc vào khoảng cách t − j. Ta có thể tạo một kernel cố định:

K_k = CĀ^kB̄

y = x * K

Đây là một lý do quan trọng khiến các SSM trước Mamba có thể huấn luyện hiệu quả: thay vì chạy recurrence tuần tự qua L bước, chúng có thể chuyển sang chế độ convolution và tận dụng các thuật toán song song hiệu quả [2].

Nhưng sự bất biến này cũng tạo ra giới hạn. Cùng một động lực được áp dụng cho mọi vị trí, nên mô hình khó thay đổi cách ghi nhớ dựa trên nội dung token. Với ngôn ngữ, một tên riêng, một dấu phân cách, một con số quan trọng và một từ ít thông tin không nên nhất thiết tác động lên memory theo cùng một quy luật.

Selection phá vỡ convolution kernel cố định

Mamba làm cho B, C và Δ trở thành hàm của input. Viết giản lược:

B_t = s_B(x_t)

C_t = s_C(x_t)

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

Sau khi rời rạc hóa, Ā và B̄ cũng thay đổi theo thời điểm:

h_t = Ā_th_t−1 + B̄_tx_t

y_t = C_th_t

Lúc này đóng góp của x_j vào một output tương lai không còn chứa lũy thừa của một ma trận Ā cố định. Nó có dạng:

C_t(Ā_tĀ_t−1 … Ā_j+1)B̄_jx_j

Hệ số này phụ thuộc vào cả chuỗi tham số nằm giữa vị trí j và t. Vì các tham số lại phụ thuộc input, không còn một kernel K cố định chỉ dựa trên khoảng cách. Theo Algorithm 2 của paper Mamba, Selective SSM vì vậy được tính ở dạng recurrence bằng scan thay vì global convolution [1].

Hình 1. Selection làm SSM chuyển từ time-invariant sang time-varying, khiến convolution kernel cố định không còn áp dụng được. Sơ đồ tái dựng dựa trên cơ chế trong [1, Fig. 1 và Algorithm 2].

Scan biến recurrence thành phép toán có thể song song hóa ra sao?

Nhìn trực tiếp vào recurrence, ta dễ kết luận rằng bước t phải chờ h_t−1, nên toàn bộ chuỗi bắt buộc chạy tuần tự. Về mặt phụ thuộc logic điều này đúng, nhưng không có nghĩa cách đánh giá toán học duy nhất là một vòng lặp tuần tự.

Ta viết recurrence ở dạng affine đơn giản:

h_t = A_th_t−1 + b_t

Trong Selective SSM, b_t có thể hiểu là phần input sau khi đi qua hệ số B̄_t. Mỗi bước t vì vậy định nghĩa một phép biến đổi affine:

T_t(h) = A_th + b_t

Xét hai bước liên tiếp:

T₂(T₁(h)) = A₂A₁h + A₂b₁ + b₂

Do đó hai bước có thể được gộp thành một cặp mới:

(A₂, b₂) ⊗ (A₁, b₁)

= (A₂A₁, A₂b₁ + b₂)

Phép hợp thành các transformation affine này có tính kết hợp (associative). Nghĩa là khi giữ nguyên thứ tự các phần tử, ta có thể nhóm phép tính theo nhiều cách khác nhau mà kết quả cuối cùng không đổi.

Đây là nền tảng của parallel prefix scan. Thay vì tính:

T₁ → T₂ → T₃ → T₄ → T₅ → T₆ → T₇ → T₈

ta có thể ghép nhiều cặp transformation đồng thời, sau đó tiếp tục ghép các kết quả theo cấu trúc cây. Một scan hoàn chỉnh còn tạo ra mọi prefix cần thiết để lấy trạng thái tại từng vị trí, không chỉ trạng thái cuối cùng.

Hình 2. Trực giác phép scan dạng cây. Tính kết hợp cho phép gộp nhiều đoạn recurrence đồng thời; prefix scan thực tế còn phục hồi trạng thái tại từng vị trí trong chuỗi.

Điều quan trọng là Selectivity không bị loại bỏ để đổi lấy song song hóa. Mỗi token vẫn có A_t, B_t, C_t hoặc Δ_t riêng theo thiết kế của Selective SSM. Scan chỉ thay đổi cách tổ chức phép tính.

Điểm dễ hiểu sai: Selective Scan không biến recurrence thành một phép tính không còn phụ thuộc thứ tự. Quan hệ nhân quả giữa các token vẫn tồn tại. Thứ được khai thác là tính kết hợp của phép hợp thành trạng thái, nhờ đó GPU có thể xử lý nhiều đoạn của phép scan song song.

Tại sao chỉ parallel scan vẫn chưa đủ nhanh?

Nếu chỉ viết một parallel scan về mặt toán học nhưng materialize toàn bộ tensor trạng thái ra bộ nhớ GPU, chi phí vẫn rất lớn. Paper Mamba xem đây là nút thắt trung tâm của Selective SSM [1].

Giả sử input có batch size B, chiều dài L, số channel D và state dimension N. Input và output có dạng:

x, y : (B, L, D)

Nếu lưu state của mọi channel tại mọi token, tensor trạng thái có kích thước hiệu dụng:

h : (B, L, D, N)

Tensor này lớn hơn input theo hệ số N. Với chuỗi dài, liên tục ghi state mở rộng xuống High Bandwidth Memory (HBM) rồi đọc lại có thể tốn nhiều thời gian hơn chính các phép nhân và cộng. Đây là ví dụ của một workload bị giới hạn bởi băng thông bộ nhớ thay vì số FLOP thuần túy.

Cách nhìn này tương đồng với nguyên tắc IO-aware đã được nhấn mạnh trong FlashAttention: trên accelerator hiện đại, giảm lượng dữ liệu di chuyển giữa HBM và bộ nhớ nhanh trên chip có thể quan trọng ngang với giảm số phép toán [3]. Mamba áp dụng tư duy tương tự cho recurrence của SSM.

Kernel fusion: không ghi mọi intermediate xuống HBM

Trong một implementation ngây thơ, ta có thể thực hiện nhiều kernel riêng:

projection → Δ,B,C → discretization → scan → output

Sau mỗi bước, tensor trung gian có thể phải được ghi xuống HBM để kernel tiếp theo đọc lại. Selective Scan giảm chuyển động dữ liệu bằng cách hợp nhất nhiều thao tác liên quan. State mở rộng được tạo và sử dụng trong những mức bộ nhớ nhanh hơn của GPU trong quá trình scan thay vì tồn tại dưới dạng tensor B × L × D × N hoàn chỉnh trong HBM [1].

Recomputation: đổi một ít compute lấy ít memory hơn

Lan truyền ngược thường cần các activation trung gian từ forward pass. Một lựa chọn trực tiếp là lưu toàn bộ state, nhưng điều đó lại tạo tensor rất lớn.

Mamba sử dụng recomputation: một số trạng thái trung gian cần cho backward được tính lại từ input và tham số thay vì giữ chúng trong bộ nhớ suốt forward pass [1]. Đây là đánh đổi có chủ đích: thực hiện thêm một số phép tính để giảm memory traffic và dung lượng activation.

Ba thành phần vì vậy bổ sung cho nhau:

• Parallel scan giải quyết tính tuần tự của recurrence.
• Kernel fusion giảm số lần đọc/ghi intermediate qua các mức bộ nhớ.
• Recomputation tránh phải lưu toàn bộ state cho backward.

Chính sự kết hợp này mới là ý nghĩa đầy đủ của cụm từ hardware-aware selective scan. Chỉ nhìn vào ký hiệu O(L) không giải thích được vì sao một implementation có thể nhanh trên GPU.

Selective Scan đạt độ phức tạp tuyến tính theo L như thế nào?

Paper Mamba phân biệt rõ chi phí của recurrence và convolution. Với batch B, chiều dài L, số channel D và state dimension N, recurrence trực tiếp cần lượng công việc theo bậc:

O(BLDN)

Khi B, D và N cố định, biểu thức này tăng tuyến tính theo L. Trong khi đó, cách convolution dựa trên FFT của các SSM trước đó có chi phí xấp xỉ O(BLD log L) theo phân tích trong Mamba [1].

Điều này thoạt nhìn có vẻ lạ: recurrence có thêm N nhưng lại có thể cạnh tranh với convolution. Lý do là N thường là một state dimension tương đối nhỏ và cố định theo kiến trúc, trong khi log L tiếp tục tăng với chiều dài chuỗi. Quan trọng hơn, recurrent scan có thể có hằng số thấp và được tối ưu trực tiếp theo memory hierarchy.

Cách tính	Đặc tính	Chi phí theo L	Hạn chế chính
LTI SSM convolution	Kernel cố định, dễ song song hóa.	Gần tuyến tính, thường chứa log L với FFT.	Không có dynamics phụ thuộc input như Mamba.
Naive selective recurrence	Giữ đầy đủ selection.	O(L) khi D, N cố định.	Vòng lặp tuần tự và materialization có thể rất chậm trên GPU.
Selective Scan	Parallel scan + IO-aware implementation.	O(L) khi các kích thước khác cố định.	Cần kernel chuyên biệt và tối ưu phần cứng.
Full self-attention	Mỗi token có thể tương tác trực tiếp với các token khác.	Phần attention tăng bậc hai theo L.	Chi phí tăng nhanh với context dài.

Bảng trên chỉ so sánh cấu trúc tính toán, không phải benchmark tốc độ thực tế. Kernel attention hiện đại, batch size, dimension, precision và phần cứng có thể làm thay đổi đáng kể wall-clock time. Vì vậy phát biểu chính xác là Selective Scan có scaling tuyến tính theo chiều dài chuỗi trong mô hình tính toán của nó, không phải “Mamba luôn nhanh hơn Transformer”.

Để xem sự khác biệt rộng hơn giữa hai họ kiến trúc, có thể tham khảo bài Mamba vs Transformer: Khi nào nên dùng State Space Model thay vì Attention?.

Ví dụ Python: vì sao recurrence có thể viết dưới dạng scan?

Đoạn code sau không triển khai CUDA Selective Scan của Mamba. Nó chỉ minh họa cấu trúc toán học quan trọng nhất: mỗi bước recurrence là một affine transformation và các transformation này có thể được hợp thành bằng một phép toán kết hợp.

import numpy as np

rng = np.random.default_rng(0)

L = 8
N = 4

x = rng.normal(size=L)

# A âm để tạo dynamics suy giảm ổn định trong ví dụ.
A = -np.exp(np.linspace(-1.0, 0.5, N))

# B_t và delta_t thay đổi theo từng input position.
B_t = rng.normal(size=(L, N))
delta_t = np.log1p(np.exp(rng.normal(size=(L, 1))))

# Zero-order hold cho A dạng đường chéo.
delta_A = delta_t * A[None, :]
A_bar = np.exp(delta_A)

B_bar = np.expm1(delta_A) / A[None, :]
B_bar = B_bar * B_t

# b_t trong h_t = A_t * h_(t-1) + b_t.
b_t = B_bar * x[:, None]

def sequential_recurrence(A_bar, b_t):
h = np.zeros(A_bar.shape[1])
states = []

```
for a, b in zip(A_bar, b_t):
    h = a * h + b
    states.append(h.copy())

return np.stack(states)
```

def compose(prefix, current): # current(prefix(h))
a1, b1 = prefix
a2, b2 = current

```
a = a2 * a1
b = a2 * b1 + b2

return a, b
```

def reference_scan(A_bar, b_t): # Reference prefix scan; vẫn dùng loop để code dễ đọc.
identity = (np.ones(A_bar.shape[1]),
np.zeros(A_bar.shape[1]))

```
prefix = identity
states = []

for a, b in zip(A_bar, b_t):
    prefix = compose(prefix, (a, b))

    # h_0 = 0 nên state prefix chính là phần b.
    states.append(prefix[1].copy())

return np.stack(states)
```

h_seq = sequential_recurrence(A_bar, b_t)
h_scan = reference_scan(A_bar, b_t)

print("Max error:", np.max(np.abs(h_seq - h_scan))) 

Hai cách tính cho cùng kết quả, sai số chỉ đến từ floating-point. Hàm reference_scan vẫn dùng vòng lặp Python để dễ đọc; nó không nhanh hơn recurrence.

Điểm quan trọng nằm ở hàm compose. Vì phép compose có tính kết hợp, một implementation GPU có thể tổ chức các phép kết hợp theo cấu trúc parallel prefix thay vì thực hiện đúng vòng lặp Python trên. Phần khó của Mamba nằm ở việc biến tính chất toán học đó thành kernel hiệu quả trong khi vẫn kiểm soát memory traffic.

Selective Scan thay đổi inference chuỗi dài như thế nào?

Trong training, Selective Scan giúp đánh giá toàn bộ chuỗi song song hiệu quả. Trong autoregressive inference, bài toán lại đơn giản hơn vì token xuất hiện từng bước. Mamba có thể trực tiếp giữ state cuối của mỗi layer và cập nhật nó khi token mới đến:

h_t ← f(h_t−1, x_t)

Không cần chạy lại prefix scan của toàn bộ lịch sử cho mỗi token mới. Mỗi bước chỉ cập nhật state hiện tại. Với kiến trúc cố định, lượng recurrent state cần giữ cho mỗi layer không tăng theo số token đã xử lý.

Đây là khác biệt quan trọng với autoregressive Transformer, nơi KV cache thường tăng theo chiều dài context. Full attention có lợi thế là token mới có thể truy cập trực tiếp biểu diễn của các token cũ; Mamba đổi khả năng truy cập trực tiếp đó lấy một state nén có kích thước cố định.

Vì vậy khả năng xử lý chuỗi dài của Mamba không đến từ việc “giữ toàn bộ chuỗi với chi phí O(L)”. Mô hình nén lịch sử vào state. Selection giúp quá trình nén phụ thuộc nội dung: thông tin nào cần tồn tại lâu, thông tin nào có thể suy giảm nhanh. Selective Scan là cách tính quá trình này hiệu quả, chứ không phải cơ chế tự nó quyết định nội dung nào quan trọng.

Những hiểu nhầm thường gặp về Selective Scan

Selective Scan và selection mechanism không phải một

Selection mechanism là thay đổi mô hình: B, C và Δ trở thành hàm của input. Selective Scan là giải pháp tính toán để chạy SSM time-varying đó hiệu quả. Không có selection, scan vẫn có thể tồn tại về mặt toán học; không có scan tối ưu, selective SSM vẫn có thể chạy nhưng hiệu năng training sẽ kém hơn nhiều.

O(L) không có nghĩa mọi phép tính chỉ tốn L thao tác

Chi phí đầy đủ phụ thuộc B, D và N. O(L) chỉ mô tả scaling theo chiều dài khi những đại lượng còn lại được coi là cố định. Tăng state dimension N hoặc model dimension D vẫn làm chi phí tăng.

Parallel Scan không xóa bỏ quan hệ nhân quả

State tại vị trí t vẫn biểu diễn kết quả tổng hợp từ các vị trí trước. GPU chỉ tận dụng tính kết hợp để thay đổi thứ tự nhóm phép toán. Đây là sự khác biệt giữa dependency của bài toán và execution schedule của thuật toán.

Linear scaling không đảm bảo nhớ hoàn hảo chuỗi cực dài

Khả năng chạy một triệu token về mặt tính toán khác với khả năng giữ chính xác mọi chi tiết trong một triệu token. Mamba nén lịch sử vào state hữu hạn; thông tin có thể bị suy giảm hoặc ghi đè. Paper gốc báo cáo khả năng mở rộng tốt trên các bài toán và modality được đánh giá [1], nhưng không nên diễn giải thành việc mọi nhiệm vụ long-context đều được giải quyết hoàn hảo.

Cách nhớ ngắn gọn: Selection làm Mamba thông minh hơn trong việc ghi và quên; parallel scan làm recurrence có thể khai thác GPU; kernel fusion giảm memory traffic; recomputation giảm activation phải lưu. Selective Scan hiệu quả vì bốn ý tưởng này phối hợp với nhau, không phải chỉ vì một công thức O(L).

Kết luận

Selective Scan giải quyết chính chi phí tính toán phát sinh khi Mamba chuyển từ SSM bất biến sang SSM phụ thuộc input. Selection làm mất khả năng dùng một convolution kernel cố định, nhưng cấu trúc recurrence vẫn có phép hợp thành kết hợp; điều này cho phép sử dụng parallel scan thay cho vòng lặp tuần tự ngây thơ.

Hiệu quả thực tế đến từ việc kết hợp scan với kernel fusion, quản lý phân cấp bộ nhớ GPU và recomputation để tránh materialize tensor state B × L × D × N trong HBM. Vì vậy, lợi thế của Mamba trên chuỗi dài là kết quả đồng thời của thiết kế toán học của Selective SSM và thiết kế thuật toán nhận biết phần cứng, chứ không chỉ từ việc loại bỏ self-attention.

Đọc thêm trên The CS Note: Giới thiệu Mamba: State Space Model cho chuỗi dài và Mamba vs Transformer: Khi nào nên dùng State Space Model thay vì Attention?.

Tài liệu tham khảo

1. [1] A. Gu and T. Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces,” in Proc. First Conference on Language Modeling (COLM), 2024. OpenReview.
2. [2] A. Gu, K. Goel, and C. Ré, “Efficiently Modeling Long Sequences with Structured State Spaces,” in Proc. 10th International Conference on Learning Representations (ICLR), 2022. OpenReview.
3. [3] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” in Advances in Neural Information Processing Systems, vol. 35, 2022, pp. 16344–16359. NeurIPS Proceedings.