小型Transformerへの算術教育

Teaching Arithmetic to Small Transformers

July 7, 2023

著者: Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, Dimitris Papailiopoulos

cs.AI

要旨

GPT-4のような大規模言語モデルは、広範なテキストデータで訓練されることで、基本的な算術などの汎用タスクにおいて創発的な能力を示します。これらのタスクは、教師なしの次トークン予測目標として明示的にエンコードされていないにもかかわらず、そのような能力が現れます。本研究では、ランダムな初期化から訓練された小さなトランスフォーマーが、次トークン予測目標を用いて、加算や乗算、平方根などの初等関数といった算術操作を効率的に学習する方法を調査します。まず、従来の訓練データが算術学習に最も効果的ではないこと、そして単純なフォーマット変更が精度を大幅に向上させ得ることを示します。これにより、訓練データの規模に応じた鋭い相転移が観察され、場合によっては低ランク行列補完との関連性を通じて説明可能です。先行研究を基に、中間ステップの結果を含むチェーン・オブ・ソート形式のデータで訓練を行います。事前訓練を全く行わない場合でも、このアプローチは精度、サンプル複雑性、収束速度を同時に大幅に向上させます。また、訓練中の算術データとテキストデータの相互作用を研究し、Few-shotプロンプティング、事前訓練、モデル規模の影響を検証します。さらに、長さの一般化に関する課題についても議論します。本研究は、次単語予測目標の特性を考慮した高品質で教育的なデータが、算術能力を迅速に引き出す上で重要であることを強調しています。

English

Large language models like GPT-4 exhibit emergent capabilities across general-purpose tasks, such as basic arithmetic, when trained on extensive text data, even though these tasks are not explicitly encoded by the unsupervised, next-token prediction objective. This study investigates how small transformers, trained from random initialization, can efficiently learn arithmetic operations such as addition, multiplication, and elementary functions like square root, using the next-token prediction objective. We first demonstrate that conventional training data is not the most effective for arithmetic learning, and simple formatting changes can significantly improve accuracy. This leads to sharp phase transitions as a function of training data scale, which, in some cases, can be explained through connections to low-rank matrix completion. Building on prior work, we then train on chain-of-thought style data that includes intermediate step results. Even in the complete absence of pretraining, this approach significantly and simultaneously improves accuracy, sample complexity, and convergence speed. We also study the interplay between arithmetic and text data during training and examine the effects of few-shot prompting, pretraining, and model scale. Additionally, we discuss length generalization challenges. Our work highlights the importance of high-quality, instructive data that considers the particular characteristics of the next-word prediction objective for rapidly eliciting arithmetic capabilities.