Longhorn: State Space Models are Amortized Online Learners

要約

大規模言語モデル (LLM) などの最新の AI 手法の最も基本的な機能は、「シーケンス モデリング」として知られる、長いトークン シーケンス内の次のトークンを予測する機能です。
Transformers モデルはシーケンス モデリングに対する現在主流のアプローチですが、シーケンスの長さに関する二次計算コストが大きな欠点です。
状態空間モデル (SSM) は、線形デコード効率とトレーニング中の高い並列性により、有望な代替手段となります。
ただし、既存の SSM は、一見アドホックな線形反復設計に依存していることがよくあります。
この研究では、オンライン学習のレンズを通して SSM の設計を探求し、特定のオンライン学習の問題に対するメタモジュールとして SSM を概念化します。
このアプローチは、SSM 設計を正確なオンライン学習目標の策定に結び付け、これらの目標の最適化から導出された状態遷移ルール​​を使用します。
この洞察に基づいて、オンライン回帰目標を最適化するための暗黙的な更新に基づく新しいディープ SSM アーキテクチャを導入します。
私たちの実験結果は、標準的なシーケンス モデリング ベンチマークと言語モデリング タスクにおいて、私たちのモデルが Mamba モデルを含む最先端の SSM よりも優れていることを示しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google