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

架構細節

RWKV 架構由一系列堆疊的殘差塊組成，每個殘差塊又由具有循環結構的時間混合和通道混合子塊組成。

循環被表示為當前輸入和前一個時間步的輸入之間的線性插值（研究者稱這種技術為時移混合或 token shift，如下圖 3 所示），該插值可以針對輸入嵌入的每個線性投影進行獨立調整（比如時間混合中的 R、K 和 V，通道混合中的 R 和 K），并作為公式 14 中形式化的 WKV 的時變更新。

類 Transformer 的并行化

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

類 RNN 的序列解碼

在循環網絡中，將狀態 t 時的輸出用作狀態 t+1 時的輸入很常見。這在語言模型的自回歸解碼推理中尤為明顯，要求每一個 token 在饋入下一步之前必須進行計算，從而使 RWKV 可以利用類 RNN 結構（即時序模式）。在這種情況下，RWKV 可以方便地循環用于推理解碼，從而利用每個輸出 token 僅依賴于最新狀態的優勢。

然后 RWKV 充當 RNN ****，在序列長度方面保持恒定速度和內存占用，從而更高效地處理更長的序列。相比之下，自注意力通常需要 KV 緩存相對于序列長度呈線性增長，這會導致效率下降，并隨序列長度增加消耗更多內存和時間。

軟件實現

RWKV 最初使用 PyTorch 深度學習庫和自定義 CUDA 內核（它用于 WKV 計算）來實現。盡管 RWKV 是一個通用循環網絡，但其當前的實現主要集中在語言建模任務（RWKV-LM）。該模型架構包含了一個嵌入層，為此研究者遵循第 4.7 節中的設置，并按照第 4.6 節中的原則依次應用幾個相同的殘差塊，具體如上圖 2 和 3 所示。

梯度穩定性和層堆疊

RWKV 架構被設計為 Transformer 和 RNN 的融合，與傳統的 RNN 相比，Transformers 具有穩定梯度和更深層次架構的優勢，同時推理效率高。

RWKV 模型具有用于更新類似注意力分數的單步過程，其中包括一個依賴于時間的 softmax 操作，該操作有助于數值穩定性并防止梯度消失（有關嚴格證明，請參見附錄 F）。直觀地說，此操作可確保梯度沿最相關的路徑傳播。Layer normalization (Ba et al., 2016) 是架構的另一個關鍵方面，它通過穩定梯度、解決梯度消失和爆炸問題來增強深度神經網絡的訓練動態。

利用時間結構進行時序數據處理

RWKV 通過三種機制的組合來捕獲和傳播時序信息：循環、時間衰減和 token shift。

RWKV 時間混合塊中的循環是模型捕獲序列元素之間復雜關系和隨時間傳播局部信息的能力的基礎。

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

token shift 或 time-shift 混合或（圖 3 中的對角線箭頭），也有助于模型適應時序數據。通過在當前輸入和前一個時間步輸入之間進行線性插值，模型自然地聚合和門控輸入通道中的信息。

實驗結果

實驗的重點是回答以下問題：

• RQ1：在參數數量和訓練 token 數量相等的情況下，RWKV 與二次 transformer 架構相比具有競爭力嗎？

• RQ2：增加參數數量時，RWKV 是否仍然具有與二次 transformer 架構相競爭的能力？

• RQ3：當 RWKV 模型被訓練用于開源二次 transformer 無法高效處理的上下文長度時，增加 RWKV 的參數是否能夠獲得更好的語言建模損失？

首先是回答 RQ1 和 RQ2 問題，從圖 4 可以看出，在六個基準測試中（Winogrande、PIQA、ARC-C、ARC-E、LAMBADA 和 SciQ），RWKV 與開源二次復雜度 transformer 模型 Pythia、OPT 和 BLOOM 具有相當的競爭力。RWKV 甚至在四個任務（PIQA、OBQA、ARC-E 和 COPA）中勝過了 Pythia 和 GPT-Neo。

對于 RQ3，圖 5 顯示，增加上下文長度會導致 Pile 上的測試損失降低，這表明 RWKV 能夠有效利用較長的上下文信息。