State-space models can learn in-context by gradient descent

要約

深い状態空間モデル (Deep SSM) は、トランスフォーマーと同様に、自己回帰タスクに関するコンテキスト内学習の機能を示しています。
ただし、リカレント ネットワークでこれを可能にするアーキテクチャ上の要件とメカニズムは依然として不明瞭です。
この研究は、状態空間モデル アーキテクチャが勾配ベースの学習を実行し、それをコンテキスト内学習に使用できることを示しています。
局所的な自己注意を強化した単一の構造化状態空間モデル層が、勾配降下法の 1 ステップ後に最小二乗損失を伴う陰的線形モデルの出力を再現できることを証明します。
私たちの重要な洞察は、対角線形再帰層が勾配アキュムレーターとして機能し、陰的回帰モデルのパラメーターに「適用」できるということです。
ランダムに初期化された拡張 SSM を単純な線形回帰タスクでトレーニングすることによって、構築を検証します。
経験的に最適化されたパラメーターは、陰的モデル構築から分析的に得られた理論的なパラメーターと一致します。
複数ステップの線形回帰および非線形回帰を拡張すると、一貫した結果が得られます。
構築された SSM には、最新の深状態空間モデルの機能が組み込まれており、一般的なタスクでもスケーラブルなトレーニングと有効性を実現できる可能性があります。
この理論的構築により、基盤モデルに特有の表現力を可能にする重要な要素として、リカレント アーキテクチャにおける局所的な自己注意と乗算的相互作用の役割が解明されます。

要約(オリジナル)

Deep state-space models (Deep SSMs) have shown capabilities for in-context learning on autoregressive tasks, similar to transformers. However, the architectural requirements and mechanisms enabling this in recurrent networks remain unclear. This study demonstrates that state-space model architectures can perform gradient-based learning and use it for in-context learning. We prove that a single structured state-space model layer, augmented with local self-attention, can reproduce the outputs of an implicit linear model with least squares loss after one step of gradient descent. Our key insight is that the diagonal linear recurrent layer can act as a gradient accumulator, which can be `applied’ to the parameters of the implicit regression model. We validate our construction by training randomly initialized augmented SSMs on simple linear regression tasks. The empirically optimized parameters match the theoretical ones, obtained analytically from the implicit model construction. Extensions to multi-step linear and non-linear regression yield consistent results. The constructed SSM encompasses features of modern deep state-space models, with the potential for scalable training and effectiveness even in general tasks. The theoretical construction elucidates the role of local self-attention and multiplicative interactions in recurrent architectures as the key ingredients for enabling the expressive power typical of foundation models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google