Longhorn: State Space Modellen zijn Geamortiseerde Online Leerders

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

Samenvatting

De meest fundamentele capaciteit van moderne AI-methoden, zoals Large Language Models (LLM's), is het vermogen om het volgende token in een lange reeks tokens te voorspellen, bekend als "sequence modeling." Hoewel het Transformers-model momenteel de dominante aanpak is voor sequence modeling, is de kwadratische rekencapaciteit ten opzichte van de sequentielengte een belangrijk nadeel. State-space modellen (SSM's) bieden een veelbelovend alternatief vanwege hun lineaire decodeerefficiëntie en hoge paralleliseerbaarheid tijdens het trainen. Bestaande SSM's vertrouwen echter vaak op ogenschijnlijk ad hoc ontworpen lineaire recurrenties. In dit werk verkennen we het ontwerp van SSM's vanuit het perspectief van online leren, waarbij we SSM's conceptualiseren als meta-modules voor specifieke online leerproblemen. Deze benadering koppelt het ontwerp van SSM's aan het formuleren van precieze online leerdoelen, waarbij de regels voor staatsovergangen worden afgeleid uit het optimaliseren van deze doelen. Op basis van dit inzicht introduceren we een nieuw diep SSM-architectuur gebaseerd op de impliciete update voor het optimaliseren van een online regressiedoel. Onze experimentele resultaten tonen aan dat onze modellen state-of-the-art SSM's, inclusief het Mamba-model, overtreffen op standaard sequence modeling benchmarks en taalmodelleringstaken.

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.