Mamba & SSMs — Selective scan thay attention, linear scaling

Bạn muốn context window 1 triệu token nhưng KV cache của Transformer chiếm hết VRAM? Mamba là câu trả lời. Thay vì lưu lại toàn bộ quá khứ như attention, Mamba nén thông tin vào một hidden state cố định, sử dụng cơ chế "gatekeeper" để quyết định gì nên nhớ, gì nên quên — tất cả với độ phức tạp O(L) thay vì O(L²).

Vấn đề

Transformer đã thống trị AI ba năm qua nhờ attention mechanism, nhưng chính attention lại là điểm nghẽn khi scale lên context dài.

Quadratic bottleneck: Attention tính similarity giữa mọi cặp token, tạo ra ma trận N×N. Với 128K tokens, bạn cần 16 tỷ giá trị attention — chỉ tính toán đã chậm, chưa nói đến KV cache chiếm 32GB+ VRAM. Inference trở thành memory-bound thay vì compute-bound.

Nỗi đau của SSM cũ (S4): Trước Mamba, State Space Models như S4 đã cho phép linear scaling bằng cách nén quá khứ vào hidden state cố định. Nhưng chúng là Linear Time Invariant (LTI) — nghĩa là B, C, Δ (các tham số state transition) cố định cho mọi token. Dù input là "the" hay "password=1234", cách cập nhật state vẫn y hệt. Kết quả: S4 có thể tóm tắt tài liệu dài nhưng thất bại thảm hại trên selective copying hay reasoning task — nó không thể "nhấp nháy" để ghi nhớ thông tin quan trọng và lướt qua filler.

Ý tưởng cốt lõi

Mamba giới thiệu Selective State Space Model: Biến B, C, và step size Δ thành hàm của input token hiện tại. Điều này cho phép model selectively propagate hoặc reset hidden state dựa trên nội dung — giống như bạn chủ động quyết định "cái này quan trọng, ghi nhớ kỹ" hay "cái này là fluff, bỏ qua".

Cơ chế "gatekeeper" Δ (delta): Đây là phần hay nhất của Mamba. Δ quyết định mức độ hidden state tiếp nhận thông tin mới:

• Δ cao → "Overwrite ngay!": Token này quan trọng, đẩy nó mạnh vào state, reset phần cũ. Giống như nghe mật khẩu, bạn ghi nhớ ngay lập tức.
• Δ thấp → "Giữ nguyên": Token này không quan trọng, state hiện tại tiếp tục flow qua. Như khi đọc "um" hay "the".

Hãy tưởng tượng Δ như aperture của máy ảnh hoặc cổng vào working memory. S4 có aperture cố định — chụp mọi thứ cùng độ sáng. Mamba điều chỉnh aperture theo từng khung hình, cho phép nó "zoom in" vào thông tin quan trọng và "lướt" qua nội dung tầm thường.

Hardware-aware scan: Vấn đề là selective mechanism làm mất tính LTI, khiến việc tính toán song song trở nên khó khăn. Mamba giải quyết bằng parallel associative scan (Blelloch scan) — thuật toán tận dụng tính kết hợp của phép nhân ma trận để tính toán recurrence trong O(log L) bước trên GPU, thay vì tuần tự.

That's it. Đó là toàn bộ bí mật: discretize hệ thống liên tục, cho phép tham số phụ thuộc input, và tính toán song song bằng scan.

Tại sao nó hoạt động

Từ liên tục đến rời rạc: Mamba bắt đầu từ phương trình vi phân liên tục: ḣ(t) = Ah(t) + Bx(t). Sau đó discretize bằng phương pháp Zero-Order Hold với step size Δ:

h_t = Āh_{t-1} + B̄x_t

Trong đó Ā = exp(AΔ) và B̄ = (ΔA)^{-1}(exp(AΔ) - I) · ΔB.

Khi Δ→0, B̄→0 và Ā→I, nghĩa là state không đổi (giữ nguyên quá khứ). Khi Δ→∞, B̄ lớn và Ā nhỏ, nghĩa là state bị overwrite bởi input mới. Bằng cách làm Δ phụ thuộc x_t, Mamba học được cách "nhấp nháy" này một cách mượt mà.

Tại sao linear scaling hoạt động: Trong Transformer, mỗi token mới phải attend với tất cả token trước đó → O(L) memory và compute per step. Trong Mamba, hidden state h_t có kích thước cố định (thường 16-64 chiều). Dù sequence dài 1K hay 1M token, chi phí cập nhật state vẫn là O(1) — bạn chỉ cần nhân ma trận cố định, không cần lưu trữ toàn bộ quá khứ như KV cache.

Associative scan: Phép cập nhật state thực chất là phép nhân ma trận liên tiếp. Vì phép nhân ma trận có tính kết hợp, ta có thể chia sequence thành các đoạn, tính song song, rồi kết hợp kết quả — tương tự cách parallel prefix sum hoạt động. Đây là lý do Mamba train nhanh dù có tính chất recurrent.

Ý nghĩa thực tế

Trade-off: Compression vs. Random Access

Transformer như photographic memory — bạn có thể tra cứu bất kỳ token nào trong quá khứ với độ chính xác 100%, nhưng phải trả giá bằng VRAM và compute. Mamba như bản tóm tắt có chọn lọc — bạn nén 100K token thành vector 64 chiều, mất khả năng tra cứu chính xác nhưng gain được linear scaling.

Điều này làm Mamba cực kỳ hiệu quả cho:

• Long-context summarization: Khi bạn chỉ cần "gist" của tài liệu dài, không cần nhớ chính xác từng từ ở đâu.
• Genomics/Audio: Sequence cực dài (1M+ tokens) nhưng pattern local.

Nhưng Mamba không thay thế hoàn toàn Attention cho task cần exact retrieval — như "copy-paste" đoạn code từ đầu file xuống cuối file hay "needle in a haystack" test. Đó là lý do hybrid architectures như Jamba ra đời — kết hợp SSM layer cho phần lớn network và Attention layer thưa thớt cho exact recall.

Benchmarks & Adoption:

• 5× higher throughput trong inference so với Transformer tương đương
• Linear memory: O(1) per step thay vì O(L), cho phép context 1M+ token trên consumer hardware
• Mistral Codestral: Sử dụng kiến trúc Mamba cho code generation
• Vision Mamba: Ứng dụng cho image processing với hiệu quả tốt hơn ViT

Hạn chế:

• Training instability: Mamba thường cần nhiều epoch hơn Transformer để hội tụ, nhạy cảm với initialization
• Ecosystem kém mature: Chưa có tối ưu hóa kernel rộng rãi như FlashAttention, debug khó hơn
• Information loss: Nén thành fixed-size state buộc model phải "quên" chi tiết — đôi khi làm suy giảm performance trên reasoning task cần tra cứu chính xác từ xa

Đào sâu hơn

Paper gốc:

• "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" (Gu & Dao, 2023) — Giới thiệu selective mechanism và hardware-aware scan

Cùng cụm New Architectures:

• Mamba-2 — Phiên bản nâng cấp dùng structured matrix multiplication, kết nối chặt chẽ với linear attention và tận dụng Tensor Cores hiệu quả hơn
• Hybrid Attention-SSM (Jamba, Zamba) — Kết hợp SSM và attention để cân bằng giữa hiệu quả memory và khả năng recall chính xác
• Linear Attention — Hướng tiếp cận khác để giảm độ phức tạp attention từ O(N²) xuống O(N) bằng kernel trick
• Diffusion Language Models — Kiến trúc non-autoregressive khác, generate song song thay vì left-to-right

Đọc tiếp:

• Transformer & Attention (Level 0) — Nền tảng quadratic attention mà Mamba thay thế
• Flash Attention (Level 1) — Cách tối ưu attention truyền thống thay vì thay thế nó
• Ring Attention & Long Context (Level 2) — Cách khác để scale context length mà vẫn giữ attention nguyên bản