Longhorn: State Space Models are Amortized Online Learners

TL;DR

This paper introduces Longhorn, a novel state-space model architecture inspired by online learning principles, achieving superior efficiency and longer context extrapolation in sequence modeling tasks compared to existing models.

Contribution

The paper presents a new SSM architecture, Longhorn, derived from online learning insights, improving sequence modeling efficiency and context length handling.

Findings

Longhorn outperforms state-of-the-art SSMs like Mamba in benchmarks.

Achieves 1.8x better sample efficiency than Mamba.

Can extrapolate contexts up to 16x longer during inference.

Abstract

Modern large language models are built on sequence modeling via next-token prediction. While the Transformer remains the dominant architecture for sequence modeling, its quadratic decoding complexity in sequence length poses a major limitation. State-space models (SSMs) present a competitive alternative, offering linear decoding efficiency while maintaining parallelism during training. However, most existing SSMs rely on linear recurrence designs that appear somewhat ad hoc. 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 solving these objectives. Based on this insight, we introduce a novel deep SSM architecture, Longhorn, whose update resembles the closed-form solution for solving the online associative recall problem. Our experimental results show that Longhorn outperforms state-of-the-art SSMs, including the Mamba model, on standard sequence modeling benchmarks, language modeling, and vision tasks. Specifically, Longhorn achieves a 1.8x improvement in sample efficiency compared to Mamba, and can extrapolate over contexts that are up to 16x longer during inference.