Longhorn: State Space Models are Amortized Online Learners

Bo Liu; Rui Wang; Lemeng Wu; Yihao Feng; Peter Stone; Qiang Liu

返回文献探索

Paper · arXiv 2407.14207

Longhorn: State Space Models are Amortized Online Learners

Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, Qiang Liu

18 upvotesJuly 19, 2024arXiv 预印本

AI 摘要

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 modelingTransformersquadratic computational coststate-space models (SSMs)linear decoding efficiencyparallelizabilityonline learningmeta-modulesstate transition rulesonline regression objectiveMamba modellanguage 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.