Raki的读paper小记：RWKV: Reinventing RNNs for the Transformer Era

Abstract&Introduction&Related Work

• 研究任务
  基础模型架构
• 已有方法和相关工作
  • RNN，CNN，Transformer
  • 稀疏注意力（Beltagy等人，2020年；Kitaev等人，2020年；Guo等人，2022年）、近似全注意力矩阵（Wang等人，2020年；Ma等人，2021年；Choromanski等人，2020年）、将分块注意力与门控结合（Ma等人，2023年）FlashAttention（Dao等人，2022a年）MLP-Mixer，Attention Free Transformer（AFT）（Zhai等人，2021年）用计算效率更高的替代方法替换了点积自注意力，可以看作是一个多头注意力，其中每个特征维度对应一个头
• 面临挑战
  • 训练以及推理开销太高
• 创新思路
  • 提出了一种新的注意力机制，改变了线性注意力的结果
  • 受到AFT的启发，RWKV采取了类似的方法，但通过修改交互权重来简化它，使其可以转换成RNN。同时，也对RNN风格的递归组件进行了修改，以增加上下文长度，例如Recurrent Memory Transformer和Linear Recurrent Units，还提出了像S4（Gu等人，2022年）及其变体这样的状态空间模型（SSM）
• 实验结论
  • RWKV利用线性注意机制、时间混合和token偏移等技术，捕捉和传播顺序信息，具有稳定梯度和较深层次的优势，具有高效处理顺序数据的能力。
  • RWKV模型能够在时间顺序模式下进行递归推理，与自注意力相比，在处理长序列时具有恒定的速度和内存占用。设计结合了注意力机制和RNN更新方式，通过时间相关的softmax运算和层归一化等技术确保梯度稳定性和学习能力。
  • 模型采用自定义CUDA内核、FFN与R门、小初始化嵌入和自定义初始化等附加优化，提高了计算效率和训练稳定性。

经典注意力机制：
$\mathrm{Attn}(Q,K,V{)}_t=\frac{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}{v}_i}{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}}.$

AFT把注意力机制改成：
${\mathrm{Attn}}^{+}(W,K,V{)}_t=\frac{{\sum }_{i=1}^t{e}^{w_{t,i}+k_i}{v}_i}{{\sum }_{i=1}^t{e}^{w_{t,i}+k_i}},$

{ w t , i } ∈ R T × T \{w_{t,i}\}\in R^{T\times T} {wt,i​}∈RT×T是学习到的位置偏移，矩阵中每个元素是一个标量

RWKV

受AFT的启发，我们让RWKV中的每个$w_t,i$成为一个channel-wise time decay vector，乘以从当前时间向后追溯的相对位置：$w_{t,i}=-(t-i)w,$ 维度为d，d是通道数

RWKV与其他模型的结构对比：

RWKV的模型结构

RWKV架构由一系列堆叠的残差块组成，每个残差块由一个时间混合和一个通道混合子块组成，具有递归结构
递归在RWKV中有两种表述方式：

1. 作为当前输入和上一个时间步骤的输入之间的线性插值（我们称之为时间偏移混合或token偏移技术，如图3中的对角线所示）这可以针对输入embedding的每个线性投影（例如时间混合中的R、K、V，以及通道混合中的R、K）进行独立调整，并且作为时间依赖的WKV的更新，其公式化在方程式14中。 WKV计算与AFT（Zhai等人，2021年）类似，但现在W是一个通道-wise向量，乘以相对位置，而不是AFT中的成对矩阵。还引入了一个向量U，用于单独关注当前token，以弥补W可能出现的退化问题。
   

在其中，WKV计算wkvt扮演了Transformer中Attn(Q, K, V)的角色，而不会产生二次成本，因为交互是在标量之间进行的。直观地说，随着时间t的增加，向量 $o_t$ 依赖于一个较长的历史，由逐渐增加的项的总和来表示。对于目标位置t，RWKV在位置间隔 $[1，t]$内执行加权求和，然后与接受度 $\sigma (r)$ 相乘。因此，交互在给定时间步内是乘性的，而在不同时间步内是加性的

time-mixing的公式与模型示意图：

各个模型的复杂度对比：

Experiments

Conclusions

RWKV是一种利用基于时间混合组件潜力的新型RNN模型。RWKV引入了几个关键策略，使其能够捕捉局部性和长程依赖，并解决当前架构的局限性：

1. 通过标量公式将二次的QK注意力替换为线性成本
2. 重新构造递归和顺序归纳偏置以实现有效的训练并行化和高效的推理
3. 使用自定义初始化增强训练动态

我们在各种NLP任务上对所提出的架构进行了基准测试，显示出与SoTA相当的性能，同时减少了成本。进一步的实验涉及表达能力、可解释性和扩展性，展示了模型的能力，并揭示了RWKV与其他LLM之间的行为相似之处。

RWKV为在顺序数据中建模复杂关系提供了一个可扩展和高效的新途径。虽然已经提出了许多替代Transformer的方法，并声称具有类似的特点，但RWKV是首个用数百亿参数预训练模型来支持这些主张的研究

Limitations

尽管提出的RWKV模型在训练和推理期间展现出了有希望的效率，但也应该意识到并解决一些限制，以便在未来的研究中加以应对。首先，RWKV的线性注意力确实带来了显著的效率提升，但它也可能限制了模型在需要回忆极其详细信息的非常长上下文任务中的性能。这是因为相比标准Transformer的二次注意力维持的完整信息，RWKV通过单个向量表示在许多时间步上传递信息。换句话说，模型的循环结构本质上限制了它“回顾”之前的标记的能力，与传统的自注意机制相对立。虽然学习的时间衰减有助于防止信息丢失，但与完整的自注意力相比，它在机制上存在一定的限制。

这项工作的另一个限制是与标准Transformer模型相比，prompt engineering的重要性增加了。RWKV中使用的线性注意力机制限制了从prompt中传递到模型继续部分的信息。因此，精心设计的提示可能对模型在任务中表现良好至关重要

Remark

是个好工作，但是跟RetNet一样，个人对其超长距离建模能力保持怀疑，希望后续工作能增加到万级别tokens来证明一下模型的真正实力？