目錄
- 概
- Mamba
- 程式碼
Gu A. and Dao T. Mamba: Linear-time sequence modeling with selective state spaces. 2023.
概
Mamba.
Mamba
- S4 和 S4D 雖然解決了 SSM 計算速度的問題, 但是有一個前提, 就是 \(A, B, C, D\) 是與時間 \(t\) 無關的. 這導致這些方法只能採取一種固定的模式取處理序列問題, 作者認為這導致 SSM 無法 text 這類強上下文關係的任務.
-
所以如上圖和上述演算法所示, \(B, C, \Delta\) 現在是與輸入有關的了, 不同的輸入會產生不同的 \(B, C, \Delta\).
-
但是, 我們知道 \(\bar{A}\) 由 \(A, \Delta t\) 共同決定, 這就導致 \(\bar{A}(x)\) 實際上也是與輸入有關的了.
-
而我們知道, S4, S4D 訓練速度快的原因就是輸出能夠透過卷積的方式實現:
\[y = \mathcal{K}_L (\bar{A}, \bar{B}, C) * u + Du, \\ \mathcal{K}_L (A, B, C) := (CB, CAB, \ldots, CA^{L-1}B). \]但是這個必須要求 \(A\) 是隨著 \(t\) 不變的, 所以我們沒法實現這一點.
-
所以作者額外設計了 scan 演算法, 如上圖所示, 這是一種 hardware-aware 的演算法, 他會把隱狀態的更新放在 GPU 中速度最快的 SRAM 位置, 我看網上大多用下面這個圖來說:
- 我對這個不太感興趣, 有興趣的同學可以找相應的部落格看看.
程式碼
[official-code]