Mamba-2 — SSM = structured matrix multiply, kết nối với linear attention

Mamba-1 đã chứng minh SSM có thể thay thế Attention với độ phức tạp tuyến tính, nhưng nó mang một "gót chân Achilles": selective scan tuần tự không thể tận dụng GPU Tensor Cores, khiến training chậm 16 lần so với matmul thông thường. Mamba-2 giải quyết điều này bằng một cú twist toán học đẹp đẽ — nó chứng minh SSM thực chất là phép nhân ma trận có cấu trúc (structured matrix multiplication), từ đó kết nối trực tiếp với Linear Attention và mở khóa 90% tính toán trên Tensor Cores.

Vấn đề

Memory-bound vs Compute-bound: Mamba-1 dùng associative scan để tính toán trạng thái ẩn (hidden state) qua thời gian. Đây là thao tác tuần tự, không thể song song hóa trên GPU hiện đại. Trong khi đó, Tensor Cores (đơn vị tính toán chuyên biệt trên A100/H100) xử lý ma trận multiplication nhanh gấp 16 lần so với các thao tác scan tuần tự. Kết quả là Mamba-1 bị kẹt ở "bottleneck băng thông nhớ" thay vì tận dụng được sức mạnh tính toán thô.

Hạn chế thứ hai là state size nhỏ (thường N=16), hạn chế khả năng lưu trữ thông tin dài hạn. Và cuối cùng, SSM tồn tại như một "hộp đen" riêng biệt so với Attention, khiến việc kết hợp hai cơ chế trở nên khó khăn trong kiến trúc hybrid như Jamba.

Ý tưởng cốt lõi

Structured State Space Duality (SSD) là cốt lõi của Mamba-2. Ý chính đơn giản đến mức gây "aha moment": SSM không phải là một loại mạng recurrent đặc biệt — nó chỉ là phép nhân ma trận tam giác dưới có cấu trúc (lower-triangular structured matrix).

Khi bạn viết lại công thức recurrent của SSM dưới dạng ma trận, bạn phát hiện ma trận này có tính chất "1-semiseparable": các khối nằm ngoài đường chéo chính (off-diagonal blocks) đều có rank-1. Tính chất này cực kỳ quan trọng vì ma trận rank-1 có thể phân tích thành outer product của hai vector, cho phép tính toán hiệu quả qua phép nhân ma trận thay vì scan tuần tự.

Điểm thứ hai là kết nối với Linear Attention. Linear Attention thay thế softmax bằng kernel function tách rời (separable kernel) để giảm độ phức tạp từ O(N²) xuống O(N). SSD chứng minh rằng khi ma trận cấu trúc của SSM có dạng exponential decay (từ ma trận A đường chéo), nó trở thành dạng tương đương toán học với Linear Attention. SSM và Linear Attention chỉ là hai góc nhìn khác nhau của cùng một cấu trúc ma trận.

4-Step Algorithm: Thay vì scan từng token, Mamba-2 chia sequence thành các chunk (Q chunks):

1. Intra-chunk: Tính output local trong mỗi chunk bằng matmul (song song hoàn toàn)
2. Final states: Tính trạng thái cuối của mỗi chunk
3. Inter-chunk scan: Chỉ scan các trạng thái cuối này (sequence ngắn hơn Q lần)
4. Aggregation: Kết hợp trạng thái đã scan với output local

Bước 1, 2, 4 đều là matmul thuần túy — chạy cực nhanh trên Tensor Cores. Chỉ bước 3 là scan tuần tự, nhưng làm trên chuỗi đã nén (length T/Q), nên overhead không đáng kể.

Multi-Head Expansion: Mamba-1 giới hạn head dimension P=1 (mỗi token chỉ có 1 scalar state). Mamba-2 tăng lên P=64, cho phép nhiều "kênh" SSM độc lập hoạt động song song, tương tự multi-head attention, tăng đáng kể expressiveness.

Đó là tất cả. Không còn ma thuật scan bí ẩn — chỉ là đại số ma trận thông minh.

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

Toán học đằng sau sự tương đương: Trong SSM, output tại position i phụ thuộc vào tất cả inputs j ≤ i qua một tích chập (convolution) ngầm. Khi viết dưới dạng ma trận L, ta có L_ij = C_i^T A_i A_{i-1}...A_{j+1} B_j. Nếu A_t là ma trận đường chéo (scalar per channel), tích này trở thành dạng rank-1, cho phép viết L = QK^T với Q, K được định nghĩa qua các tích lũy (cumulative products). Đây chính là dạng kernel của Linear Attention.

Tận dụng phần cứng: GPU hiện đại (Ampere/Ada/Hopper) có kiến trúc chia rõ ràng giữa Tensor Cores (xử lý ma trận 16x16 hoặc lớn hơn cực nhanh) và CUDA cores (xử lý tổng hợp/scalar). Mamba-2 đẩy 90% FLOPs vào Tensor Cores, trong khi Mamba-1 phải dùng CUDA cores cho scan tuần tự.

Online Softmax tái xuất: Cơ chế gộp (aggregation) trong Mamba-2 sử dụng "online softmax" tương tự FlashAttention — tính toán normalization incrementally để tránh materialize ma trận attention N×N. Điều này cho phép training với sequence dài mà không bị OOM.

Ý nghĩa thực tế

Tốc độ training: Mamba-2 nhanh hơn 2–8 lần so với Mamba-1, và cạnh tranh được với Transformer ở quy mô 8B parameters khi train trên 3.5T tokens. Điều này biến SSM từ "thí nghiệm thú vị" thành "lựa chọn production thực sự".

State lớn hơn: Với P=64, Mamba-2 có khả năng lưu trữ và truy xuất thông tin tốt hơn Mamba-1, giảm bớt "amnesia" trên các task đòi hỏi recall chính xác từ xa.

Trade-offs rõ ràng:

• Short sequences (< 1k tokens): Vẫn chậm hơn Transformer do overhead khởi tạo kernel và chia chunk. Mamba-2 chỉ "bùng nổ" ở context dài.
• Exact recall: Dù cải thiện, pure SSM vẫn thua kém Attention thuần túy trên "needle-in-haystack" tasks — đây là lý do Jamba kết hợp cả hai.

Ai đang dùng: Các implementation nghiên cứu (state-spaces), mô hình vision (Vision Mamba variants), và các kiến trúc multimodal mới. Tuy nhiên, việc áp dụng rộng vẫn thận trọng do độ nhạy với training recipe (hyperparameters, initialization) cao hơn Transformer.

Đào sâu hơn

Paper gốc: "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" (Dao & Gu, 2024) — trình bày lý thuyết SSD và chứng minh tương đương với Linear Attention.

Bài đọc thêm: Blog series của Tri Dao giải thích chi tiết 4-step algorithm và cách triển khai trên CUDA.

Bài liên quan TroiSinh:

Cùng cụm (new-architectures):

• Mamba & SSMs — Nền tảng selective scan và linear scaling từ Mamba-1, prerequisite để hiểu sự cải tiến của Mamba-2
• Hybrid Attention-SSM (Jamba, Zamba) — Cách kết hợp SSM với Attention layers để giữ precision recall khi cần
• Linear Attention — Hiểu rõ mối liên hệ toán học mà Mamba-2 khai thác
• Diffusion LM — Một hướng kiến trúc khác bên ngoài Transformer autoregressive

Đọc tiếp:

• Flash Attention — Cơ chế online softmax và tiling mà Mamba-2 mượn để tối ưu memory
• KV Cache — So sánh với cách Transformer lưu trữ history, hiểu tại sao SSM state nhỏ gọn hơn
• Ring Attention — Nếu cần context triệu tokens, đây là cách Transformer vượt qua giới hạn bộ nhớ GPU