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

但實(shí)際上，SSM 和 Transformer 并不是非此即彼的兩種架構(gòu)，它們完全可以組合起來(lái)！

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

但在我們介紹這種新方法前，先簡(jiǎn)單說(shuō)說(shuō) Transformer。在許多不同的自然語(yǔ)言處理（NLP）任務(wù)上，Transformer 的表現(xiàn)都非常出色?？梢哉f(shuō) Transformer 已經(jīng)很大相當(dāng)程度上替代了循環(huán)神經(jīng)網(wǎng)絡(luò)。不僅如此，它也正在圖像和視頻等 NLP 之外的領(lǐng)域大展拳腳。

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

盡管 Transformer 相比 RNN 有很多優(yōu)勢(shì)，但它在輸入序列長(zhǎng)度的擴(kuò)展上依然存在問(wèn)題，其中涉及計(jì)算性能和質(zhì)量等方面的原因。更進(jìn)一步說(shuō)，Transformer 的運(yùn)行時(shí)間會(huì)隨輸入序列長(zhǎng)度的增長(zhǎng)成二次方增長(zhǎng)，這會(huì)讓訓(xùn)練這些模型的成本越來(lái)越高。

此外，眾所周知使用注意力的 Transformer 在長(zhǎng)輸入分類(lèi)任務(wù)上表現(xiàn)不佳。最基本的 Transformer 在長(zhǎng)序列上訓(xùn)練時(shí)可能不穩(wěn)定，而且其 token 重要度聚焦在當(dāng)前時(shí)間步驟周?chē)s 50 個(gè) token 的局部感受野中。

近來(lái)，越來(lái)越多的研究表明狀態(tài)空間模型（SSM）可以替代 Transformer，因?yàn)?SSM 可以捕獲極長(zhǎng)序列之中的依賴(lài)關(guān)系，同時(shí)還有更高的計(jì)算效率和更好的并行化能力。

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

盡管 SSM 在長(zhǎng)程分類(lèi)任務(wù)上很成功，但如果要用作通用語(yǔ)言建模的現(xiàn)成可用序列模型，SSM 還完全趕不上 Transformer。

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

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

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

據(jù)介紹，該模型能在處理長(zhǎng)輸入序列的同時(shí)整合注意力機(jī)制來(lái)預(yù)測(cè)下一個(gè) token。相比于基于 Transformer 的層，BST 是完全可并行化的，能擴(kuò)展用于更長(zhǎng)得多的序列，同時(shí)速度還能快 10 倍。

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

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

最后的運(yùn)行時(shí)間復(fù)雜度可以表示成一個(gè)和：??(??2)+??(?? log ??)，其中前一項(xiàng)表示 Transformer 子層的時(shí)間復(fù)雜度，后一項(xiàng)是 SSM 子層的時(shí)間復(fù)雜度。

只要有支持并行計(jì)算的硬件，相較于 Block-Recurrent Transformer 的 ??(????)，這是一個(gè)重大提升。此外，由于硬件施加的限制，SSM 在完整序列上的運(yùn)行時(shí)間復(fù)雜度與 Block Transformer 在 token 塊上的運(yùn)行時(shí)間復(fù)雜度相當(dāng)，這進(jìn)一步意味著 BST 層不存在速度瓶頸。該團(tuán)隊(duì)使用包含數(shù)十萬(wàn) token 的序列通過(guò)實(shí)驗(yàn)驗(yàn)證了這一點(diǎn)。

方法

這里研究的是通過(guò)僅解碼器語(yǔ)言模型實(shí)現(xiàn)下一 token 預(yù)測(cè)的問(wèn)題。

對(duì)狀態(tài)空間的前置說(shuō)明

狀態(tài)空間模型可以分為兩大類(lèi)：

狀態(tài)空間：結(jié)構(gòu)化核S4、S5、S4D、DSS遵循卷積核的一種結(jié)構(gòu)化初始化，方式是展開(kāi)一種線性時(shí)間不變式（LTI）動(dòng)態(tài)系統(tǒng)，如下所示：

其中的參數(shù)包括狀態(tài)矩陣 ??∈?^{N×N}，向量 ??∈?^{N×1}、??∈?^{1×N}、??∈?^{1×1}。SSM 會(huì)將一維的輸入信號(hào) u_k 映射成一維的輸出信號(hào) y_k。

顯式參數(shù)化的過(guò)濾器。不同于結(jié)構(gòu)化核，還可以將卷積核參數(shù)化為可訓(xùn)練的權(quán)重并優(yōu)化它們。但是，這會(huì)導(dǎo)致性能很差，除非對(duì)這些核使用特定類(lèi)型的正則化方法。替代 Transformer 的無(wú)注意力模型中也有使用可訓(xùn)練核的，比如 Hyena 涉及到沿核對(duì)權(quán)重進(jìn)行指數(shù)衰減。

Block-State Transformer（BST）層

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

這個(gè)輸出經(jīng)過(guò)了上下文編碼，也就是說(shuō)每個(gè)時(shí)間步驟的項(xiàng)目都可能包含有關(guān)該序列中元素之前的所有時(shí)間步驟的信息。他們從上下文序列收集一定數(shù)量 S 的「上下文狀態(tài)」，并使得 S ? L。

這些上下文狀態(tài)會(huì)被饋送給 Block Transformer，以替代 Block-Recurrent Transformer 中的「循環(huán)狀態(tài)向量」。如圖 1 右側(cè)所示，后續(xù)操作保持不變，只是無(wú)需再運(yùn)行 BRECT 單元的循環(huán)單元，因?yàn)楝F(xiàn)在是通過(guò) SSM 來(lái)維護(hù)上下文。除了上下文狀態(tài)，Block Transformer 的輸入中還有長(zhǎng)度 W 的 token 嵌入的塊/窗口；然后在這個(gè)窗口與上下文狀態(tài)上使用交叉注意力。然后將這個(gè)交叉注意力操作的輸出與自注意力在輸入嵌入上的輸出連接起來(lái)，之后是一個(gè)簡(jiǎn)單的投影。

SSM 不僅能在更長(zhǎng)時(shí)間尺度上保留信息，而且使用 SSM 來(lái)維持上下文狀態(tài)以替代循環(huán)單元，可以得到計(jì)算效率更高的層。通過(guò)將 SSM 整合進(jìn) Transformer 層，可以移除循環(huán)部分，從而讓 Block-State Transformer 層可以完全并行化。

上下文狀態(tài)

盡管從技術(shù)上看，最新的 SSM 輸出包含有關(guān)整個(gè)序列的信息，但僅從最后的狀態(tài)檢索單個(gè) token 可能是不可行的。為了彌補(bǔ)這一點(diǎn)，該團(tuán)隊(duì)將一系列狀態(tài)連接了起來(lái)，對(duì)應(yīng)于最新的 token 塊。這與 BRECT 采用的方法類(lèi)似。這種表征可以通過(guò)冗余來(lái)確保可檢索性和易訪問(wèn)性。

在新提出的方法中，上下文狀態(tài)是使用 SSM 的輸出構(gòu)建的，并會(huì)被饋送給 Transformer 的注意力頭。這些上下文狀態(tài)的構(gòu)建方式有很多。為了引導(dǎo)設(shè)計(jì)決策，該團(tuán)隊(duì)考慮了多種設(shè)計(jì)方案，包括使用單頭（Single-Head）、多頭（Multi-Head）或多過(guò)濾器（Multi-Filter）。其中單頭設(shè)計(jì)見(jiàn)圖 1。下圖 2 則展示了多頭和多過(guò)濾器的設(shè)計(jì)方案。

比較下來(lái)，多過(guò)濾器的記憶狀態(tài)的冗余最少，多頭次之，單頭的冗余最大。

結(jié)果

該團(tuán)隊(duì)在 PG19、GitHub 和 arXiv 三個(gè)數(shù)據(jù)集上進(jìn)行了實(shí)驗(yàn)，檢驗(yàn)了新提出的 BST 在不同長(zhǎng)度的英語(yǔ)文本、latex 科學(xué)文章和源代碼上的效果。下表 1 總結(jié)了實(shí)驗(yàn)結(jié)果。

下圖 3 則給出了長(zhǎng)度泛化分析并報(bào)告了困惑度。實(shí)驗(yàn)中，新模型和基準(zhǔn)模型的參數(shù)數(shù)量都約為 4 億，訓(xùn)練時(shí)的序列長(zhǎng)度為 4k，測(cè)試中的序列長(zhǎng)度為 {512, 16k, 65k}。

可以看到，在 PG19、GitHub 和 arXiv 上，當(dāng)序列長(zhǎng)度為 65k 時(shí)，BST:SH:S4-L 的困惑度最好。

在效率方面，下圖 4 左給出了 BST 層在 GPU 上的基準(zhǔn)測(cè)試結(jié)果。

可以看到 SSM 帶來(lái)了非常顯著的增長(zhǎng)——比包含循環(huán)單元的 Block-Recurrent Transformer 快 6-11 倍；即使在序列長(zhǎng)度達(dá)到 65k token 時(shí)，還依然能有 6 倍的提升，而這時(shí)候硬件就已經(jīng)開(kāi)始飽和了。當(dāng)使用結(jié)構(gòu)化的 SSM 時(shí)，計(jì)算復(fù)雜度與 SSM 的內(nèi)部記憶狀態(tài)大小 N 緊密相關(guān)。對(duì)于報(bào)告的性能，N = 16。

研究者表示，如果使用其它自動(dòng)微分框架中近期引入的更快的針對(duì)硬件的 I/O 感知型實(shí)現(xiàn)，BST 方法的速度還能更快。

更多技術(shù)細(xì)節(jié)和實(shí)驗(yàn)結(jié)果參閱原論文。