線性複雜度語言模型的前景
處理更長且更複雜序列的需求不斷增長，這也推動了大型語言模型 (LLMs) 的發展邊界。傳統的基於 Transformer 的架構雖然強大，但由於其複雜度與序列長度呈二次方關係，因此面臨著顯著的擴展問題。當處理擴展的上下文輸入時，這種限制變得尤為明顯，阻礙了它們有效地捕獲和利用來自序列遠端部分的信息的能力。為了應對這一挑戰，湧現出了一波創新方法，旨在實現處理長序列的線性複雜度。
這些方法包括線性注意力模型、狀態空間模型 (如 Mamba)、線性 RNN (如 DeltaNet) 和 RWKV。這些架構中的每一種都為二次複雜度問題提供了一種獨特的解決方案，從而能夠更有效地處理長序列。然而，這些線性架構經常遇到難以完全理解和利用長上下文信息的問題。
例如，RWKV-7（一個 29 億參數模型）在密鑰檢索任務中表現出高達 28K tokens 的高準確性。但是，超過此閾值後，其性能會迅速下降。即使使用 128K 長度的數據進行持續預訓練，長上下文限制仍然存在。這個問題並非 RWKV 獨有；它延伸到像 Mamba 這樣的其他架構，代表了這類模型的一個根本挑戰。在擴展上下文中保持性能的困難凸顯了線性複雜度語言模型中一個需要改進的關鍵領域。
線性複雜度語言模型的概況
線性複雜度語言模型已成為基於 Transformer 架構的有吸引力的替代方案，避開了處理長序列中固有的二次方計算負擔。RWKV 模型系列在這個領域中脫穎而出，巧妙地將訓練期間 Transformer 的並行性與類似 RNN 的循環狀態表示相結合。
RWKV 的發展歷程跨越了幾個迭代，從基礎的 RWKV-4 開始，發展到 RWKV-5、RWKV-6，最終達到 RWKV-7。每次迭代都帶來了改進和完善，增強了模型的功能並解決了限制。此外，像 Jamba、Zamba 和 MiniMax 這樣的混合語言模型，通過引入獨特的混合設計，也留下了它們的印記，進一步豐富了線性複雜度模型的概況。
對高效長上下文處理的追求也促使了創新注意力機制的發展。例如，原生稀疏注意力將 tokens 組織成時間塊，採用三種不同的注意力路徑：用於全局上下文的壓縮粗粒度 tokens、用於局部細節的選擇性保留細粒度 tokens，以及用於捕獲局部上下文信息的滑動窗口。其他值得注意的注意力機制包括 SeerAttention 和 Block Attention (MoBA)，每種機制都提供了獨特的策略來關注長序列中的相關信息。
RWKV-X：用於增強遠程上下文建模的混合架構
來自廣東人工智慧與數字經濟實驗室（深圳）、深圳、河海大學、南京、深圳大學和青海大學的研究人員推出了一種名為 RWKV-X 的新型混合架構。這種架構巧妙地結合了 RWKV 在建模短程依賴關係方面的效率，以及專門設計用於捕獲遠程上下文的稀疏注意力機制。
與之前的混合方法不同，RWKV-X 在訓練期間實現了線性時間複雜度，在推理解碼期間實現了常數時間複雜度。這使其在處理長序列方面非常高效。該模型在持續預訓練 64K-token 序列時，在 64K 密鑰檢索基準測試中表現出接近完美的準確性。它在長上下文基準測試中始終優於之前的 RWKV-7 模型，同時在短上下文任務中保持強勁的性能。
RWKV-X 的創新代表了在解決長上下文語言建模挑戰方面的一個重大進步。通過結合循環模型和稀疏注意力機制的優勢，RWKV-X 在效率和準確性之間取得了平衡，為更有效地處理擴展序列鋪平了道路。
RWKV-X：架構和訓練
RWKV-X 體現了一種混合架構，將 RWKV-7 塊與稀疏注意力塊集成在一起，以利用這兩種方法的優勢。RWKV-X 並非從頭開始訓練，而是基於現有模型，使用受 LLaMA Pro 啟發的交錯塊擴展方法和零初始化機制。
訓練過程包括兩個階段，經過精心設計，旨在優化模型在短上下文和長上下文中的性能：
- 短上下文預訓練： 最初，該模型在從 MiniPile 數據集中提取的短 1024-token 上下文中進行訓練。在此階段，除了新添加塊中的參數外，所有參數都被凍結，確保保留來自基本 RWKV-7 模型的預訓練知識。這允許新添加的塊適應現有架構，而不會破壞預訓練的表示。
- 長上下文持續預訓練： 第二階段涉及使用 ProLong-64K 數據集和 64K tokens 的上下文長度進行長上下文持續預訓練，總共處理大約 10 億個 tokens。在此階段，所有參數都被解凍並聯合優化，允許模型微調其表示並學習遠程依賴關係。訓練採用長上下文交叉熵 (LongCE) 損失，該損失根據 tokens 的重要性動態地對其進行加權。這種損失函數有助於模型專注於序列中最相關的部分，從而提高其捕獲遠程關係的能力。
兩階段訓練過程使 RWKV-X 能夠有效地結合 RWKV-7 在短程建模方面的效率，以及稀疏注意力機制的長程上下文感知能力。通過首先在短上下文上進行預訓練，然後在長上下文上進行微調，該模型學會有效地整合來自序列不同部分的信息。
RWKV-X：評估與性能
短上下文評估顯示，RWKV-X 在標準基準測試中保持了具有競爭力的性能，證明了其有效處理較短序列的能力。較小的 RWKV-X (0.22B) 取得了 51.0 的平均分，與 RWKV-7 的 51.8 相當。在更大的規模上，RWKV-X (3.6B) 達到 71.9，與 RWKV-7 (2.9B, 72.8) 和 Qwen2.5-3B (71.4) 非常接近，同時超過了 LLaMA3.2-3B (69.7)。這些結果證實了 RWKV-X 作為通用 LLM 主幹的有效性，而不會犧牲在較短上下文中的性能。
此外，效率分析表明 RWKV-X 在長序列方面具有卓越的擴展特性。在 128K tokens 時，RWKV-X 比 Flash-Attention v3 實現了 1.37 倍的加速，並且隨著上下文長度的增加，這種優勢會不斷擴大。這表明，隨著序列長度的增長，與其他注意力機制相比，RWKV-X 變得越來越高效。
RWKV-X 在短上下文和長上下文中的強勁性能凸顯了其作為語言模型的多功能性和效率。它既能在較短序列上保持具有競爭力的性能，又能在較長序列上實現顯著的加速，使其成為適用於各種應用前景廣闊的架構。
RWKV-X：局限性與未來方向
RWKV-X 作為一種混合語言模型脫穎而出，它成功地結合了 RWKV 在建模短程依賴關係方面的效率，以及專門為長程上下文建模設計的新型稀疏注意力機制。儘管 RWKV-X 在長上下文語言建模中表現出強勁的性能和效率，但仍然存在一些局限性。
首先，其稀疏注意力機制依賴於 top-k chunk 選擇，採用了一種啟發式方法，可能會忽略語義上相關的依賴關係。top-k 選擇策略可能並不總能捕獲序列中最重要的信息，從而可能導致次優性能。
其次，當前實現顯示稀疏注意力解碼的運行速度慢於原始 RWKV，這表明需要進一步的工程努力來優化性能。儘管 RWKV-X 與長序列上的其他注意力機制相比實現了顯著的加速，但其稀疏注意力解碼仍然慢於原始 RWKV，這表明其實現還有改進的空間。
未來的研究可以側重於通過探索更複雜的稀疏注意力機制、優化稀疏注意力解碼的實現以及研究替代訓練策略來解決這些局限性。通過克服這些挑戰，RWKV-X 有可能成為一種更強大、更高效的長上下文應用語言模型。