大型語言模型（LLMs）之所以能夠在語言理解與生成上取得巨大成功，Transformer架構(gòu)是其強(qiáng)大的支撐。從開源的Mistral，到OpenAI開發(fā)的閉源模型ChatGPT，都采用了這一架構(gòu)。

然而，技術(shù)的探索從未止步。為進(jìn)一步提升LLMs的性能，學(xué)界正在研發(fā)能夠超越Transformer的新架構(gòu)。其中，Mamba模型以其創(chuàng)新的狀態(tài)空間模型（State Space Model）成為研究的焦點(diǎn)。

本文介紹Mamba模型及其在語言建模領(lǐng)域的應(yīng)用，逐步解析狀態(tài)空間模型的基本概念，并通過豐富的可視化內(nèi)容，讓讀者直觀地理解這一技術(shù)如何有望挑戰(zhàn)現(xiàn)有的Transformer架構(gòu)。

1 Transformer架構(gòu)的挑戰(zhàn)

我們先對(duì)Transformer架構(gòu)做一個(gè)快速回顧，并指出其存在的一個(gè)主要缺陷。

Transformer架構(gòu)將文本輸入視為由一系列token構(gòu)成的序列。

其核心優(yōu)勢在于，無論面對(duì)何種輸入，都能追溯到序列中的早期token，以此來推導(dǎo)出其深層的語義表示。

1.1 Transformer的核心組件

Transformer由架構(gòu)兩部分組成：編碼器和解碼器。編碼器負(fù)責(zé)解析文本，而解碼器則負(fù)責(zé)生成文本。這種結(jié)構(gòu)的結(jié)合，使之能夠勝任從文本翻譯到內(nèi)容創(chuàng)作的多種任務(wù)。

進(jìn)一步地，我們可以僅利用解碼器部分來創(chuàng)建生成式模型。這種基于Transformer的模型，即生成預(yù)訓(xùn)練Transformer（GPT），通過解碼器來續(xù)寫或補(bǔ)全輸入的文本，展現(xiàn)出其在文本生成方面的強(qiáng)大能力。

1.2 自注意力的高效訓(xùn)練

單個(gè)解碼器塊由兩個(gè)主要部分組成，即“掩蔽自注意力機(jī)制（Masked Self-attention）”和“前饋神經(jīng)網(wǎng)絡(luò)（Feedforward Neural Network）”。

自注意力是這些模型運(yùn)行如此良好的主要原因。它提供了整個(gè)序列的未壓縮視圖，并加快了訓(xùn)練速度。

具體來說，自注意力機(jī)制的工作原理是通過創(chuàng)建一個(gè)矩陣，該矩陣對(duì)序列中的每個(gè)token與之前所有token進(jìn)行比較，并通過計(jì)算它們之間的相關(guān)性來確定權(quán)重。

在訓(xùn)練過程中，自注意力矩陣是一次性整體構(gòu)建的，這表示不需要依次等待每個(gè)token的注意力計(jì)算完成，而是可以同時(shí)進(jìn)行整個(gè)序列的注意力計(jì)算。例如，在處理“我”和“名字”的關(guān)聯(lián)之前，無需先完成“名字”和“是”的關(guān)聯(lián)計(jì)算。

這種設(shè)計(jì)實(shí)現(xiàn)了訓(xùn)練過程的并行化，極大地提升了訓(xùn)練速度，使Transformer架構(gòu)在處理大規(guī)模數(shù)據(jù)集時(shí)更加高效。

1.3 訓(xùn)練與推理的矛盾

然而，Transformer架構(gòu)也有其局限性。每當(dāng)生成新的token時(shí)，必須對(duì)整個(gè)序列的注意力權(quán)重重新進(jìn)行計(jì)算，哪怕此前已經(jīng)生成了若干token。

生成長度為L的序列需要大約L2次計(jì)算，隨著序列的延長，計(jì)算成本會(huì)急劇上升。

這種對(duì)序列全面重新計(jì)算的需求，是Transformer架構(gòu)的一個(gè)主要瓶頸。

接下來，讓我們看看傳統(tǒng)的遞歸神經(jīng)網(wǎng)絡(luò)（RNN）是如何克服這一推理過程中的效率問題。

1.4 RNN的潛力

遞歸神經(jīng)網(wǎng)絡(luò)（RNN）是一種處理序列數(shù)據(jù)的網(wǎng)絡(luò)結(jié)構(gòu)。在序列的每個(gè)時(shí)間點(diǎn)，RNN接收兩個(gè)輸入：當(dāng)前時(shí)間點(diǎn)t的輸入數(shù)據(jù)和上一個(gè)時(shí)間點(diǎn)t-1的隱藏狀態(tài)，以此來計(jì)算下一個(gè)隱藏狀態(tài)并預(yù)測輸出結(jié)果。

RNN具有循環(huán)機(jī)制，能夠?qū)v史信息傳遞至下一步，類似于將每一步的信息“串聯(lián)”起來。這種機(jī)制可以通過可視化的方式“展開”，以便更清晰地理解其工作原理。

在生成輸出時(shí)，RNN僅依賴于前一步驟的隱藏狀態(tài)和當(dāng)前的輸入數(shù)據(jù)，避免了像Transformer那樣需要重新計(jì)算整個(gè)序列的歷史隱藏狀態(tài)。

正因如此，RNN在進(jìn)行推理時(shí)速度較快，因?yàn)樗挠?jì)算量與序列長度呈線性關(guān)系，理論上能夠處理無限長的上下文。

舉例來說，當(dāng)我們將RNN應(yīng)用于之前的輸入文本。

每個(gè)隱藏狀態(tài)都是對(duì)之前所有狀態(tài)的壓縮匯總。

但這里存在一個(gè)問題：隨著時(shí)間的推移，比如在生成名字"Maarten"時(shí)，最后一個(gè)隱藏狀態(tài)可能已經(jīng)丟失了對(duì)"Hello"的記憶，因?yàn)镽NN在每一步只考慮了前一個(gè)狀態(tài)的信息。

此外，RNN的這種順序依賴性也導(dǎo)致了另一個(gè)問題：它的訓(xùn)練過程無法并行化，必須按順序逐步進(jìn)行。

與Transformer相比，RNN在推理速度上具有優(yōu)勢，但在訓(xùn)練并行化方面卻存在不足。

這就引出了一個(gè)問題：能否找到一種結(jié)合了Transformer訓(xùn)練并行化優(yōu)勢和RNN線性推理優(yōu)勢的架構(gòu)呢？

答案是肯定的，Mamba模型就是。在深入了解Mamba架構(gòu)之前，先來了解狀態(tài)空間模型的世界。

2 狀態(tài)空間模型（SSM）

狀態(tài)空間模型（SSM），像Transformer和RNN一樣，處理信息序列，如文本和信號(hào)。在這一部分中，我們將了解SSM的基礎(chǔ)知識(shí)以及其與文本數(shù)據(jù)的關(guān)系。

2.1 什么是狀態(tài)空間

狀態(tài)空間模型是一種通過數(shù)學(xué)方法全面描述系統(tǒng)狀態(tài)的方式，包含了描述系統(tǒng)所必需的全部最小變量。

簡單來說，就像我們?cè)诿詫m中尋找路徑，狀態(tài)空間就是那張展示所有可能位置（即狀態(tài)）的地圖。在這張地圖上，每個(gè)點(diǎn)都代表一個(gè)獨(dú)特的位置，并且攜帶了如距離出口遠(yuǎn)近等具體信息。

進(jìn)一步簡化這個(gè)概念，可以將“狀態(tài)空間表示”理解為這張地圖的提煉，它不僅告訴我們當(dāng)前所在的位置（即當(dāng)前狀態(tài)），還展示了可能的目的地（未來狀態(tài)），以及如何通過特定的行動(dòng)（比如右轉(zhuǎn)或左轉(zhuǎn)）達(dá)到下一個(gè)狀態(tài)。

雖然狀態(tài)空間模型依賴方程和矩陣來捕捉系統(tǒng)的行為，但其核心目標(biāo)是追蹤系統(tǒng)的位置、可能的移動(dòng)方向及其變化路徑。

在這個(gè)模型中，用以描述狀態(tài)的變量，如示例中的X和Y坐標(biāo)或者到出口的距離，統(tǒng)稱為“狀態(tài)向量”。

這聽起來有點(diǎn)熟悉，因?yàn)樵谡Z言模型中，類似的嵌入或向量經(jīng)常用來描述輸入序列的“狀態(tài)”。例如，你當(dāng)前的位置狀態(tài)就可以通過一個(gè)向量來表示：

在神經(jīng)網(wǎng)絡(luò)的語境下，系統(tǒng)的“狀態(tài)”通常指的是其隱藏狀態(tài)，這在生成新token的過程中，尤其是在大型語言模型的背景下，扮演著至關(guān)重要的角色。

本文轉(zhuǎn)載自 ??AI科技論談??，作者： AI科技論談