RWKV-LM
We propose the RWKV language model, with alternating time-mix and channel-mix layers:
-
The R, K, V are generated by linear transforms of input, and W is parameter. The idea of RWKV is to decompose attention into R(target) * W(src, target) * K(src). So we can call R "receptance", and sigmoid means it's in 0~1 range.
-
The Time-mix is similar to AFT (https://arxiv.org/abs/2105.14103). There are two differences.
(1) We changed the normalization (denominator). For masked language models, we define:
(2) We decompose W_{t,u,c} and introduce multi-head W (here h is the corresponding head of c):
(3) You don't need LayerNorm for Time-mix. In fact, the model converges faster when LayerNorm is removed.
Moreover we multiply the final output of Time-mix layer by γ(t). The reason for the α β γ factors, is because the context size is smaller when t is small, and this can be compensated using the α β γ factors.
-
The Channel-mix is similar to GeGLU (https://arxiv.org/abs/2002.05202) with an extra R factor.
-
Finally, we add extra time-mixing as in (https://github.com/BlinkDL/minGPT-tuned). You can try reducing the amt of time-mixing in upper layers of deep models.
We also propose a new sampling method (as in src/utils.py):
(1) Find the max probability p_max after softmax.
(2) Remove all entries whose probability is lower than 0.02 * pow(p_max, 2)
(3) Feel free to tune the 0.02 and 2 factor.
Training loss, RWKV vs MHA+Rotary+GeGLU:
(this is character-level loss with simplebooks-92 dataset https://dldata-public.s3.us-east-2.amazonaws.com/simplebooks.zip)
2 Nov 20, 2022
4 Jul 01, 2022
132 Dec 14, 2022
2 Apr 24, 2022
1 Dec 23, 2022
3.5k Dec 30, 2022
33 Oct 13, 2022
0 Jan 05, 2022
338 Dec 02, 2022
6 Feb 10, 2022
197 Dec 25, 2022
42 Sep 21, 2022
1 Dec 22, 2021
0 Nov 12, 2022
3 Aug 08, 2021
3 Jan 04, 2022
375 Dec 28, 2022
116 Dec 27, 2022
37 Dec 14, 2022