Fixed-Point RNNs: Interpolating from Diagonal to Dense

We thank the reviewer for their positive and detailed feedback and we’re glad that they found our paper easy to follow, describing a novel idea that would benefit the research community and could unlock a plethora of models. We address their suggestions as follows:

1. Reconcile results with Koopman theory.

Thank you for asking for clarification on this important point! Indeed, according to Koopman theory any nonlinear dynamical system can be approximated by an SSM with a large enough hidden state $d$ — as long as it is followed by an MLP tasked with computing eigenfunctions. It is important to note, however, that the form of the Koopman operator depends on the nature of the dynamical system. In the case of a scalar-valued dynamical system the state-to-state transition acts as an element-wise linear operator (like Mamba). But the Koopman operator of a vector-valued dynamical system is generally a dense object. It can be diagonalized into an element-wise operator only if there exists an eigenbasis which decouples the vector-valued dynamical system. This is where the expressivity results due to Merrill et al (2024) and Cirone et al (2024) come in: certain problems such as state-tracking cannot be expressed by a decoupled dynamical system. In the case of state-tracking, it is easy to see that the time-varying permutation matrices which represent the linear operator cannot be co-diagonalized. We will try to make this clearer in the introduction section, and point to the Discussion in Cirone et al (2024), who directly addressed this issue mathematically in their section 4.3.

Since Mamba can express progressively more coupled dynamical systems with depth, we empirically investigated scaling it’s the depth (instead of width $d$) to operate under the same runtime as FP-Mamba. We hope that answer we gave to the other reviewers addresses your question:

2. Make compute trade-offs more transparent.

We agree that the main submission provided insufficient guidance on the compute trade-offs of our FP-RNN approach when restricting the number of iterations. Therefore, we had already included plot to relate compute time to generalisation on the state-tracking task $A\_5$ in Appendix E.4, Figure 11 and now provide an analysis here:

”Wall clock time is plotted against the longest test sequence length with > 90% accuracy for every model. While baselines of increasing depth $\\in \\{1, 2, 4, 6, 8\\}$ cannot generalize beyond the training sequence length 16 (horizontal pink line), our proposed framework allows to achieve much higher generalization by scaling training time through the number of fixed-point iterations $\\ell\_\\text{max} \\in \\{2, 4, 8, 16\\}$. There is further room to improve efficiency, as suggested by a simple randomization scheme (gray stars) where $\\ell\_{\\text{max}} \\sim\\Gamma(4,1)$ is sampled from a Gamma distribution with mean $4$ for every batch.”

We will include the above analysis in Sec 4.5 and move Figure 11 to the main text (replacing Figure 7).

3. Scaling of $\\ell^*$ with progressively more challenging tasks

We would love to provide benchmarking analysis, but unfortunately the NeurIPS rebuttal format does not allow for rich media. Therefore, we will resort to explaining the theoretical cost of executing a single layer of our method vs. the underlying sequence mixer used (Mamba) instead. Denoting the sequence length by $T$, the state size by $N$, and the number of iterations to convergence by $\\ell^*$, the FLOPS are:

Note that the overall cost of FP-Mamba is $\\ell^\*$ times the cost of channel+sequence mixing for the forward pass while no additional cost for the iteration occurs for the backward pass. While the sequence mixing cost is inherited from Mamba, the channel mixing incurs an additional cost that is sub-quadratic in state size and linear in sequence length in all three variants of the mixers investigated in this paper.

Figure 3c depicts the value of $\\ell^\*$ in a progressively more challenging task. In this state tracking task $A\_5$, the difficulty corresponds to the number of state transitions, which is equal to the sequence length $T$. From the figure, we observe a linear relationship between $\\ell^\*$ and $T$, indicating that, at least in the state tracking problem, $\\ell^\*$ scales linearly with the difficulty of the task. Note, however, that on simpler tasks (Figures 5b, 7, 10) the required number of fixed-point iterations $\\ell^\*$ remains well below the sequence length $T$. Finally, the overall FLOPs are scaled by the task-specific $\\ell^\*$ according to the cost in the table above. We will update the paper with this analysis.

4. Effect of normalization factor $\\mathbf{(I-\\Lambda\_t)}$ on class of matrices $A\_t$.

Thank you for pointing out the missing description of the representable dense matrices $\\mathbf{A}\_t$ in the presence of the normalization factor $(\\mathbf{I}-\\mathbf{\\Lambda}\_t)$. To drive $\\mathbf{A}\_t$, let us start by assuming a fixed-point was found according to Thm 3.1:

$h\_t^* = \\mathbf{\\Lambda}\_t h\_{t-1}^* + (\\mathbf{I}-\\mathbf{\\Lambda}\_t) (\\mathbf{Q}\_t\\mathbf{B}\_t x\_t + (\\mathbf{I} -\\mathbf{Q}\_t)h\_t^*)$

Rearranging the terms allows to move $h\_t^*$ to the other side:

$(\\mathbf{I} - (\\mathbf{I} - \\mathbf{\\Lambda}\_t)(\\mathbf{I} - \\mathbf{Q}\_t))h\_t^* = \\mathbf{\\Lambda}\_t h\_{t-1}^* + (\\mathbf{I}-\\mathbf{\\Lambda}\_t) \\mathbf{Q}\_t\\mathbf{B}\_t x\_t$

Therefore, $\\mathbf{A}\_t = (\\mathbf{I} - (\\mathbf{I} - \\mathbf{\\Lambda}\_t)(\\mathbf{I} - \\mathbf{Q}\_t))^{-1} \\mathbf{\\Lambda}\_t = (\\mathbf{\\Lambda}\_t^{-1}(\\mathbf{I} - (\\mathbf{I} - \\mathbf{\\Lambda}\_t)(\\mathbf{I} - \\mathbf{Q}\_t)))^{-1} = (\\mathbf{\\Lambda}\_t^{-1} - (\\mathbf{\\Lambda}\_t^{-1} - \\mathbf{I})(\\mathbf{I} - \\mathbf{Q}\_t))^{-1} =(\\mathbf{I} + (\\mathbf{\\Lambda}\_t^{-1} - \\mathbf{I})\\mathbf{Q}\_t)^{-1}$

Following the standard assumptions that $0\\preceq \\mathbf{\\Lambda}\_t, \\mathbf{Q}\_t\\preceq \\mathbf{I}$, the matrix $\\left(\\mathbf{I}+\\left(\\mathbf{\\Lambda}\_t^{-1} - \\mathbf{I}\\right)\\mathbf{Q}\_t\\right)$ is full rank and $\\succeq \mathbf{I}$. Therefore its inverse $\\mathbf{A}\_t$ exists and is contractive. The expressivity of $\\mathbf{A}\_t$ is only limited if $\\mathbf{\\Lambda}\_t\\approx \\mathbf{I}$. This however, would also be problematic for a diagonal SSM and therefore the Mamba initialization is biased towards $\\mathbf{\\Lambda}\_t \\prec \\mathbf{I}$. Thus, we argue that in practice, these modifications do not pose a significant problem. We will make sure to refer to this analysis when we introduce normalization in Sec. 3.1.

Questions:

1. Are FP-RNNs unstable during training? Although a large portion of our theoretical contributions focus on this issue, we notice that under certain conditions we can still face instability in training. Specifically, we notice that for long sequences in the state tracking tasks ($A\_5$, $S\_5$), where the difficulty of the task scales linearly with the sequence length, training can become unstable. In this case, we notice that incorporating off-the-shelf stability techniques, such as weight normalization, can significantly help with this issue (Gent et al, 2023).
2. Are there trade-offs between the number of parameters and the required number of fixed-point iterations? In our experiments, we indeed observe a relationship between the parameter numbers and the required number of fixed-point iterations for finding the fixed point. However, this only happens when the model is modified in depth, and not in width. Specifically, we observe that when increasing the depth of the model for more challenging tasks, the number of iterations indeed reduces.
3. Is there a way to reduce the number of fixed-point iterations? We have several potential solutions to reducing the number of fixed-point iterations. Firstly, we observe that off-the-shelf optimization techniques do not perform as well as fixed-point iterations near the initial point of the hidden state, but are much faster to find the exact solution once the hidden state is in the vicinity of the fixed-point. Consequently, we believe a hybrid approach would be possible to significantly convergence. We also introduce some techniques in Appendix B.4 for this purpose. Finally, as we noted in our response in the previous question, we observe that for deeper networks, the fixed-point becomes much easier to find. This may indicate the existence of a structure in the mixer (e.g. positive semi-definiteness) that can significantly improve the cost of finding the fixed-point.

We thank the reviewer for the list of minor issues which we will fix in the camera ready version. We hope that our response addresses the reviewer’s concerns and would be interested to hear if they are considering raising their score to help increase the visibility of our work during the conference. In case there are any remaining questions, we are of course happy to resolve them during the discussion period.

Merrill et al (2024), ”The Illusion of State in State-Space Models”, https://arxiv.org/abs/2404.08819

Cirone et al. (2024), ”Theoretical Foundations of Deep Selective State-Space Models”, https://arxiv.org/abs/2402.19047

Geng et al. (2023). "Torchdeq: A library for deep equilibrium models." https://arxiv.org/abs/2310.18605

We thank the reviewer for their positive and detailed feedback and we’re glad that they found our paper a pleasure to read, describing a theoretically grounded, novel, and clean method that clearly points to new research directions. We address their suggestions as follows:

1. Make compute trade-offs more transparent.

We agree that the main submission provided insufficient guidance on the compute trade-offs of our FP-RNN approach when restricting the number of iterations. Therefore, we had already included plot to relate compute time to generalisation on the state-tracking task $A\_5$ in Appendix E.4, Figure 11 and now provide an analysis here:

”Wall clock time is plotted against the longest test sequence length with > 90% accuracy for every model. While baselines of increasing depth $\\in \\{1, 2, 4, 6, 8\\}$ cannot generalize beyond the training sequence length 16 (horizontal pink line), our proposed framework allows to achieve much higher generalization by scaling training time through the number of fixed-point iterations $\\ell\_\\text{max} \\in \\{2, 4, 8, 16\\}$. There is further room to improve efficiency, as suggested by a simple randomization scheme (gray stars) where $\\ell\_{\\text{max}} \\sim\\Gamma(4,1)$ is sampled from a Gamma distribution with mean $4$ for every batch.”

We will include the above analysis in Sec 4.5 and move Figure 11 to the main text (replacing Figure 7). Thank you for this excellent suggestion which makes the paper stronger!

2. More evaluations on real-world benchmarks.

Thank you for sharing this perspective! It’s true that the main goal of the paper is to introduce a framework which solves the theoretically motivated state-tracking task and we therefore focus on this aspect in the evaluation. Nevertheless, we evaluate a small ~160M FP-Mamba language model trained on 2B tokens of Fine-Web (Penedo et al, 2024) in Appendix E.1. In addition, we now provide new results of FP-Mamba on the reasoning benchmark CatbAbi (Schlag et al, 2021):

In this experiment, we report the test accuracy, using the settings of (Schlag et al, 2021). We optimize the learning rates on Mamba, and use the same learning rate to train FP-Mamba. The vanilla Mamba model hits at the diminishing returns point with more layers, while our proposed dense variant further improves the accuracy with more dense layers. Furthermore, we can see that in this case, the Householder mixer performs better than the Kronecker mixer, indicating that the Householder mixer is more suitable for text data. We observed a similar pattern in language modeling, where FPMamba achieves a perplexity of 26.00 with the Kronecker mixer on the test set, while with a single Householder component it achieves the perplexity of 25.24. The vanilla Mamba achieves a perplexity of 27.01 in this case.

We agree with the reviewer that these initial results on real-world tasks are indeed an important signal to the community and we will thus place them more prominently in the paper. We believe that our old language and new reasoning results point into a promising direction and we would like to include more benchmarks in the camera-ready version if the reviewer has some concrete suggestions. However, we believe that a thorough, large-scale study to investigate which application domains benefit from our verified state-tracking capabilities is beyond the scope of this work.

3. Further discuss required number of iterations.

Yes, we do observe a complexity of $O(T^2)$ on state-tracking in Figure 3c. However, as we discuss in L365-365, we believe that “this is not necessarily a disadvantage if the model is capable of adapting its sequential steps to the difficulty of the task, with negligible cost for the least demanding tasks”. Indeed, on copying (Figure 5b), language modelling (Figure 6), and modular arithmetic (Figure 10) we observe the required number of fixed-point iterations $\ell^*$ is well below the sequence length $T$ suggesting that the model adapts to $O(T)$ complexity on simpler tasks. Finally, a new theoretical study suggests that logarithmic depth might be both necessary and sufficient to achieve state tracking from an expressivity perspective (Merrill et al 2025). In future work, we will investigate how the number of fixed-point iterations could be reduced to this theoretical lower bound in practice. We will make sure to discuss this better in Sec. 4.5 of the paper, thank you for raising this concern!

4. Ablate iteration count, tolerance, and mixer depth. Thank you for the suggestion. We first note that the iteration count is not controllable by us, since it is the model that determines when it reaches the fixed-point. So we assume you are referring to the maximum number of fixed-point iterations. We agree these ablation studies are valuable, and the results in Figure 1 provide a preliminary analysis of these parameters. Our work ultimately focuses on ablating core architectural components, such as the choice of mixer, as we view iteration count and tolerance as solver-specific hyperparameters. But we are happy to expand on these initial findings with a more detailed study for the camera-ready version if needed.

Regarding the mixer-depth choice, we would like to kindly ask you to further elaborate on what you are exactly interested in. In case your question is with regards to the number of Householder (H) components, we have investigate the effect of more Householder components in the paper (up to 4).

5. Highlight state-of-the-art results in the paper. Thank you for your attention to detail. Our claim on state-of-the-art results in the paper specifically refer to the A5/S5 task, for which we believe our reported result for figure 5-a and 5-b - achieving near perfect $>\\times 3$ length generalization - is the best performance reported for these tasks for training in sequence length 16. Furthermore, using the setting proposed for the DeltaProduct method (Siems et al, 2025) for training in sequence length 128, we achieve better $\\times 4$ length generalization than their proposed method on both the $A\_5$ and $S\_5$ tasks (average of 3 separate runs), the result of which are available in the following table:

We will make the distinction in the abstract of the paper about this point.

Questions:

1. What are theoretical guarantees once the recurrence depends on the hidden state? While we did not establish a formal proof, we’re happy to share some ideas on how to prove convergence in this richer setting. To start, we agree that a state-dependent iteration can admit multiple solutions. However, note that in the time dimension, sequence mixing occurs only uni-directionally. This allows for an inductive argument in both dimensions: Since $Q\_0^{\\ell}$ is contractive, depends only on the input and no (varying) hidden state, the iteration in depth on the first time step $h\_0^{\\ell}$ converges (similar to Thm 3.1). Once $h\_0^{\\ell}$ converged and does not change anymore, $Q\_1^{\\ell}$ will not change anymore and allow for $h\_1^{\\ell}$ to converge. This way, the iterations $h\_t^{\\ell}$ converge sequentially in time and the system cannot iterate forever.

2. Is there any guidance on which structure performs best on different tasks? Since the submission, we have ablated the mixers on all tasks and conclude that the Kronecker structure seems to be a good default. The language experiment, however, is the only one where Householder mixers perform slightly better than Kronecker mixers. We will add these new results to the camera-ready version of the paper.

We thank the reviewer for the list of minor issues which we will fix in the camera ready version. We hope that our response addresses the reviewer’s concerns and would be interested to hear if they are considering raising their score to help increase the visibility of our work during the conference. In case there are any remaining questions, we are of course happy to resolve them during the discussion period.

Penedo et al. (2024), "The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale", https://arxiv.org/abs/2406.17557

Schlag et al. (2021), "Learning associative inference using fast weight memory.", https://arxiv.org/abs/2011.07831

Merrill et al. (2025), "A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers.", https://arxiv.org/abs/2503.03961

Siems et al. (2025), "Deltaproduct: Improving state-tracking in linear rnns via Householder products." https://arxiv.org/abs/2502.10297

Thank you for taking the time to explain in so much detail the experimental setup you would have liked to see. This is very valuable feedback since there are many possible configurations for ablations and we believed that we had covered the most important aspects. We particularly liked your suggestion to relate maximum iteration budget $\\ell\_{\\text{max}}$ at training time to the effective number of iterations $\\ell^*$ at test time and ran the corresponding experiment at least for $A\_5$ in the past days. We will create a new ablation section in the Appendix and add results depicted in the table below as a plot. We plan to add similar plots to elucidate the effect of the tolerance $\epsilon$ at test-time for models with tolerance $\epsilon=0.1$ at training-time. Although we might not see your response until the final reviews are released, we would still be interested to hear if this is the ablation you meant or what you would deem more meaningful.

Finally, we want to apologize for the confusion caused by the legend in Fig.1, where $r\\in\{1,2,4\}$ denotes the number of reflections r of the Householder (H) sequence mixer $Q\_t$. We updated all the legends now to make this clearer. We believe that investigating the effect of different mixer structures is a crucial result of our work beyond just ablations. Therefore we had added more results on all the tasks, including rank r of the Diagonal Plus Low-Rank (DPLR) and the hyperparameter-free Kronecker (K) sequence mixer $Q\_t$ after the submission. For the Householder, we observed in preliminary experiments that increasing the number of reflections to $r>4$ yields diminishing returns especially since increasing the rank also increases the computational cost. This is again visible in Fig. 11 of Appendix E.4, where the best performing model ($\\ell_{\\text{max}}=16$ at $r=4$) is 25% slower than the next best model ($\\ell_{\\text{max}}=16$ at $r=2$) but generalizes only to 10% longer sequences.

Ablation train-time $\\ell\_{\\text{max}}$ vs performance and test-time $\\ell^*$

Effect of train-time iteration budget on performance and effective test-time iterations for a 1-layer FP-Mamba with $r=1$ Householder reflections on the state-tracking task $A\_5$. Each cell contains the test accuracy in % and the corresponding test-time number of iterations $l^*$. Reported are averages over 4 runs.

We thank the reviewer for their valuable feedback and we appreciate their assessment of our paper as well written and easy to follow, describing a novel, elegant method that opens a path towards improved sequence models. We address their suggestions as follows:

1. Make compute trade-offs more transparent.

We agree that the main submission provided insufficient guidance on the compute trade-offs of our FP-RNN approach when restricting the number of iterations. Therefore, we had already included plot to relate compute time to generalisation on the state-tracking task $A\_5$ in Appendix E.4, Figure 11 and now provide an analysis here:

”Wall clock time is plotted against the longest test sequence length with > 90% accuracy for every model. While baselines of increasing depth $\\in \\{1, 2, 4, 6, 8\\}$ cannot generalize beyond the training sequence length 16 (horizontal pink line), our proposed framework allows to achieve much higher generalization by scaling training time through the number of fixed-point iterations $\\ell\_\\text{max} \\in \\{2, 4, 8, 16\\}$. There is further room to improve efficiency, as suggested by a simple randomization scheme (gray stars) where $\\ell\_{\\text{max}} \\sim\\Gamma(4,1)$ is sampled from a Gamma distribution with mean $4$ for every batch.”

We will include the above analysis in Sec 4.5 and move Figure 11 to the main text (replacing Figure 7). Thank you for this excellent suggestion which makes the paper stronger!

2. More evaluations on real-world benchmarks.

Thank you for sharing this perspective! It’s true that the main goal of the paper is to introduce a framework which solves the theoretically motivated state-tracking task and we therefore focus on this aspect in the evaluation. Nevertheless, we evaluate a small ~160M FP-Mamba language model trained on 2B tokens of Fine-Web (Penedo et al, 2024) in Appendix E.1. In addition, we now provide new results of FP-Mamba on the reasoning benchmark CatbAbi (Schlag et al, 2021):

In this experiment, we report the test accuracy, using the settings of (Schlag et al, 2021). We optimize the learning rates on Mamba, and use the same learning rate to train FP-Mamba. The vanilla Mamba model hits the diminishing returns point with more layers, while our proposed dense variant further improves the accuracy with more dense layers. Furthermore, we can see that in this case, the Householder mixer performs better than the Kronecker mixer, indicating that the Householder mixer is more suitable for text data. We observed a similar pattern in language modeling, where FPMamba achieves a perplexity of 26.00 with the Kronecker mixer on the test set, while with a single Householder component it achieves the perplexity of 25.24. The vanilla Mamba achieves a perplexity of 27.01 in this case.

We agree with the reviewer that these initial results on real-world tasks are indeed an important signal to the community and we will thus place them more prominently in the paper. We believe that our old language and new reasoning results point into a promising direction and we would like to include more benchmarks in the camera-ready version if the reviewer has some concrete suggestions. However, we believe that a thorough, large-scale study to investigate which application domains benefit from our verified state-tracking capabilities is beyond the scope of this work.

3. Provide intuition/heuristics for choosing the number of fixed-point iterations.

As described in L222 and L298, the fixed-iteration stops once the relative distance between iterates falls below a threshold of 0.1, so the number of fixed-point iterations is chosen by the model itself and not the user. In order to avoid excessive computation on outlier samples during training, we usually set the maximum number of iterations $\\ell\_{max}$ to the sequence length $T$ and apply other heuristics as described in Appendix B.4.

Questions:

1. Are there examples where the convergence of the FP-RNN is very slow? We observe different convergence speeds depending on the application. For example state-tracking tasks are very hard to solve, whereas language modelling is easier and converges usually in a few iterations as depicted in Figure 6.

2. How should the mixer structure be selected in practice? We suggest empirically evaluating the best performing mixer structure depending on the specific application. In the meantime, we have ablated the mixers on all tasks and conclude that the Kronecker structure seems to be a good default for state tracking, and the Householder structure more suitable for language and copying. However, we encourage the community to investigate more channel mixer structures under the FP-RNN framework.

3. Could you provide concrete cost measurements comparing FP-RNNs to existing models? As discussed before, we compare the runtime of FP-RNN to existing models in Appendix E.4, Figure 11 and will move this discussion into the main text.

We hope that our response addresses the reviewer’s concerns and would be interested to hear if they are considering raising their score to help increase the visibility of our work during the conference. In case there are any remaining questions, we are of course happy to resolve them during the discussion period.

Penedo et al. (2024), "The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale", https://arxiv.org/abs/2406.17557

Schlag et al. (2021), "Learning associative inference using fast weight memory.", https://arxiv.org/abs/2011.07831