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）のような現代のAI手法の最も基本的な能力は、長いトークン列における次のトークンを予測する能力であり、これは「シーケンスモデリング」として知られています。Transformerモデルは現在、シーケンスモデリングにおける主流のアプローチですが、シーケンス長に対して二次的な計算コストがかかるという重大な欠点があります。状態空間モデル（SSM）は、線形デコード効率とトレーニング中の高い並列化可能性から、有望な代替手段を提供します。しかし、既存のSSMはしばしば一見アドホックな線形再帰設計に依存しています。本研究では、オンライン学習の視点からSSM設計を探求し、SSMを特定のオンライン学習問題のためのメタモジュールとして概念化します。このアプローチにより、SSM設計は正確なオンライン学習目標の定式化と結びつき、状態遷移規則はこれらの目標を最適化することから導出されます。この洞察に基づいて、オンライン回帰目標を最適化するための暗黙的更新に基づく新しい深層SSMアーキテクチャを導入します。実験結果は、我々のモデルが標準的なシーケンスモデリングベンチマークと言語モデリングタスクにおいて、Mambaモデルを含む最先端のSSMを上回ることを示しています。

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.