Longhorn: State Space Models are Amortized Online Learners

Authors: Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, qiang liu
Primary area: foundation or frontier models, including LLMs
Keywords: Deep State Space Models, Linear Attention Models, Online Learning, Language Modeling

Abstract:
The most fundamental capability of modern AI methods such as Large Language Models (LLMs) is the ability to predict the next token in a long sequence of tokens, known as “sequence modeling.”  Although the Transformers model is the current dominant approach to sequence modeling, its quadratic computational cost with respect to sequence length is a significant drawback. State-space models (SSMs) offer a promising alternative due to their linear decoding efficiency and high parallelizability during training. However, existing SSMs often rely on seemingly ad hoc linear recurrence designs.
In this work, we explore SSM design through the lens of online learning, conceptualizing SSMs as meta-modules for specific online learning problems. This approach links SSM design to formulating precise online learning objectives, with state transition rules derived from optimizing these objectives.
Based on this insight, we introduce a novel deep SSM architecture based on the implicit update for optimizing an online regression objective. Our experimental results show that our models outperform state-of-the-art SSMs, including the Mamba model, on standard sequence modeling benchmarks and language modeling tasks.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730518598624

thread_type:
official_review

<thread_object>
creator: reviewer_dvh4

<content>
summary:
This paper presents a novel perspective on SSM models through the lens of online learning, offering a fresh analytical framework. The learning process consists of two losses: one ensuring minimal state updates, and another optimizing the reconstruction of current x from the state using key k. Within this framework, the authors propose Longhorn, which introduces differential importance weighting for various dimensions during reconstruction. Experimental results across different scales, tasks, and sequence lengths demonstrate Longhorn's superior compression rate on PassKey tasks and better length generalization, while maintaining comparable performance with baseline SSMs on other tasks. While the theoretical framework shows significant value, the paper makes an approximation step without thoroughly discussing its impact on the overall theory, which intuitively could lead to substantial differences.

strengths:
- Novel theoretical perspective analyzing SSMs through online learning
- Proposed theoretical framework shows strong potential for length generalization
- Comprehensive empirical validation across various settings

weaknesses:
- Insufficient discussion of the approximation's impact, creating a gap between theory and practice
- Limited comparison with contemporary methods (e.g., Mamba2 and DeltaNet); given the concurrent timing with DeltaNet, would appreciate author response on this comparison

questions:
- 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?
- Could the authors provide analysis on why DeltaNet struggles with extrapolation, while Longhorn demonstrates superior extrapolation capabilities, especially given their similarities in the update equation?
- Is there an ablation study on the beta parameters in the OCP formula? What guidelines exist for its optimal selection?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696455754

thread_type:
official_review

<thread_object>
creator: reviewer_ftdh

<content>
summary:
The paper introduces Longhorn, a novel state-space model (SSM) architecture designed as a meta-module that effectively handles sequence modeling problems. It described a theoretical framework based on online learning principles to derive the closed-form solutions for the online associative recall problem. 

The empirical results convincingly demonstrate that Longhorn surpasses other state-of-the-art SSMs in performance, particularly highlighted by its impressive recall capabilities on the Multi-Query Associative Recall (MQAR) benchmark.

strengths:
- the online learning framework provides a fair theoretical underpinning for understanding the Linear attention model / SSMs. This approach not only supports the conceptual innovations presented but also enhances the interpretability of SSM behaviors in practical applications.

- emirpical results: Longhorn has good sample efficiency compared to STOA models such as Mamba and GLA. This advantage is critical in scenarios where computational resources are limited.

weaknesses:
Appoximation: 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.

questions:
1. can you provide more details on the sample efficiency experiments? Say, what kinds of hyper-parameters did you try? Can you do an abaltion study? 

2. Echoing the weakness of the paper, it is unclear to me that after using such an appoximation, is the theory framework still well aligned to the experiments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702534933

thread_type:
official_review

<thread_object>
creator: reviewer_lqbo

<content>
summary:
This paper introduces Longhorn, a new state-space model (SSM) architecture. By adopting an online learning optimization perspective, the authors unify several popular SSM architectures, bringing clarity to the structural differences between them. The explanation in Appendix A and Table 4 are especially helpful for understanding these nuances. Building on this unified approach, the authors propose a simplified SSM structure through a novel value retrieval mechanism based on key structures, offering insightful explanations of their method. The paper concludes by deriving a closed-form update formula for the state S in the SSM, supported by effective empirical results.

strengths:
1. The paper is well-written and easy to follow. The clarity of explanation makes complex ideas accessible, particularly in sections like Appendices A and B, which provide valuable insights into the nuances of different approaches.
2. The novelty of the new formulation for the Longhorn approach is impressive. The retrieval-based perspective is both innovative and elegantly presented, offering a fresh solution that enhances the field.
3. The exploration of SSM structure variances through online learning optimization is also intriguing and adds depth to the paper’s contribution.

weaknesses:
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.
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.
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.

questions:
see weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710200602

thread_type:
official_review

<thread_object>
creator: reviewer_dxpc

<content>
summary:
Longhorn formulates state-space models by solving an online regression problem. By designing various online learning objectives, it can induce different linear recurrent models, providing a unified framework and principled approach for developing new models. This work adopts an objective that encourages the hidden state to remain close to its previous state (i.e., \( ||S_{t+1} - S_t||_F \)) and includes a term that promotes key-input association (i.e., \( ||S k_t - x_t||_{\text{diag}(\beta_t)^2} \)). While the optimal solution is derived, it presents computational challenges. To address this, a diagonal approximation is used, resulting in a model computationally similar to Mamba (and thus similarly efficient) but with improved empirical performance, as demonstrated on synthetic datasets and medium-scale language modeling and image modeling.

strengths:
1. The paper is well-written and easy to follow. 
2. The formulation via online learning for solving in-context associative recall is interesting and elegant. It explains why Longhorn (and also DeltaNet) performs well in MQAR tasks.
3. Empirical results look good.

weaknesses:
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.
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?
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.
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.

questions:
Are there any actual wall-time comparisons in terms of training & inference?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper offers a novel online-learning perspective on SSM design, introducing a novel architecture called Longhorn and demonstrating improved performance over strong baselines like Mamba. All reviewers appreciated the clarity of exposition, the elegance of the theoretical framework, and the thorough empirical comparisons showing Longhorn’s effectiveness in tasks requiring long-context processing. While the reviewers initially expressed some concerns (e.g., diagonal approximation), they later found the rebuttal and additional experiments convincing.