少数のトークンで小型ベース言語モデルを事前学習する

Pre-training Small Base LMs with Fewer Tokens

April 12, 2024

著者: Sunny Sanyal, Sujay Sanghavi, Alexandros G. Dimakis

cs.AI

要旨

既存の大規模なベース言語モデル（LM）から出発して、小さなベースLMを開発するためのシンプルなアプローチの有効性を研究します。まず、大規模LMからいくつかのTransformerブロックを継承し、その後、この小さなモデルを大規模モデルの生の事前学習データのごく一部（0.1％）で訓練します。このシンプルなレシピを「Inheritune」と呼び、3Bパラメータの大規模LMの最初の数層を使用して、1Bトークンで1.5Bパラメータの小さなベースLMを構築するために最初に実証します。これを単一のA6000 GPUで半日未満で行います。9つの多様な評価データセットおよびMMLUベンチマークにおいて、結果として得られたモデルは、1B-2Bサイズの公開されているベースモデルと比較して有利であり、そのうちいくつかは50〜1000倍以上のトークンを使用して訓練されています。 Inherituneを、大規模LMとその完全な事前学習データセットを活用して小さなLMを訓練する、やや異なる設定で調査します。ここでは、GPT2-medium（355M）およびGPT-2-large（770M）のいくつかの層を活用して訓練された小さなLMが、OpenWebTextデータセットの9Bトークンで同じ訓練ステップ数でゼロから訓練された場合、それらの大きな対応モデルのval損失に効果的に匹敵することを示します。広範な実験を通じてレシピを分析し、多様な設定での有効性を実証します。私たちのコードはhttps://github.com/sanyalsunny111/LLM-Inherituneで利用可能です。

English

We study the effectiveness of a simple approach to develop a small base language model (LM) starting from an existing large base LM: first inherit a few transformer blocks from the larger LM, and then train this smaller model on a very small subset (0.1\%) of the raw pretraining data of the larger model. We call our simple recipe Inheritune and first demonstrate it for building a small base LM with 1.5B parameters using 1B tokens (and a starting few layers of larger LM of 3B parameters); we do this using a single A6000 GPU for less than half a day. Across 9 diverse evaluation datasets as well as the MMLU benchmark, the resulting model compares favorably to publicly available base models of 1B-2B size, some of which have been trained using 50-1000 times more tokens. We investigate Inheritune in a slightly different setting where we train small LMs utilizing larger LMs and their full pre-training dataset. Here we show that smaller LMs trained utilizing some of the layers of GPT2-medium (355M) and GPT-2-large (770M) can effectively match the val loss of their bigger counterparts when trained from scratch for the same number of training steps on OpenWebText dataset with 9B tokens. We analyze our recipe with extensive experiments and demonstrate it efficacy on diverse settings. Our code is available at https://github.com/sanyalsunny111/LLM-Inheritune.