Transformer、RNN和SSM的相似性探究：揭示看似不相關的LLM架構之間的聯系

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

盡管Mamba等線性循環神經網絡(RNN)和狀態空間模型(SSM)近來備受關注，Transformer架構仍然是LLM的主要支柱。這種格局可能即將發生變化：像Jamba、Samba和Griffin這樣的混合架構展現出了巨大的潛力。這些模型在時間和內存效率方面明顯優于Transformer，同時在能力上與基于注意力的LLM相比并未顯著下降。

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

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

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

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

這些聯系不僅有趣，還可能對未來的模型設計產生深遠影響。

LLM中的掩碼自注意力機制

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

更詳細的結構如下：

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

● 將查詢矩陣Q和鍵矩陣K相乘，得到一個L×L的矩陣，包含查詢和鍵的標量積。

● 對結果矩陣進行歸一化。

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

● 對結果應用softmax函數。

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

這意味著第i個值是通過"第t個查詢對第i個鍵的注意力權重"來加權的。

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

線性化注意力

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

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

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

隨著上下文長度L的增加,左側公式的計算成本rapidly become prohibitively非常的高。為了解決這個問題，我們可以考慮移除Softmax。詳細展開帶有Softmax的公式:

其中：

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

那么歸一化因子

也隨之消失。

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

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

這種注意力機制的變體被稱為線性化注意力。它的一個重要優勢是允許我們利用結合律:

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

如果M是一個因果掩碼，即對角線及以下為1，對角線以上為0：

那么計算可以進一步簡化：

這可以通過一種簡單的遞歸方式計算:

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

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

線性化注意力具有兩個主要優勢：

● 作為遞歸機制，它在推理時相對于序列長度L具有線性復雜度。

● 作為Transformer模型，它可以高效地并行訓練。

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

進一步探索

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

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

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

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

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

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

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

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

注意力掩碼

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

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

系數m_ts不再相同，也不存在將y_3與y_2關聯的簡單遞歸公式。因此，對于每個t我們都需要從頭開始計算總和，這使得計算復雜度再次變為L的二次方而不是線性的。

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

首先明確這個表達式的含義：

如果深入到單個索引級別：

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

現在我們可以提出一個四步算法：

步驟1. 利用K和V創建一個三維張量Z，其中：

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

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

這正好得到：

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

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

這將得到：

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

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

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

在這個過程中，最關鍵的是第二步，我們故意省略了其復雜度分析。一個簡單的估計是：

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

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

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

在下一節中將討論這種矩陣類型的一個重要例子——半可分離矩陣，它與狀態空間模型有著密切的聯系。

半可分離矩陣與狀態空間模型

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

在離散形式中，SSM本質上是一個帶有跳躍連接的復雜線性RNN。為了簡化后續討論，我們甚至可以通過設置D_t = 0來忽略跳躍連接。

讓我們將SSM表示為單個矩陣乘法：

其中

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

這種類型的矩陣具有一個重要的優勢：

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

這意味著每個狀態空間模型都對應一個結構化的注意力掩碼M，可以在具有線性化注意力的高效Transformer模型中使用。

即使沒有周圍的查詢-鍵-值機制，半可分離矩陣M本身已經相當復雜和富有表現力。它本身可能就是一個掩碼注意力機制。我們將在下一節中詳細探討這一點。

狀態空間對偶性

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

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

這里的

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

現在我們還需要定義矩陣

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

這個表達式是否看起來很熟悉？這實際上是一個掩碼注意力機制，其中：

● G作為掩碼

● C作為查詢矩陣Q

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

● u作為值矩陣V

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

進一步探索

使用矩陣混合器而不是更復雜的架構并不是一個全新的idea。一個早期的例子是是MLP-Mixer，它在計算機視覺任務中使用MLP而不是卷積或注意力來進行空間混合。

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

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

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

總結

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