Longhorn: I Modelli di Spazio degli Stati sono Apprenditori Online Ammortizzati

Longhorn: State Space Models are Amortized Online Learners

July 19, 2024

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

cs.AI

Abstract

La capacità più fondamentale dei metodi moderni di intelligenza artificiale, come i Large Language Models (LLM), è quella di prevedere il token successivo in una lunga sequenza di token, nota come "modellazione di sequenza". Sebbene il modello Transformers rappresenti l'approccio attualmente dominante per la modellazione di sequenza, il suo costo computazionale quadratico rispetto alla lunghezza della sequenza costituisce uno svantaggio significativo. I modelli a spazio di stato (SSM) offrono un'alternativa promettente grazie alla loro efficienza lineare nel decoding e all'elevata parallelizzabilità durante l'addestramento. Tuttavia, gli SSM esistenti spesso si basano su progetti di ricorrenza lineare apparentemente ad hoc. In questo lavoro, esploriamo la progettazione degli SSM attraverso la lente dell'apprendimento online, concettualizzando gli SSM come meta-moduli per specifici problemi di apprendimento online. Questo approccio collega la progettazione degli SSM alla formulazione di obiettivi precisi di apprendimento online, con regole di transizione di stato derivate dall'ottimizzazione di tali obiettivi. Sulla base di questa intuizione, introduciamo una nuova architettura di SSM profonda basata sull'aggiornamento implicito per ottimizzare un obiettivo di regressione online. I nostri risultati sperimentali dimostrano che i nostri modelli superano gli SSM all'avanguardia, incluso il modello Mamba, nei benchmark standard di modellazione di sequenza e nei task di modellazione del linguaggio.

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.