Longhorn: State Space Models are Amortized Online Learners

Author Response

We sincerely thank the reviewer for the valuable suggestions and thoughtful questions. Below, we address each of your concerns in detail. **1. Insufficient discussion of the approximation's impact, creating a gap between theory and practice.** We acknowledge that the diagonal approximation introduces some discrepancies. However, to the best of our knowledge, there is currently no efficient method to implement the exact closed-form solution. Notably, the diagonal approximation aligns with the implementation of the Mamba model, ensuring direct comparability and avoiding additional overhead. In principle, an efficient matrix pscan algorithm could enable the exact closed-form implementation, which we leave as a direction for future work. **2. Limited comparison with contemporary methods (e.g., Mamba2 and DeltaNet); given the concurrent timing with DeltaNet, would appreciate author's response on this comparison.** We acknowledge that Mamba2 and DeltaNet are concurrent. Due to the limited resources we have and the short period of time in rebuttal, we cannot train DeltaNet. But please consider the following two remarks: - Mamba2’s main change to Mamba is to reduce the richness of the recurrence, and in return make the state size larger. Given this, if the state sizes are the same, Mamba should be strictly better than Mamba2. In addition, take Figure 9 in Mamba2 as an example, one can see that even though Mamba’s state size is much smaller than Mamba2, its scaling performance is quite the same as Mamba2 (indicating the importance of the richness of the recurrence form). Moreover, recent works like Jamba-1.5 [1] observes that in hybrid architecture that interleaves self-attention and SSMs, Mamba’s performance is much better than Mamba2. So we can expect that the A (state transition matrix) needs to be rich enough to bring extra benefit to self-attention. - Regarding DeltaNet, we find that it shares a similar state size as Mamba2, hence is much larger than Mamba’s state size. In the meantime, as the 100B subset of Slimpajama dataset DeltaNet uses and we use are different, it is hard to compare against DeltaNet fairly. Note that our Mamba model’s performance is better than the results in DeltaNet’s paper, which was taken from the GLA paper. From the method perspective, we want to argue that the contributions of DeltaNet and ours are different. DeltaNet focuses on computation efficient training (parallelization) of the original DeltaNet work. Both the original one and the more efficient DeltaNet require certain normalization techniques to make sure the A matrix is stable. In addition, as seen in the last paragraph of Section 5.3 in the more recent DeltaNet paper [2], the model does not extrapolate well beyond the training context. In comparison, Longhorn successfully extrapolates beyond its training context, which closely aligns with the original motivation that it does online learning. And Longhorn’s update ensures that the A matrix is always stable (absolute value of the eigenvalues is smaller than 1). Moreover, Longhorn is not only an architecture, but is also a framework. Future works can explore different online learning objectives, instead of directly working on the recurrence form. **3. There appears to be a discrepancy between the significant improvements in PPL versus the modest gains in downstream task metrics. Could the authors elaborate on this phenomenon?** Thanks for raising this point. The downstream evaluation metrics essentially are not what the model optimizes for, hence lower loss doesn’t always mean better metrics. But Longhorn indeed performs better in terms of both. On the other hand, as the 1B model is still small, the downstream evaluation metrics are not comparable to larger models, we expect those metrics will be more meaningful when actual large models are trained. In that regime, often the validation loss aligns better with the downstream metrics, and we expect that Longhorn’s advantage will be more significant.

Author Response

We sincerely thank the reviewer for the valuable suggestions and thoughtful questions. Below, we address each of your concerns in detail. **1. Approximation: while the diagonal approximation is a key aspect of Longhorn's implementation, its impact on the theoretical framework's alignment with empirical results remains unclear to me. I would expect a deeper exploration into how this approximation influences model performance could bridge the gap between theoretical predictions and observed outcomes.** We appreciate the reviewer’s observation. However, to the best of our knowledge, the exact closed-form SSM update cannot be efficiently implemented, preventing us from directly evaluating its modeling performance. The current approximation is the most efficient implementation we could identify that closely aligns with the original formulation. It also retains the same parallel structure as the Mamba architecture, ensuring a fair comparison. Despite this approximation, we observed that Longhorn improves sample efficiency in terms of perplexity compared to Mamba, suggesting that its inductive bias benefits language modeling. Implementing the exact closed form would require matrix pscan, which we have explored; however, it is currently numerically unstable and significantly slower. **2. Can you provide more details on the sample efficiency experiments? Say, what kinds of hyper-parameters did you try? Can you do an ablation study?** Given the size of the model and training, we did **not** have the resources to do a hyperparameter sweep or sensitivity analysis. We use the same hyperparameters as Mamba. **3. Echoing the weakness of the paper, it is unclear to me that after using such an approximation, is the theory framework still well aligned to the experiments?** It is true that there is this discrepancy due to the approximation, however, this is the closest approximation we can find that has an efficient parallel form. As diagonal approximations of matrices have been widely used in optimization and other fields of deep learning for efficiency, we think this diagonal approximation is reasonable. The experiments also suggest improved sample efficiency even using the diagonal approximation, hence we find that it indicates that Longhorn benefits from the inductive bias. One way of thinking of it is that Longhorn is a practical method that is inspired by a theoretical formulation. This discrepancy between theory and practice is not uncommon, and indeed is often inevitable as in this case.

Author Response

We sincerely thank the reviewer for the valuable suggestions and thoughtful questions. Below, we address each of your concerns in detail. **1. While this paper presents a focused study on architecture, the data and model scale seem limited. Expanding the experimental scale and providing a more comprehensive analysis would significantly enhance the paper's impact.** We agree with the reviewer on this. However, due to the limited computation resources we have in our lab, the current experiments have already taken us several months to run. We are happy to experiment on larger models/dataset once we have more compute in the future. **2. The reduction in perplexity compared to Mamba is notable. However, the results in Table 2 appear mixed, which could benefit from further clarification or exploration.** Yes, the reduction in perplexity indicates that Longhorn benefits from its inductive bias (in terms of language modeling). We think the downstream evaluations might not directly reflect this as they are noisy evaluation metrics of a given model, and since we are only comparing 1B size models, they might not differ too much from each other. **3. Including additional experiments, such as MMLU, GSM-8K, and more extensive long-context benchmarks, would strengthen the findings and provide a more robust evaluation of the model's capabilities.** We appreciate the reviewer’s suggestion to include evaluations on benchmarks like MMLU and GSM-8K. After investigating, we found that all 1B models in our experiments yield near-random performance on these benchmarks. For reference, even LLaMA 7B, trained on 1–2T tokens, achieves only ~25% accuracy on MMLU in the 0-shot setting (and 35% in 5-shot) (see Table 9 of [1]), which aligns with random guessing. Given that our models are 1B models trained on 100B tokens, their similar near-random performance is expected. Similarly, GSM-8K results are limited, as the SlimPajama dataset lacks sufficient high-quality mathematical reasoning data. However, we emphasize the compact design of Mamba and Longhorn, which have a state size of only 16—significantly smaller than that of GLA or self-attention. Despite this, Longhorn achieves superior performance on 4K context length tasks and outperforms Mamba on both synthetic and large-scale language modeling benchmarks. We believe Longhorn's direct approach to solving the online associative recall problem will demonstrate even greater advantages as context lengths increase (just as what is shown in the MQAR example). [1] LLaMA: Open and Efficient Foundation Language Models. https://arxiv.org/pdf/2302.13971

Author Response

We sincerely thank the reviewer for the valuable suggestions and thoughtful questions. Below, we address each of your concerns in detail. **1. The main issue with this work is that the implementation does not fully align with the theory. Using a diagonal matrix to approximate an “identity-plus-low-rank” dense matrix is coarse, and it’s unclear if the theoretical advantage translates to this setting.** We acknowledge that the diagonal matrix is an approximation. However, to the best of our knowledge, there is currently no efficient method to implement the exact closed-form solution. Notably, the diagonal approximation aligns with the implementation of the Mamba model, ensuring direct comparability and avoiding additional overhead. In principle, an efficient matrix pscan algorithm could enable the exact closed-form implementation, which we leave as a direction for future work. **2. In Eq5, the norm (\text{diag}(\beta_t)) appears unusual and is not well-motivated or empirically validated. Why is a vector-valued (\epsilon) necessary? If not, the DeltaNet structure could leverage the kernel from Yang et al. (2024, https://arxiv.org/abs/2406.06484), which would easily scale up the head dimension and likely benefit recall-intensive tasks requiring a large state size. Longhorn, as it stands, cannot be expressed in matmul form, leading to similar challenges as in Mamba. Would Mamba2-like optimization, potentially resulting in a DeltaNet-like model with scalar (\epsilon), be preferable?** Having $\beta$ as a vector is reasonable because, like Mamba, we treat each dimension of $x$ as a separate state space model. This approach enhances the sequence model's representational power. Since each dimension of $x$ can have its own $\beta$, it is natural to represent $\beta$ as a vector, leading to the $(\text{diag}(\beta_t))$ norm. However, we agree with the reviewer that this introduces the limitation that, similar to Mamba, Longhorn cannot be expressed in matmul form. If $\beta$ were instead shared across all dimensions of $x$, Longhorn could be implemented in the same way as DeltaNet. **3. MQAR is a synthetic dataset and insufficient to demonstrate Longhorn’s advantages in recall-intensive tasks. Results on real-world recall-intensive tasks proposed in Arora 2024 [https://arxiv.org/abs/2402.18668] would provide a stronger case. Could you report zero-shot accuracy on these tasks? A table similar to [Table 1, https://arxiv.org/abs/2407.05483] would be very useful and necessary.** We thank the reviewer for the suggestion. However, as Table 1 in the paper mentioned, the authors trained 1B Mamba over 300B tokens but we train Longhorn only using 100B tokens, it would not make a fair comparison. Given the short period of time during rebuttal, we could not retrain another Longhorn model over 300B tokens (it might take several weeks). Additionally, we would like to point out that to compare different state space models’ recall ability really fairly, one would make the state size the same. From this perspective, Mamba and Longhorn both have the smallest state sizes across all SSMs in most recall comparison experiments (even including MQAR). We hypothesize that they would have even better recall rates if Mamba and Longhorn were scaled to larger state sizes. We will add this evaluation as an interesting direction for future work. **4. This work lacks several ablation studies. For instance, the "value" projection is removed compared to standard models, yet this change is not analyzed. Additionally, the model does not clarify the benefits of parameter tying.** We did not omit the value projection, as we aimed for a fair comparison with Mamba, which uses the same parameterization except for the updating functions. To ensure consistency, we followed the Mamba architecture exactly: in Mamba, $v$ is first projected, and $q$ and $k$ are linearly projected from $v$. We adopted the same approach, modifying only the SSM component, as shown in Figure 3. Additionally, we clarified in Line 9 of Algorithm 1 (Lines 123–125) that $x_t$ is preprocessed through a linear projection followed by a 1D convolution, where the resulting $x_t$ serves as the value. **5. Wall-clock time comparison.** Unfortunately, we did not keep track of the wall-clock training time. But we did some offline kernel speed tests and Mamba and Longhorn achieved almost the same speed (0.122s for Mamba and 0.124s for Longhorn, which differs within 2% in performance). And the Mamba and Longhorn architecture is exactly the same except for the SSM kernel, hence the overall training/inference speeds are also nearly the same.

Response to Reviewer's Followup Question

We thank the reviewer for their follow-up on the concerns and feel sorry for the late response as we were preparing the experiment results. Here are our latest results and response and hope those could address your concern. **1. Is vector-valued $\beta$ useful in practice?** To verify that the vector value $\beta$ is useful, we train a 120M size Longhorn but keep the $\beta$ term a learnable scalar. So everything is the same, including the update form. The only change is that we broadcast the scalar $\beta$ to all dimensions. The experiment setting follows exactly as in Section 5.2 of the paper, where we use a 1024 context length. The results are the following: | Model | Validation Loss | |------------------------------|-----------------| | Longhorn (with scalar $\beta$) | 3.262 | | Longhorn (with vector $\beta$) | 3.225 | | Mamba | 3.238 | | LLaMA | 3.247 | | GLA | 3.381 | From this, we conclude that the performance of Longhorn with scalar $\beta$ is much worse than that of Longhorn with vector $\beta$. **2. Real-world recall intensive benchmark.** We thank the reviewer for the clarification of the concern. We have evaluated the 1B models we trained on the suggested benchmark. Here are the results we got for Longhorn-1B and Mamba-1B. | Model | FDA | SWDE | NQ | SQuAD | TriviaQA | Drop | AVG | |-----------|------------|-----------|-----------|-----------|-----------|-----------|------------| | Mamba | 33.2/40.6 | 35.0/36.2 | 26.6/32.6 | 37.4/52.7 | 56.3/56.9 | 20.4/31.5 | 34.82/41.75 | | Longhorn | 40.2/50.4 | 33.2/42.3 | 27.6/33.2 | 35.0/55.0 | 58.5/55.9 | 21.3/33.3 | 35.97/45.0 | From the results, we observe that Longhorn achieved a significant 1.15 (4.25 if read twice) improvement over Mamba. Note that there is a gap between the Mamba 1B in the suggested paper and ours here. The reason might be due to: 1. we trained for 100B instead of 300B tokens; 2. we used SlimPajama instead of Pile. We will include this additional result in our paper in the final version. Thanks for suggesting this benchmark. **3. Longhorn saves parameters.** For Mamba’s SSM module, it requires two linear projections to the B and C matrices (in our case, the K and Q matrices). It also requires one (bottlenecked) linear projection to compute the $\Delta \in \mathbb{R}^d$ term (the step size), this is equivalent to our $\beta \in \mathbb{R}^d$. However, Mamba additional requires an input-independent A matrix $A \in \mathbb{R}^{m \times d}$, which Longhorn does not require. Therefore, we say that Longhorn has slightly fewer parameters than Mamba, and the gap is proportional to $\mathcal{O}(md)$. To further convince the reviewer that the inductive bias from Longhorn is useful, in one of our old experiments, we compared Longhorn against a Mamba’s variant that is more similar to Longhorn, where we made Mamba’s $A$ matrix, instead of a constant learnable matrix, also an input dependent vector. So $A = W_A x_t \in \mathbb{R}^m$ and $W_a \in \mathbb{R}^{m \times d}$. We call this Mamba (with learnable $A$). The final update rule of Mamba (with learnable A) is: $$S_t = \exp( -\Delta \otimes A_t ) \odot S_{t-1} + (\Delta \odot x_t) \otimes B_t$$ In comparison, Longhorn’s update is $$S_t = (1_{d \times m} - \epsilon_t \otimes k_t^{\odot 2}) \odot S_{t-1} + (\epsilon_t \odot x_t) \otimes k_t $$ Here, $S_t \in \mathbb{R}^{d \times m}$. $k_t$ matches the size of $B_t$ in Mamba, $\epsilon_t$ matches the size of $\Delta_t$ in Mamba. Note that the main difference is that in Longhorn, the forget gate $k_t^{\odot 2}$ is linked to the input gate $k_t$, while in Mamba (with learnable $A$) the two are separate. The old experiment setting used a 350M model and 1024 context length (also the same experiment setting as in Section 5.2), the results are the following: | Model | Validation Loss | |------------------------------|-----------------| | Longhorn (with vector $\beta$) | 2.888 | | Mamba | 2.902 | | LLaMA | 2.891 | | GLA | 3.018 | | Mamba (with learnable $A$) | 2.922 | From this, we see that with this learnable forgetting gate $A = W_A x_t$, the performance of Mamba (with learnable $A$) is even worse than the original Mamba. We do not know the exact reason behind this, but given the similarity between the recurrent update rule of Mamba (with learnable $A$) and that of Longhorn, we believe it indicates the importance of the inductive bias in Longhorn. Namely, the forgetting is roughly about the square of the input might be beneficial (i.e., when $k$’s magnitude < 1, which is usually the case, it means Longhorn forgets less but inputs more).