Transformer、RNN和SSM的相似性探究：揭示看似不相關(guān)的LLM架構之間的聯(lián)系

通過(guò)探索看似不相關(guān)的大語(yǔ)言模型(LLM)架構之間的潛在聯(lián)系，我們可能為促進(jìn)不同模型間的思想交流和提高整體效率開(kāi)辟新的途徑。

盡管Mamba等線(xiàn)性循環(huán)神經(jīng)網(wǎng)絡(luò )(RNN)和狀態(tài)空間模型(SSM)近來(lái)備受關(guān)注，Transformer架構仍然是LLM的主要支柱。這種格局可能即將發(fā)生變化：像Jamba、Samba和Griffin這樣的混合架構展現出了巨大的潛力。這些模型在時(shí)間和內存效率方面明顯優(yōu)于Transformer，同時(shí)在能力上與基于注意力的LLM相比并未顯著(zhù)下降。

近期研究揭示了不同架構選擇之間的深層聯(lián)系，包括Transformer、RNN、SSM和matrix mixers，這一發(fā)現具有重要意義，因為它為不同架構間的思想遷移提供了可能。本文將深入探討Transformer、RNN和Mamba 2，通過(guò)詳細的代數分析來(lái)理解以下幾點(diǎn):

● Transformer在某些情況下可以視為RNN(第2節)

● 狀態(tài)空間模型可能隱藏在自注意力機制的掩碼中(第4節)

● Mamba在特定條件下可以重寫(xiě)為掩碼自注意力(第5節)

這些聯(lián)系不僅有趣，還可能對未來(lái)的模型設計產(chǎn)生深遠影響。

LLM中的掩碼自注意力機制

首先,讓我們回顧一下經(jīng)典的LLM自注意力層的結構：

更詳細的結構如下：

自注意力層的工作流程如下：

● 將查詢(xún)矩陣Q和鍵矩陣K相乘，得到一個(gè)L×L的矩陣，包含查詢(xún)和鍵的標量積。

● 對結果矩陣進(jìn)行歸一化。

● 將歸一化后的矩陣與L×L的注意力掩碼進(jìn)行元素級乘法。圖中展示了默認的因果掩碼——左側的0-1矩陣。這一步驟將較早查詢(xún)與較晚鍵的乘積置零，防止注意力機制"看到未來(lái)"。

● 對結果應用softmax函數。

● 最后，將注意力權重矩陣A與值矩陣V相乘。輸出的第t行可表示為：

這意味著(zhù)第i個(gè)值是通過(guò)"第t個(gè)查詢(xún)對第i個(gè)鍵的注意力權重"來(lái)加權的。

這種架構中的多個(gè)設計選擇都可能被修改。接下來(lái)我們將探討一些可能的變體。

線(xiàn)性化注意力

注意力公式中的Softmax函數確保了值是以和為1的正系數混合的。這種設計保持了某些統計特性，但同時(shí)也帶來(lái)了限制。例如即使我們希望利用結合律，如(QK^T)V = Q(K^TV),也無(wú)法突破Softmax的限制。

為什么結合律如此重要？因為改變乘法順序可能顯著(zhù)影響計算復雜度：

左側公式需要計算一個(gè)L×L矩陣，如果這個(gè)矩陣完全顯現在內存中，復雜度為O(L2d),內存消耗為O(L2)。右側公式需要計算一個(gè)d×d矩陣，復雜度為O(Ld2)，內存消耗為O(d2)。

隨著(zhù)上下文長(cháng)度L的增加,左側公式的計算成本rapidly become prohibitively非常的高。為了解決這個(gè)問(wèn)題，我們可以考慮移除Softmax。詳細展開(kāi)帶有Softmax的公式:

其中：

是Softmax函數。指數函數是主要的障礙,它阻止了我們從中提取任何項。如果我們直接移除指數函數：

那么歸一化因子

也隨之消失。

這個(gè)簡(jiǎn)化后的公式存在一個(gè)問(wèn)題:q_t^T k_s不能保證為正,這可能導致值以不同符號的系數混合,這在理論上是不合理的。更糟糕的是,分母可能為零,會(huì )導致計算崩潰。為了緩解這個(gè)問(wèn)題,我們可以引入一個(gè)"良好的"元素級函數φ(稱(chēng)為核函數):

原始研究建議使用φ(x) = 1 + elu(x)作為核函數。

這種注意力機制的變體被稱(chēng)為線(xiàn)性化注意力。它的一個(gè)重要優(yōu)勢是允許我們利用結合律:

括號中M, K^T和V之間的關(guān)系現在變得相當復雜，不再僅僅是普通的矩陣乘法和元素級乘法。我們將在下一節詳細討論這個(gè)計算單元。

如果M是一個(gè)因果掩碼，即對角線(xiàn)及以下為1，對角線(xiàn)以上為0：

那么計算可以進(jìn)一步簡(jiǎn)化：

這可以通過(guò)一種簡(jiǎn)單的遞歸方式計算:

這是在2020年ICML上首次提出線(xiàn)性化注意力的論文"Transformers are RNNs"。在這個(gè)公式中，我們有兩個(gè)隱藏狀態(tài)：向量z_t和矩陣h_t(φ(k_t)^T v_t是列向量乘以行向量，得到一個(gè)d×d矩陣。

而近期的研究often以更簡(jiǎn)化的形式呈現線(xiàn)性化注意力，去除了φ函數和分母:

線(xiàn)性化注意力具有兩個(gè)主要優(yōu)勢：

● 作為遞歸機制，它在推理時(shí)相對于序列長(cháng)度L具有線(xiàn)性復雜度。

● 作為T(mén)ransformer模型，它可以高效地并行訓練。

但是你可能會(huì )問(wèn)：如果線(xiàn)性化注意力如此優(yōu)秀，為什么它沒(méi)有在所有LLM中廣泛應用？我們在討論注意力的二次復雜度問(wèn)題？實(shí)際上基于線(xiàn)性化注意力的LLM在訓練過(guò)程中stability較低，且capability略遜于標準自注意力。這可能是因為固定的d×d形狀的瓶頸比可調整的L×L形狀的瓶頸能傳遞的信息更少。

進(jìn)一步探索

RNN和線(xiàn)性化注意力之間的聯(lián)系在近期的多項研究中得到了重新發(fā)現和深入探討。一個(gè)common pattern是使用具有如下更新規則的矩陣隱藏狀態(tài)：

其中k_t和v_t可以視為某種"鍵"和"值"，RNN層的輸出形式為：

這本質(zhì)上等同于線(xiàn)性注意力。下面兩篇論文提供了有趣的一些樣例：

1、xLSTM (2024年5月)：該論文提出了對著(zhù)名的LSTM遞歸架構的改進(jìn)。其mLSTM塊包含一個(gè)矩陣隱藏狀態(tài)，更新方式如下：

輸出通過(guò)將這個(gè)狀態(tài)與一個(gè)"查詢(xún)"相乘得到。(注意：該論文的線(xiàn)性代數設置與我們的相反，查詢(xún)、鍵和值是列向量而非行向量，因此v_t k_t^T的順序看起來(lái)可能有些奇怪。)

2、Learning to (learn at test time) (2024年7月)：這是另一種具有矩陣隱藏狀態(tài)的RNN架構，它的隱藏狀態(tài)W是一個(gè)函數的參數，在t的迭代過(guò)程中通過(guò)梯度下降優(yōu)化：

這里的設置也是轉置的，因此順序看起來(lái)有些不同。盡管數學(xué)表達比W_t = W_{t-1} + v_t k_t^T更復雜，但可以簡(jiǎn)化為這種形式。

以上兩篇論文我們都詳細介紹過(guò)，有興趣的可以自行搜索。

注意力掩碼

在簡(jiǎn)化了掩碼注意力機制后，我們可以開(kāi)始探索其潛在的發(fā)展方向。一個(gè)明顯的研究方向是選擇不同的下三角矩陣（確保不會(huì )"看到未來(lái)"）作為掩碼M，而不是簡(jiǎn)單的0-1因果掩碼。在進(jìn)行這種探索之前，我們需要解決由此帶來(lái)的效率問(wèn)題。

在前一節中，我們使用了一個(gè)簡(jiǎn)單的0-1因果掩碼M，這使得遞歸計算成為可能。但在一般情況下，這種遞歸技巧不再適用：

系數m_ts不再相同，也不存在將y_3與y_2關(guān)聯(lián)的簡(jiǎn)單遞歸公式。因此，對于每個(gè)t我們都需要從頭開(kāi)始計算總和，這使得計算復雜度再次變?yōu)長(cháng)的二次方而不是線(xiàn)性的。

解決這個(gè)問(wèn)題的關(guān)鍵在于我們不能使用任意的掩碼M，而應該選擇特殊的、"良好"的掩碼。我們需要那些可以快速與其他矩陣相乘（注意不是元素級乘法）的掩碼。為了理解如何從這種特性中獲益，讓我們詳細分析如何高效計算：

首先明確這個(gè)表達式的含義：

如果深入到單個(gè)索引級別：

為了便于后續討論，可以用不同的顏色標記索引，而不是塊：

現在我們可以提出一個(gè)四步算法：

步驟1. 利用K和V創(chuàng )建一個(gè)三維張量Z，其中：

（每個(gè)軸都標注了其長(cháng)度。）這一步驟需要O(Ld2)的時(shí)間和內存復雜度。值得注意的是，如果我們在洋紅色軸t上對這個(gè)張量求和，我們將得到矩陣乘積K^T V：

步驟2. 將M乘以這個(gè)張量（注意不是元素級乘法）。M乘以Z沿著(zhù)洋紅色軸t的每個(gè)"列"。

這正好得到：

將這個(gè)結果記為H。接下來(lái)只需要將所有內容乘以q，這將在接下來(lái)的兩個(gè)步驟中完成。

步驟3a. 取Q并與H的每個(gè)j = const層進(jìn)行元素級乘法：

這將得到：

這一步驟需要O(Ld2)的時(shí)間和內存復雜度。

步驟3b. 沿i軸對結果張量求和：

這一步驟同樣需要O(Ld2)的時(shí)間和內存復雜度。最終得到了所需的結果：

在這個(gè)過(guò)程中，最關(guān)鍵的是第二步，我們故意省略了其復雜度分析。一個(gè)簡(jiǎn)單的估計是：

每次矩陣乘法需要O(L2)的復雜度，重復d2次

這將導致一個(gè)巨大的O(L2d2)復雜度。但是我們的目標是選擇特殊的M，使得將M乘以一個(gè)向量的復雜度為O(RL)，其中R是某個(gè)不太大的常數。

例如如果M是0-1因果矩陣，那么與它相乘實(shí)際上就是計算累積和，這可以在O(L)時(shí)間內完成。但還存在許多其他具有快速向量乘法特性的結構化矩陣選項。

在下一節中將討論這種矩陣類(lèi)型的一個(gè)重要例子——半可分離矩陣，它與狀態(tài)空間模型有著(zhù)密切的聯(lián)系。

半可分離矩陣與狀態(tài)空間模型

讓我們回顧一下（離散化的）狀態(tài)空間模型（SSM）的定義。SSM是一類(lèi)連接1維輸入x_t、r維隱藏狀態(tài)h_t和1維輸出u_t的序列模型，其數學(xué)表達式如下：

在離散形式中，SSM本質(zhì)上是一個(gè)帶有跳躍連接的復雜線(xiàn)性RNN。為了簡(jiǎn)化后續討論，我們甚至可以通過(guò)設置D_t = 0來(lái)忽略跳躍連接。

讓我們將SSM表示為單個(gè)矩陣乘法：

其中

M是一個(gè)下三角矩陣，類(lèi)似于我們之前討論的注意力掩碼。

這種類(lèi)型的矩陣具有一個(gè)重要的優(yōu)勢：

一個(gè)L × L的下三角矩陣，如果其元素可以以這種方式表示，則可以使用O(rL)的內存存儲，并且具有O(rL)的矩陣-向量乘法復雜度，而不是默認的O(L2)。

這意味著(zhù)每個(gè)狀態(tài)空間模型都對應一個(gè)結構化的注意力掩碼M，可以在具有線(xiàn)性化注意力的高效Transformer模型中使用。

即使沒(méi)有周?chē)牟樵?xún)-鍵-值機制，半可分離矩陣M本身已經(jīng)相當復雜和富有表現力。它本身可能就是一個(gè)掩碼注意力機制。我們將在下一節中詳細探討這一點(diǎn)。

狀態(tài)空間對偶性

在這里，我們將介紹Mamba 2論文中的一個(gè)核心結果。

讓我們再次考慮y = Mu，其中u = u(x)是輸入的函數，M是一個(gè)可分離矩陣。如果我們考慮一個(gè)非常特殊的情況，其中每個(gè)A_t都是一個(gè)標量矩陣：A_t = a_t I。在這種情況下公式變得特別簡(jiǎn)單：

這里的

只是一個(gè)標量。還可以將C_i和B_i堆疊成矩陣B和C，使得：

現在我們還需要定義矩陣

然后就可以很容易地驗證：

這個(gè)表達式是否看起來(lái)很熟悉？這實(shí)際上是一個(gè)掩碼注意力機制，其中：

● G作為掩碼

● C作為查詢(xún)矩陣Q

● B作為轉置的鍵矩陣K^T

● u作為值矩陣V

在經(jīng)典的SSM中，B和C是常量。但在Mamba模型中，它們被設計為依賴(lài)于數據，這進(jìn)一步強化了與注意力機制的對應關(guān)系。這種特定狀態(tài)空間模型與掩碼注意力之間的對應關(guān)系在Mamba 2論文中被稱(chēng)為狀態(tài)空間對偶性。

進(jìn)一步探索

使用矩陣混合器而不是更復雜的架構并不是一個(gè)全新的idea。一個(gè)早期的例子是是MLP-Mixer，它在計算機視覺(jué)任務(wù)中使用MLP而不是卷積或注意力來(lái)進(jìn)行空間混合。

盡管當前研究主要集中在大語(yǔ)言模型（LLM）上，但也有一些論文提出了用于編碼器模型的非Transformer、矩陣混合架構。例如：

● 來(lái)自Google研究的FNet，其矩陣混合器M基于傅里葉變換。

● Hydra，除了其他創(chuàng )新外，還提出了半可分離矩陣在非因果（非三角）工作模式下的適應性方案。

總結

本文深入探討了Transformer、循環(huán)神經(jīng)網(wǎng)絡(luò )（RNN）和狀態(tài)空間模型（SSM）之間的潛在聯(lián)系。文章首先回顧了傳統的掩碼自注意力機制，然后引入了線(xiàn)性化注意力的概念，解釋了其計算效率優(yōu)勢。接著(zhù)探討了注意力掩碼的優(yōu)化，引入了半可分離矩陣的概念，并闡述了其與狀態(tài)空間模型的關(guān)系。最后介紹了狀態(tài)空間對偶性，揭示了特定狀態(tài)空間模型與掩碼注意力之間的對應關(guān)系。通過(guò)這些分析，展示了看似不同的模型架構之間存在深層聯(lián)系，為未來(lái)模型設計和跨架構思想交流提供了新的視角和可能性。