在Transformer時(shí)代重塑RNN，RWKV將非Transformer架構(gòu)擴(kuò)展到數(shù)百億參數(shù)（1）

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ì)算需求上呈線(xiàn)性增長(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ì)話(huà)式人工智能、時(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)更有效的線(xiàn)性擴(kuò)展，同時(shí)保留了使 Transformer 在這個(gè)領(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

本文利用線(xiàn)性注意力機(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)行了重新闡述，引入了線(xià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 中，線(xiàn)性注意力的實(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)而提出線(xià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é)