與以往方法的聯(lián)系和區(qū)別

表 1 從不同角度對(duì) RetNet 與以往的方法進(jìn)行了比較。對(duì)比結(jié)果與圖 2 所示的「不可能三角」相呼應(yīng)。此外，RetNet 對(duì)于長(zhǎng)序列具有線性記憶復(fù)雜性，因?yàn)樗捎昧朔謮K循環(huán)表示。

Transformer：retention 的并行表示與 Transformers [VSP^+17] 有著相似的思路。最相關(guān)的 Transformer 變體是 Lex Transformer [SDP^+22]，它實(shí)現(xiàn)了 xPos 作為位置嵌入。如式 (3) 所示，retention 的推導(dǎo)與 xPos 一致。與注意力相比，retention 消除了 softmax 并使循環(huán)公式成為可能，這非常有利于推理。

S4：與式 (2) 不同，如果 Q_n 和 K_n 是 content-unaware 的，則公式可簡(jiǎn)并為 S4 [GGR21]，其中

Linear Attention：變體通常使用各種 kernel來取代 softmax 函數(shù)。然而，線性注意力難以有效地編碼位置信息，導(dǎo)致模型性能下降。此外，研究者從頭開始重新檢查序列建模，而不是以近似 softmax 為目標(biāo)。

AFT/RWKV：Attention Free Transformer (AFT) 簡(jiǎn)化了點(diǎn)積對(duì)元素運(yùn)算的關(guān)注，并將 softmax 移動(dòng)到關(guān)鍵向量。RWKV 用指數(shù)衰減取代 AFT 的位置嵌入，并循環(huán)運(yùn)行模型進(jìn)行訓(xùn)練和推理。相比之下，retention 保留了高維狀態(tài)來編碼序列信息，有助于提高表達(dá)能力和性能。

xPos/RoPE：與為 Transformers 提出的相對(duì)位置嵌入方法相比，公式（3）呈現(xiàn)出與 xPos [SDP^+22] 和 RoPE [SLP^+21] 類似的表達(dá)式。

Sub-LayerNorm：如公式（8）所示，retention 層使用 Sub-LayerNorm [WMH^+22] 對(duì)輸出進(jìn)行歸一化。由于多尺度建模導(dǎo)致不同頭的方差不同，研究者將原始的 LayerNorm 替換為 GroupNorm。

實(shí)驗(yàn)結(jié)果

該研究進(jìn)行了大量的實(shí)驗(yàn)來評(píng)估 RetNet，包括語言建模任務(wù)、下游任務(wù)上零樣本、少樣本學(xué)習(xí)性能，此外，研究者還比較了 RetNet 訓(xùn)練和推理的速度、內(nèi)存消耗和延遲等指標(biāo)。

與 Transformer 的比較

語言建模任務(wù)。圖 5 報(bào)告了基于 Transformer 和 RetNet 的語言模型在驗(yàn)證集上的困惑度（perplexity）結(jié)果。實(shí)驗(yàn)給出了 13 b、2.7B 和 6.7B 三種模型尺寸的縮放曲線。表明，RetNet 取得了與 Transformer 可比較的結(jié)果。

更重要的是，這一結(jié)果還表明了 RetNet 在大小擴(kuò)展方面更具優(yōu)勢(shì)。除了性能優(yōu)勢(shì)外，實(shí)驗(yàn)中 RetNet 的訓(xùn)練也非常穩(wěn)定。RetNet 是 Transformer 的有力競(jìng)爭(zhēng)對(duì)手。研究者根據(jù)經(jīng)驗(yàn)發(fā)現(xiàn)，當(dāng)模型規(guī)模大于 2B 時(shí)，RetNet 開始超越 Transformer。

該研究還在各種下游任務(wù)上對(duì)語言模型進(jìn)行了比較。他們使用 6.7B 大小的模型進(jìn)行了零樣本和 4 個(gè)樣本學(xué)習(xí)的評(píng)估，如表 3 所示。表中展示的關(guān)于準(zhǔn)確率的數(shù)字與圖 5 中呈現(xiàn)的語言建模困惑度一致。在零樣本學(xué)習(xí)和上下文學(xué)習(xí)設(shè)置中，RetNet 在性能上與 Transformer 相當(dāng)。

訓(xùn)練成本

表 4 比較了 Transformer 和 RetNet 在訓(xùn)練速度和內(nèi)存開銷方面的結(jié)果，其中訓(xùn)練序列長(zhǎng)度為 8192。此外，該研究還將其與 FlashAttention 進(jìn)行了比較。

實(shí)驗(yàn)結(jié)果表明，在訓(xùn)練過程中，RetNet 比 Transformer 更節(jié)省內(nèi)存，并且具有更高的吞吐量。即使與 FlashAttention 相比，RetNet 在速度和內(nèi)存成本方面仍然具有競(jìng)爭(zhēng)力。此外，由于不依賴于特定的內(nèi)核，用戶可以輕松高效地在其他平臺(tái)上訓(xùn)練 RetNet。例如，研究者可以在具有良好吞吐量的 AMD MI200 集群上訓(xùn)練 RetNet 模型。

推理成本

圖 6 比較了 Transformer 和 RetNet 在推理過程中的內(nèi)存成本、吞吐量和延遲。實(shí)驗(yàn)中使用了 A100-80GB GPU 評(píng)估了 6.7B 模型。圖 6 顯示，RetNet 在推理成本方面優(yōu)于 Transformer。

內(nèi)存：如圖 6a 所示，由于 KV（鍵和值）緩存，Transformer 的內(nèi)存成本呈線性增長(zhǎng)。相比之下，RetNet 的內(nèi)存消耗即使對(duì)于長(zhǎng)序列也保持一致。

吞吐量：如圖 6b 所示，隨著解碼長(zhǎng)度的增加，Transformer 的吞吐量開始下降。相比之下，RetNet 通過利用 Retention 的循環(huán)表征，在解碼過程中具有更高的吞吐量，并且與長(zhǎng)度無關(guān)。

延遲：延遲是部署中的重要指標(biāo)，它極大地影響用戶體驗(yàn)。圖 6c 報(bào)告了解碼延遲。實(shí)驗(yàn)結(jié)果顯示，增加批次大小會(huì)使 Transformer 的延遲變大。此外，Transformer 的延遲隨著輸入長(zhǎng)度的增加而增加得更快。為了使延遲可接受，研究者不得不限制批次大小，這會(huì)損害 Transformer 的整體推理吞吐量。相比之下，RetNet 的解碼延遲優(yōu)于 Transformer，并且在不同的批次大小和輸入長(zhǎng)度下幾乎保持不變。

與 Transformer 變體比較

下表表明，RetNet 在不同的數(shù)據(jù)集上優(yōu)于先前的方法。RetNet 不僅在領(lǐng)域內(nèi)語料庫上取得更好的評(píng)估結(jié)果，還在幾個(gè)領(lǐng)域外數(shù)據(jù)集上獲得更低的困惑度。這種優(yōu)越的性能使得 RetNet 成為 Transformer 的有力繼任者。

消融實(shí)驗(yàn)

下表列出了 RetNet 的各種設(shè)計(jì)選擇，并在表 6 中報(bào)告了語言建模結(jié)果。