[論文レビュー] Efficient Long Sequence Modeling via State Space Augmented Transformer

Simiao Zuo, Xiaodong Liu+5 more|arXiv (Cornell University)|Dec 15, 2022

Topic Modeling被引用数 7

ひとこと要約

本稿では、長期間の依存関係を捉えるために最初の層に状態空間モデル（SSM）を統合し、以降の層では効率的な局所的アテンションを用いることで、長期間のシーケンスモデリングを向上させるハイブリッドなTransformerアーキテクチャ、SPADEを提案する。この手法は線形の計算量を達成し、Long Range Arenaおよび言語モデリングのベンチマークでベースラインを上回り、NLUおよび生成タスクのための大規模事前学習モデルへのスケーリングも効果的である。

ABSTRACT

Transformer models have achieved superior performance in various natural language processing tasks. However, the quadratic computational cost of the attention mechanism limits its practicality for long sequences. There are existing attention variants that improve the computational efficiency, but they have limited ability to effectively compute global information. In parallel to Transformer models, state space models (SSMs) are tailored for long sequences, but they are not flexible enough to capture complicated local information. We propose SPADE, short for $\underline{ extbf{S}}$tate s$\underline{ extbf{P}}$ace $\underline{ extbf{A}}$ugmente$\underline{ extbf{D}}$ Transform$\underline{ extbf{E}}$r. Specifically, we augment a SSM into the bottom layer of SPADE, and we employ efficient local attention methods for the other layers. The SSM augments global information, which complements the lack of long-range dependency issue in local attention methods. Experimental results on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method. To further demonstrate the scalability of SPADE, we pre-train large encoder-decoder models and present fine-tuning results on natural language understanding and natural language generation tasks.

研究の動機と目的

• 標準のTransformerが長期間のシーケンスモデリングにおいて二次的な計算コストを抱える問題に対処すること。
• 局所的アテンション機構が長距離依存関係を捉える能力に制限を受ける問題を克服すること。
• 状態空間モデル（SSM）のグローバルモデリングの強みと、効率的アテンションの局所的インダクティブバイアスを統合すること。
• 長期間のNLPタスクに適したスケーラブルで効率的かつ効果的なアーキテクチャを設計すること。
• 下流のNLUおよび生成タスクにおける事前学習と微調整を通じて、この手法の有効性を実証すること。

提案手法

• SPADEは、最初の層が状態空間モデル（SSM）である多層Transformerを用いる。SSM層により、シーケンスに粗いグローバル情報を注入する。
• SSM層の後には、複数の効率的局所アテンション層（例：窓ベースのアテンション）が続き、局所的依存関係を精緻化する。
• SSMは線形時間および空間計算量を有し、長期間のシーケンスの効率的処理を可能にする。
• この手法は、SSMのグローバルなインダクティブバイアスとアテンションの局所的モデリングの柔軟性を組み合わせ、両者の欠点を補完する。
• アーキテクチャは、自己回帰的および非自己回帰的シーケンスモデリングの両方をサポートしており、エンコーダ・デコーダ構成も含む。
• モデルは標準的な目的関数で訓練され、GLUEおよび要約生成ベンチマークで微調整される。

実験結果

リサーチクエスチョン

• RQ1 SSMと局所的アテンションを組み合わせたハイブリッドアーキテクチャは、Transformerにおける長距離依存関係モデリングを改善できるか？
• RQ2 最初の層にSSMを統合することで、Long Range Arenaのような長期間シーケンスベンチマークでの性能が向上するか？
• RQ3 SPADEは、二次的アテンションおよびスパースアテンションのベースラインを上回りながら、線形計算量を維持できるか？
• RQ4 大規模言語モデリングで事前学習し、下流タスクで微調整した場合、SPADEはどの程度スケーリングできるか？
• RQ5 SSM層と局所的アテンション層のどちらが全体のパフォーマンス向上に寄与しているか？

主な発見

• SPADEは、Long Range Arenaベンチマークで、完全アテンション、スパースアテンション、低ランクアテンションを用いたモデルを上回る最先端のパフォーマンスを達成した。
• 自己回帰的言語モデリングにおいて、SPADEはネイティブなTransformerよりも著しく高速であり、より良いパープレキシティスコアを達成した。
• モデルは強力なスケーラビリティを示し、事前学習済みのエンコーダ・デコーダバージョンがGLUEおよび要約タスクで優れた結果を得た。
• アブレーションスタディにより、SSM層が長距離依存関係を捉える上で不可欠であることが確認され、局所的アテンション層が局所的パターンを精緻化していることが示された。
• SPADEは、8k長のシーケンスに対してもGPUメモリ消費量を80GB未満に抑えることができ、長期間のシーケンスモデリングを実用的なものにした。
• 構造的バイアスが限定的であっても、この手法は位置エンコーディングやスパースパターンに依存するモデルを上回る高いパフォーマンスを維持した。