[論文レビュー] Simplifying and Understanding State Space Models with Diagonal Linear RNNs

Ankit Gupta, Harsh Mehta+1 more|arXiv (Cornell University)|Dec 1, 2022

Topic Modeling被引用数 4

ひとこと要約

本稿では、離散化を不要とする、状態空間モデルの簡素化・連続時間形式である対角線形RNN（DLR）を提案する。これにより、より洗練され、解釈性の高いアーキテクチャが実現される。単純であるにもかかわらず、DLRは長距離シーケンスタスク、特にLong Range Arenaおよび生音声分類において、先行するSSMおよびアテンションモデルと同等またはそれ以上の性能を発揮し、O(L)の推論複雑度を維持する。

ABSTRACT

Sequence models based on linear state spaces (SSMs) have recently emerged as a promising choice of architecture for modeling long range dependencies across various modalities. However, they invariably rely on discretization of a continuous state space, which complicates their presentation and understanding. In this work, we dispose of the discretization step, and propose a model based on vanilla Diagonal Linear RNNs ($\mathrm{DLR}$). We empirically show that, despite being conceptually much simpler, $\mathrm{DLR}$ is as performant as previously-proposed SSMs on a variety of tasks and benchmarks including Long Range Arena and raw speech classification. Moreover, we characterize the expressivity of SSMs (including $\mathrm{DLR}$) and attention-based models via a suite of $13$ synthetic sequence-to-sequence tasks involving interactions over tens of thousands of tokens, ranging from simple operations, such as shifting an input sequence, to detecting co-dependent visual features over long spatial ranges in flattened images. We find that while SSMs report near-perfect performance on tasks that can be modeled via $ extit{few}$ convolutional kernels, they struggle on tasks requiring $ extit{many}$ such kernels and especially when the desired sequence manipulation is $ extit{context-dependent}$. Despite these limitations, $\mathrm{DLR}$ reaches high performance on two higher-order reasoning tasks $\mathrm{ListOpsSubTrees}$ and $\mathrm{PathfinderSegmentation} ext{-}\mathrm{256}$ with input lengths $8K$ and $65K$ respectively, and gives encouraging performance on $\mathrm{PathfinderSegmentation} ext{-}\mathrm{512}$ with input length $262K$ for which attention is not a viable choice.

研究の動機と目的

• 離散化ステップを排除することで、状態空間モデルの理解および実装の複雑さを軽減すること。
• 長距離シーケンスモデリングの出発点として、バニラ対角線形RNN（DLR）に基づく直接的かつ解釈可能な定式化を提案すること。
• 合成および実世界の長シーケンスタスクにおけるDLRの表現力と性能を実験的に評価し、SSMおよびアテンションベースのモデルと比較すること。
• SSMが文脈依存的および高カーネルタスクを処理する際の限界を調査し、アテンション機構との相対的な表現力を特徴づけること。
• 13の合成的シーケンス・ツー・シーケンスタスクを用いた、長距離モデリング能力を評価する洗練された、アクセスしやすいベンチマークフレームワークを提供すること。

提案手法

• 連続的状態空間の離散化を回避する、対角線形RNNの直接的パラメータ化としてDLRを提案する。
• ユニット円周上に固有値を初期化した複素数行列を用い、数百万ステップにわたる長距離信号伝播を可能にする。
• 長さNの任意の畳み込みカーネルを表現できるように、DLRサイズNの周期的初期化スキームを導入する。
• 単純なシーケンスシフトから、フラット化された画像における複雑な空間的推論までをカバーする13の合成タスクのスイートでDLRを評価する。
• Long Range Arena、音声分類、因果的言語モデリングにおいて、S4、S4D、Transformer、LocalAttentionとDLRをベンチマーク比較する。
• 入力長を最大262Kトークンまで拡大し、極めて長いシーケンスにおけるスケーラビリティと性能をテストする。

実験結果

リサーチクエスチョン

• RQ1 連続時間離散化の複雑さを回避する直接的対角線形RNN（DLR）定式化は、離散化されたSSMと同等またはそれ以上の性能を達成できるか？
• RQ2 長距離依存性および複雑な相互作用を要するタスクにおいて、DLRの表現力はSSMおよびアテンションベースのモデルと比べてどうか？
• RQ3 現在のSSMは、文脈依存的シーケンス操作や多数の畳み込みカーネルを必要とするタスクを処理する際に、どのような限界を示すか？
• RQ4 DLRは、生音声分類や因果的言語モデリングといった実世界の長シーケンスタスクに、競争力のある性能と効率性を発揮できるか？
• RQ5 DLRのパラメータ化における周期性は、任意の畳み込みカーネル表現を可能にする原理的初期化を可能にするか？

主な発見

• DLRは、Path-XやListOpsを含むLong Range Arenaの全タスクで、S4およびS4Dと同等またはそれ以上の性能を発揮し、より単純で解釈性の高い定式化を実現した。
• 262Kの入力長を持つPathfinderSegmentation-512タスクにおいて、アテンションはこのスケールで計算的に非現実的であるにもかかわらず、DLRは強力な性能を発揮した。
• ListOps-SubTreesおよびPathfinderSegmentation-256において、DLRはLocalAttentionを上回り、特に長シーケンス長で顕著な性能向上を示し、高いスループットを維持した。
• 10クラスの音声コマンドデータセットでは、DLRはS4およびS4Dと同等またはそれ以上の性能を発揮し、O(L)のデコード複雑度で97.1%の正確度を達成した。
• PG-19における因果的言語モデリングでは、1.2Bパラメータと1.0×スループットで、Transformerと同等のテスト交差エントロピー損失（2.44–2.88）を達成した。
• Transformerがこのようなアーキテクチャ的選択に極めて感受的であるのに対し、DLRはレイヤーノーマライゼーションの配置に対してより頑健であった。