Longhorn: Zustandsraummodelle sind amortisierte Online-Lerner.

Longhorn: State Space Models are Amortized Online Learners

July 19, 2024

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

cs.AI

Zusammenfassung

Die grundlegendste Fähigkeit moderner KI-Methoden wie Großer Sprachmodelle (LLMs) ist die Fähigkeit, das nächste Token in einer langen Sequenz von Tokens vorherzusagen, bekannt als "Sequenzmodellierung". Obwohl das Transformer-Modell derzeit der dominante Ansatz für die Sequenzmodellierung ist, ist sein quadratischer Rechenaufwand in Bezug auf die Sequenzlänge ein bedeutender Nachteil. Zustandsraummodelle (SSMs) bieten aufgrund ihrer linearen Dekodierungseffizienz und hohen Parallelisierbarkeit während des Trainings eine vielversprechende Alternative. Allerdings stützen sich bestehende SSMs oft auf scheinbar ad hoc lineare Rekurrenzdesigns. In dieser Arbeit erforschen wir das Design von SSMs durch die Linse des Online-Lernens und konzeptualisieren SSMs als Meta-Module für spezifische Online-Lernprobleme. Dieser Ansatz verknüpft das SSM-Design mit der Formulierung präziser Online-Lernziele, wobei die Zustandsübergangsregeln aus der Optimierung dieser Ziele abgeleitet werden. Basierend auf dieser Erkenntnis stellen wir eine neuartige tiefe SSM-Architektur vor, die auf dem impliziten Update zur Optimierung eines Online-Regression-Ziels basiert. Unsere experimentellen Ergebnisse zeigen, dass unsere Modelle auf Standard-Sequenzmodellierungs-Benchmarks und Sprachmodellierungsaufgaben die modernsten SSMs, einschließlich des Mamba-Modells, übertreffen.

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.