MicroRWKV
Efficient language model combining RWKV linear attention with MoE. Constant per-token latency, linear memory, CPU-trainable.
Architecture
Four components in sequence: RWKV time-mixing for temporal dependencies, RWKV channel-mixing for feature transformation, sliding window attention for local context, and Mixture of Experts for capacity scaling.
RWKV TimeMix
Linear complexity attention
RWKV ChannelMix
Feature transformation
Sliding Window
Local context capture
TinyMoE
4 experts, top-1 routing
Configuration
130M parameters total. Trained on the same data as GPT-2 124M for direct comparison.
| Parameter | Value |
|---|---|
| Layers | 8 |
| Attention heads | 8 |
| Embedding dimension | 768 |
| MoE experts | 4 |
| Window size | 1024 |
| Vocabulary | 50304 |
| Total parameters | 130M |
Loss comparison
Training and validation loss on the same splits GPT-2 used. MicroRWKV reaches competitive performance with the efficiency benefits of linear attention.
| Model | Parameters | Train loss | Val loss |
|---|---|---|---|
| MicroRWKV (ours) | 130M | 2.85 | 2.88 |
| GPT-2 | 124M | 2.82 | 2.86 |
The small gap in training loss reflects the tradeoff made for efficiency. Linear attention cannot attend to all previous tokens equally, but the constant-time inference and linear memory scaling make deployment feasible in resource-constrained environments where transformer inference is not.
Efficiency characteristics
Constant per-token latency
RWKV linear attention processes each token independently. No quadratic dependency on sequence length.
Linear total latency
Processing N tokens takes N × constant time. Transformer inference is O(N²) in the attention layers.
Constant memory
Hidden state is a fixed-size recurrent state. No KV cache that grows with sequence length.
No custom CUDA kernel
Pure PyTorch operations. Runs anywhere PyTorch runs, including CPU-only environments.
CPU-trainable
The model trained to convergence on CPU. Slow, but possible, which is not true for many architectures at this scale.
Visualizations
Loss curves and attention mechanism internals from the training run.
