スケーリング則を超えて：連想メモリを用いたTransformerの性能理解

Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

May 14, 2024

著者: Xueyan Niu, Bo Bai, Lei Deng, Wei Han

cs.AI

要旨

Transformerモデルのサイズを増大させても、必ずしも性能が向上するわけではない。この現象は、経験的なスケーリング則では説明できない。さらに、モデルが訓練サンプルを記憶するにつれて、汎化能力が向上するという現象が観察される。本論文では、Transformerベースの言語モデルの記憶プロセスと性能ダイナミクスを解明する理論的フレームワークを提示する。我々は、Hopfieldネットワークを用いてTransformerの振る舞いを連想記憶としてモデル化し、各Transformerブロックが近似的な最近傍探索を効果的に実行するように定式化する。これに基づいて、現代の連続Hopfieldネットワークと類似したエネルギー関数を設計し、Attentionメカニズムに対する洞察に富む説明を提供する。Majorization-minimization手法を用いて、Transformerの階層的アーキテクチャを捉えたグローバルなエネルギー関数を構築する。特定の条件下では、達成可能な最小の交差エントロピー損失が約1に等しい定数によって下から制限されることを示す。我々の理論的結果を検証するため、GPT-2を用いて様々なデータサイズで実験を行い、また200万トークンのデータセットでVanilla Transformerを訓練する実験を実施する。

English

Increasing the size of a Transformer model does not always lead to enhanced performance. This phenomenon cannot be explained by the empirical scaling laws. Furthermore, improved generalization ability occurs as the model memorizes the training samples. We present a theoretical framework that sheds light on the memorization process and performance dynamics of transformer-based language models. We model the behavior of Transformers with associative memories using Hopfield networks, such that each transformer block effectively conducts an approximate nearest-neighbor search. Based on this, we design an energy function analogous to that in the modern continuous Hopfield network which provides an insightful explanation for the attention mechanism. Using the majorization-minimization technique, we construct a global energy function that captures the layered architecture of the Transformer. Under specific conditions, we show that the minimum achievable cross-entropy loss is bounded from below by a constant approximately equal to 1. We substantiate our theoretical results by conducting experiments with GPT-2 on various data sizes, as well as training vanilla Transformers on a dataset of 2M tokens.