Sequential-Parallel Duality in Prefix Scannable Models

要約

最新のニューラルシーケンスモデルは、並列化可能なトレーニングと高速順次推論の二重の任務を満たすように設計されています。
最近の開発は、そのような「シーケンシャル平行二重性」を達成するゲート線形注意（GLA）やマンバなどのさまざまなモデルを生み出しています。これは自然な疑問を提起します。
まず、このようなモデルの広範なクラス（状態空間モデル）を説明することから、カスタム連想集合体演算子を使用した古典的な並列プレフィックススキャンアルゴリズムを使用して状態更新を計算できるものとして説明することから始めます。
次に、SoftMaxの注意などの任意の（潜在的に非共同的）機能を可能にするために、状態集約演算子を緩和することにより、より一般的なクラスであるプレフィックススカンド可能モデル（PSM）を定義します。
この一般化は、要素ごとのRNN（MAMBAなど）や線形変圧器（例：GLA、MAMBA2、MLSTM）を含む多くの既存のアーキテクチャを統合し、O（1）償却型の断片化されたメモリを実現するO（1）償却型のメモリを達成するSoftMax様オペレーターを持つ新しいモデルを導入し、そのようなモデルを評価します。
州の追跡と連想リコールを含む合成タスク。
経験的には、PSMは、状態空間モデルの推論効率と一致しながら、変圧器ベースのアーキテクチャの表現性を保持していることがわかります。

要約(オリジナル)

Modern neural sequence models are designed to meet the dual mandate of parallelizable training and fast sequential inference. Recent developments have given rise to various models, such as Gated Linear Attention (GLA) and Mamba, that achieve such “sequential-parallel duality.” This raises a natural question: can we characterize the full class of neural sequence models that support near-constant-time parallel evaluation and linear-time, constant-space sequential inference? We begin by describing a broad class of such models — state space models — as those whose state updates can be computed using the classic parallel prefix scan algorithm with a custom associative aggregation operator. We then define a more general class, Prefix-Scannable Models (PSMs), by relaxing the state aggregation operator to allow arbitrary (potentially non-associative) functions such as softmax attention. This generalization unifies many existing architectures, including element-wise RNNs (e.g., Mamba) and linear transformers (e.g., GLA, Mamba2, mLSTM), while also introducing new models with softmax-like operators that achieve O(1) amortized compute per token and log(N) memory for sequence length N. We empirically evaluate such models on illustrative small-scale language modeling and canonical synthetic tasks, including state tracking and associative recall. Empirically, we find that PSMs retain the expressivity of transformer-based architectures while matching the inference efficiency of state space models — in some cases exhibiting better length generalization than either.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google