Longhorn: Модели пространства состояний являются амортизированными онлайн-обучающимися.

Longhorn: State Space Models are Amortized Online Learners

July 19, 2024

Авторы: Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, Qiang Liu

cs.AI

Аннотация

Самая фундаментальная способность современных методов искусственного интеллекта, таких как модели больших языковых моделей (LLM), заключается в возможности предсказать следующий токен в длинной последовательности токенов, известной как "моделирование последовательности". Хотя модель Transformers является текущим доминирующим подходом к моделированию последовательностей, ее квадратичная вычислительная сложность относительно длины последовательности является значительным недостатком. Модели пространства состояний (SSM) предлагают многообещающую альтернативу благодаря своей линейной эффективности декодирования и высокой параллелизуемости во время обучения. Однако существующие SSM часто опираются на кажущиеся произвольными линейные рекуррентные конструкции. В данной работе мы исследуем проектирование SSM через призму онлайн-обучения, концептуализируя SSM как мета-модули для конкретных задач онлайн-обучения. Этот подход связывает проектирование SSM с формулировкой точных целей онлайн-обучения, с правилами перехода состояний, выведенными из оптимизации этих целей. Основываясь на этом понимании, мы представляем новую глубокую архитектуру SSM на основе неявного обновления для оптимизации цели онлайн-регрессии. Наши экспериментальные результаты показывают, что наши модели превосходят современные SSM, включая модель Mamba, на стандартных бенчмарках моделирования последовательностей и задачах языкового моделирования.

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.