Longhorn: State Space Models are Amortized Online Learners
Authors: name: Bo Liu; name: Rui Wang; name: Lemeng Wu; name: Yihao Feng; name: Peter Stone; name: Qiang Liu
Published: 2024-07-19 11:12:08

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.