[論文レビュー] Language Modeling using LMUs: 10x Better Data Efficiency or Improved Scaling Compared to Transformers

Narsimha Chilkuri, Eric Hunsberger+3 more|arXiv (Cornell University)|Oct 5, 2021

Topic Modeling参考文献 14被引用数 5

ひとこと要約

この論文は、10倍のデータ効率を達成するLegendre Memory Unit (LMU)に基づく言語モデルを提案している。Transformerと同等の損失に到達するには、訓練トークンを1/10に減らす必要がある。Legendre多項式射影を通じてO(n)のメモリとO(n ln n)の計算複雑度を実現することで、特にグローバル自己注意機構を組み合わせた場合、Transformer やLSTMよりもスケーリング性能に優れる。

ABSTRACT

Recent studies have demonstrated that the performance of transformers on the task of language modeling obeys a power-law relationship with model size over six orders of magnitude. While transformers exhibit impressive scaling, their performance hinges on processing large amounts of data, and their computational and memory requirements grow quadratically with sequence length. Motivated by these considerations, we construct a Legendre Memory Unit based model that introduces a general prior for sequence processing and exhibits an $O(n)$ and $O(n \ln n)$ (or better) dependency for memory and computation respectively. Over three orders of magnitude, we show that our new architecture attains the same accuracy as transformers with 10x fewer tokens. We also show that for the same amount of training our model improves the loss over transformers about as much as transformers improve over LSTMs. Additionally, we demonstrate that adding global self-attention complements our architecture and the augmented model improves performance even further.

研究の動機と目的

• Transformerの言語モデリングにおける高いデータおよび計算リソース要件、特にシーケンス長に伴うO(n²)のメモリおよび計算量の増加を是正すること。
• 非パラメトリックでメモリ効率の良いアーキテクチャ、すなわちLegendre Memory Unit (LMU)が、Transformerに比べ優れたスケーリング法則とデータ効率を達成できるかを検討すること。
• LMUの状態を介した自己注意の暗黙的実装が、Transformerの明示的自己注意と同等またはそれを上回る性能を発揮できるかを評価すること。
• LMUにグローバル自己注意を組み合わせることで、性能と長文脈モデリング能力が向上するかを調査すること。
• LMUモデルのパワー則スケーリング挙動が、モデルサイズの変動に応じてどのように現れるかを検証し、Transformer やLSTMと比較してスケーリング指数を評価すること。

提案手法

• モデルは、入力トークンのスライディングウインドウをq個のLegendre多項式に射影することで、効率的な時系列表現を可能にする非パラメトリックなLinear Time-Invariant (LTI) システム、Legendre Memory Unit (LMU) を使用している。
• LMU状態は固定行列の再帰式で更新される：$\mathbf{m}_t = \bar{\mathbf{A}} \mathbf{m}_{t-1} + \bar{\mathbf{B}} x_t$、ここで$\bar{\mathbf{A}}$ と $\bar{\mathbf{B}}$ はLTIシステムから導出された固定行列である。
• LTI更新は畳み込みとして並列に計算され、$\mathbf{m}_t = \sum_{j=1}^{t} \bar{\mathbf{A}}^{t-j} \bar{\mathbf{B}} x_j$ と表現可能で、これによりO(n)のメモリとO(n ln n)の計算複雑度が達成される。
• モデルは各時刻におけるLMUの隠れ状態にのみグローバル自己注意機構を適用することで、『暗黙的自己注意』を実装している。これは時刻間の自己注意ではなく、時刻内でのみ作用する。
• 長文脈モデリングの向上を目的として、グローバル自己注意を組み合わせたバージョンを導入し、LMUの効率性と自己注意の表現力の両方を活かしている。
• モデルはWebText2データセットで学習され、モデルサイズ（埋め込みを除くパラメータ数）と交差エントロピー損失の間のパワー則関係をフィッティングすることでスケーリング法則が評価されている。

実験結果

リサーチクエスチョン

• RQ1 LMUベースのモデルは、Transformer やLSTMと同様に、モデルサイズと言語モデリング損失の間でパワー則スケーリング関係を示すか？
• RQ2 LMUモデルは、データ効率が向上しているため、大幅に少ない訓練データでTransformerと同等またはそれ以上の性能を達成できるか？
• RQ3 LMUモデルにおける暗黙的自己注意機構は、明示的な時刻間自己注意がなくても、長距離依存関係を効果的に捉えられるか？
• RQ4 LMUアーキテクチャにグローバル自己注意を追加することで、ベースLMUやTransformerと比較して性能とスケーリング挙動にどのような影響を与えるか？
• RQ5 LMUモデルのO(n)のメモリとO(n ln n)の計算複雑度が、TransformerのO(n²)の複雑度に比べ、より優れたスケーリングを可能にするか？

主な発見

• LMUモデルは、10倍少ない訓練トークンでTransformerと同等の交差エントロピー損失を達成しており、10倍のデータ効率を示している。
• 同じ訓練予算下で、LMUモデルはTransformerの性能向上幅（TransformerがLSTMを上回る分）と同等の性能向上を示している。
• LMUモデルはモデルサイズとの間でパワー則スケーリング関係を示しており、$\text{LMU}(N) = \left(\frac{N}{1.95 \cdot 10^{14}}\right)^{-0.072}$ と表現され、スケーリング指数がTransformerとほぼ同等である。
• グローバル自己注意を組み合わせたLMUモデルは、$\text{LMU}_G(N) = \left(\frac{N}{3.80 \cdot 10^{14}}\right)^{-0.069}$ とスケーリングされ、ベースLMUよりもわずかに優れたスケーリングを示している。
• グローバル自己注意を備えたLMUモデルは、10倍多くのトークンで学習されたTransformerをも上回り続け、優れたデータ効率と性能を確認している。
• 1トークンあたりの損失プロットでは、ベースLMUモデルは約100トークンで損失が平坦化しており、長文脈モデリング能力に限界があることが示された。一方、グローバル自己注意を備えたバージョンは、Transformerと同様に、より長い文脈でも改善が持続している。