Longhorn: State Space Models are Amortized Online Learners
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.
66a2706c66ead78120db36e3
Bo Liu
claimed_verified
2024-07-26T07:40:27.764Z
64375a5ab8ea86e50329d2da
/avatars/b3dcef9e2e0234156a22d699a30b6f1f.svg
Bo Liu
user
cranial-xix
66a2706c66ead78120db36e4
Rui Wang
66a2706c66ead78120db36e5
Lemeng Wu
claimed_verified
2024-07-26T07:40:23.961Z
6364acf56fa33d2f59397c59
/avatars/f642a7dcc098f50fef865f30f93a25c3.svg
Lemeng Wu
user
klight
66a2706c66ead78120db36e6
Yihao Feng
66a2706c66ead78120db36e7
Peter Stone
66a2706c66ead78120db36e8
Qiang Liu
Longhorn: State Space Models are Amortized Online Learners
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.
A novel deep state-space model based on implicit update rules outperforms existing models in sequence and language modeling tasks.
Large Language Models (LLMs) sequence modeling Transformers quadratic computational cost state-space models (SSMs) linear decoding efficiency parallelizability online learning meta-modules state transition rules online regression objective Mamba model language modeling