3effecc), trained in this repository as the reference (“the fork”). “Ours” is the small model from the factory: the w64 preset (64 channels, 4 heads of 16, 13 layers). It reads the same 92 inputs as Alex’s model (with input 22, the review state, zeroed) and has 331,511 parameters. Besides the shape, its training differs: knowledge distillation from the fork model, long review windows (up to 32,768 reviews), a final learning-rate decay, and int8 weights plus int8 state at scoring time. The older iter36 config (32 channels, 64 fixed curves plus a correction head) is shown separately in the output heads.How close ours gets to Alex’s model
LogLoss, lower is better. Curve scores the predicted recall probability at the next review; instant scores the predicted grade. Δ is relative to Alex’s model (the fork) on the same reviews. Read live from the factory data, so it follows the latest result.
The 92 numbers each review becomes
One square per input number, in the order the model receives them. Both Alex’s model and ours use exactly this vector; ours sets input 22 (review state) to 0 in training and scoring. Striped squares are blanked on the query copy of a review, the row that has to predict the grade without seeing it. Hover a square or a group.
From review to prediction
Data flows top to bottom. Each stream re-sorts the reviews by its own ID, so the card stream only sees one card's history while the user stream sees every review. Bars are drawn to scale: bar width is the hidden width, segments are heads, the thinner bar under each layer is its feed-forward (channel-mixer) width.
Where the parameters go
Each square is 4,000 parameters, coloured by component, on the same scale for both models. Hover a legend entry to highlight it.
Share of each model's own budget
Inside one RWKV layer
How a single layer splits its parameters. 100 squares = the whole layer (so the two grids are not to the same scale; the totals are). The four mixing matrices grow with width squared, so narrowing from 128 to 64 channels shrinks them 4×. Ours keeps a channel mixer in only four of its 13 layers, and the first layer of each stream has no value-residual LoRA.
- R, K, V, O: receptance (what to read), key (where to write), value (what to write) and output projections of the time mixer.
- Decay LoRA sets how fast each channel of the memory fades; in-context rate LoRA sets how strongly a new review overwrites what the memory already holds at that key.
- Gate LoRA scales the output; value-residual LoRA mixes in the first layer's value so deep layers keep a direct path to the input.
- Channel mixer is a small squared-ReLU feed-forward network, applied after each time mixer.
What has to be remembered
At inference time every card, deck, note, preset and user carries its own recurrent state. Each small grid is one head’s state matrix in one layer (32×32 for Alex, 16×16 for ours); per layer the model also keeps a token-shift vector as wide as the hidden layer for the time mixer, and another for each channel mixer that is still present. Sizes are fp32; ours is scored with int8 state, which takes a quarter of that.
Alex
Ours
Hidden vector between layers
One square per channel, coloured by head.
The output heads
Curve predicts the forgetting curve. Alex mixes 128 fixed-stability basis curves and adds a learned correction sampled at fixed time points; ours uses 3 learned power-law curves and no correction. Instant predicts the four grades. Each tick below is one basis curve's stability (time to fall to 90%), one correction point, or (ours) a curve's starting stability, on a log time axis.
Our head is the single three-tick rug: its three curves have their weight, stability and decay predicted for every review, so the ticks only show where the stabilities start (1 hour, about 4 days, 1 year; the initialisation in the head code). The two rugs under the iter36 label are this repository’s older config; there the -parameter correction head is active and the whole model has parameters. Ours has no such head.
Our head: a few learned power-law curves
Our head (gru_head=3) is modelled on the GRU model in srs-benchmark. Instead of weighting dozens of curves with fixed stabilities, three small linear layers on the curve head’s hidden output predict, for every review, the weight wi, stability Si and decay di of N curves. The Instant head is unchanged.
R(t) = Σi=1..N wi · (1 + t / Si)−di Si, di > 0, weights sum to 1
This is a power-law shape like FSRS’s forgetting curve. The correction head is switched off , so every curve, and their mix, can only fall over time. Our parameter total is computed from the w64 configuration and matches the count the training code reports. Try both kinds of head in the curve explainer on the home page.
Side by side
Sources: rwkv/architecture.py, the iter36 and w64 configurations, rwkv/model/srs_model.py, rwkv/model/rwkv_model.py, rwkv/data_processing.py, rwkv/prepare_batch.py and rwkv/config.py. Parameter counts are computed in the page from the configurations and match the totals in the repository README.