[論文レビュー] Long Range Language Modeling via Gated State Spaces

Harsh Mehta, Ankit Gupta+2 more|arXiv (Cornell University)|Jun 27, 2022

Topic Modeling被引用数 44

ひとこと要約

この論文は自己回帰言語モデリングのためのGated State Space (GSS) 層を導入し、DSSより2–3×速い学習時間を実現しつつ、Transformerベースラインと競争力を維持し、65kトークンまでの長さ一般化が強いことを示す。GSS-Transformer ハイブリッドはさらに結果を向上させる。

ABSTRACT

State space models have shown to be effective at modeling long range dependencies, specially on sequence classification tasks. In this work we focus on autoregressive sequence modeling over English books, Github source code and ArXiv mathematics articles. Based on recent developments around the effectiveness of gated activation functions, we propose a new layer named Gated State Space (GSS) and show that it trains significantly faster than the diagonal version of S4 (i.e. DSS) on TPUs, is fairly competitive with several well-tuned Transformer-based baselines and exhibits zero-shot generalization to longer inputs while being straightforward to implement. Finally, we show that leveraging self-attention to model local dependencies improves the performance of GSS even further.

研究の動機と目的

• 言語データにおける固定長ウィンドウを超えた長距離依存のモデリングを動機づける。
• ゲーティング状態空間に基づく高速でスケーラブルな自己回帰層を開発し、計算のボトルネックを削減する。
• GSSがDSSより学習が速い一方で、長距離LMベンチマークで競争力のあるパープレキシティを維持することを示す。
• 局所自己注意またはTransformerブロックの少量の組み込みが性能をさらに向上させることを示す。
• トレーニング時に見られた長さを超えるシーケンスで評価することにより、長さ一般化を評価する。）

提案手法

• ゲーティングユニットをDSS似の状態空間モジュールと組み合わせるGated State Space (GSS)層を導入する。
• GELU活性化とゲーティング機構を用いて、FFT中心の演算の実効次元を削減する。
• Deltaを1に固定し、Lambdaの初期化をランダム化して、厳密なHiPPOベースの初期化を避けることでDSSを簡略化する。
• XW1, XW2に対してGELUを適用したU, Vを介してDSSコンテキストをゲーティッドパスで計算し、Y = DSS(U) および O = (YW3 * V)W4 に残差接続を追加する。
• オプションとして、GSSスタックにTransformerブロックを組み合わせたGSS-Transformer-Hybridを導入し、Transformer層の入力を長いシーケンスで区切って処理する（長さ512の非重複チャンク）。
• tokenレベルのパープレキシティを用いて、長距離LMベンチマーク（LM1B、PG-19、ArXiv LaTeX、Githubコード）で訓練・評価し、DSSおよび適切に調整されたBlock Recurrent Transformersと比較する。

実験結果

リサーチクエスチョン

• RQ1 GSSは長距離テキストでパープレキシティを維持または改善しつつ、DSSよりも速く自己回帰言語モデルを訓練できるか。
• RQ2 GSSと軽量な局所アテンションを持つTransformerハイブリッドを統合すると、過度な訓練コストなく追加の利得が得られるか。
• RQ3 GSSは訓練時に見られた長さを超える長さのシーケンスへどの程度一般化するか（長さ一般化）。
• RQ4 初期化とゲーティングが、状態空間LM層の訓練安定性と性能にどのような影響を与えるか。

主な発見

• GSSは複数の長距離LMベンチマークでDSSより2–3×速く訓練される。
• GSSは LM1B, PG-19, ArXiv, Github などの複数データセットにおいて、固定計算・固定パラメータの比較で十分競争力がある。
• GSSは長さ一般化が強く、訓練長を超えるシーケンス長（最大65k）で性能が劣化せず、むしろ改善することが多い。
• GSS-Transformer-Hybridは、状態空間層の強みと局所アテンションの強みを組み合わせて性能をさらに向上させる。
• Delta=1固定とLambdaのランダムなガウス様初期化のような簡略化は、パフォーマンスを維持しつつ訓練ボトルネックを削減する。