Griffin
Griffin ：將門控線性循環與局部注意力混合以實現高效的語言模型
arXiv
模型架構
我們所有的模型都包含以下組件：（i）殘差塊，（ii）MLP 塊，以及（iii）時間混合塊。雖然 (i) 和 (ii) 在所有模型中都是相同的，但我們考慮三個時間混合區塊：全局多查詢注意力 (MQA)、局部（滑動視窗）MQA 和我們提出的循環區塊。作為循環塊的一部分，我們使用實門線性循環單元（RG-LRU）——一種受線性循環單元 Orvieto 等人，2023b 啟發的新型循環層。
如圖 2(a) 所示，殘差塊定義了我們模型的全域結構，其靈感來自於前範數 Transformers（Xiong 等人，2020）。嵌入輸入序列後，我們將其傳遞給 $N$這樣的塊（ $N$表示模型深度），然後我們應用 RMSNormZhang 和 Sennrich，2019 來產生最終的活化。為了計算 token 機率，我們應用了最後一個線性層，後面跟著一個 softmax。此層的權重與輸入嵌入層共享。
殘塊
圖 2：a）我們的模式架構的主要支柱是殘差塊，它是堆疊的 $N$次。 b) 我們使用的門控 MLP 模組。 c）我們提出的循環塊作為多查詢注意力（MQA）的替代方案。它使用我們提出的 RG-LRU 層，在 2.4 節中定義。
殘差塊包含兩個依序應用的元件。第一個組件採用隱藏狀態 $chi$並套用 RMSNorm Zhuang 和 Sennrich，2019，然後是時間混合區塊。然後我們將輸出與跳躍連接合併 $chi$透過加法。類似地，第二個元件應用 RMSNorm，後面跟著 MLP 模組，然後將其輸出與來自 RMSNorm 輸入的跳躍連接合併。該塊如圖 2 (a) 所示。
MLP塊
我們使用門控 MLP 模組 Dauphin et al., 2017（如圖 2(b) 所示），它根據維度輸入創建兩個分支 $D$ 。我們應用具有輸出維度的線性層 $MD$在每個分支上，其中 $M$表示擴展因子。為了簡單起見，我們使用 $M=3$貫穿整個這項工作。我們在其中一個分支上應用 GeLU 非線性 Hendrycks 和 Gimpel，2016，然後透過逐元素乘法合併它們，類似於 GeGeLU Shazeer，2020。 $D$在 GeGeLU 層的輸出上。
時間混合區塊
時間混合區塊是我們模型的元件，它聚合序列中不同時間位置的隱藏層活化。我們考慮三個時間混合區塊：全域 MQA Shazeer，2019、局部 MQA Beltagy 等人，2020 以及我們提出的 Recurrent 區塊。
全域多查詢注意力
除非另有說明，我們使用 MQA 而不是 MHA 來提高 Transformer 基線 Shazeer，2019 的推理速度。 $D_{頭}=128$ ，我們固定注意力頭的數量 $H$這樣 $HD_{頭}=D$ 。這需要模型尺寸 $D$是 128 的倍數。
局部滑動視窗注意力
使用全域注意力的主要缺點之一是其計算複雜度隨序列長度呈二次方增長。為了解決這個問題，一些工作已經開始採用局部注意力 Beltagy et al., 2020，也稱為滑動視窗注意力。它允許每個位置在過去只參與固定數量的代幣。這不僅減少了計算的 FLOP，而且還將 KV 快取的大小限制為視窗的大小，使其不再是序列長度的二次方。所有其他細節與全局 MQA 相同。
循環區塊
我們的循環區塊（圖 2(c)）類似於 GSS 區塊 Mehta et al., 2022 以及 Mamba Gu 和 Dao, 2023 所使用的區塊。 $D$並應用兩個具有輸出維度的線性層 $D_{RNN}$並行地建立兩個分支。在第一個分支上，我們應用了一個小的可分離的Conv1D 層，其靈感來自於H3 Dao 等人，2022b 中的Shift-SSM，時間濾波器維度為4。 ，只有 $4D$參數。我們在 Conv1D 層後面加上我們提出的 RG-LRU 層（定義如下）。然後我們應用具有輸出維度的最終線性層 $D$ 。
實門線性循環單元 (RG-LRU)
我們提出的 RG-LRU 層具有受線性循環單元 (LRU) Orvieto 等人，2023b 啟發的簡單循環，但結合了受非線性 RNN 文獻啟發的門控機制，特別是 LSTM Hochreiter 和 Schmidhuber，1997 和GRUs Chung et al., 2014。
$$begin{align} r_t &= sigma(W_{a} x_t + b_a), & text{遞歸閘} \ i_t &= sigma(W_{x} x_t + b_x), & text{輸入閘} \ a_t &= a^{cr_t}, & text{} \ h_t &= a_t odot h_{t-1} + sqrt{1 - a_t^2} odot (i_t odot x_t ）。 & text{} end{對齊}$$
該層的輸出為 $y_t=h_t$ ，和非線性 $西格瑪$方程式中是 sigmoid 函數。經常性體重 $a$等式（4）中是對角線。因此，所有操作都是逐元素的。我們參數化 $a$等式（3）中為 $a=sigma(Lambda)$ ， 在哪裡 $Lambda$是一個可學習的參數。這保證了 $0 <= a <= 1$ ，確保復發穩定。變數 $c$是一個標量值常數，設定為 8。 $a^{cr_t}$在日誌空間中（請參閱附錄 A）。此層在兩個輸入上都有門 $x$和經常性體重 $a$ 。然而，這兩個門都不依賴循環狀態 $h_{t-1}$ ，這確保了計算可以在設備上有效執行。我們初始化兩者 $W_{a}$和 $W_{b}$使用 LeCun init LeCun et al., 2002。 $Lambda$這樣 $a^c$均勻分佈在 $0.9$和 $0.999$在訓練開始時，類似於（Orvieto et al., 2023b.）。與 SSM 文獻中的許多最新作品不同，RG-LRU 不使用受正交多項式理論啟發的初始化 Gu 等人，2020，並且它也沒有定義為底層連續系統的離散化 Gu 等人，2021a。與原始 LRU 層不同，我們在遞歸中不使用複雜代數。雖然使用複雜的遞歸會導致更具表現力的層Orvieto et al., 2023a，但我們發現複雜的遞歸對實踐中的語言建模不利，正如Gu 和Dao, 2023 所觀察到的那樣。 B）
門行為
輸入門 $i_t$類似於LSTM中的，可以對輸入進行過濾（或按比例縮小） $x_t$ 。然而，據我們所知，我們的遞歸門 $r_t$與文獻中的其他門控機制不同。例如，Mamba Gu and Dao, 2023 中提出的選擇機制相當於 GRU 的更新門，它插值 $x_t$ 。它對隱藏狀態的影響使其能夠重置其狀態並忘記過去保存的任何訊息，類似於 LSTM 中的遺忘門。相較之下，我們的遞歸閘門可以在Orvieto 等人2023a 的標準LRU 更新和先前的隱藏狀態之間進行近似插值，這使得它能夠有效地丟棄輸入並保留先前歷史記錄中的所有資訊（有關更多詳細信息，請參閱附錄A） ）。我們認為這個門的關鍵作用是透過減少無資訊輸入的影響來使模型實現超指數記憶。
