在Transformer時代重塑RNN，RWKV將非Transformer架構擴展到數百億參數（1）

Transformer 模型在幾乎所有自然語言處理（NLP）任務中都帶來了革命，但其在序列長度上的內存和計算復雜性呈二次方增長。相比之下，循環神經網絡（RNNs）在內存和計算需求上呈線性增長，但由于并行化和可擴展性的限制，很難達到與 Transformer 相同的性能水平。本文提出了一種新穎的模型架構，Receptance Weighted Key Value（RWKV），將 Transformer 的高效可并行訓練與 RNN 的高效推理相結合。實驗證明，RWKV 的性能與相同規模的 Transformer 相當。

深度學習技術在人工智能領域取得了重大進展，在各種科學和工業應用中發揮了關鍵作用。這些應用通常涉及復雜的序列數據處理任務，包括自然語言理解、對話式人工智能、時間序列分析等，其中用到的技術主要包括循環神經網絡（RNNs）、卷積神經網絡（CNNs）和 Transformer 等。

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

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

為了應對這些挑戰，一個由 27 所大學、研究機構組成的開源研究團隊，聯合發表論文《 RWKV: Reinventing RNNs for the Transformer Era 》，文中介紹了一種新型模型：RWKV（Receptance Weighted Key Value），這是一種新穎的架構，有效地結合了 RNN 和 Transformer 的優點，同時規避了兩者的缺點。RWKV 設計精良，能夠緩解 Transformer 所帶來的內存瓶頸和二次方擴展問題，實現更有效的線性擴展，同時保留了使 Transformer 在這個領域占主導的一些性質。

• 論文地址：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

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

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

該研究表示，開發 RWKV 的主要動機是彌補神經網絡架構在計算效率和表達能力之間的差距。它為處理涉及數十億參數的大規模模型的任務提供了一個有希望且可行的解決方案，以極低的計算成本展現出強有力的競爭性。

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

總結而言，本文的貢獻如下：

• 引入了 RWKV 網絡架構，該架構結合了 RNN 和 Transformer 的優點，同時減輕了它們已知的限制。

• 本文提出了一個新的注意力機制重構，進而提出線性注意力，避開了與標準 Transformer 模型相關的二次復雜性。

• 本文在基準數據集上進行了一系列全面的實驗，展示了 RWKV 在處理涉及大規模模型和長距離依賴任務上的性能、效率和可擴展性。

• 發布了預訓練模型，其大小從 1.69 億到 140 億的參數不等，這些模型是在 Pile 上訓練的。

值得注意的是，論文參與機構之一的 EleutherAI 表示：這篇論文還不是最終版本，后續會不斷完善。

RWKV 模型

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

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

• W：權重（weight）是位置權重衰減向量，是可訓練的模型參數；

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

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

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

架構細節