[論文レビュー] When Attention Meets Fast Recurrence: Training Language Models with Reduced Compute

Tao Lei|arXiv (Cornell University)|Feb 24, 2021

Topic Modeling被引用数 6

ひとこと要約

この論文では、高速な再帰と最小限のアテンションを組み合わせることで、訓練計算コストを3倍から10倍削減しながら、最先端の言語モデリング性能を達成する高効率なシーケンスモデル、SRU++を提案する。SRUの入力投影を自己アテンションコンポーネントに置き換えることで、並列処理が維持され、enwik8、Wiki-103、Billion Word データセットにおいて、トップクラスのTransformerモデルと同等またはそれを上回るビット/文字数とパープレキシティを達成する。GPU使用時間も大幅に削減されている。

ABSTRACT

Large language models have become increasingly difficult to train because of the growing computation time and cost. In this work, we present SRU++, a highly-efficient architecture that combines fast recurrence and attention for sequence modeling. SRU++ exhibits strong modeling capacity and training efficiency. On standard language modeling tasks such as Enwik8, Wiki-103 and Billion Word datasets, our model obtains better bits-per-character and perplexity while using 3x-10x less training cost compared to top-performing Transformer models. For instance, our model achieves a state-of-the-art result on the Enwik8 dataset using 1.6 days of training on an 8-GPU machine. We further demonstrate that SRU++ requires minimal attention for near state-of-the-art performance. Our results suggest jointly leveraging fast recurrence with little attention as a promising direction for accelerating model training and inference.

研究の動機と目的

• 純粋なアテンションベースのTransformerに代わる効率的なアーキテクチャの代替案を検討し、大規模言語モデルの訓練コストの増大に対処すること。
• 高速な再帰と最小限のアテンションを組み合わせることで、高いモデリング能力を維持しながら、訓練および推論コストを大幅に削減できるかどうかを調査すること。
• 再帰とアテンションがシーケンスモデリングにおいて補完的であることを示し、再帰により低い計算コストで優れた性能を達成できることを明らかにすること。
• 最小限のアテンションと位置エンコーディングを一切不要としないが、ほぼ最先端の結果を達成する計算効率の高いTransformerの代替案を提供すること。

提案手法

• SRU++は、単純再帰ユニット（SRU）を拡張し、その入力投影を自己アテンション機構に置き換えることで、アーキテクチャの並列処理性を保ったまま実装する。
• モデルは、SRUのゲート付き再帰メカニズムを適用する前に、文脈に適した入力表現を計算するために1層の自己アテンションコンポーネントを使用する。
• 位置エンコーディングや多頭アテンションを回避し、計算オーバーヘッドを低減するために、単一ヘッドのアテンションコンponentに依存する。
• モデルはSRUが持つ効率的なパラメータ化と計算を維持しており、高速なフォワードパスとメモリ使用量の削減を実現する。
• 訓練は標準的な言語モデリング目的に基づき、混合精度訓練と標準的な最適化設定を用いて実施される。
• 実装はPyTorchでオープンソース化されており、再現性と今後の拡張が可能である。

実験結果

リサーチクエスチョン

• RQ1 高速な再帰と最小限のアテンションを組み合わせたモデルが、標準的な言語モデリングベンチマークで最先端の性能を達成しながら、訓練コストを削減できるか？
• RQ2 限定的アテンションを備えた再帰ベースのモデルは、完全なアテンションを備えたTransformerと比較して、パープレキシティとビット/文字数の点でどの程度の性能を示すか？
• RQ3 シーケンスモデルにおいて、モデリング能力を損なわずにアテンションをどの程度まで最小限にできるか？
• RQ4 位置エンコーディングと多頭アテンションを欠如させても、高速な再帰と組み合わせた場合、性能に悪影響を及えるか？

主な発見

• enwik8データセットでは、SRU++は8-GPUマシンでたった1.6日間の訓練で、1.017ビット/文字という最先端の性能を達成し、訓練コストが1/8に抑えられたTransformer-XLを上回った。
• Wiki-103およびBillion Wordデータセットにおいて、SRU++はトップクラスのTransformerモデルと同等または優れたパープレキシティを達成したが、訓練計算コストは3倍から10倍も削減された。
• わずか2層のアテンションを備えたSRU++は、高速な再帰と組み合わせることで、最小限のアテンションでも十分な性能を発揮することを示した。
• 行列乗算やテンソル転置のオーバーヘッドが少ないため、バッチあたりの処理速度がTransformer-XLの4〜5倍にのぼり、特に相対アテンションコンponentで顕著だった。
• 位置エンコーディングや多頭アテンションを必要としないが、強力な性能を維持しており、再帰が長距離依存関係を効果的にモデル化できることを示している。
• SRU++のPyTorch実装はオープンソース化されており、再現性とさらに効率的なシーケンスモデリングアーキテクチャの研究を促進している。