The RWKV language mannequin: An RNN with the benefits of a transformer
For some time, I’ve been following and contributing to the RWKV language mannequin, an open supply massive language mannequin with nice potential. As ChatGPT and huge language fashions usually have gotten a variety of consideration just lately, I believe it’s time to jot down about RWKV. On this publish, I’ll attempt to clarify what’s so particular about RWKV in comparison with most language fashions (transformers). The other RWKV post is extra technical, exhibiting intimately how RWKV truly works (with a ~100 line minimal implementation).
At a excessive stage, the RWKV mannequin is a intelligent RNN structure that allows it to be educated like a transformer. So to elucidate RWKV, I want to elucidate RNNs and transformers first.
Classically, the neural networks used for sequence (resembling textual content) processing have been RNNs (like LSTMs). RNNs take two inputs: a state vector and a token. It goes by the enter sequence one token at a time, every token updating the state. We could for instance use an RNN to course of a textual content right into a single state vector. This will then be used to categorise the textual content into “optimistic” or “unfavourable”. Or we could use the ultimate state to foretell the subsequent token, which is how RNNs are used to generate textual content.
Due to the sequential nature of RNNs, they’re laborious to massively parallelize throughout many GPUs. This motivated utilizing an “consideration” mechanism as an alternative of sequential processing, leading to an structure known as a transformer. A transformer processes all tokens on the identical time, evaluating every token to all earlier tokens in parallel. Particularly, the eye calculates “key”, “worth” and “question” vectors for every token, then contributions between all pairs of tokens are computed utilizing these.
Along with having the ability to pace up coaching by large parallelization, massive transformers usually rating higher than RNNs on benchmarks.
Nonetheless, the eye mechanism scales quadratically with the size of the sequence to be processed. This successfully limits the mannequin’s enter dimension (or “context size”). Moreover, due to the eye mechanism, when producing textual content, we have to maintain consideration vectors for all earlier tokens in reminiscence. This requires way more reminiscence than an RNN which solely shops a single state.
RWKV combines one of the best options of RNNs and transformers. Throughout coaching, we use the transformer sort formulation of the structure, which permits large parallelization (with a type of consideration which scales linearly with the variety of tokens). For inference, we use an equal formulation which works like an RNN with a state. This enables us to get one of the best of each worlds.
So we principally have a mannequin which trains like a transformer, besides that lengthy context size is just not costly. And through inference, we want considerably much less reminiscence and may implicitly deal with “infinite” context size (although in apply, the mannequin might need a tough time generalizing to for much longer context lengths than it noticed throughout coaching).
OK, however what concerning the efficiency? Since RWKV an RNN, it’s pure to suppose that it may well’t carry out in addition to a transformer on benchmarks. Additionally, this simply seems like linear consideration. Not one of the many earlier linear time consideration transformer architectures (like “Linformer”, “Nystromformer”, “Longformer”, “Performer”) appeared to take off.
Properly, RWKV appears to scale in addition to SOTA transformers. No less than as much as 14 billion parameters.
RWKV is an open supply group mission. Be a part of the Discord and contribute (or ask questions or no matter).
When taking a look at RWKV 14B (14 billion parameters), it’s straightforward to ask what occurs once we scale to 175B like GPT-3. Nonetheless, coaching a 175B mannequin is dear. Calculating the approximate coaching value of a transformer-like structure is definitely simple.
The bottleneck for coaching is basically multiplying by all of the parameters, after which including that collectively, for every enter token. With automated differentiation, we are able to calculate the gradient with about one other 2x that, for a complete of 6 FLOPs per parameter per token. So a 14B mannequin educated on 300 billion tokens takes about (14B occasions 300B occasions 6 = 2.5 occasions 10^{22}) FLOPs. We use A100 GPUs for coaching. Utilizing 16-bit floating level numbers, an A100 can theoretically do as much as 312 TFLOPS, or about (1.1times 10^{18}) FLOPs per hour. So we theoretically want no less than 22 436 hours of A100 time to coach. In apply, RWKV 14B was educated on 64 A100s in parallel, sacrificing a little bit of efficiency for numerous causes. RWKV 14B took about 3 months (approx 140 160) A100 hours to coach, thus attaining about 20% theoretical effectivity (because it took about 5x longer than the theoretical minimal). Current variations can prepare RWKV 14B at round 50% theoretical efficency.
As a tough value estimate, on the time of writing, the most affordable A100 value at cloud-gpus.com was $0.79/h. Coaching the unique 14B RWKV there would therefore value round $100k, however with the current coaching code enhancements we may scale back this to $40k. In apply, there are different concerns like ease of use, timeouts, multi-gpu communication pace, and so on. Thus, one may need extra high-end choices like AWS at $4.096/h. RWKV was educated on compute donated by Stability and EleutherAI.
Now you may think about that coaching 10x extra parameters and 10x extra information will value 100x extra, making it prohibitively costly.