This is a custom architecture for the nanoRWKV project from RWKV-v4neo. The architecture is based on the original nanoRWKV architecture, but with some modifications.
nanoRWKV
This is RWKV "x051a" which does not require custom CUDA kernel to train, so it works for any GPU / CPU.
RWKV is essentially an RNN with unrivaled advantage when doing inference. Here we benchmark the speed and space occupation of RWKV, along with its Transformer counterpart (code could be found here). We could easily find:
single token generation latency of RWKV is an constant.
overall latency of RWKV is linear with respect to context length.
overall memory occupation of RWKV is an constant.
benchmark
Prerequisites
Before kicking off this project, make sure you are familiar with the following concepts:
RNN: RNN stands for Recurrent Neural Network. It is a type of artificial neural network designed to work with sequential data or time-series data. Check this tutorial about RNN.
Transformer: A Transformer is a type of deep learning model introduced in the paper Attention is All You Need. It is specifically designed for handling sequential data, like natural language processing tasks, by using a mechanism called self-attention. Check this post to know more about Transformer.
LLM: LLM, short for Large Language Model, has taken the world by storm. Check this Awesome-LLM repo and State of GPT.
nanoGPT: the simplest, fastest repository for training/finetuning medium-sized GPTs by great Andrej Karpathy. Here you could find the code and the teaching video.
RWKV Language Model: an RNN with GPT-level LLM performance, which can also be directly trained like a GPT transformer (parallelizable). The model is created by an independent researcher Bo Peng. Get more information here.
RWKV_TimeMix: This component applies a time-based mixing operation to the input, which helps the model capture temporal dependencies.
RWKV_ChannelMix: The channel-based mixing operation is performed in this module, allowing the model to learn better representations across different feature channels.
Sliding Window Attention: This attention mechanism operates on a sliding window of the input, enabling the model to efficiently capture local and global dependencies.
GroupedQAttention: This attention module applies a grouped approach to the query, key, and value computations, improving the model's ability to capture multi-headed attention.
TinyMoE: The Tiny Mixture of Experts (TinyMoE) layer is a lightweight and efficient implementation of a Mixture of Experts (MoE) mechanism, which can help the model learn specialized representations.
Detailed Explanation
RWKV_TimeMix:
This module applies a time-based mixing operation to the input, which helps the model capture temporal dependencies.
It uses several learnable parameters, such as time_maa_k, time_maa_v, time_maa_r, and time_maa_g, to control the mixing process.
The module also applies a time-decay mechanism using the time_decay parameter, which allows the model to give more importance to recent inputs.
The output of this module is then passed through a series of linear layers, including the receptance, key, value, and gate layers.
Time Mixing
RWKV_ChannelMix:
This module performs a channel-based mixing operation on the input, allowing the model to learn better representations across different feature channels.
It uses a time-shift operation and learnable parameters, such as time_maa_k and time_maa_r, to control the mixing process.
The module applies a key, value, and receptance linear layers to the mixed input, and the output is then passed through a sigmoid activation function.
Channel Mixing
Sliding Window Attention:
This attention mechanism operates on a sliding window of the input, enabling the model to efficiently capture both local and global dependencies.
The module computes the query, key, and value matrices using a linear layer, and then applies a sliding window attention operation to the input.
The output of the sliding window attention is then passed through a final linear layer to produce the final output.
GroupedQAttention:
This attention module applies a grouped approach to the query, key, and value computations, improving the model's ability to capture multi-headed attention.
The module first computes the query, key, value, and weight matrices using a single linear layer, and then splits these matrices into groups.
The attention computation is then performed on each group, and the results are concatenated and passed through a final linear layer.
TinyMoE:
The Tiny Mixture of Experts (TinyMoE) layer is a lightweight and efficient implementation of a Mixture of Experts (MoE) mechanism, which can help the model learn specialized representations.
The module computes attention scores using a linear layer, and then applies these scores to a set of expert networks to produce the final output.
The module also includes an auxiliary loss term that encourages the experts to learn diverse representations, improving the overall performance of the model.
Usage (Inference)
To use this model for inference, you can follow these steps:
Download and paste model weights in the out directory.
Copy and paste the values like: block_size, vocab_size, etc from the table into the class GPTConfig in generate.py.
Explain:
This command will generate text based on the input prompt "One day" using the model weights stored in the out directory. The max_num_tokens parameter specifies the maximum number of tokens to generate, and the model_name parameter specifies the name of the model weights file to load. For model_name, you can specify the name of the model weights file without the extension, like "ckpt-500" or "ckpt-1000" or only "ckpt".
Tables
name_model
BLOCK_SIZE
VOCAB_SIZE
N_LAYER
N_HEAD
N_EMBD
NUM_EXPERTS
NUM_ACTIVE_EXPERTS
EXPERT_DIM
DIM
DROPOUT
BIAS
DATASET
ckpt-500.pth
1024
50304
8
8
768
4
4
512
768
0.0
False
tinystories_15k
Results
Prompt: One day
Generated text: One day: Sharing positive bought Isabel a rainbow hug. Her name was an vitamins, so only one favorite thing to cheer she were.
Lily picked up a hay and proudly went to a small portion. She was very happened. When Tommy said it
Generated text length: 227 | Inference time: 3 seconds
We got the results as follows
image
model
params
train loss
val loss
GPT-2
124M
2.82
2.86
RWKV
130M
2.85
2.88
baselines
Existing OpenAI GPT-2 checkpoints and RWKV checkpoints allow us to get some baselines in place for openwebtext. We can get the numbers as follows:
Notice that both models are not trained in the openwebtext (RWKV in The Pile and OpenAI GPT-2 in private WebText), so they could be further improved due to dataset domain gap.
Dependencies
torch
numpy
tiktoken
Conclusion
The MicroRWKV model is a custom neural network architecture that combines several cutting-edge techniques, such as time-based and channel-based mixing, sliding window attention, grouped attention, and a Tiny Mixture of Experts (TinyMoE) layer. These components work together to enhance the model's ability to capture both local and global dependencies, as well as to learn specialized representations. The combination of these techniques results in a powerful and efficient model that can be used for a variety of natural language processing tasks.
Reference
Here are some useful references (offering my sincerest gratitude):
nanoGPT - The original nanoGPT implementation by [Andrej Karpathy]
