经典注意力机制:
Attn ( Q , K , V ) t = ∑ i = 1 T e q t ⊤ k i v i ∑ i = 1 T e q t ⊤ k i . \operatorname{Attn}(Q,K,V)_t=\large\frac{\sum_{i=1}^Te^{q_t^\top k_i}v_i}{\sum_{i=1}^Te^{q_t^\top k_i}}. Attn(Q,K,V)t=∑i=1Teqt⊤ki∑i=1Teqt⊤kivi.
AFT把注意力机制改成:
Attn + ( W , K , V ) t = ∑ i = 1 t e w t , i + k i v i ∑ i = 1 t e w t , i + k i , \operatorname{Attn}^+(W,K,V)_t=\large\frac{\sum_{i=1}^te^{w_{t,i}+k_i}v_i}{\sum_{i=1}^te^{w_{t,i}+k_i}}, Attn+(W,K,V)t=∑i=1tewt,i+ki∑i=1tewt,i+kivi,
{ w t , i } ∈ R T × T \{w_{t,i}\}\in R^{T\times T} {wt,i}∈RT×T是学习到的位置偏移,矩阵中每个元素是一个标量
受AFT的启发,我们让RWKV中的每个 w t , i w_t, i wt,i成为一个channel-wise time decay vector,乘以从当前时间向后追溯的相对位置: w t , i = − ( t − i ) w , w_{t,i}=-(t-i)w, wt,i=−(t−i)w, 维度为d,d是通道数
RWKV与其他模型的结构对比:
RWKV的模型结构
RWKV架构由一系列堆叠的残差块组成,每个残差块由一个时间混合和一个通道混合子块组成,具有递归结构
递归在RWKV中有两种表述方式:
在其中,WKV计算wkvt扮演了Transformer中Attn(Q, K, V)的角色,而不会产生二次成本,因为交互是在标量之间进行的。直观地说,随着时间t的增加,向量 o t o_t ot 依赖于一个较长的历史,由逐渐增加的项的总和来表示。对于目标位置t,RWKV在位置间隔 [ 1 , t ] [1,t] [1,t]内执行加权求和,然后与接受度 σ ( r ) σ(r) σ(r) 相乘。因此,交互在给定时间步内是乘性的,而在不同时间步内是加性的
time-mixing的公式与模型示意图:
各个模型的复杂度对比:
RWKV是一种利用基于时间混合组件潜力的新型RNN模型。RWKV引入了几个关键策略,使其能够捕捉局部性和长程依赖,并解决当前架构的局限性:
我们在各种NLP任务上对所提出的架构进行了基准测试,显示出与SoTA相当的性能,同时减少了成本。进一步的实验涉及表达能力、可解释性和扩展性,展示了模型的能力,并揭示了RWKV与其他LLM之间的行为相似之处。
RWKV为在顺序数据中建模复杂关系提供了一个可扩展和高效的新途径。虽然已经提出了许多替代Transformer的方法,并声称具有类似的特点,但RWKV是首个用数百亿参数预训练模型来支持这些主张的研究
尽管提出的RWKV模型在训练和推理期间展现出了有希望的效率,但也应该意识到并解决一些限制,以便在未来的研究中加以应对。首先,RWKV的线性注意力确实带来了显著的效率提升,但它也可能限制了模型在需要回忆极其详细信息的非常长上下文任务中的性能。这是因为相比标准Transformer的二次注意力维持的完整信息,RWKV通过单个向量表示在许多时间步上传递信息。换句话说,模型的循环结构本质上限制了它“回顾”之前的标记的能力,与传统的自注意机制相对立。虽然学习的时间衰减有助于防止信息丢失,但与完整的自注意力相比,它在机制上存在一定的限制。
这项工作的另一个限制是与标准Transformer模型相比,prompt engineering的重要性增加了。RWKV中使用的线性注意力机制限制了从prompt中传递到模型继续部分的信息。因此,精心设计的提示可能对模型在任务中表现良好至关重要
是个好工作,但是跟RetNet一样,个人对其超长距离建模能力保持怀疑,希望后续工作能增加到万级别tokens来证明一下模型的真正实力?