[論文レビュー] Toeplitz Neural Network for Sequence Modeling

Zhen Qin, Xiaodong Han+7 more|arXiv (Cornell University)|May 8, 2023

Topic Modeling被引用数 5

ひとこと要約

この論文では、自己注意機構を学習可能なトーペリッツ行列に置き換えることで、対数線形の空間的・時間的計算量を達成する、系列モデリングアーキテクチャであるToeplitz Neural Network (TNN) を提案する。軽量な相対的位置エンコーダーと指数関数的減衰バイアスを用いることで、14Kトークンに達するような長系列長にわたり一貫した性能を維持し、Long-Range Arena などの長文脈タスクで競合モデルを上回りながらも、著しく高速である。

ABSTRACT

Sequence modeling has important applications in natural language processing and computer vision. Recently, the transformer-based models have shown strong performance on various sequence modeling tasks, which rely on attention to capture pairwise token relations, and position embedding to inject positional information. While showing good performance, the transformer models are inefficient to scale to long input sequences, mainly due to the quadratic space-time complexity of attention. To overcome this inefficiency, we propose to model sequences with a relative position encoded Toeplitz matrix and use a Toeplitz matrix-vector production trick to reduce the space-time complexity of the sequence modeling to log linear. A lightweight sub-network called relative position encoder is proposed to generate relative position coefficients with a fixed budget of parameters, enabling the proposed Toeplitz neural network to deal with varying sequence lengths. In addition, despite being trained on 512-token sequences, our model can extrapolate input sequence length up to 14K tokens in inference with consistent performance. Extensive experiments on autoregressive and bidirectional language modeling, image modeling, and the challenging Long-Range Arena benchmark show that our method achieves better performance than its competitors in most downstream tasks while being significantly faster. The code is available at https://github.com/OpenNLPLab/Tnn.

研究の動機と目的

• 長系列モデリングにおけるTransformerの自己注意機構の2次的計算量のボトル neck を解消すること。
• コンテンツ依存の注意機構に依存せずに、相対的位置情報のみで有効な系列モデリングが可能かどうかを検証すること。
• 再訓練なしに可変な系列長に一般化可能な、パラメータ効率の良いアーキテクチャを設計すること。
• Long-Range Arena などの長文脈ベンチマークで、強力な性能と外挿能力を達成すること。

提案手法

• トークン間の相対的位置関係を符号化する学習可能なトーペリッツ行列に、標準的な注意行列を置き換える。
• 高速フーリエ変換（FFT）に基づくトーペリッツ行列-ベクトル乗算を活用し、計算量を O(n²) から O(n log n) に削減する。
• 系列長に依存せず、固定されたパラメータ予算内でトーペリッツ係数を生成する軽量な相対的位置エンコーダー（RPE）を導入する。
• 推論時に長系列に一般化できるように、トーペリッツ行列に指数関数的減衰バイアスを直接適用する。
• トークン混合ブロックにおける表現能力を向上させるために、ゲーテッド線形ユニット（GLU）とゲーテッド・ティード・ユニット（GTU）を組み合わせる。
• TNN を、Transformers や CNN や状態空間モデルを特別なケースとして含む統一的フレームワークとして定式化する。

実験結果

リサーチクエスチョン

• RQ1 コンテンツベースの注意機構に依存せず、相対的位置エンコーディングのみで強力な性能を達成できるか？
• RQ2 対数線形計算量のモデルが、2次的計算量のTransformerを長文脈系列モデリングタスクで上回れるか？
• RQ3 再トレーニングなしに、14Kトークンのような、学習時とは大きく異なる系列長に一般化できるパラメータ効率の良いアーキテクチャは可能か？
• RQ4 長文脈ベンチマークにおいて、提案されたTNNは、速度、メモリ、精度の観点で最先端モデルと比べてどうか？
• RQ5 TNN は、Transformers や CNN といった既存のアーキテクチャを包含する統一的フレームワークであるか？

主な発見

• WikiText-103 語彙モデルベンチマークにおいて、TNN はテストパープレキシティが 23.98 に達し、ベースラインのTransformer や他の効率的注意機構の変種を上回った。
• Long-Range Arena ベンチマークでは、1K系列長で 25.72 ステップ/秒の最高の推論速度を記録し、テストした全系列長にわたり一貫した性能を維持した。
• 512トークンの系列で学習したにもかかわらず、推論時に14Kトークンの系列に一般化でき、平均パープレキシティは 23.70 であった。これは、強い外挿能力を示している。
• 相対的位置エンコーダー（RPE）を導入したことで、RPE を持たないTNNの変種と比べて2.47 PPL の性能向上が確認され、位置に依存する表現を学習する有効性が裏付けられた。
• 減衰率 0.99 の指数関数的減衰が、安定した外挿を可能にした。一方、減衰なし、または学習可能な減衰では、性能が著しく低下した。
• 数学的に、TNN が、Transformers、CNN、状態空間モデルを特別なケースとして包含することを示し、系列モデリングの統一的視点を確立した。