Transformer 模型在幾乎所有自然語(yǔ)言處理（NLP）任務(wù)中都帶來(lái)了革命，但其在序列長(zhǎng)度上的內(nèi)存和計(jì)算復(fù)雜性呈二次方增長(zhǎng)。相比之下，循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）在內(nèi)存和計(jì)算需求上呈線性增長(zhǎng)，但由于并行化和可擴(kuò)展性的限制，很難達(dá)到與 Transformer 相同的性能水平。本文提出了一種新穎的模型架構(gòu)，Receptance Weighted Key Value（RWKV），將 Transformer 的高效可并行訓(xùn)練與 RNN 的高效推理相結(jié)合。實(shí)驗(yàn)證明，RWKV 的性能與相同規(guī)模的 Transformer 相當(dāng)。

深度學(xué)習(xí)技術(shù)在人工智能領(lǐng)域取得了重大進(jìn)展，在各種科學(xué)和工業(yè)應(yīng)用中發(fā)揮了關(guān)鍵作用。這些應(yīng)用通常涉及復(fù)雜的序列數(shù)據(jù)處理任務(wù)，包括自然語(yǔ)言理解、對(duì)話式人工智能、時(shí)間序列分析等，其中用到的技術(shù)主要包括循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）、卷積神經(jīng)網(wǎng)絡(luò)（CNNs）和 Transformer 等。

不過(guò)，這些方法各自存在不同的缺點(diǎn)，從而限制了它們?cè)谀承﹫?chǎng)景下的效率。循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）面臨著梯度消失的問(wèn)題，使得它們難以對(duì)長(zhǎng)序列進(jìn)行訓(xùn)練。此外，在訓(xùn)練過(guò)程中無(wú)法在時(shí)間維度上并行化，進(jìn)而限制了其可擴(kuò)展性。另一方面，卷積神經(jīng)網(wǎng)絡(luò)（CNNs）只擅長(zhǎng)捕捉局部模式，在處理長(zhǎng)程依賴(lài)方面還很欠缺，而這對(duì)于許多序列處理任務(wù)至關(guān)重要。

Transformer 模型由于其處理局部和長(zhǎng)程依賴(lài)關(guān)系的能力以及可并行化訓(xùn)練的特點(diǎn)而成為一個(gè)強(qiáng)大的替代方案，如 GPT-3、ChatGPT、GPT-4、LLaMA 和 Chinchilla 等都展示了這種架構(gòu)的能力，推動(dòng)了自然語(yǔ)言處理領(lǐng)域的前沿。盡管取得了這些重大進(jìn)展，Transformer 中固有的自注意力機(jī)制帶來(lái)了獨(dú)特的挑戰(zhàn)，主要是由于其二次復(fù)雜度造成的。這種復(fù)雜性使得該架構(gòu)在涉及長(zhǎng)輸入序列或資源受限情況下計(jì)算成本高昂且占用內(nèi)存。這也促使了大量研究的發(fā)布，旨在改善 Transformer 的擴(kuò)展性，但往往以犧牲一些特性為代價(jià)。

為了應(yīng)對(duì)這些挑戰(zhàn)，一個(gè)由 27 所大學(xué)、研究機(jī)構(gòu)組成的開(kāi)源研究團(tuán)隊(duì)，聯(lián)合發(fā)表論文《 RWKV: Reinventing RNNs for the Transformer Era 》，文中介紹了一種新型模型：RWKV（Receptance Weighted Key Value），這是一種新穎的架構(gòu)，有效地結(jié)合了 RNN 和 Transformer 的優(yōu)點(diǎn)，同時(shí)規(guī)避了兩者的缺點(diǎn)。RWKV 設(shè)計(jì)精良，能夠緩解 Transformer 所帶來(lái)的內(nèi)存瓶頸和二次方擴(kuò)展問(wèn)題，實(shí)現(xiàn)更有效的線性擴(kuò)展，同時(shí)保留了使 Transformer 在這個(gè)領(lǐng)域占主導(dǎo)的一些性質(zhì)。

• 論文地址：https:///pdf/2305.13048.pdf

• RWKV 模型下載：https:///BlinkDL/rwkv-4-raven

• Demo 地址：https://www./i/app/BlinkDL/ChatRWKV/RWKV-4-Raven-7B

本文利用線性注意力機(jī)制，允許將模型定義為 Transformer 或 RNN，從而在訓(xùn)練期間并行化計(jì)算，并在推理過(guò)程中保持恒定的計(jì)算和內(nèi)存復(fù)雜性，使其成為第一個(gè)可擴(kuò)展到數(shù)百億參數(shù)的非 Transformer 架構(gòu)。

RWKV 其中的一個(gè)特征是它能夠提供并行訓(xùn)練和強(qiáng)大的可擴(kuò)展性，類(lèi)似于 Transformer。此外，該研究對(duì) RWKV 中的注意力機(jī)制進(jìn)行了重新闡述，引入了線性注意力的一個(gè)變體，避開(kāi)了傳統(tǒng)點(diǎn)積（dot-product）token 交互，轉(zhuǎn)而采用更有效的通道導(dǎo)向注意力（ channel directed attention ）。這種方法與傳統(tǒng)的 Transformer 架構(gòu)形成了鮮明的對(duì)比，其中特定的 token 交互主導(dǎo)了注意力。在 RWKV 中，線性注意力的實(shí)施是無(wú)需近似的，這在效率上提供了顯著的改進(jìn)，并增強(qiáng)了可擴(kuò)展性，詳見(jiàn)表 1。

該研究表示，開(kāi)發(fā) RWKV 的主要?jiǎng)訖C(jī)是彌補(bǔ)神經(jīng)網(wǎng)絡(luò)架構(gòu)在計(jì)算效率和表達(dá)能力之間的差距。它為處理涉及數(shù)十億參數(shù)的大規(guī)模模型的任務(wù)提供了一個(gè)有希望且可行的解決方案，以極低的計(jì)算成本展現(xiàn)出強(qiáng)有力的競(jìng)爭(zhēng)性。

實(shí)驗(yàn)結(jié)果表明，RWKV 可以成為一個(gè)有價(jià)值的工具，用于解決各個(gè)領(lǐng)域擴(kuò)展和部署人工智能模型的各種挑戰(zhàn)，特別是那些涉及序列數(shù)據(jù)處理的領(lǐng)域。RWKV 為下一代更可持續(xù)、計(jì)算效率更高的序列處理任務(wù)的 AI 模型鋪平了道路。

總結(jié)而言，本文的貢獻(xiàn)如下：

• 引入了 RWKV 網(wǎng)絡(luò)架構(gòu)，該架構(gòu)結(jié)合了 RNN 和 Transformer 的優(yōu)點(diǎn)，同時(shí)減輕了它們已知的限制。

• 本文提出了一個(gè)新的注意力機(jī)制重構(gòu)，進(jìn)而提出線性注意力，避開(kāi)了與標(biāo)準(zhǔn) Transformer 模型相關(guān)的二次復(fù)雜性。

• 本文在基準(zhǔn)數(shù)據(jù)集上進(jìn)行了一系列全面的實(shí)驗(yàn)，展示了 RWKV 在處理涉及大規(guī)模模型和長(zhǎng)距離依賴(lài)任務(wù)上的性能、效率和可擴(kuò)展性。

• 發(fā)布了預(yù)訓(xùn)練模型，其大小從 1.69 億到 140 億的參數(shù)不等，這些模型是在 Pile 上訓(xùn)練的。

值得注意的是，論文參與機(jī)構(gòu)之一的 EleutherAI 表示：這篇論文還不是最終版本，后續(xù)會(huì)不斷完善。

RWKV 模型

RWKV 架構(gòu)的名稱(chēng)來(lái)源于時(shí)間混合和通道混合塊中使用的四個(gè)主要模型元素，分別如下：

• R：Receptance 向量，用于接收以往信息；

• W：權(quán)重（weight）是位置權(quán)重衰減向量，是可訓(xùn)練的模型參數(shù)；

• K：鍵（Key）是類(lèi)似于傳統(tǒng)注意力中 K 的向量；

• V：值（Value）是類(lèi)似于傳統(tǒng)注意力中 V 的向量。

每一時(shí)間步的主要元素之間的交互是相乘增加的，具體如下圖 2 所示。

架構(gòu)細(xì)節(jié)

RWKV 架構(gòu)由一系列堆疊的殘差塊組成，每個(gè)殘差塊又由具有循環(huán)結(jié)構(gòu)的時(shí)間混合和通道混合子塊組成。

循環(huán)被表示為當(dāng)前輸入和前一個(gè)時(shí)間步的輸入之間的線性插值（研究者稱(chēng)這種技術(shù)為時(shí)移混合或 token shift，如下圖 3 所示），該插值可以針對(duì)輸入嵌入的每個(gè)線性投影進(jìn)行獨(dú)立調(diào)整（比如時(shí)間混合中的 R、K 和 V，通道混合中的 R 和 K），并作為公式 14 中形式化的 WKV 的時(shí)變更新。

類(lèi) Transformer 的并行化

RWKV 可以在時(shí)間并行模式下進(jìn)行高效地并行化，讓人聯(lián)想到 Transformer。單個(gè)層中一個(gè) batch 序列的時(shí)間復(fù)雜度為 O (BTd^2 )，它主要由矩陣乘法 W_□,  □ ∈ {r, k, v, o}（假設(shè) B 個(gè)序列、T 個(gè)最大 token 和 d 個(gè)通道）。同時(shí)更新注意力分?jǐn)?shù) wkv_t 需要串行掃描，并且復(fù)雜度為 O (BTd)。

類(lèi) RNN 的序列解碼

在循環(huán)網(wǎng)絡(luò)中，將狀態(tài) t 時(shí)的輸出用作狀態(tài) t+1 時(shí)的輸入很常見(jiàn)。這在語(yǔ)言模型的自回歸解碼推理中尤為明顯，要求每一個(gè) token 在饋入下一步之前必須進(jìn)行計(jì)算，從而使 RWKV 可以利用類(lèi) RNN 結(jié)構(gòu)（即時(shí)序模式）。在這種情況下，RWKV 可以方便地循環(huán)用于推理解碼，從而利用每個(gè)輸出 token 僅依賴(lài)于最新?tīng)顟B(tài)的優(yōu)勢(shì)。

然后 RWKV 充當(dāng) RNN 解碼器，在序列長(zhǎng)度方面保持恒定速度和內(nèi)存占用，從而更高效地處理更長(zhǎng)的序列。相比之下，自注意力通常需要 KV 緩存相對(duì)于序列長(zhǎng)度呈線性增長(zhǎng)，這會(huì)導(dǎo)致效率下降，并隨序列長(zhǎng)度增加消耗更多內(nèi)存和時(shí)間。

軟件實(shí)現(xiàn)

RWKV 最初使用 PyTorch 深度學(xué)習(xí)庫(kù)和自定義 CUDA 內(nèi)核（它用于 WKV 計(jì)算）來(lái)實(shí)現(xiàn)。盡管 RWKV 是一個(gè)通用循環(huán)網(wǎng)絡(luò)，但其當(dāng)前的實(shí)現(xiàn)主要集中在語(yǔ)言建模任務(wù)（RWKV-LM）。該模型架構(gòu)包含了一個(gè)嵌入層，為此研究者遵循第 4.7 節(jié)中的設(shè)置，并按照第 4.6 節(jié)中的原則依次應(yīng)用幾個(gè)相同的殘差塊，具體如上圖 2 和 3 所示。

梯度穩(wěn)定性和層堆疊

RWKV 架構(gòu)被設(shè)計(jì)為 Transformer 和 RNN 的融合，與傳統(tǒng)的 RNN 相比，Transformers 具有穩(wěn)定梯度和更深層次架構(gòu)的優(yōu)勢(shì)，同時(shí)推理效率高。

RWKV 模型具有用于更新類(lèi)似注意力分?jǐn)?shù)的單步過(guò)程，其中包括一個(gè)依賴(lài)于時(shí)間的 softmax 操作，該操作有助于數(shù)值穩(wěn)定性并防止梯度消失（有關(guān)嚴(yán)格證明，請(qǐng)參見(jiàn)附錄 F）。直觀地說(shuō)，此操作可確保梯度沿最相關(guān)的路徑傳播。Layer normalization (Ba et al., 2016) 是架構(gòu)的另一個(gè)關(guān)鍵方面，它通過(guò)穩(wěn)定梯度、解決梯度消失和爆炸問(wèn)題來(lái)增強(qiáng)深度神經(jīng)網(wǎng)絡(luò)的訓(xùn)練動(dòng)態(tài)。

利用時(shí)間結(jié)構(gòu)進(jìn)行時(shí)序數(shù)據(jù)處理

RWKV 通過(guò)三種機(jī)制的組合來(lái)捕獲和傳播時(shí)序信息：循環(huán)、時(shí)間衰減和 token shift。

RWKV 時(shí)間混合塊中的循環(huán)是模型捕獲序列元素之間復(fù)雜關(guān)系和隨時(shí)間傳播局部信息的能力的基礎(chǔ)。

時(shí)間衰減機(jī)制（等式 14 中的 e^?w 和 e^u）保持了對(duì)序列元素之間位置關(guān)系的敏感性。通過(guò)逐漸減少以往信息隨時(shí)間的影響，該模型保留了時(shí)間局部性和進(jìn)展感，這對(duì)于時(shí)序處理至關(guān)重要。

token shift 或 time-shift 混合或（圖 3 中的對(duì)角線箭頭），也有助于模型適應(yīng)時(shí)序數(shù)據(jù)。通過(guò)在當(dāng)前輸入和前一個(gè)時(shí)間步輸入之間進(jìn)行線性插值，模型自然地聚合和門(mén)控輸入通道中的信息。

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

實(shí)驗(yàn)的重點(diǎn)是回答以下問(wèn)題：

• RQ1：在參數(shù)數(shù)量和訓(xùn)練 token 數(shù)量相等的情況下，RWKV 與二次 transformer 架構(gòu)相比具有競(jìng)爭(zhēng)力嗎？

• RQ2：增加參數(shù)數(shù)量時(shí)，RWKV 是否仍然具有與二次 transformer 架構(gòu)相競(jìng)爭(zhēng)的能力？

• RQ3：當(dāng) RWKV 模型被訓(xùn)練用于開(kāi)源二次 transformer 無(wú)法高效處理的上下文長(zhǎng)度時(shí)，增加 RWKV 的參數(shù)是否能夠獲得更好的語(yǔ)言建模損失？

首先是回答 RQ1 和 RQ2 問(wèn)題，從圖 4 可以看出，在六個(gè)基準(zhǔn)測(cè)試中（Winogrande、PIQA、ARC-C、ARC-E、LAMBADA 和 SciQ），RWKV 與開(kāi)源二次復(fù)雜度 transformer 模型 Pythia、OPT 和 BLOOM 具有相當(dāng)?shù)母?jìng)爭(zhēng)力。RWKV 甚至在四個(gè)任務(wù)（PIQA、OBQA、ARC-E 和 COPA）中勝過(guò)了 Pythia 和 GPT-Neo。

對(duì)于 RQ3，圖 5 顯示，增加上下文長(zhǎng)度會(huì)導(dǎo)致 Pile 上的測(cè)試損失降低，這表明 RWKV 能夠有效利用較長(zhǎng)的上下文信息。