状態-予測分離仮説

要旨

Transformerは、次のトークンを予測するためと、将来のトークン予測に有用な状態を保持するために、同一の前方計算ストリームを使用しています。本稿では、状態-予測分離仮説を定式化します。すなわち、これら二つの役割を切り離すことで、より優れた言語モデリング性能が得られるというものです。この二つの機能を分離するために二つの計算ストリームを用いるTransformerの変種を設計し、様々な規模で事前学習実験を行いました。実験の結果、状態-予測分離はデータ効率と計算効率の両方において一貫して優れており、検証損失を改善し、下流タスクにおいて標準的なTransformerを平均2〜3パーセントポイント上回る性能を示しました。また、潜在的な交絡因子を排除し、本設計に固有の勾配の根本的な差異を実証する広範な実証分析も実施しました。

English

Transformers use the same forward computation stream to both predict the next token and store useful state for future token predictions. We formulate the state-prediction separation hypothesis: disentangling the two roles yields better language modeling performance. We design a Transformer variant that uses two computation streams to separate the two functions, and conduct pretraining experiments across various scales. Our experiments show that state-prediction separation consistently offers better data and compute efficiencies, improving validation loss and outperforming standard Transformers by 2--3 percentage points on average on downstream tasks. We also conduct extensive empirical analysis that rules out potential confounders and demonstrates the fundamental difference in the gradients our design entails.