大規模な日本語の事前学習言語モデルSarashina1・2の公開

概要

SB Intuitionsの清野舜,李凌寒,高瀬翔です。

弊社では高品質な日本語大規模言語モデルの構築に取り組んでいます。 大規模言語モデルの構築は理論的にも経験的にもまだまだ発展途上の領域であり,学術界・産業界の日本語大規模言語モデルの研究開発促進のために,7B,13B,65BパラメータのモデルからなるSarashina1および,7B,13BパラメータのモデルからなるSarashina2を公開しました。

この記事ではSarashina1・2の詳細について解説します。

使い方

Sarashina1・2共にHuggingFace Hubにて,商用利用も可能なライセンス(MIT)のもと公開しています。

以下のようにtransformersライブラリからお試し頂けます。

なお,下記にあるように,tokenizerを呼び出す際にuse_fast=Falseを設定することで想定通りの挙動となります。

2024/06/27追記: use_fast=False を指定しなくても所望の動作が可能となりました。use_fast=False を用いる場合はHuggingFace Hubの各モデルのREADMEをご参照ください。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, set_seed
model = AutoModelForCausalLM.from_pretrained("sbintuitions/sarashina2-7b", torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("sbintuitions/sarashina2-7b")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
set_seed(123)
text = generator(
"おはようございます、今日の天気は",
max_length=30,
do_sample=True,
pad_token_id=tokenizer.pad_token_id,
num_return_sequences=3,
)

for t in text:
print(t)

# These examples are generated by sarashina2-7b parameters model
# {'generated_text': 'おはようございます、今日の天気は晴れです。ちょっと風が強い。\n昨日は、久しぶりにゆっくりとしていました。\n2週間位間があいてしまったかも、でもその間に'}
# {'generated_text': 'おはようございます、今日の天気は曇。朝は曇っていてどんよりしていましたね。昼からは晴れそうですが。気温は徐々に上昇しています。昨日は春らしい陽気でした。'}
# {'generated_text': 'おはようございます、今日の天気はくもり、少し寒気がします。 この土日に、家族で一泊二日で旅行に行ってきました。といっても、100キロ'}

Sarashina1・2の仕様

Sarashina1 Sarashina2
訓練データ 日・英・コード
訓練トークン数 1T 2.1T
語彙サイズ 51,200 102,400
フレームワーク GPT-NeoX Megatron-LM
モデル構造 GPT-NeoX Llama
計算精度 float16 bfloat16
学習率 2e-4, 1e-4, 6e-5 2.5e-4
スケジューラ コサイン コサイン
Weight decay 0.01 0.1
Auxiliary z-loss 1e-4 1e-4
Scaled embed

訓練データ

Sarashina1は日本語のデータ,Sarashina2は日本語に加えて英語とプログラミングコードのデータを用いて学習させました。 日本語にはCommonCrawlから独自に構築したデータを用いているほか,英語にはSlimPajama*1,プログラミングコードにはstarcoderdataを用いました。 Sarashina2において,日本語・英語・プログラミングコードの割合は5:4:1に設定しました。

Sarashina1・2はともに日本語データを約1Tトークン学習しており,これは国外の大規模言語モデルと比較しても圧倒的に大きいです。 例えばMeta社の公開しているLlama2において,学習データ中の日本語は約0.10%と報告されており,これは約2Bトークンに相当すると考えられます。 つまり,SarashinaはLlama2の数百倍規模の日本語データを学習していると言えます。

Sarashinaの日本語性能を高めるための工夫として,日本語データの構築には特に力を注いでいます。 まず,CommonCrawlから取得したダンプデータにCCNetパイプラインを適用し,大規模な日本語のテキストデータを獲得しました。 その後,OSSライブラリであるHojiCharを用いてクリーニング処理をおこない,最後にデータセット中に含まれる重複した文書をMinHashLSHを用いて除去しました。

クリーニング処理や重複文書の除去の効果については網羅的な比較実験によって検証しています。 例えば,実験を通して,URLを用いた重複除去がLLMの性能に悪影響を及ぼすことが分かったため,同処理はSarashina2の日本語データには適用していません。 一連の実験とその結果については言語処理学会第30回年次大会にて以下のタイトルで発表済みですので,ぜひ論文もご一読ください。

https://www.anlp.jp/proceedings/annual_meeting/2024/pdf_dir/A3-2.pdf

日本語LLM構築におけるコーパスクリーニングの網羅的評価
新里顕大 (SB Intuitions/京大), 高瀬翔, 清野舜, 李凌寒, 加藤卓也, 水本智也, 小林滉河, 佐藤潤一, 柴田知秀 (SB Intuitions)

トークナイザと語彙サイズ

Sarashina1・2にはそれぞれ独自に構築したトークナイザを用いました。 具体的には,sentencepieceの提供するユニグラム言語モデルのアルゴリズムを用いてサブワードを構築しました。 その際,Sarashina2の語彙サイズ(102,400)はSarashina1の語彙サイズ(51,200)の倍に設定しました。 これにより,日本語・英語・プログラミングコードといった複数種類のデータを効率的に学習可能です。 また社内での研究では,語彙サイズが大きいほうが最終的なLLMの性能が高いことを確認しています。 この結果については以下の論文をご参照ください。

arxiv.org

Large Vocabulary Size Improves Large Language Models
Sho Takase, Ryokan Ri, Shun Kiyono, Takuya Kato (SB Intuitions)

学習を安定させつつ良いモデルを得るために

大規模言語モデルのフルスクラッチの学習には多大な計算コストを要するため,学習途中でモデルの損失値が跳ね上がってしまうなどの失敗を避けつつ,なるべく良い性能を達成したいという要求があります。 これを実現するために学習率や正則化のハイパーパラメータ調整といくつかの工夫を加えました。 それぞれについて説明します。

学習率

学習率は高いほど性能が良くなりますが,学習も不安定になると報告されています。 Sarashina1の7Bパラメータでは少々高い値を用いましたが,13Bと65Bでは学習が失敗しないことを優先し,GPT-3の論文で13B,175Bパラメータのモデルに用いられている値を採用しました。

Sarashina2では学習率が2.5e-4とSarashina1よりも高くなっていますが,これはモデル構造をLlamaにしたため,Llamaの論文で報告されている値までなら学習率を上げられると考えたからです。 ただ,学習の失敗をできる限り回避するため,Llama論文での学習率である3.0e-4よりは小さな値に設定しました。

下記はSarashina2-7B・13Bについて,開発データでの損失値をプロットした図です。 このような高い学習率でも損失値が跳ね上がることなく,滑らかに下がっていっていることがわかります。

Weight decay

Sarashina1では0.01としていましたが,Sarashina2では0.1としました。 学習初期では0.01としていた方が性能が高いのですが,学習後期に性能差が逆転します。 下記はSarashina2の7Bパラメータモデルについて,Weight decayを0.01,0.1とした際の,更新回数に対する開発データでの損失値をプロットした図です。 更新回数がおおよそ40万回の時点でWeight decayを0.1としたモデルの性能が上回っていることがわかります。 また,0.1の方が学習も安定します。

Auxiliary z-loss

トークンを出力する際に,確率分布を計算するためソフトマックス関数を用います。 このソフトマックス関数の分母の値が小さくなるような損失項がz-lossであり,これを活用することで学習の安定性が向上することがPaLMの論文で報告されています。 実際に,我々のモデル構築においてもz-lossを用いることで学習が安定することを確認しています。 z-loss項への重みはPaLMの論文と同様に1e-4としました。

Scale embed

大規模言語モデルの学習では各層の出力するベクトルの分散が小さい値になるよう初期化を行います。 このような状況では埋め込み表現の分散を1に近づけることで各層の勾配のノルムが安定し,学習も安定すると報告されています。 例えばPaLMでは平均0,分散1の正規分布から埋め込み表現を構築しています。 我々はTransformerの論文Gemmaで用いられている,埋め込み表現を次元数の平方根倍することで分散を1に近づけるという手法を採用しました。 詳細は下記の論文をご参照ください。

arxiv.org

Spike No More: Stabilizing the Pre-training of Large Language Models
Sho Takase, Shun Kiyono (SB Intuitions), Sosuke Kobayashi, Jun Suzuki (Tohoku University)

おわりに

本記事ではSarashina1・2の仕様や構築過程の詳細を解説しました。性能については次の記事で紹介致します。 冒頭にも記したように,Sarashina1・2は商用利用も可能なMITライセンスで公開していますので,ぜひ言語モデルの研究開発にご利用いただければと思います。

*1:bookコーパスは著作権侵害の恐れがあるため,SlimPajamaから取り除いています。