RWKV-7 "Goose" with Expressive Dynamic State Evolution

1. One possible concern ...

We agree about the importance of separating out entangled effects. To address this specific issue we provided a set of ablation studies in which we train all models on the Pile. Please refer to Appendix L and Table 19 for these ablation study details. A concise plot of some of these results for RWKV-7 and Mamba-2 can also be found in the Pile models shown in Figure 7.

2. There lacks the evaluation on real-world generation tasks (e.g., story generation, stylist generation, role playing, etc.) which requires annotations from human evaluators.

While we agree that human evaluations and annotations are valuable for real-world generation tasks, they are unfortunately beyond the scope of our research focusing on developing new modeling architectures.

1. Could the dataset proposed RWKV-world3 dataset be used on other architectures? What would be the results?

This would be a fascinating study, and we would love to see other models adopt the multilingual RWKV-World3 dataset. Our belief is that it would be beneficial for any architecture, and could help improve accessibility for worldwide adoption of open source models. We have more open datasets coming soon that perform even better in our internal testing, and hopefully these will be a valuable resource to other model authors.

2. Figure 3 could be improved with better resolution and alignment.

Thank you for noticing this oversight. We have updated the image in the manuscript accordingly for camera ready.

1. The state update mechanism and overall architecture are quite complex, making the model potentially difficult to implement and understand for those not already familiar with linear attention variants.

Although we cannot simplify the actual architecture further without reducing performance, we may be able to provide a more succinct description. Please let us know if you think it would be helpful for us to include the following 30 line numpy implementation in an appendix:

group_norm = lambda x, w, b : ((x - x.mean(axis=1,keepdims=1)) / (x.var(axis=1,keepdims=1) + 64e-5)**0.5).flatten() * w + b
sigmoid = lambda x : 1/(1+np.exp(-x))
def time_mixing(x, v0, last_x, S, params):   
    mr,mw,mk,mv,ma,mg, w_bias, r_k, Ww1,Ww2, Wa1,Wa2,a_bias, Wg1,Wg2 = params[:15]
    k_k,k_a, Wr,Wk,Wv,Wo, ln_w,ln_b = params[-8:]

    xr,xw,xk,xv,xa,xg = [x + m * (last_x - x) for m in [mr,mw,mk,mv,ma,mg]]

    r = Wr @ xr
    w = np.exp(-sigmoid(np.tanh(xw @ Ww1) @ Ww2 + w_bias)/np.e**0.5)
    k = Wk @ xk
    v = Wv @ xv
    if v0 is None:
        v0 = v
    else:
        Wv2,Wv1,v_bias = params[15:18]
        v += (v0 - v) * sigmoid(xv @ Wv1 @ Wv2 + v_bias)
    a = sigmoid(xa @ Wa1 @ Wa2 + a_bias)
    g = sigmoid(xg @ Wg1) @ Wg2
    kk = k * k_k
    k += k * (a-1) * k_a

    r,w,k,v,kk,a,r_k = [i.reshape(N_HEAD, HEAD_SIZE, 1) for i in [r,w,k,v,kk,a,r_k]]
    kk /= np.maximum(np.linalg.norm(kk, axis=1,keepdims=1), 1e-12)

    S = S * w.mT - S @ kk * (kk*a).mT + v * k.mT
    y = S @ r

    y = group_norm(y, ln_w, ln_b)
    y += ((r * k * r_k).sum(axis=1,keepdims=1) * v).flatten()
    return Wo @ (y * g), v0, x, S

2. While some ablations are provided, a more comprehensive analysis separating the contributions of each architectural component would strengthen the paper.

Thank you for bringing this to our attention. We will consider if there are other ablations we can run that might help illustrate the specific role that each component plays.

Questions To Authors:

When comparing architectures trained on identical datasets (as in the Pile experiments), what are the specific speed

Training optimization can vary widely across different model architectures, sizes, and devices, so such comparisons may unfortunately be beyond the scope of this paper.

"we expand the in-context learning rate" <- What is the ICL rate? Addendum after reading: please add a forward reference to where the ICL rate is defined, i.e., "we expand the in-context learning rate, described in Section..."

Thank you for the suggestion. We have updated the manuscript to reflect this change for camera-ready.

Can the authors speak to both the range of values and the types of instabilities encountered? Might these instabilities be side-stepped and the range of negative values increased during fine-tuning?

You raise a very good point, and we supply some details about this in Theorem 1 in Appendix B. It is worth noting that, although all the eigenvalues are bounded in $[-1,1]$, this does not always bound the largest singular value (the spectral norm) below one, which is a sufficient (but not necessary) condition for overall stability. In practice, in some timesteps the largest singular value does exceed 1, but further inspections (Appendix M) show that it does not affect general stability of state values. We will continue to investigate this issue in our future studies.

What model is Flash Attention 3 being run

In Figure 3 we show results for isolated kernels, rather than full models.

Why is English Lambada perplexity not reported

Thank you for catching this oversight, which may have occurred due to concerns that the text in the table had gotten too small. We will add these results back into the paper for the camera ready.

it would have been interesting to see mamba-2.8b-slimpj

Sure! Here are the results:

To verify, the non-instruct-tuned versions of relevant models are used for comparison in Table 2 and 3, correct? If possible, could the authors include an appendix table of the competitor HF checkpoints?

Thank you for the excellent suggestion - we have updated the tables in the manuscript to add hyperlinks to all listed models. All models listed in Table 2 and 3 are base models without instruction fine tuning.

What are the parameter counts for the models in Table 4?

According to the default settings listed in https://github.com/athms/mad-lab/blob/main/benchmark.py, the models have two layers with dimension 128. This accounts for approximately 2.5M parameters.

Could you please add implementation details for the compression rate metric? Also, to make the paper self-contained, could you add a description of the compression rate metric used in Table 15.

Thank you for the suggestion. We have added the following to the caption under Table 15 for camera-ready: We define compression rate as: $\frac{\log_2(e) * \text{average loss over document} * \text{document tokens}}{8 * \text{document length in bytes}}.$ This definition is also consistent with Shannon’s information theory.

We sincerely hope these clarifications could strengthen the manuscript. Please don’t hesitate to reach out if there are any further questions or suggestions.

The paper demonstrates that RWKV-7 has expressivity that exceeds TC complexity class under standard complexity conjectures, which is a theoretical advantage over Transformers. But I am not so sure or confident about this conclusion, which I would encourage authors to elaborate further.

Thank you for this point about needing further elaboration on the actual advantages that come from the complexity class improvements of RWKV-7 over Transformers. Transformers aggregate history through commutative sums, which limits their ability to solve sequential state tracking problems where the precise order of actions matters for the result (formally, these problems are outside TC0 under standard conjectures). A key example of state tracking is computing the state of a chess board given a sequence of moves, each specified in notation like ‘e2e4’, indicating the source and target square. Transformers, S4 and Mamba cannot solve this chess state tracking problem beyond short input lengths unless extra 'thinking' tokens are given (See https://arxiv.org/abs/2404.08819 and https://arxiv.org/abs/2310.07923). In contrast, our result shows that RWKV-7 can solve the chess state tracking task for any sequence length as long as the model has at least four layers with sufficient width, and a head size that is linear with respect to the number of board positions. We prove this result theoretically in Appendix D. We will add further description and examples like this for camera-ready.

Numerical Precision Issues: The paper notes sensitivity to numerical precision in certain operators, particularly the WKV7 kernel, which may complicate deployment. But I don't think this belies the strength of the paper.

Thank you for raising this important concern about numerical precision. Currently, we have multiple implementations of WKV7 kernel for model training, some written in CUDA and some in Triton, some optimized for higher precision and some for balanced performance under different sequence lengths. For example, a "backstepping" implementation computes gradients in FP32, achieving a numerical error of 9e-5, which is generally satisfactory. On the other hand, a faster "chunked" implementation uses BF16 precision, reducing computation time by 2x (saving ~8ms per pass in our 0.1B model) but introduces a slightly higher error of 5e-3. Unlike Flash Attention v3, our kernels have not yet undergone GPU architecture-specific optimizations (like Nvidia Hopper or AMD MI300 series), and we are actively working to optimize these implementations further.

Lack of Instruction Tuning: The presented models are base models without instruction tuning or alignment, limiting direct comparability to fully-tuned models in practical applications. But this would be helpful to further understand how this model can be actually used in the production environment.

Our pre-trained RWKV-7 base models do not feature any true instruction tuning or alignment, having been exposed to only the limited amount of open-source instruction data in the pretraining stage contained in the RWKV World v3 corpus. Specific instruction tuning and alignment processes can vary widely, and are unfortunately beyond the scope of our architecture paper.

Questions To Authors: In my view, the paper would benefit from: Moving some of the key theoretical results from the appendix to the main paper Providing more intuitive explanations of why these complexity results matter Including concrete examples that demonstrate how these theoretical capabilities translate to practical advantages Better connecting the theoretical guarantees to the empirical results observed in benchmarks

These are great ideas and seem feasible to implement. While we were limited by space concerns for the submission version, we hope that the camera ready might allow us one more page that we can use for this purpose. We can give examples, including experiments shown in the paper, of the kinds of problems these complexity improvements allow the model to solve without requiring extra tokens.