Transformer 很厲害，但并不完美，尤其是在處理長序列方面。而狀態空間模型（SSM）則在長序列上的表現相當不俗。早在去年就有研究者提出可使用 SSM 替代 Transformer，參見文章《預訓練無需注意力，擴展到4096個token不成問題，與BERT相當》，前些天基于 SSM 方法的 Mamba 更是異軍突起，推理吞吐量達到了 Transformer 的五倍之多，參閱《五倍吞吐量，性能全面包圍Transformer：新架構Mamba引爆AI圈》。

但實際上，SSM 和 Transformer 并不是非此即彼的兩種架構，它們完全可以組合起來！

近日公布的一篇 NeurIPS 2023 論文《Block-State Transformers》就采用了這種做法，其不僅能輕松支持 65k token 長度的超長輸入，而且計算效率還非常高，速度相比使用循環單元的 Transformer 足可提升十倍之多！這篇論文也得到了 Mamba 作者 Tri Dao 的點贊，他表示：「SSM 和Transformer 似乎可以互補。」

但在我們介紹這種新方法前，先簡單說說 Transformer。在許多不同的自然語言處理（NLP）任務上，Transformer 的表現都非常出色。可以說 Transformer 已經很大相當程度上替代了循環神經網絡。不僅如此，它也正在圖像和視頻等 NLP 之外的領域大展拳腳。

其成功的原因有很多，包括計算效率和架構層面的歸納偏差，這讓它們非常適合在自然語言任務進行大規模訓練。在計算方面，Transformer 能以并行方式處理輸入序列的 token，從而使其能充分利用現代加速器硬件。此外，注意力機制讓 Transformer 可以找到更長序列之間的關系，其方式是在推斷下一個 token 時讀取從過去 token 提取的所有信息。相比于 RNN 和 LSTM，自注意力有兩個優勢：(1) 存儲信息以及將這些信息直接用作上下文的能力得到了極大提升，(2) 在更長序列上能更穩定地訓練。

盡管 Transformer 相比 RNN 有很多優勢，但它在輸入序列長度的擴展上依然存在問題，其中涉及計算性能和質量等方面的原因。更進一步說，Transformer 的運行時間會隨輸入序列長度的增長成二次方增長，這會讓訓練這些模型的成本越來越高。

此外，眾所周知使用注意力的 Transformer 在長輸入分類任務上表現不佳。最基本的 Transformer 在長序列上訓練時可能不穩定，而且其 token 重要度聚焦在當前時間步驟周圍約 50 個 token 的局部感受野中。

近來，越來越多的研究表明狀態空間模型（SSM）可以替代 Transformer，因為 SSM 可以捕獲極長序列之中的依賴關系，同時還有更高的計算效率和更好的并行化能力。

盡管 SSM 依然屬于自回歸序列模型，但其底層的線性時間不變式動態系統可使用基于快速傅立葉變換（FFT）的可并行化卷積算子來高效地處理序列，而且這個過程的復雜度僅為 ??(?? log ??)，其中 ?? 是序列的長度。此外，借用在線函數近似的方法，通過推導循環更新規則，可以確保在長序列上保留過去的信息，甚至可達成千上萬個時間步驟。在 Long-Range Arena 基準上，SSM 甚至超過了 Transformer 一大截，參閱機器之心報道《六項任務、多種數據類型，谷歌、DeepMind提出高效Transformer評估基準》。

盡管 SSM 在長程分類任務上很成功，但如果要用作通用語言建模的現成可用序列模型，SSM 還完全趕不上 Transformer。

近期又有研究《Long Range Language Modeling via Gated State Spaces》認為 Transformer 和 SSM 完全可以互補。

DeepMind 等機構提出的新架構 Block-State Transformer（BST）將強大的基于局部注意力的歸納偏差與長期上下文建模能力組合到了一起，做成了單一層。

論文地址：https://arxiv.org/pdf/2306.09539.pdf

據介紹，該模型能在處理長輸入序列的同時整合注意力機制來預測下一個 token。相比于基于 Transformer 的層，BST 是完全可并行化的，能擴展用于更長得多的序列，同時速度還能快 10 倍。

在每一層 BST 中，有一個 SSM 將輸入的整個序列映射進一個同樣長度的「上下文」序列。這個 SSM 子層使用基于 FFT 的卷積。然后將這個上下文序列分成大小相等的上下文塊，這個大小即為窗口長度 W；然后再將每個上下文塊輸入一個 Transformer 層，其注意力關注的是大小為 W 的子序列。之后對輸入 token 嵌入塊與對應的上下文狀態塊使用交叉注意力，如圖 1 所示。

注意，通過將 SSM 用作一種上下文化的方法，就可以完全不需要序列循環，這樣一來就能以完全并行的方式運行這種 SSM-Transformer 混合層。

最后的運行時間復雜度可以表示成一個和：??(??2)+??(?? log ??)，其中前一項表示 Transformer 子層的時間復雜度，后一項是 SSM 子層的時間復雜度。

只要有支持并行計算的硬件，相較于 Block-Recurrent Transformer 的 ??(????)，這是一個重大提升。此外，由于硬件施加的限制，SSM 在完整序列上的運行時間復雜度與 Block Transformer 在 token 塊上的運行時間復雜度相當，這進一步意味著 BST 層不存在速度瓶頸。該團隊使用包含數十萬 token 的序列通過實驗驗證了這一點。

方法

這里研究的是通過僅解碼器語言模型實現下一 token 預測的問題。

對狀態空間的前置說明

狀態空間模型可以分為兩大類：

狀態空間：結構化核S4、S5、S4D、DSS遵循卷積核的一種結構化初始化，方式是展開一種線性時間不變式（LTI）動態系統，如下所示：

其中的參數包括狀態矩陣 ??∈?^{N×N}，向量 ??∈?^{N×1}、??∈?^{1×N}、??∈?^{1×1}。SSM 會將一維的輸入信號 u_k 映射成一維的輸出信號 y_k。

顯式參數化的過濾器。不同于結構化核，還可以將卷積核參數化為可訓練的權重并優化它們。但是，這會導致性能很差，除非對這些核使用特定類型的正則化方法。替代 Transformer 的無注意力模型中也有使用可訓練核的，比如 Hyena 涉及到沿核對權重進行指數衰減。

Block-State Transformer（BST）層

Block-State Transformer 層將 SSM 與 Block Transformer 組合到了一起。在每一次訓練迭代中，都會從一個長文檔采樣一個包含 L 個 token 的序列。然后嵌入該 token 并將其饋送給模型。這個模型由堆疊的 Block-State Transformer 層構成。每一層 BST 都會選擇性地包含一個 SSM 子層，其負責為 Block Transformer 層提供長程上下文，這與 Block-Recurrent Transformer（BRECT）單元的工作方式類似。這個 SSM 子層的輸入是前一層的 token 嵌入序列，輸出則是一個長度同樣為 L 的序列。

這個輸出經過了上下文編碼，也就是說每個時間步驟的項目都可能包含有關該序列中元素之前的所有時間步驟的信息。他們從上下文序列收集一定數量 S 的「上下文狀態」，并使得 S ? L。

這些上下文狀態會被饋送給 Block Transformer，以替代 Block-Recurrent Transformer 中的「循環狀態向量」。如圖 1 右側所示，后續操作保持不變，只是無需再運行 BRECT 單元的循環單元，因為現在是通過 SSM 來維護上下文。除了上下文狀態，Block Transformer 的輸入中還有長度 W 的 token 嵌入的塊/窗口；然后在這個窗口與上下文狀態上使用交叉注意力。然后將這個交叉注意力操作的輸出與自注意力在輸入嵌入上的輸出連接起來，之后是一個簡單的投影。

SSM 不僅能在更長時間尺度上保留信息，而且使用 SSM 來維持上下文狀態以替代循環單元，可以得到計算效率更高的層。通過將 SSM 整合進 Transformer 層，可以移除循環部分，從而讓 Block-State Transformer 層可以完全并行化。

上下文狀態

盡管從技術上看，最新的 SSM 輸出包含有關整個序列的信息，但僅從最后的狀態檢索單個 token 可能是不可行的。為了彌補這一點，該團隊將一系列狀態連接了起來，對應于最新的 token 塊。這與 BRECT 采用的方法類似。這種表征可以通過冗余來確保可檢索性和易訪問性。

在新提出的方法中，上下文狀態是使用 SSM 的輸出構建的，并會被饋送給 Transformer 的注意力頭。這些上下文狀態的構建方式有很多。為了引導設計決策，該團隊考慮了多種設計方案，包括使用單頭（Single-Head）、多頭（Multi-Head）或多過濾器（Multi-Filter）。其中單頭設計見圖 1。下圖 2 則展示了多頭和多過濾器的設計方案。

比較下來，多過濾器的記憶狀態的冗余最少，多頭次之，單頭的冗余最大。

結果

該團隊在 PG19、GitHub 和 arXiv 三個數據集上進行了實驗，檢驗了新提出的 BST 在不同長度的英語文本、latex 科學文章和源代碼上的效果。下表 1 總結了實驗結果。

下圖 3 則給出了長度泛化分析并報告了困惑度。實驗中，新模型和基準模型的參數數量都約為 4 億，訓練時的序列長度為 4k，測試中的序列長度為 {512, 16k, 65k}。

可以看到，在 PG19、GitHub 和 arXiv 上，當序列長度為 65k 時，BST:SH:S4-L 的困惑度最好。

在效率方面，下圖 4 左給出了 BST 層在 GPU 上的基準測試結果。

可以看到 SSM 帶來了非常顯著的增長——比包含循環單元的 Block-Recurrent Transformer 快 6-11 倍；即使在序列長度達到 65k token 時，還依然能有 6 倍的提升，而這時候硬件就已經開始飽和了。當使用結構化的 SSM 時，計算復雜度與 SSM 的內部記憶狀態大小 N 緊密相關。對于報告的性能，N = 16。

研究者表示，如果使用其它自動微分框架中近期引入的更快的針對硬件的 I/O 感知型實現，BST 方法的速度還能更快。

更多技術細節和實驗結果參閱原論文。