Longhorn: Los Modelos de Espacio de Estados son Aprendices en Línea Amortizados

Longhorn: State Space Models are Amortized Online Learners

July 19, 2024

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

cs.AI

Resumen

La capacidad más fundamental de los métodos de IA modernos como los Modelos de Lenguaje Extensos (MLE) es la habilidad de predecir el siguiente token en una larga secuencia de tokens, conocida como "modelado de secuencias". Aunque el modelo Transformers es el enfoque dominante actual para el modelado de secuencias, su costo computacional cuadrático con respecto a la longitud de la secuencia es una desventaja significativa. Los modelos de espacio de estados (MEE) ofrecen una alternativa prometedora debido a su eficiencia de decodificación lineal y alta capacidad de paralelización durante el entrenamiento. Sin embargo, los MEE existentes a menudo se basan en diseños de recurrencia lineal aparentemente ad hoc. En este trabajo, exploramos el diseño de MEE a través del prisma del aprendizaje en línea, conceptualizando los MEE como meta-módulos para problemas específicos de aprendizaje en línea. Este enfoque vincula el diseño de MEE a la formulación de objetivos precisos de aprendizaje en línea, con reglas de transición de estado derivadas de la optimización de estos objetivos. Basándonos en esta idea, introducimos una novedosa arquitectura profunda de MEE basada en la actualización implícita para optimizar un objetivo de regresión en línea. Nuestros resultados experimentales muestran que nuestros modelos superan a los MEE de última generación, incluido el modelo Mamba, en pruebas estándar de modelado de secuencias y tareas de modelado de lenguaje.

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.