201
个编辑
更改
→二、RWKV 模型结构（以 v4/5 为例）
----
----
== 二、RWKV 模型结构（以 v4/5 为例） ==
== 二、RWKV 模型结构 ==
=== RWKV Block（基本结构） ===
=== RWKV Block（基本结构） ===
<math>\mathrm{output} = \mathrm{ReLU}(W_1 \times x + b_1) \times W_2 + b_2</math>
<math>\mathrm{output} = \mathrm{ReLU}(W_1 \times x + b_1) \times W_2 + b_2</math>
* '''自注意力层简化表达'''：
* '''注意力层简化表达'''
<math>x_{t+1} = \sigma(Rx_t) \cdot \frac{\exp(Kx_t) \cdot (Vx_t) + \exp(W) \cdot a_t}{\exp(Kx_t) + \exp(W) \cdot b_t} </math>
<math>x_{t+1} = \sigma(Rx_t) \cdot \frac{\exp(Kx_t) \cdot (Vx_t) + \exp(W) \cdot a_t}{\exp(Kx_t) + \exp(W) \cdot b_t} </math>
* '''核心机制代码'''
<syntaxhighlight lang="python3">
def ref_fwd(r, w, k, v, a, b):
r = r.view(B, T, H, N)
k = k.view(B, T, H, N)
v = v.view(B, T, H, N)
a = a.view(B, T, H, N)
b = b.view(B, T, H, N)
w = torch.exp(-torch.exp(w.view(B, T, H, N)))
out = torch.zeros((B, T, H, N), device=DEVICE)
state = torch.zeros((B, H, N, N), device=DEVICE)
for t in range(T):
kk = k[:, t, :]
rr = r[:, t, :]
vv = v[:, t, :]
aa = a[:, t, :]
bb = b[:, t, :]
sab = torch.einsum('bhik,bhk,bhj->bhij', state, aa, bb)
state = state * w[: , t, :, None, :] + sab + torch.einsum('bhj,bhi->bhij', kk, vv)
out[:, t, :] = torch.einsum('bhj,bhij->bhi', rr, state)
return out.view((B, T, C))
</syntaxhighlight>
----
----