Transformer 模型在幾乎所有自然語言處理（NLP）任務(wù)中都帶來了革命，但其在序列長度上的內(nèi)存和計算復(fù)雜性呈二次方增長。相比之下，循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）在內(nèi)存和計算需求上呈線性增長，但由于并行化和可擴(kuò)展性的限制，很難達(dá)到與 Transformer 相同的性能水平。本文提出了一種新穎的模型架構(gòu)，Receptance Weighted Key Value（RWKV），將 Transformer 的高效可并行訓(xùn)練與 RNN 的高效推理相結(jié)合。實(shí)驗證明，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ù)，包括自然語言理解、對話式人工智能、時間序列分析等，其中用到的技術(shù)主要包括循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）、卷積神經(jīng)網(wǎng)絡(luò)（CNNs）和 Transformer 等。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

RWKV 模型

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

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

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

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

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

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

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