Longhorn : Les modèles d'état spatial sont des apprenants en ligne amortis

Longhorn: State Space Models are Amortized Online Learners

July 19, 2024

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

cs.AI

Résumé

La capacité la plus fondamentale des méthodes modernes d'IA, telles que les modèles de langage à grande échelle (LLMs), est la prédiction du token suivant dans une longue séquence de tokens, connue sous le nom de « modélisation de séquence ». Bien que le modèle Transformers soit actuellement l'approche dominante pour la modélisation de séquence, son coût computationnel quadratique par rapport à la longueur de la séquence constitue un inconvénient majeur. Les modèles à espace d'états (SSMs) offrent une alternative prometteuse grâce à leur efficacité de décodage linéaire et à leur grande parallélisabilité pendant l'entraînement. Cependant, les SSMs existants reposent souvent sur des conceptions de récurrence linéaire qui semblent ad hoc. Dans ce travail, nous explorons la conception des SSMs à travers le prisme de l'apprentissage en ligne, en conceptualisant les SSMs comme des méta-modules pour des problèmes spécifiques d'apprentissage en ligne. Cette approche relie la conception des SSMs à la formulation d'objectifs précis d'apprentissage en ligne, avec des règles de transition d'état dérivées de l'optimisation de ces objectifs. Sur la base de cette intuition, nous introduisons une nouvelle architecture de SSM profond basée sur la mise à jour implicite pour optimiser un objectif de régression en ligne. Nos résultats expérimentaux montrent que nos modèles surpassent les SSMs de pointe, y compris le modèle Mamba, sur des benchmarks standards de modélisation de séquence et des tâches de modélisation du langage.

English

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.