Introduction

Transformer 架構自 2017 年推出以來，一直主宰住深度學習嘅世界。由 GPT 到 Claude，幾乎所有 state-of-the-art 嘅 language models 都係基於 Transformer。但係，Transformer 有一個致命弱點：quadratic complexity。當 sequence length 增加，計算成本同記憶體需求會以平方倍數增長。

呢個限制催生咗一個新嘅研究方向：State Space Models (SSMs)。而喺 2023 年底，由 Carnegie Mellon University 嘅 Albert Gu 同 Princeton University 嘅 Tri Dao 推出嘅 Mamba 架構，終於證明咗 SSM 可以喺 language modeling 上面同 Transformer 平起平坐，甚至喺某啲任務上表現更好。[1]

今日，我哋會深入探討呢場由 Transformer 到 Mamba 嘅演變史。

💡 生活化類比：諗像你要整理一年嘅電郵。Transformer 就好似每次要搵資料，都要打開所有電郵逐封睇（超慢！）；而 Mamba 就好似你有個智能助手，會自動將重要資訊摘錄落筆記本，你只需要翻筆記本就搵到想要嘅嘢（快好多！）。

📚 目錄

點解 Transformer 需要被挑戰？
State Space Models (SSMs) 係乜嘢？
S4：SSM 嘅突破
Mamba：Selective State Space Models
Mamba vs Transformer：Performance Comparison
演變時間軸總結
技術深入：點樣實現 Selective SSM？
未來展望
Conclusion

點解 Transformer 需要被挑戰？

Transformer 嘅核心問題

Transformer 嘅核心係 self-attention mechanism。每個 token 都需要同序列入面所有其他 tokens 計算 attention scores，所以計算複雜度係 O(n²)，其中 n 係 sequence length。

呢個設計帶嚟幾個問題：

記憶體瓶頸：處理 1M tokens 嘅序列需要存儲一個 1M × 1M 嘅 attention matrix
計算成本：訓練同推理速度隨住 sequence length 急劇下降
長序列限制：實際應用上好難處理超過幾萬 tokens 嘅文檔

雖然有 Flash Attention、Linear Attention 等優化方法，但始終未能完全解決呢個 quadratic bottleneck。

🏢 現實例子：公司開會
想像你公司有個會議，需要每個人都同其他所有人傾偈一次（類似 Transformer 嘅 all-to-all attention）：
• 10 個人：需要 45 次對話（10×9/2）
• 100 個人：需要 4,950 次對話
• 1000 個人：需要 499,500 次對話！

你會發現人數一倍，對話次數就變四倍。呢個就係 O(n²) complexity 嘅問題！當序列越長，Transformer 就好似要組織一個超大型會議，所有人都要互相傾偈，根本冇可能。

早期嘅替代方案

喺 Mamba 出現之前，研究者嘗試過好多方法去挑戰 Transformer：

Linear Attention：將 attention 機制簡化到 linear complexity，但效果通常比唔上標準 Transformer
Recurrent models：例如 LSTM、GRU，但有 vanishing gradient 問題同難以並行化
Gated convolution：效率較高，但喺 language modeling 表現麻麻

而 State Space Models 就係喺呢個背景下開始受到關注。

Loading diagram...

State Space Models (SSMs) 係乜嘢？

數學基礎

State Space Models 源自控制理論 (control theory)，用嚟描述動態系統嘅演變。一個 continuous-time SSM 可以用以下方程式表示：

\begin{align} h'(t) &= \mathbf{A}h(t) + \mathbf{B}x(t) \\ y(t) &= \mathbf{C}h(t) + \mathbf{D}x(t) \end{align}

其中：

$h(t)$ 係 hidden state（隱藏狀態）
$x(t)$ 係 input（輸入）
$y(t)$ 係 output（輸出）
$\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D}$ 係 learnable parameters

但係喺 deep learning 入面，我哋需要 discretize 呢個 continuous system。透過 discretization（例如用 zero-order hold），我哋可以得到：

\begin{align} h_t &= \bar{\mathbf{A}}h_{t-1} + \bar{\mathbf{B}}x_t \\ y_t &= \mathbf{C}h_t \end{align}

SSM 嘅兩種計算模式

SSM 有個好特別嘅地方：佢可以用兩種等價嘅方式去計算！

1. Recurrent Mode（循環模式）

喺推理 (inference) 時，我哋可以逐步更新 hidden state：

h = initial_state
for x_t in sequence:
    h = A @ h + B @ x_t
    y_t = C @ h

呢種方式係 O(n) complexity，而且只需要 constant memory（因為只需要保存當前嘅 state）。

2. Convolution Mode（卷積模式）

📱 生活類比：兩種方式睇 Instagram
Recurrent Mode（逐條睇）：好似你逐條 post 咁睇 Instagram feed，每次只需要記住上一條 post 嘅印象。記憶體需求低（只記住當前印象），但要逐條睇，唔能夠跳過。

Convolution Mode（一次過睇曬）：好似你將成個 feed 嘅所有 posts 一次過 load 曬出嚟，用 grid view 同時睇。速度快（可以並行處理），但需要將所有內容一次過 load 入記憶體。

SSM 嘅神奇之處：**呢兩種方式計算出嚟嘅結果係一樣嘅！**訓練時用 Convolution Mode（快），推理時用 Recurrent Mode（省記憶體）。

喺訓練時，我哋可以將成個 SSM 展開成一個 convolution kernel：

\mathbf{K} = (\mathbf{C}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}^2\bar{\mathbf{B}}, \ldots, \mathbf{C}\bar{\mathbf{A}}^{L-1}\bar{\mathbf{B}})

然後用 Fast Fourier Transform (FFT) 計算 convolution：

y = \mathbf{K} \ast x

呢種方式可以完全並行化，訓練效率好高！

Loading diagram...

S4：SSM 嘅突破

Structured State Space Sequence Model (S4)

2021 年，Albert Gu 等人提出咗 S4 (Structured State Space Sequence Model)，係 SSM 應用喺 deep learning 嘅一個重大突破。[2]

S4 嘅核心創新：

1. HiPPO Matrix 初始化

S4 使用 HiPPO (High-order Polynomial Projection Operators) 嚟初始化矩陣 $mathbf{A}$ 。HiPPO matrix 有個特性：佢可以記住長期歷史信息，有效解決咗 long-range dependency 問題。

2. Diagonal Plus Low-Rank (DPLR) 結構

為咗高效計算，S4 將 $\mathbf{A}$ 參數化為 diagonal plus low-rank 形式：

\mathbf{A} = \Lambda - PQ^*

呢種結構化設計令到計算 convolution kernel 可以用 Cauchy kernel 快速完成。

3. 表現

S4 喺多個 long-range sequence modeling benchmarks 上面達到 SOTA，包括：

Long Range Arena (LRA)：專門測試長序列處理能力
Speech recognition
Time-series forecasting

但係，S4 喺 language modeling 上面嘅表現始終比唔上 Transformer。點解？

S4 嘅限制

S4 同其他 SSM 架構有個共同弱點：time-invariance（時不變性）。

因為 parameters $\mathbf{A}, \mathbf{B}, \mathbf{C}$ 係 fixed（唔會隨住 input 改變），所以 SSM 難以做到 content-based reasoning。簡單啲講，佢唔能夠根據當前嘅 token 內容去決定應該記住乜嘢、遺忘乜嘢。

呢個限制對於 language modeling 特別致命，因為語言理解需要好強嘅 context-dependent 處理能力。

📝 生活例子：記筆記嘅兩種方式
傳統 SSM (S4) = 機械式抄寫員：
• 有固定嘅抄寫規則，唔理內容重唔重要都照抄
• 就好似一個機器人，見到乜就記乜，完全唔會思考
• 老師講「今日天氣好好」同「下星期考試範圍係第三章」都用同一種方式記錄

Mamba (Selective SSM) = 聰明學生：
• 會判斷內容重要性，決定應該詳細記定係簡單帶過
• 聽到「下星期考試」就會特別留神，用螢光筆標記
• 聽到閒話家常就簡單記一句或者唔記
• 會根據 context 調整筆記策略

呢個就係 time-invariance 同 selectivity 嘅分別！S4 唔能夠根據內容調整，而 Mamba 可以。

Loading diagram...

Mamba：Selective State Space Models

核心創新：Selectivity

2023 年 12 月，Albert Gu 同 Tri Dao 推出咗 Mamba，徹底解決咗 SSM 嘅 content-based reasoning 問題。[1]

Mamba 嘅關鍵創新係引入 selectivity（選擇性）：

\begin{align} \mathbf{B}_t &= s_B(x_t) \\ \mathbf{C}_t &= s_C(x_t) \\ \Delta_t &= \tau_{\Delta}(s_{\Delta}(x_t)) \end{align}

即係話，parameters $\mathbf{B}, \mathbf{C}$ 同 discretization step size $\Delta$ 都係 input-dependent！模型可以根據當前 token 嘅內容動態調整：

記住重要信息：當遇到重要 token，增大 $\mathbf{B}_t$ 將信息存入 state
遺忘無關信息：減少唔重要信息對 state 嘅影響
控制時間尺度：透過調整 $\Delta_t$ 控制信息衰減速度

Selection Mechanism 嘅直覺

🎬 現實例子：睇電影記劇情
想像你睇緊一套兩個鐘嘅電影，之後要同朋友講返個故事：

傳統 SSM（固定規則）：
• 每分鐘都記同樣份量嘅內容
• 主角表白嘅重要場景：記一句
• 路人甲行過嘅無關場景：都記一句
• 結果：重要情節記唔夠，無聊片段佔咗好多記憶空間

Mamba（選擇性記憶）：
• $\mathbf{B}_t$ （記憶閘門）：遇到重要情節（表白、轉折、伏筆），開大個閘，詳細記低
• $\Delta_t$ （遺忘速度）：無關場景（路人、風景）就加快遺忘，唔佔空間
• $\mathbf{C}_t$ （提取選擇）：朋友問起劇情時，選擇性提取相關記憶

結果：你可以用有限嘅記憶空間，記住成套電影嘅精華！呢個就係 selectivity 嘅威力。

諗返 attention mechanism：佢可以根據 query 同 key 嘅相似度去決定應該 attend 邊啲 tokens。Mamba 嘅 selection mechanism 做緊類似嘅嘢，但係用一種更高效嘅方式！

Hardware-Aware Algorithm

但係有個問題：一旦 parameters 變成 input-dependent，我哋就冇辦法預先計算 convolution kernel，即係冇辦法用 FFT 加速訓練！

Mamba 嘅第二個創新係設計咗一個 hardware-aware parallel algorithm：

核心思想：避免 Materialization

傳統做法會將 expanded state 完整寫入 GPU memory (HBM)，但呢個操作好慢。Mamba 嘅做法係：

Kernel fusion：將 discretization、recurrence、output 嘅計算 fuse 成一個 CUDA kernel
Recomputation：喺 backward pass 重新計算 intermediate states，而唔係存儲佢哋
SRAM 優化：盡量將計算保持喺 GPU 嘅 fast memory (SRAM) 入面

呢種設計令到 Mamba 即使冇咗 convolution 嘅優勢，都可以做到：

Linear scaling：O(n) complexity
Fast inference：比 Transformer 快 5× throughput
Memory efficient：可以處理 million-length sequences

Mamba 架構設計

Mamba 將 selective SSM 整合入一個極簡嘅架構：

Input
  ↓
Linear Projection (expand dimension)
  ↓
[Selective SSM Block] ← 核心！
  ↓
Linear Projection (reduce dimension)
  ↓
Residual Connection
  ↓
Layer Norm
  ↓
Output

冇 attention，甚至冇 MLP blocks！整個架構就係 stacked selective SSM layers。

Mamba Block 詳細結構解析

下圖展示咗一個完整嘅 Mamba block 內部結構。我哋逐個部分講解：

📥 1. Input Tokens

頂部嘅 "Hello ! My name is" (tokens 1-5) 係輸入序列。每個 token 經過 embedding 後成為一個 hidden dimension 嘅 vector。

🔄 2. RMS Norm (Root Mean Square Normalization)

位置：Block 最頂部
功能：將 input 正規化，穩定訓練過程
類似於：LayerNorm，但計算更高效
公式： $text{RMS}(x) = frac{x}{sqrt{frac{1}{n}sum x_i^2 + epsilon}}$

⚡ 3. Skip Connection（殘差連接）

位置：由 input 直接跳到 output 嘅虛線
功能：Residual connection！將原始 input 直接加到 block output
重要性：
- 解決 gradient vanishing 問題
- 允許梯度直接流過，訓練更深嘅網絡
- 類似 Transformer 嘅 residual connections
公式： $text{output} = text{SSM_block}(x) + x$

🟢 4. Projection Branches（左右兩個綠色分支）

位置：分成左右兩條並行路徑
功能：
- 左分支：包含 convolution 同 SSM 嘅主要處理路徑
- 右分支：產生 gating signal，控制信息流動
設計理念：類似 Gated Linear Units (GLU)

左分支：SSM 路徑

a) Projection（投影）

將 input dimension 投影到 expanded dimension
通常會 expand 2× (例如 512 → 1024)

b) 1D Convolution

Kernel size：通常 4
功能：捕捉 local patterns，類似 n-gram
點解需要：SSM 主要處理 global dependencies，convolution 補充 local modeling

c) SiLU Activation

公式： $text{SiLU}(x) = x cdot sigma(x)$
特性：Smooth, non-monotonic activation function

d) Selective SSM 核心

呢個就係 Mamba 嘅靈魂！
執行 selective state space computation
動態調整 $B_t, C_t, \Delta_t$ 根據 input content

右分支：Gating 路徑

a) Projection

同左分支一樣 expand dimension

b) SiLU Activation

產生 smooth gating values

🔵 5. Element-wise Multiplication（逐元素相乘）

位置：兩條分支匯合嘅地方
功能：將 SSM output 同 gating signal 相乘
公式： $text{output} = text{SSM_output} odot text{gate}$
作用：Gate 控制邊啲信息應該通過，邊啲應該被抑制

🔷 6. Final Projection（藍色投影）

功能：將 expanded dimension 投影返回 original dimension
例如：1024 → 512
確保 output 同 input 有相同 dimension，先可以做 residual addition

➕ 7. Residual Addition（⊕）

位置：Skip connection 同 processed output 相加
功能：將 Mamba block 嘅輸出同原始 input 相加
呢個就係 residual connection 嘅實現！

🔄 8. Final RMS Norm

位置：Block 最底部
功能：再次正規化 output
準備餵入下一個 Mamba block 或 output layer

📤 9. Output Head

Linear + Softmax：最終預測下一個 token
Example output："Maarten" (token 6)

🔑 關鍵設計特點

Residual Connection 嘅重要性

Mamba 使用 pre-norm residual structure：

# Pre-norm residual (Mamba 用呢種)
x = x + MambaBlock(RMSNorm(x))

# Post-norm residual (舊式 Transformer)
x = RMSNorm(x + MambaBlock(x))

Pre-norm 有幾個優勢：

更穩定嘅訓練：Normalization 喺 residual block 之前
更深嘅網絡：可以堆疊更多 layers 而唔會 gradient explosion
唔需要 warmup：訓練初期已經好穩定

Gated Architecture 嘅作用

雙分支設計（SSM path + Gating path）類似：

LSTM gates：控制信息流動
GLU (Gated Linear Units)：選擇性地通過信息
Transformer FFN with gating：SwiGLU, GeGLU

呢種設計令到模型可以：

動態選擇：根據 context 決定邊啲 features 重要
增強表達能力：兩條路徑可以學習唔同嘅 representations
提升非線性：Gating 引入額外嘅非線性變換

點解 Mamba 咁高效？

No Attention Matrix：唔需要計算同存儲 $n \times n$ attention matrix
Linear Complexity：SSM 係 O(n)，唔係 O(n²)
Hardware-Optimized：CUDA kernels 針對 GPU memory hierarchy 優化
Minimal Components：冇 MLP blocks，簡化咗架構

對比 Transformer Block

Component	Transformer	Mamba
Normalization	LayerNorm	RMS Norm
Main Operation	Multi-Head Attention (O(n²))	Selective SSM (O(n))
FFN	2-layer MLP	冇！（已包含喺 SSM）
Gating	Optional (SwiGLU)	Built-in (dual branches)
Residual	Yes	Yes
Position Encoding	Required	唔需要！（implicit in SSM）

Mamba 嘅設計證明：Simpler can be better！透過精心設計嘅 SSM + gating mechanism，可以達到同 Transformer 相當嘅效果，但效率更高。🚀

Loading diagram...

SSM vs Transformer：核心差異

喺深入 Mamba 嘅 performance 之前，我哋先講清楚 State Space Models (SSMs) 同 Transformer 嘅根本分別。

🎯 核心機制對比

Transformer：Attention-Based

Transformer 嘅核心係 self-attention mechanism：

# Simplified attention
Q = input @ W_q  # Query
K = input @ W_k  # Key  
V = input @ W_v  # Value

attention_scores = softmax(Q @ K.T / sqrt(d_k))  # n × n matrix!
output = attention_scores @ V

工作原理：

每個 token 同所有其他 tokens計算相似度
產生一個 $n \times n$ 嘅 attention matrix
根據相似度 weight 唔同 tokens 嘅信息

優點：

✅ 強大嘅 content-based reasoning
✅ 可以直接 attend 到任何位置嘅 token
✅ 完全並行化訓練
✅ 成熟嘅生態系統同工具

缺點：

❌ O(n²) complexity：sequence 越長，成本指數上升
❌ 記憶體密集：需要存儲 n × n attention matrix
❌ 長序列瓶頸：實際上好難處理超過 100K tokens

SSM：Recurrence-Based

SSM 嘅核心係 state space equations：

# Simplified SSM
h_t = A @ h_{t-1} + B @ x_t  # Update hidden state
y_t = C @ h_t                 # Generate output

工作原理：

維持一個 fixed-size hidden state $h_t$
逐步更新 state，將信息壓縮入去
唔需要同所有 previous tokens 直接交互

優點：

✅ O(n) complexity：Linear scaling with sequence length
✅ Constant memory：只需要保存當前 state
✅ Long sequences：可以處理 million-token sequences
✅ Fast inference：Recurrent mode 逐步生成

缺點：

❌ Fixed-size state bottleneck：所有信息必須壓縮入 fixed-size state
❌ Weaker content-based reasoning（傳統 SSM）
❌ 較新嘅架構：生態未夠成熟

📊 詳細對比表

特性	Transformer	SSM (S4/Mamba)
核心機制	Self-Attention	State Space Equations
信息流動	每個 token attend 所有 tokens	信息透過 hidden state 流動
時間複雜度（訓練）	O(n²)	O(n) 或 O(n log n)
空間複雜度	O(n²) attention matrix	O(d) hidden state (d = state size)
推理速度	較慢（需要整個 context）	快 5× (recurrent mode)
最大序列長度	~10K-100K tokens (實際限制)	1M+ tokens
並行化（訓練）	完全並行	可並行（convolution mode）
Position Encoding	必需（explicit）	唔需要（implicit in dynamics）
Content-based Reasoning	極強（attention 機制）	傳統 SSM 較弱；Mamba 透過 selectivity 解決
Long-range Dependencies	理論上可以，但受限於 context window	天生擅長（通過 state propagation）
Interpretability	Attention weights 較易解釋	Hidden state 較難解釋

🔄 資訊處理方式嘅根本分別

Transformer："Look at Everything"

想像你讀緊一本書：

Transformer 就好似每次讀到一個詞，你都翻返去頭睇返所有之前嘅詞
然後決定邊啲詞對理解當前詞最重要
呢個過程雖然強大，但當書越來越長，翻返去頭嘅成本越來越高

Loading diagram...

每個 token 都同所有其他 tokens 計算 attention！

SSM："Compress and Flow"

SSM 就好似你有個筆記本（hidden state）
每次讀到新嘅詞，你更新筆記本，將重要信息記低
你唔需要成日翻返去頭，只需要睇住個筆記本
當書越來越長，你只需要繼續更新筆記本，成本保持不變

Loading diagram...

信息透過 state 逐步流動，無需回頭！

⚖️ Trade-offs 總結

點解 Transformer 咁成功？

Content-based reasoning 強大：可以根據內容動態決定關注乜嘢
訓練高效：完全並行化
靈活性高：Attention 可以學到複雜嘅 dependencies
生態成熟：工具、優化、預訓練模型豐富

但係，當 sequence 長到一定程度，quadratic complexity 就成為致命傷。

點解 SSM 值得關注？

效率驚人：Linear complexity 令長序列處理成為可能
記憶體友好：唔需要存儲巨大嘅 attention matrix
推理快：Recurrent mode 特別適合 autoregressive generation
理論優雅：源自控制理論，數學基礎紮實

但係，傳統 SSM（如 S4）喺 language modeling 表現唔夠好。

Mamba 嘅突破：兩者兼得

Mamba 透過 selective mechanism 解決咗傳統 SSM 嘅弱點：

特性	傳統 SSM (S4)	Mamba
Content-based Reasoning	❌ 弱（time-invariant）	✅ 強（selective）
效率	✅ O(n)	✅ O(n)
Language Modeling	❌ 比唔上 Transformer	✅ Match Transformer

Loading diagram...

🎓 直觀類比

Transformer = 圖書館

你可以隨時翻查任何一本書（任何 token）
但係書越多，搵書越慢
需要好大嘅空間存放所有書

SSM = 筆記系統

你將重要信息摘錄落筆記
只需要睇住筆記，唔需要翻返所有書
但係筆記空間有限，可能會遺漏細節

Mamba = 智能筆記系統

會自動判斷乜嘢信息重要
動態調整筆記內容
既快速又唔會遺漏關鍵信息

Mamba vs Transformer：Performance Comparison

Language Modeling

喺 language modeling 任務上，Mamba 達到咗驚人嘅結果：

Mamba-3B outperforms 同 size 嘅 Transformers
Mamba-3B matches Transformer-6B 嘅表現（兩倍 size！）
喺 The Pile dataset 上，Mamba 嘅 perplexity 同 Transformer 相當或更好

Inference Speed

5× higher throughput than Transformers
Linear scaling in sequence length（Transformer 係 quadratic）
可以處理 1 million tokens 嘅序列

Long-Range Dependencies

💼 實際應用例子：分析年度報告
想像你要分析一份 500 頁嘅公司年度報告（~100K tokens）：

Transformer 嘅困境：
• 第 1 頁講「我哋今年推出新產品 X」
• 第 450 頁講「產品 X 嘅銷售額」
• 但 Transformer 嘅 context window 只有 32K tokens，根本睇唔到咁遠！
• 就算勉強處理，運算成本會爆炸（O(n²)）

Mamba 嘅優勢：
• 用 linear complexity 輕鬆處理成份報告
• 透過 hidden state 將第 1 頁嘅信息「記住」到第 450 頁
• 5× 更快：處理長文檔唔會拖慢速度
• 可以處理 1M tokens：相當於一本幾千頁嘅書！

呢個對於 legal documents、scientific papers、code repositories 等長文檔分析超有用。

Mamba 喺處理長序列任務表現優異：

Genomics：DNA sequence modeling
Audio：long audio generation
Document understanding：超長文檔處理

限制

當然，Mamba 都有佢嘅限制：

Copying tasks：有研究指出 Transformer 喺簡單嘅 copying 任務上表現更好[3]
Fixed-size state：雖然可以處理長序列，但 hidden state size 係 fixed，可能限制咗某啲類型嘅推理能力
Interpretability：SSM 嘅內部機制比 attention 更難解釋

演變時間軸總結

2017: Transformer 時代開始

"Attention is All You Need" 論文發布
Self-attention mechanism 成為主流

2020-2021: SSM 嘅探索

HiPPO: 提出可以記住長期歷史嘅 projection operators
S4 (2021)：第一個喺 long-range tasks 達到 competitive performance 嘅 SSM

2022-2023: SSM 嘅改進

S5 (Simplified S4)：簡化 S4 架構
各種 SSM variants：探索唔同嘅結構化方法

2023: Mamba 嘅突破

Mamba (Dec 2023)：引入 selective mechanism，首次喺 language modeling 上 match Transformer

2024-2025: SSM 嘅未來

Mamba-2：進一步優化
Hybrid models：結合 SSM 同 Transformer 嘅優勢（例如 Jamba）
更多應用領域：vision、multimodal、time-series

Loading diagram...

技術深入：點樣實現 Selective SSM？

如果你想自己 implement 一個簡化版嘅 selective SSM，核心步驟係：

import torch
import torch.nn as nn

class SelectiveSSM(nn.Module):
    def __init__(self, d_model, d_state):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # Selection projections
        self.proj_B = nn.Linear(d_model, d_state)
        self.proj_C = nn.Linear(d_model, d_state)
        self.proj_delta = nn.Linear(d_model, d_model)
        
        # SSM parameters (simplified)
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        batch, seq_len, _ = x.shape
        
        # Selection mechanism
        B = self.proj_B(x)  # (batch, seq_len, d_state)
        C = self.proj_C(x)  # (batch, seq_len, d_state)
        delta = torch.sigmoid(self.proj_delta(x))  # (batch, seq_len, d_model)
        
        # Discretization (simplified)
        A_bar = torch.exp(delta.unsqueeze(-1) * self.A.unsqueeze(0).unsqueeze(0))
        
        # Recurrent computation
        h = torch.zeros(batch, self.d_state, device=x.device)
        outputs = []
        
        for t in range(seq_len):
            h = A_bar[:, t] * h + B[:, t]
            y = (C[:, t] * h).sum(dim=-1)
            outputs.append(y)
        
        return torch.stack(outputs, dim=1)

當然，實際嘅 Mamba implementation 會複雜好多，包括 CUDA kernels、parallel scan algorithms 等等。

未來展望

Mamba 會取代 Transformer 嗎？

短期內唔會。但係 Mamba 證明咗：

SSM 係 viable alternative：唔係 Transformer 係唯一選擇
Efficiency matters：linear complexity 喺長序列處理上有巨大優勢
Architecture diversity：唔同架構適合唔同應用場景

Hybrid Approaches

最有前景嘅方向可能係 hybrid models：

用 Transformer layers 處理需要 content-based reasoning 嘅部分
用 SSM layers 處理長序列同 efficient inference
例如 Jamba model 就係呢種設計

其他應用領域

SSM 嘅優勢唔止 language modeling：

Vision：long video understanding
Audio：speech synthesis with long context
Time-series：高效處理連續數據流
Edge computing：低功耗設備上嘅 AI[4]

Conclusion

由 Transformer 到 Mamba 嘅演變，唔係一個簡單嘅架構替代，而係對 sequence modeling 根本問題嘅重新思考：

Transformer 用 quadratic attention 換取強大嘅 content-based reasoning
SSM 用 linear recurrence 換取高效嘅長序列處理
Mamba 透過 selective mechanism 將兩者嘅優勢結合

呢個演變仲未結束。隨住 hardware 進步、算法優化、同新嘅架構設計，我哋可能會見到更多突破。但有一點係肯定嘅：Smarter architecture > Bigger models。

Loading diagram...

喺 AI 發展嘅下一個階段，效率同性能嘅平衡會越嚟越重要。而 Mamba 同 SSM 正正指向呢個方向。🚀

參考論文

必讀論文：

Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023)

Albert Gu, Tri Dao

https://arxiv.org/abs/2312.00752
Efficiently Modeling Long Sequences with Structured State Spaces (S4) (2021)

Albert Gu et al.

https://openreview.net/forum?id=uYLFoz1vlAC
From S4 to Mamba: A Comprehensive Survey on Structured State Space Models (2025)

https://arxiv.org/abs/2503.18970

延伸閱讀：

Repeat After Me: Transformers are Better than State Space Models at Copying (2024)

https://arxiv.org/pdf/2402.01032
A Visual Guide to Mamba and State Space Models

https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state
Mamba Explained (The Gradient)

https://thegradient.pub/mamba-explained/

About the Author

Billy Tse 係一位 AI 研究者同 SmartQuest 嘅創辦人，專注於 deep learning、NLP 同 education technology。對 transformer architectures、state space models 同 AI 喺教育領域嘅應用有深入研究。

Introduction

今日，我哋會深入探討呢場由 Transformer 到 Mamba 嘅演變史。

💡 生活化類比：諗像你要整理一年嘅電郵。Transformer 就好似每次要搵資料，都要打開所有電郵逐封睇（超慢！）；而 Mamba 就好似你有個智能助手，會自動將重要資訊摘錄落筆記本，你只需要翻筆記本就搵到想要嘅嘢（快好多！）。

📚 目錄

點解 Transformer 需要被挑戰？
State Space Models (SSMs) 係乜嘢？
S4：SSM 嘅突破
Mamba：Selective State Space Models
Mamba vs Transformer：Performance Comparison
演變時間軸總結
技術深入：點樣實現 Selective SSM？
未來展望
Conclusion

點解 Transformer 需要被挑戰？

Transformer 嘅核心問題

呢個設計帶嚟幾個問題：

記憶體瓶頸：處理 1M tokens 嘅序列需要存儲一個 1M × 1M 嘅 attention matrix
計算成本：訓練同推理速度隨住 sequence length 急劇下降
長序列限制：實際應用上好難處理超過幾萬 tokens 嘅文檔

雖然有 Flash Attention、Linear Attention 等優化方法，但始終未能完全解決呢個 quadratic bottleneck。

🏢 現實例子：公司開會
想像你公司有個會議，需要每個人都同其他所有人傾偈一次（類似 Transformer 嘅 all-to-all attention）：
• 10 個人：需要 45 次對話（10×9/2）
• 100 個人：需要 4,950 次對話
• 1000 個人：需要 499,500 次對話！

你會發現人數一倍，對話次數就變四倍。呢個就係 O(n²) complexity 嘅問題！當序列越長，Transformer 就好似要組織一個超大型會議，所有人都要互相傾偈，根本冇可能。

早期嘅替代方案

喺 Mamba 出現之前，研究者嘗試過好多方法去挑戰 Transformer：

Linear Attention：將 attention 機制簡化到 linear complexity，但效果通常比唔上標準 Transformer
Recurrent models：例如 LSTM、GRU，但有 vanishing gradient 問題同難以並行化
Gated convolution：效率較高，但喺 language modeling 表現麻麻

而 State Space Models 就係喺呢個背景下開始受到關注。

Loading diagram...

State Space Models (SSMs) 係乜嘢？

數學基礎

State Space Models 源自控制理論 (control theory)，用嚟描述動態系統嘅演變。一個 continuous-time SSM 可以用以下方程式表示：

\begin{align} h'(t) &= \mathbf{A}h(t) + \mathbf{B}x(t) \\ y(t) &= \mathbf{C}h(t) + \mathbf{D}x(t) \end{align}

其中：

$h(t)$ 係 hidden state（隱藏狀態）
$x(t)$ 係 input（輸入）
$y(t)$ 係 output（輸出）
$\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D}$ 係 learnable parameters

但係喺 deep learning 入面，我哋需要 discretize 呢個 continuous system。透過 discretization（例如用 zero-order hold），我哋可以得到：

\begin{align} h_t &= \bar{\mathbf{A}}h_{t-1} + \bar{\mathbf{B}}x_t \\ y_t &= \mathbf{C}h_t \end{align}

SSM 嘅兩種計算模式

SSM 有個好特別嘅地方：佢可以用兩種等價嘅方式去計算！

1. Recurrent Mode（循環模式）

喺推理 (inference) 時，我哋可以逐步更新 hidden state：

h = initial_state
for x_t in sequence:
    h = A @ h + B @ x_t
    y_t = C @ h

呢種方式係 O(n) complexity，而且只需要 constant memory（因為只需要保存當前嘅 state）。

2. Convolution Mode（卷積模式）

📱 生活類比：兩種方式睇 Instagram
Recurrent Mode（逐條睇）：好似你逐條 post 咁睇 Instagram feed，每次只需要記住上一條 post 嘅印象。記憶體需求低（只記住當前印象），但要逐條睇，唔能夠跳過。

Convolution Mode（一次過睇曬）：好似你將成個 feed 嘅所有 posts 一次過 load 曬出嚟，用 grid view 同時睇。速度快（可以並行處理），但需要將所有內容一次過 load 入記憶體。

SSM 嘅神奇之處：**呢兩種方式計算出嚟嘅結果係一樣嘅！**訓練時用 Convolution Mode（快），推理時用 Recurrent Mode（省記憶體）。

喺訓練時，我哋可以將成個 SSM 展開成一個 convolution kernel：

\mathbf{K} = (\mathbf{C}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}^2\bar{\mathbf{B}}, \ldots, \mathbf{C}\bar{\mathbf{A}}^{L-1}\bar{\mathbf{B}})

然後用 Fast Fourier Transform (FFT) 計算 convolution：

y = \mathbf{K} \ast x

呢種方式可以完全並行化，訓練效率好高！

Loading diagram...

S4：SSM 嘅突破

Structured State Space Sequence Model (S4)

2021 年，Albert Gu 等人提出咗 S4 (Structured State Space Sequence Model)，係 SSM 應用喺 deep learning 嘅一個重大突破。[2]

S4 嘅核心創新：

1. HiPPO Matrix 初始化

2. Diagonal Plus Low-Rank (DPLR) 結構

為咗高效計算，S4 將 $\mathbf{A}$ 參數化為 diagonal plus low-rank 形式：

\mathbf{A} = \Lambda - PQ^*

呢種結構化設計令到計算 convolution kernel 可以用 Cauchy kernel 快速完成。

3. 表現

S4 喺多個 long-range sequence modeling benchmarks 上面達到 SOTA，包括：

Long Range Arena (LRA)：專門測試長序列處理能力
Speech recognition
Time-series forecasting

但係，S4 喺 language modeling 上面嘅表現始終比唔上 Transformer。點解？

S4 嘅限制

S4 同其他 SSM 架構有個共同弱點：time-invariance（時不變性）。

呢個限制對於 language modeling 特別致命，因為語言理解需要好強嘅 context-dependent 處理能力。

📝 生活例子：記筆記嘅兩種方式
傳統 SSM (S4) = 機械式抄寫員：
• 有固定嘅抄寫規則，唔理內容重唔重要都照抄
• 就好似一個機器人，見到乜就記乜，完全唔會思考
• 老師講「今日天氣好好」同「下星期考試範圍係第三章」都用同一種方式記錄

Mamba (Selective SSM) = 聰明學生：
• 會判斷內容重要性，決定應該詳細記定係簡單帶過
• 聽到「下星期考試」就會特別留神，用螢光筆標記
• 聽到閒話家常就簡單記一句或者唔記
• 會根據 context 調整筆記策略

呢個就係 time-invariance 同 selectivity 嘅分別！S4 唔能夠根據內容調整，而 Mamba 可以。

Loading diagram...

Mamba：Selective State Space Models

核心創新：Selectivity

2023 年 12 月，Albert Gu 同 Tri Dao 推出咗 Mamba，徹底解決咗 SSM 嘅 content-based reasoning 問題。[1]

Mamba 嘅關鍵創新係引入 selectivity（選擇性）：

\begin{align} \mathbf{B}_t &= s_B(x_t) \\ \mathbf{C}_t &= s_C(x_t) \\ \Delta_t &= \tau_{\Delta}(s_{\Delta}(x_t)) \end{align}

即係話，parameters $\mathbf{B}, \mathbf{C}$ 同 discretization step size $\Delta$ 都係 input-dependent！模型可以根據當前 token 嘅內容動態調整：

記住重要信息：當遇到重要 token，增大 $\mathbf{B}_t$ 將信息存入 state
遺忘無關信息：減少唔重要信息對 state 嘅影響
控制時間尺度：透過調整 $\Delta_t$ 控制信息衰減速度

Selection Mechanism 嘅直覺

🎬 現實例子：睇電影記劇情
想像你睇緊一套兩個鐘嘅電影，之後要同朋友講返個故事：

傳統 SSM（固定規則）：
• 每分鐘都記同樣份量嘅內容
• 主角表白嘅重要場景：記一句
• 路人甲行過嘅無關場景：都記一句
• 結果：重要情節記唔夠，無聊片段佔咗好多記憶空間

Mamba（選擇性記憶）：
• $\mathbf{B}_t$ （記憶閘門）：遇到重要情節（表白、轉折、伏筆），開大個閘，詳細記低
• $\Delta_t$ （遺忘速度）：無關場景（路人、風景）就加快遺忘，唔佔空間
• $\mathbf{C}_t$ （提取選擇）：朋友問起劇情時，選擇性提取相關記憶

結果：你可以用有限嘅記憶空間，記住成套電影嘅精華！呢個就係 selectivity 嘅威力。

諗返 attention mechanism：佢可以根據 query 同 key 嘅相似度去決定應該 attend 邊啲 tokens。Mamba 嘅 selection mechanism 做緊類似嘅嘢，但係用一種更高效嘅方式！

Hardware-Aware Algorithm

但係有個問題：一旦 parameters 變成 input-dependent，我哋就冇辦法預先計算 convolution kernel，即係冇辦法用 FFT 加速訓練！

Mamba 嘅第二個創新係設計咗一個 hardware-aware parallel algorithm：

核心思想：避免 Materialization

傳統做法會將 expanded state 完整寫入 GPU memory (HBM)，但呢個操作好慢。Mamba 嘅做法係：

Kernel fusion：將 discretization、recurrence、output 嘅計算 fuse 成一個 CUDA kernel
Recomputation：喺 backward pass 重新計算 intermediate states，而唔係存儲佢哋
SRAM 優化：盡量將計算保持喺 GPU 嘅 fast memory (SRAM) 入面

呢種設計令到 Mamba 即使冇咗 convolution 嘅優勢，都可以做到：

Linear scaling：O(n) complexity
Fast inference：比 Transformer 快 5× throughput
Memory efficient：可以處理 million-length sequences

Mamba 架構設計

Mamba 將 selective SSM 整合入一個極簡嘅架構：

Input
  ↓
Linear Projection (expand dimension)
  ↓
[Selective SSM Block] ← 核心！
  ↓
Linear Projection (reduce dimension)
  ↓
Residual Connection
  ↓
Layer Norm
  ↓
Output

冇 attention，甚至冇 MLP blocks！整個架構就係 stacked selective SSM layers。

Mamba Block 詳細結構解析

下圖展示咗一個完整嘅 Mamba block 內部結構。我哋逐個部分講解：

📥 1. Input Tokens

頂部嘅 "Hello ! My name is" (tokens 1-5) 係輸入序列。每個 token 經過 embedding 後成為一個 hidden dimension 嘅 vector。

🔄 2. RMS Norm (Root Mean Square Normalization)

位置：Block 最頂部
功能：將 input 正規化，穩定訓練過程
類似於：LayerNorm，但計算更高效
公式： $text{RMS}(x) = frac{x}{sqrt{frac{1}{n}sum x_i^2 + epsilon}}$

⚡ 3. Skip Connection（殘差連接）

位置：由 input 直接跳到 output 嘅虛線
功能：Residual connection！將原始 input 直接加到 block output
重要性：
- 解決 gradient vanishing 問題
- 允許梯度直接流過，訓練更深嘅網絡
- 類似 Transformer 嘅 residual connections
公式： $text{output} = text{SSM_block}(x) + x$

🟢 4. Projection Branches（左右兩個綠色分支）

位置：分成左右兩條並行路徑
功能：
- 左分支：包含 convolution 同 SSM 嘅主要處理路徑
- 右分支：產生 gating signal，控制信息流動
設計理念：類似 Gated Linear Units (GLU)

左分支：SSM 路徑

a) Projection（投影）

將 input dimension 投影到 expanded dimension
通常會 expand 2× (例如 512 → 1024)

b) 1D Convolution

Kernel size：通常 4
功能：捕捉 local patterns，類似 n-gram
點解需要：SSM 主要處理 global dependencies，convolution 補充 local modeling

c) SiLU Activation

公式： $text{SiLU}(x) = x cdot sigma(x)$
特性：Smooth, non-monotonic activation function

d) Selective SSM 核心

呢個就係 Mamba 嘅靈魂！
執行 selective state space computation
動態調整 $B_t, C_t, \Delta_t$ 根據 input content

右分支：Gating 路徑

a) Projection

同左分支一樣 expand dimension

b) SiLU Activation

產生 smooth gating values

🔵 5. Element-wise Multiplication（逐元素相乘）

位置：兩條分支匯合嘅地方
功能：將 SSM output 同 gating signal 相乘
公式： $text{output} = text{SSM_output} odot text{gate}$
作用：Gate 控制邊啲信息應該通過，邊啲應該被抑制

🔷 6. Final Projection（藍色投影）

功能：將 expanded dimension 投影返回 original dimension
例如：1024 → 512
確保 output 同 input 有相同 dimension，先可以做 residual addition

➕ 7. Residual Addition（⊕）

位置：Skip connection 同 processed output 相加
功能：將 Mamba block 嘅輸出同原始 input 相加
呢個就係 residual connection 嘅實現！

🔄 8. Final RMS Norm

位置：Block 最底部
功能：再次正規化 output
準備餵入下一個 Mamba block 或 output layer

📤 9. Output Head

Linear + Softmax：最終預測下一個 token
Example output："Maarten" (token 6)

🔑 關鍵設計特點

Residual Connection 嘅重要性

Mamba 使用 pre-norm residual structure：

# Pre-norm residual (Mamba 用呢種)
x = x + MambaBlock(RMSNorm(x))

# Post-norm residual (舊式 Transformer)
x = RMSNorm(x + MambaBlock(x))

Pre-norm 有幾個優勢：

更穩定嘅訓練：Normalization 喺 residual block 之前
更深嘅網絡：可以堆疊更多 layers 而唔會 gradient explosion
唔需要 warmup：訓練初期已經好穩定

Gated Architecture 嘅作用

雙分支設計（SSM path + Gating path）類似：

LSTM gates：控制信息流動
GLU (Gated Linear Units)：選擇性地通過信息
Transformer FFN with gating：SwiGLU, GeGLU

呢種設計令到模型可以：

動態選擇：根據 context 決定邊啲 features 重要
增強表達能力：兩條路徑可以學習唔同嘅 representations
提升非線性：Gating 引入額外嘅非線性變換

點解 Mamba 咁高效？

No Attention Matrix：唔需要計算同存儲 $n \times n$ attention matrix
Linear Complexity：SSM 係 O(n)，唔係 O(n²)
Hardware-Optimized：CUDA kernels 針對 GPU memory hierarchy 優化
Minimal Components：冇 MLP blocks，簡化咗架構

對比 Transformer Block

Component	Transformer	Mamba
Normalization	LayerNorm	RMS Norm
Main Operation	Multi-Head Attention (O(n²))	Selective SSM (O(n))
FFN	2-layer MLP	冇！（已包含喺 SSM）
Gating	Optional (SwiGLU)	Built-in (dual branches)
Residual	Yes	Yes
Position Encoding	Required	唔需要！（implicit in SSM）

Mamba 嘅設計證明：Simpler can be better！透過精心設計嘅 SSM + gating mechanism，可以達到同 Transformer 相當嘅效果，但效率更高。🚀

Loading diagram...

SSM vs Transformer：核心差異

喺深入 Mamba 嘅 performance 之前，我哋先講清楚 State Space Models (SSMs) 同 Transformer 嘅根本分別。

🎯 核心機制對比

Transformer：Attention-Based

Transformer 嘅核心係 self-attention mechanism：

# Simplified attention
Q = input @ W_q  # Query
K = input @ W_k  # Key  
V = input @ W_v  # Value

attention_scores = softmax(Q @ K.T / sqrt(d_k))  # n × n matrix!
output = attention_scores @ V

工作原理：

每個 token 同所有其他 tokens計算相似度
產生一個 $n \times n$ 嘅 attention matrix
根據相似度 weight 唔同 tokens 嘅信息

優點：

✅ 強大嘅 content-based reasoning
✅ 可以直接 attend 到任何位置嘅 token
✅ 完全並行化訓練
✅ 成熟嘅生態系統同工具

缺點：

❌ O(n²) complexity：sequence 越長，成本指數上升
❌ 記憶體密集：需要存儲 n × n attention matrix
❌ 長序列瓶頸：實際上好難處理超過 100K tokens

SSM：Recurrence-Based

SSM 嘅核心係 state space equations：

# Simplified SSM
h_t = A @ h_{t-1} + B @ x_t  # Update hidden state
y_t = C @ h_t                 # Generate output

工作原理：

維持一個 fixed-size hidden state $h_t$
逐步更新 state，將信息壓縮入去
唔需要同所有 previous tokens 直接交互

優點：

✅ O(n) complexity：Linear scaling with sequence length
✅ Constant memory：只需要保存當前 state
✅ Long sequences：可以處理 million-token sequences
✅ Fast inference：Recurrent mode 逐步生成

缺點：

❌ Fixed-size state bottleneck：所有信息必須壓縮入 fixed-size state
❌ Weaker content-based reasoning（傳統 SSM）
❌ 較新嘅架構：生態未夠成熟

📊 詳細對比表

特性	Transformer	SSM (S4/Mamba)
核心機制	Self-Attention	State Space Equations
信息流動	每個 token attend 所有 tokens	信息透過 hidden state 流動
時間複雜度（訓練）	O(n²)	O(n) 或 O(n log n)
空間複雜度	O(n²) attention matrix	O(d) hidden state (d = state size)
推理速度	較慢（需要整個 context）	快 5× (recurrent mode)
最大序列長度	~10K-100K tokens (實際限制)	1M+ tokens
並行化（訓練）	完全並行	可並行（convolution mode）
Position Encoding	必需（explicit）	唔需要（implicit in dynamics）
Content-based Reasoning	極強（attention 機制）	傳統 SSM 較弱；Mamba 透過 selectivity 解決
Long-range Dependencies	理論上可以，但受限於 context window	天生擅長（通過 state propagation）
Interpretability	Attention weights 較易解釋	Hidden state 較難解釋

🔄 資訊處理方式嘅根本分別

Transformer："Look at Everything"

想像你讀緊一本書：

Transformer 就好似每次讀到一個詞，你都翻返去頭睇返所有之前嘅詞
然後決定邊啲詞對理解當前詞最重要
呢個過程雖然強大，但當書越來越長，翻返去頭嘅成本越來越高

Loading diagram...

每個 token 都同所有其他 tokens 計算 attention！

SSM："Compress and Flow"

SSM 就好似你有個筆記本（hidden state）
每次讀到新嘅詞，你更新筆記本，將重要信息記低
你唔需要成日翻返去頭，只需要睇住個筆記本
當書越來越長，你只需要繼續更新筆記本，成本保持不變

Loading diagram...

信息透過 state 逐步流動，無需回頭！

⚖️ Trade-offs 總結

點解 Transformer 咁成功？

Content-based reasoning 強大：可以根據內容動態決定關注乜嘢
訓練高效：完全並行化
靈活性高：Attention 可以學到複雜嘅 dependencies
生態成熟：工具、優化、預訓練模型豐富

但係，當 sequence 長到一定程度，quadratic complexity 就成為致命傷。

點解 SSM 值得關注？

效率驚人：Linear complexity 令長序列處理成為可能
記憶體友好：唔需要存儲巨大嘅 attention matrix
推理快：Recurrent mode 特別適合 autoregressive generation
理論優雅：源自控制理論，數學基礎紮實

但係，傳統 SSM（如 S4）喺 language modeling 表現唔夠好。

Mamba 嘅突破：兩者兼得

Mamba 透過 selective mechanism 解決咗傳統 SSM 嘅弱點：

特性	傳統 SSM (S4)	Mamba
Content-based Reasoning	❌ 弱（time-invariant）	✅ 強（selective）
效率	✅ O(n)	✅ O(n)
Language Modeling	❌ 比唔上 Transformer	✅ Match Transformer

Loading diagram...

🎓 直觀類比

Transformer = 圖書館

你可以隨時翻查任何一本書（任何 token）
但係書越多，搵書越慢
需要好大嘅空間存放所有書

SSM = 筆記系統

你將重要信息摘錄落筆記
只需要睇住筆記，唔需要翻返所有書
但係筆記空間有限，可能會遺漏細節

Mamba = 智能筆記系統

會自動判斷乜嘢信息重要
動態調整筆記內容
既快速又唔會遺漏關鍵信息

Mamba vs Transformer：Performance Comparison

Language Modeling

喺 language modeling 任務上，Mamba 達到咗驚人嘅結果：

Mamba-3B outperforms 同 size 嘅 Transformers
Mamba-3B matches Transformer-6B 嘅表現（兩倍 size！）
喺 The Pile dataset 上，Mamba 嘅 perplexity 同 Transformer 相當或更好

Inference Speed

5× higher throughput than Transformers
Linear scaling in sequence length（Transformer 係 quadratic）
可以處理 1 million tokens 嘅序列

Long-Range Dependencies

💼 實際應用例子：分析年度報告
想像你要分析一份 500 頁嘅公司年度報告（~100K tokens）：

Transformer 嘅困境：
• 第 1 頁講「我哋今年推出新產品 X」
• 第 450 頁講「產品 X 嘅銷售額」
• 但 Transformer 嘅 context window 只有 32K tokens，根本睇唔到咁遠！
• 就算勉強處理，運算成本會爆炸（O(n²)）

Mamba 嘅優勢：
• 用 linear complexity 輕鬆處理成份報告
• 透過 hidden state 將第 1 頁嘅信息「記住」到第 450 頁
• 5× 更快：處理長文檔唔會拖慢速度
• 可以處理 1M tokens：相當於一本幾千頁嘅書！

呢個對於 legal documents、scientific papers、code repositories 等長文檔分析超有用。

Mamba 喺處理長序列任務表現優異：

Genomics：DNA sequence modeling
Audio：long audio generation
Document understanding：超長文檔處理

限制

當然，Mamba 都有佢嘅限制：

Copying tasks：有研究指出 Transformer 喺簡單嘅 copying 任務上表現更好[3]
Fixed-size state：雖然可以處理長序列，但 hidden state size 係 fixed，可能限制咗某啲類型嘅推理能力
Interpretability：SSM 嘅內部機制比 attention 更難解釋

演變時間軸總結

2017: Transformer 時代開始

"Attention is All You Need" 論文發布
Self-attention mechanism 成為主流

2020-2021: SSM 嘅探索

HiPPO: 提出可以記住長期歷史嘅 projection operators
S4 (2021)：第一個喺 long-range tasks 達到 competitive performance 嘅 SSM

2022-2023: SSM 嘅改進

S5 (Simplified S4)：簡化 S4 架構
各種 SSM variants：探索唔同嘅結構化方法

2023: Mamba 嘅突破

Mamba (Dec 2023)：引入 selective mechanism，首次喺 language modeling 上 match Transformer

2024-2025: SSM 嘅未來

Mamba-2：進一步優化
Hybrid models：結合 SSM 同 Transformer 嘅優勢（例如 Jamba）
更多應用領域：vision、multimodal、time-series

Loading diagram...

技術深入：點樣實現 Selective SSM？

如果你想自己 implement 一個簡化版嘅 selective SSM，核心步驟係：

import torch
import torch.nn as nn

class SelectiveSSM(nn.Module):
    def __init__(self, d_model, d_state):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # Selection projections
        self.proj_B = nn.Linear(d_model, d_state)
        self.proj_C = nn.Linear(d_model, d_state)
        self.proj_delta = nn.Linear(d_model, d_model)
        
        # SSM parameters (simplified)
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        batch, seq_len, _ = x.shape
        
        # Selection mechanism
        B = self.proj_B(x)  # (batch, seq_len, d_state)
        C = self.proj_C(x)  # (batch, seq_len, d_state)
        delta = torch.sigmoid(self.proj_delta(x))  # (batch, seq_len, d_model)
        
        # Discretization (simplified)
        A_bar = torch.exp(delta.unsqueeze(-1) * self.A.unsqueeze(0).unsqueeze(0))
        
        # Recurrent computation
        h = torch.zeros(batch, self.d_state, device=x.device)
        outputs = []
        
        for t in range(seq_len):
            h = A_bar[:, t] * h + B[:, t]
            y = (C[:, t] * h).sum(dim=-1)
            outputs.append(y)
        
        return torch.stack(outputs, dim=1)

當然，實際嘅 Mamba implementation 會複雜好多，包括 CUDA kernels、parallel scan algorithms 等等。

未來展望

Mamba 會取代 Transformer 嗎？

短期內唔會。但係 Mamba 證明咗：

SSM 係 viable alternative：唔係 Transformer 係唯一選擇
Efficiency matters：linear complexity 喺長序列處理上有巨大優勢
Architecture diversity：唔同架構適合唔同應用場景

Hybrid Approaches

最有前景嘅方向可能係 hybrid models：

用 Transformer layers 處理需要 content-based reasoning 嘅部分
用 SSM layers 處理長序列同 efficient inference
例如 Jamba model 就係呢種設計

其他應用領域

SSM 嘅優勢唔止 language modeling：

Vision：long video understanding
Audio：speech synthesis with long context
Time-series：高效處理連續數據流
Edge computing：低功耗設備上嘅 AI[4]

Conclusion

由 Transformer 到 Mamba 嘅演變，唔係一個簡單嘅架構替代，而係對 sequence modeling 根本問題嘅重新思考：

Transformer 用 quadratic attention 換取強大嘅 content-based reasoning
SSM 用 linear recurrence 換取高效嘅長序列處理
Mamba 透過 selective mechanism 將兩者嘅優勢結合

呢個演變仲未結束。隨住 hardware 進步、算法優化、同新嘅架構設計，我哋可能會見到更多突破。但有一點係肯定嘅：Smarter architecture > Bigger models。

Loading diagram...

喺 AI 發展嘅下一個階段，效率同性能嘅平衡會越嚟越重要。而 Mamba 同 SSM 正正指向呢個方向。🚀

參考論文

必讀論文：

Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023)

Albert Gu, Tri Dao

https://arxiv.org/abs/2312.00752
Efficiently Modeling Long Sequences with Structured State Spaces (S4) (2021)

Albert Gu et al.

https://openreview.net/forum?id=uYLFoz1vlAC
From S4 to Mamba: A Comprehensive Survey on Structured State Space Models (2025)

https://arxiv.org/abs/2503.18970

延伸閱讀：

Repeat After Me: Transformers are Better than State Space Models at Copying (2024)

https://arxiv.org/pdf/2402.01032
A Visual Guide to Mamba and State Space Models

https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state
Mamba Explained (The Gradient)

https://thegradient.pub/mamba-explained/

Introduction

📚 目錄

點解 Transformer 需要被挑戰？

Transformer 嘅核心問題

早期嘅替代方案

State Space Models (SSMs) 係乜嘢？

數學基礎

SSM 嘅兩種計算模式

1. Recurrent Mode（循環模式）

2. Convolution Mode（卷積模式）

2. Convolution Mode（卷積模式）

S4：SSM 嘅突破

Structured State Space Sequence Model (S4)

1. HiPPO Matrix 初始化

2. Diagonal Plus Low-Rank (DPLR) 結構

3. 表現

S4 嘅限制

Mamba：Selective State Space Models

核心創新：Selectivity

Selection Mechanism 嘅直覺

Selection Mechanism 嘅直覺

Hardware-Aware Algorithm

核心思想：避免 Materialization

Mamba 架構設計

Mamba Block 詳細結構解析

🔑 關鍵設計特點

對比 Transformer Block

SSM vs Transformer：核心差異

🎯 核心機制對比

Transformer：Attention-Based

SSM：Recurrence-Based

📊 詳細對比表

🔄 資訊處理方式嘅根本分別

Transformer："Look at Everything"

SSM："Compress and Flow"

⚖️ Trade-offs 總結

點解 Transformer 咁成功？

點解 SSM 值得關注？

Mamba 嘅突破：兩者兼得

🎓 直觀類比

Mamba vs Transformer：Performance Comparison

Language Modeling

Inference Speed

Long-Range Dependencies

Long-Range Dependencies

限制

演變時間軸總結

2017: Transformer 時代開始

2020-2021: SSM 嘅探索

2022-2023: SSM 嘅改進

2023: Mamba 嘅突破

2024-2025: SSM 嘅未來

技術深入：點樣實現 Selective SSM？

未來展望

Mamba 會取代 Transformer 嗎？

Hybrid Approaches

其他應用領域

Conclusion

參考論文

相關資源

About the Author

Introduction

📚 目錄

點解 Transformer 需要被挑戰？

Transformer 嘅核心問題

早期嘅替代方案

State Space Models (SSMs) 係乜嘢？

數學基礎

SSM 嘅兩種計算模式

1. Recurrent Mode（循環模式）

2. Convolution Mode（卷積模式）

2. Convolution Mode（卷積模式）

S4：SSM 嘅突破

Structured State Space Sequence Model (S4)

1. HiPPO Matrix 初始化

2. Diagonal Plus Low-Rank (DPLR) 結構

3. 表現

S4 嘅限制

Mamba：Selective State Space Models

核心創新：Selectivity