Revisiting Bi-Linear State Transitions in Recurrent Neural Networks

1. How do bi-linear RNNs compare to the other RNNs in Figure 1 on "real" sequential processing tasks? E.g., sequential MNIST/Long Range Arena?

The primary focus of our work is to analyze the state-tracking capabilities of bilinear recurrent networks and identify the architectural principles that govern them. For this reason, we concentrated on synthetic tasks that allow for controlled and precise evaluation. That said, we agree that scaling these models for natural language pretraining is a crucial next step. We believe that robust state-tracking is a significant advantage for complex reasoning, but scaling up introduces challenges, including improving recall and managing gradient flow during parallelizable training on long context. These challenges could potentially be addressed by integrating gating mechanisms, specialized convolutions or memory layers, as done also for SSMs in the past. We'll add a discussion on this to the revised manuscript.

---

2. Is there any connection to associative memory models, which also have multiplicative connections (albeit between hidden units, and not between inputs). For example, see Figure 1 here. There is also a (very) recent paper which contains multiplicative interactions between hidden states and external inputs. It would be interesting to know if there is some theoretical connection between bi-linear RNNs and associative memory.

We thank the reviewer for this insightful question and for the pointers to the related work (which we will include in the paper). There is indeed a fundamental theoretical connection between the bilinear formulation of recurrence and associative memory networks, particularly with modern dense associative memory models (modern Hopfield networks). The foundational link is the principle of multiplicative interactions modulated by an external signal. While classic associative memories focused on hidden-to-hidden connections, modern Hopfield networks (variants with continuous state, or the IDP Hopfield model as pointed out by the reviewer) have an update rule reminiscent of attention mechanisms, where an input query dynamically weights the stored memory patterns. This is conceptually similar to the bilinear recurrence formulation, where the input modulates the hidden-state transition dynamics of the system in a multiplicative manner. We acknowledge that deeper understanding and study of this connection could lead to interesting future directions and appreciate the point raised by the reviewer.

Several theoretical, conceptual, and experimental similarities with prior art...

Thank you for pointing us to the work by Terzic et al. (2025). It is a well-written paper with interesting connections to our work. We will add a detailed discussion in the related work section and highlight similarities and differences in other relevant subsections. We acknowledge the similarities: both papers analyze the state-tracking capabilities of linear RNNs, advocate for dense transitions to improve state-tracking, and show that complex diagonal transition matrices can only learn commutative operations. However, our work also differs in several key aspects. The model in Terzic et al. constructs its transition matrix as a convex combination of a fixed set of learnable dense matrices. This can be viewed as a specific, parameter-efficient approximation of a more general bilinear model, where the transition matrix is fully parameterized by the input. Our approach to stability also differs: while they use additive terms in the recurrence and normalize the columns of the transition matrix, our bilinear recurrence formulation does not inherently include additive terms. This enables scale-invariance and allows for direct normalization of the hidden states at arbitrary points in time. Furthermore, our work contributes a more detailed hierarchy of state-tracking by showing, both theoretically and experimentally, that a real-diagonal model can only learn parity, but with extreme data efficiency. Finally, we introduce and study two parameter-efficient variants: the factorized and block-diagonal bilinear models.

---

Q1) The empirical evaluations in Table 1 are limited to only three synthetic tasks. How about the non-commutative automata with solvable transformation semigroups, such as the dihedral group? A pointer for these tasks: B. Liu, et al.

We thank the reviewer for this suggestion and the pointer. Following the suggestion, we benchmarked all models on the dihedral group simulation task in the same setup as discussed in the paper (training length 10, testing length 500, best of 3 learning rates). As the results below demonstrate, bilinear RNN models successfully learn a length-generalizable solution to this task:

---

Q2) In Table 1, what is the reason for the drop in the OOD accuracy of RNNs as the Module increases? RNNs have sufficient expressiveness to learn any regular language.

We hypothesize that the observed drop in accuracy could be due to the learning problem becoming harder as the modulus (m) grows. A larger modulus means a larger state space and a more complex transition function for the model to learn. This could make the optimization more difficult and amplify the well-known challenges of propagating gradients through time in standard RNNs. Therefore, while the required solution is representable within the RNN's capacity, the model may struggle to converge to it as the task complexity increases.

---

Q3) The best-performing model is reported (out of three using different learning rates). It is also good to report the min-max, or average-max, performance across models.

We thank the reviewer for this suggestion. We will add a table to the appendix showing the minimum and maximum accuracy for each model across the learning rates, complementing the best-performing results presented in the main text.

---

Q4) Table 2 (right) shows the detrimental effect of additive terms on the parity automata, which is consistent with the prior art. How about benchmarking when learning a non-commutative automaton (e.g., D_30)? Could the model provide better length generalization once the B matrix is introduced? See Table 4 of the prior art.

The mentioned work shows that the complex diagonal model can generalize to longer sequences on a commutative automaton, but only when the additive term is removed and a linear readout is used. However, the same model fails to learn a non-commutative automaton. In that case, the best performance is achieved when additive terms are included and a non-linear readout layer is used. Even then, generalization remains limited, though the additive terms offer a slight improvement.

We performed an ablation study on the non-commutative state machine simulation task. We report results for the diagonal, 2D block-diagonal, and full bilinear models on parity, modular addition, and state machine simulation task (as a reminder, please note that accuracy numbers are scaled so that 0.0 represents random chance performance):

Validation Accuracy (Length 2-10)

Test Accuracy (Length 500)

As observed, and in line with the proposed hierarchy, the diagonal model is only capable of learning the parity task (only when the additive terms are excluded), the 2D block-diagonal model can learn modular addition (again only without the additive terms). In contrast, the full bilinear model is able to learn both modular addition and the non-commutative task, regardless of whether the additive terms are included. Consistent with prior work, we find that additive terms improve the models' performance on the non-commutative task, although the overall performance remains low in both cases. We will include this discussion in the revised version of the paper, along with citations to the relevant work.

---

Q5) Is there any sample efficiency comparison with RNNs in Figure 2?

Based on your question, we also performed the data-efficiency experiments for the RNN using the same setup described in Section 3.2, and we report the results below:

Modular Addition

State Machine

Modular Arithmetic

Compared to the LSTM, we found that the simple RNN shows slightly lower data efficiency on the modular addition and state-machine tasks, but slightly better efficiency on the modular arithmetic task when using larger moduli. We will include this in Figure 2.

---

Q6) What are the potential limitations and implications of applying this method towards making LLMs on real-world datasets?

The benefits of improved state-tracking are particularly evident in math and code benchmarks. There is supporting evidence in the literature, for example, Grazzi et al. (2025) (figure 10), which shows that enabling even the weakest form of state-tracking in models like Mamba and DeltaNet leads to improved perplexity in math and code. Therefore, we believe that models with stronger state-tracking capabilities will have a significant advantage in reasoning more generally.

Currently, almost all large-scale pretrained language models have limited internal state-tracking capabilities. Although they rely on chain of thought to make the state observable in the token space, this becomes impractical when the state is large or continuous.

Although bilinear models demonstrate strong state-tracking, we identify two potential limitations for scaling them to large-scale pretraining on natural language. First is their recall performance, which is crucial in the natural language domain. A second potential challenge is managing gradient flow over very long sequences while maintaining parallelizable training. Similar to recent developments in SSMs, limitations like these could be overcome by utilizing additional gating mechanisms, skip connections, 1D convolutions, or even hybrid layers that incorporate memory modules. While this is out of scope for this work, we do consider it important future work.

Grazzi, Riccardo, et al. "Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues." The Thirteenth International Conference on Learning Representations.

While the authors do address the parameter-inefficient nature of bilinear RNNs by considering low-rank or factorized parameterizations, their experiments do not compare different model classes in a parameter-matched way; instead, hidden dimension is held fixed, giving a potentially massive parameter count advantage to the bilinear RNNs. A more complete set of experiments comparing bilinear RNNs and their factorized variants to other architectures with conserved parameter count would be needed to understand whether bilinear RNNs are actually practically useful on state-tracking tasks, and would be a useful addition to round out the paper regardless of what the results end up showing. I am open to bumping up my score to "Accept" if this can be done.

We thank the reviewer for their constructive suggestion. First, we would like to emphasize that we chose to use a consistent hidden dimension size of 256 for all models in Table 1 which resulted in very large parameter counts for some of the models. As suggested, we performed a new experiment to compare models in a parameter-matched setting instead, by adjusting the hidden state size of the models to approximately match parameter counts:

We report the OOD test accuracy on sequences of length 500, with training performed on sequences up to length 10. The results show the superior performance of bilinear models in most state-tracking tasks (but a mild degradation for modular arithmetic), even with matched parameter counts:

Thank you again for the suggestion. We will add these results to the revised version of the paper.

---

Can the authors comment further on why including biases degrades performance of bilinear RNNs? Why can't the model learn to simply ignore the biases, and why does this lead to instabilities?

Our hypothesis is that the presence of additive terms may alter the loss landscape such that, while a zero-bias solution exists, it may not be stable or easily reachable via gradient descent. This is similar to how different initialization schemes can lead optimization toward different local minima, even when the networks can represent better solutions. We would also like to highlight an additional benefit of removing bias terms from the recurrence: the model becomes scale-invariant, which enables normalization of the hidden states to maintain stability over long sequences. This is demonstrated in our length generalization setup with sequences of length 500.

The paper mentions that the stability problems of bi-linear recurrent models were the main reason for them not being used, and seems to solve this issue by dropping any additive terms in the state update. However, further (details on) experiments on this issue are missing.

Yes, one benefit of removing additive terms from the recurrence is that the model becomes scale-invariant, which allows for normalization of the hidden states to ensure stability. The evidence for this stability is our strong length generalization result: models trained on sequences up to length 10 performed successfully on sequences of length 500, which is made possible by performing iterative normalization during inference. We thank the reviewer for raising this point; we will clarify the link between scale-invariance, normalization, and our OOD results in the updated version of the paper.

---

There are no „real world“ experiments, e.g. on language or vision. It could be at least explained if these models are applicable to a more large scale training setup?

The primary focus of our work is to analyze the state-tracking capabilities of bilinear recurrent networks and identify the architectural principles that govern them. For this reason, we concentrated on synthetic tasks that allow for controlled and precise evaluation. That said, we agree that scaling these models for natural language pretraining is a crucial next step. This would be possible in principle, because unlike non-linear RNNs (but like linear RNNs/SSMs) all relevant computations could be parallelized for bilinear RNNs. We believe that robust state-tracking is a significant advantage for complex reasoning, but scaling up introduces challenges, including improving recall and managing gradient flow during parallelizable training on long context. These challenges can be effectively addressed by integrating gating mechanisms, specialized convolutions or memory layers. We'll add a discussion on these future directions to the revised manuscript.

---

What are the exact parameter counts of the models? Since it is mentioned in the discussion section that the parameter count is larger due to the cubic scaling with embedding dim, it would make this point more clear if the parameter counts are added to Table 1, too.

Thank you for this helpful suggestion. While this information is already provided for a subset of models in Figure 2, we will add the parameter counts additionally for all models to Table 1 for added clarity:

• Full Bilinear: ~16.7M
• Block Diagonal, block size 1: ~69K, block size 2: ~135K, block size 8: ~528K, block size 64: ~ 4.2M
• Factored Bilinear rank 1: ~4K, rank 2: ~5K, rank 4: ~7K, rank 8: ~10K, rank 16: ~16K, rank 64: ~53K, rank 128: ~102K, rank 256: ~200K, rank 512: ~397K, rank 1024: ~790K, rank 2048: ~1.5M
• 2D Rotation: ~38K
• LSTM: 531K
• RNN: 135K
• Transformer 1-layer: ~1M, 2-layer: ~1.8M, 4-layer: ~3.4M
• Mamba 1-layer: 473K, 2-layer: 943K, 4-layer: 1.8M

However, we also want to clarify a key point regarding these numbers: the large parameter counts for some of our models (e.g., the full bilinear variant) are a direct result of matching hidden dimension (256) rather than parameters in these experiments. To address the valid concern about parameter efficiency, we conducted a new set of experiments where models are matched in parameter count by adjusting their hidden dimension:

We report the OOD test accuracy on sequences of length 500, with training performed on sequences up to length 10. These new results still show superior state-tracking performance of bilinear models in most tasks (a slight degradation for modular arithmetic). We will add them to the updated version of the paper.

---

L 271: Why do you need to normalize during inference and not during training?

We did not apply normalization during training simply because the input sequences were short enough. It is entirely possible to apply hidden state normalization during both training and inference, and we will clarify in the revised text that its omission during our training runs was for simplicity, not a requirement of the model.