Longhorn: Modelos de Espaço de Estados são Aprendizes Online 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

Resumo

A capacidade mais fundamental dos métodos de IA modernos, como os Modelos de Linguagem Grandes (LLMs), é a capacidade de prever o próximo token em uma longa sequência de tokens, conhecida como "modelagem de sequência". Embora o modelo Transformers seja a abordagem dominante atual para modelagem de sequência, seu custo computacional quadrático em relação ao comprimento da sequência é uma desvantagem significativa. Modelos de espaço de estados (SSMs) oferecem uma alternativa promissora devido à sua eficiência de decodificação linear e alta capacidade de paralelização durante o treinamento. No entanto, os SSMs existentes frequentemente dependem de designs de recorrência linear aparentemente ad hoc. Neste trabalho, exploramos o design de SSMs através da ótica da aprendizagem online, conceitualizando SSMs como meta-módulos para problemas específicos de aprendizagem online. Esta abordagem vincula o design de SSM à formulação de objetivos precisos de aprendizagem online, com regras de transição de estado derivadas da otimização desses objetivos. Com base nessa percepção, introduzimos uma nova arquitetura profunda de SSM baseada na atualização implícita para otimizar um objetivo de regressão online. Nossos resultados experimentais mostram que nossos modelos superam os SSMs de ponta, incluindo o modelo Mamba, em benchmarks padrão de modelagem de sequência e tarefas de modelagem de linguagem.

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.