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

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

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

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

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

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

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

1.1 Transformer的核心組件

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

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

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

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

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

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

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

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

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

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

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

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

接下來，讓我們看看傳統(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)。在序列的每個時間點，RNN接收兩個輸入：當前時間點t的輸入數(shù)據(jù)和上一個時間點t-1的隱藏狀態(tài)，以此來計算下一個隱藏狀態(tài)并預(yù)測輸出結(jié)果。

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

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

正因如此，RNN在進行推理時速度較快，因為它的計算量與序列長度呈線性關(guān)系，理論上能夠處理無限長的上下文。

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

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

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

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

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

這就引出了一個問題：能否找到一種結(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一樣，處理信息序列，如文本和信號。在這一部分中，我們將了解SSM的基礎(chǔ)知識以及其與文本數(shù)據(jù)的關(guān)系。

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

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

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

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

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

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

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

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

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