LLM 勉強会 LLM-jp

Blog

ブログ

LLM-jpモデルに対するOLMo2ベースの中間学習の検討

Date: 2025.8.27
Author: Koshi EGUCHI, Sosuke Hosokawa, Kouta NAKAYAMA


0. はじめに

LLM-jpではオープンかつ日本語に強い大規模言語モデルの開発を進めています。LLM-jp-3シリーズの公開に続き、現在新しいモデルシリーズの公開に向けて活動しています。

本記事はこの新しいモデルシリーズの中間学習1データセット構築にあたり、Allen Institute for AI (以下、Ai2)が公開するOLMo2シリーズの中間学習データセットを検証した記事です。OLMo2では、中間学習を適用することにより、事前学習済みモデルの性能を大幅に向上させています。そして、Ai2はOLMo2シリーズのモデルの重みだけでなく、学習データセット、学習設定等を全て公開しています。

このAi2が公開する中間学習データセットを学習率やサイズを変更して、ベースラインとする事前学習済みチェックポイントに適用した時、どれだけ性能向上するかを検証しました。その結果、OLMo2の中間学習データセットを事前学習済みベースラインモデルの1.3B、7.7Bサイズのモデルに適用することで、gsm8kの正答率が、1.3Bモデルでは1.17%から30.17%に、7.7Bモデルでは14.86%から59.06%へと顕著に向上しました。

1. OLMo2中間学習データセットを用いたLLM-jp-v4モデルの中間学習検証

1.1. OLMo2の紹介

先述した通りOLMo2はAi2が透明性の高いオープンなプロセスで開発した第2世代の大規模言語モデルファミリーです。モデル重みだけでなく学習データセット(OLMo Mix 1124/ Dolmino Mix 1124)、訓練コード、評価ログまで、関連する成果物が包括的に公開されています。詳細は論文をご参照ください。

1.2. OLMo2 中間学習データセット

表1はOLMo2で使用される数学、論文、コーディングなどのデータセット群で構成される中間学習用の高品質データセットの構成を示しています。OLMo2では、この高品質データセットをDolmino Mix 1124と呼んでいます。

表1) OLMo2の中間学習データセット(Dolmino Mix 1124)の構成2

カテゴリデータソースデータタイプトークン数ワード数バイト数 (B)ドキュメント数
Mid-Training Dolmino High Quality SubsetDCLM-Baseline FastText top 7% Fine Web >=2High quality web752B670B4.56T606M
FLAN from Balma 17 decontaminatedInstruction data17.0B14.4B98.2B57.3M
peS2o from Dolma 17Academic papers58.6B51.1B413B38.8M
Wikipedia & Wikibooks from Dolma 1.7Encyclopedic3.7B3.16B16.2B6.17M
Stack Exchange 09/30/2024 dump curated Q&A dataQ&A1.26B1.14B7.72B2.48M
High quality total832.6B739.8B5.09T710.8M
Mid-training Dolmino Math MixTuluMathSynthetic math230M222M1.03B220K
Dolmino SynthMathSynthetic math28.7M35.1M163M725K
TinyGSM-MINDSynthetic math6.48B5.68B25.52B17M
MathCoder2 Synthetic Ajibawa-2023 M-A-P MatrixSynthetic Math3.87B3.71B18.4B2.83M
Metamath OWM-filteredMath84.2M76.6M741M383K
CodeSearchNet OWM-filteredCode1.78M1.41M29.8M7.27K
GSM8K Train splitMath2.74M3.00M25.3M17.6K
Math total10.7B9.73B45.9B21.37M

そして、Dolmino Mix 1124を表2のようにサンプリングし50B、100B、300Bからなる三種類の中間学習データセットが作成されました。

表2) Dolmino Mix 1124の内訳。Source(%)は元のDolmino Mix 1124に含まれる各ソースからどの程度サンプリングしたかを、Mix(%)は作成された中間学習データのうちどの程度の割合を占めるかを示します。(Mix(%)の各列の合計値は100%になります。)

データソーストークン数50B100B300B
Source (%)Mix (%)Source (%)Mix (%)Source (%)Mix (%)
Filtered DCLM752B3.2347.26.8550.220.7851.9
Decontam. FLAN17.0B50.016.610016.720011.3
StakExchange Q&A1.26B1002.452002.474001.68
peS2o58.6B5.155.8516.79.5210019.4
Wikipedia/Wikibooks3.7B1007.111003.574004.86
Dolmino Math10.7B10020.820017.540010.8

3つのDolmino Mix 1124を使用し、論文のOlmo2 7B、13Bモデルは以下のように作成されました。

  • Olmo2 7B: 50Bデータセットを異なるデータ順序で中間学習を3回実行。それぞれのチェックポイントにModel Soup (複数の学習済みモデルの重みを単純平均するアンサンブル手法)を適用し、最終モデルとした。
  • Olmo2 13B: 100Bデータセット3回、300Bデータセット1回の中間学習を実行。それら4つのチェックポイントをModel Soupを適用し、最終モデルとした。

これらのModel Soupingにより、個別実行よりも一貫して性能が向上することが論文では確認されています。

これらの異なるサイズの中間学習データセットを、事前学習済みベースラインの1.3B、7.7Bのモデルに適用し、性能がどのように向上するかを検証しました。また、Model Soupの検証も実施しました。前学習済みベースラインの1.3B、7.7Bのモデルに適用し、性能がどのように向上するかを検証しました。また、Model Soupの検証も実施しました。

2. 評価

2.1. 実験設定

2.1.1. 中間学習データセット

検証にあたり、ベースラインとするモデルの事前学習で使用したものと同じトークナイザーでOLMo2の中間学習データをトークナイズしたLLM-jp Dolmino Mix 1124を作成しました。そして、表2のSource(%)と同じ割合でLLM-jp Dolmino Mix 1124をサンプリングし、表3の3つのサイズのデータセットを作成しました。表2と比較するとトークナイザーの違いによって、Mix(%)が変化していることがわかります。

表3) 本検証で使用した中間学習データセット(LLM-jp Dolmino Mix)の内訳。Source(%)を元のDolmino Mixと合うように調整しました。そのため、各データセットのMix(%)の割合は表2とはわずかに異なっています。また、50B、100B、300Bトークンのデータセットサイズは厳密にはそれぞれ55.8B、114.9B、337.7Bトークンですが、以下では便宜上50B、100B、300Bと呼ぶことにします。

データソーストークン数50B100B300B
Source (%)Mix (%)Source (%)Mix (%)Source (%)Mix (%)
Filtered DCLM821B3.2347.576.8548.9820.7850.57
Decontam. FLAN18.5B50.016.5610016.0820010.95
StakExchange Q&A1.46B1002.632002.554001.74
peS2o62.9B5.155.8016.79.1310018.61
Wikipedia/Wikibooks3.9B1006.981003.394004.62
Dolmino Math11.4B10020.4620019.8740013.52

2.1.2. その他学習設定

ベースラインとする1.3B、7.7Bサイズのモデルは英語、中国語、日本語、韓国語から構成される19.5T トークンで学習しました。


中間学習の学習率設定は以下の2通りを試しました。

  1. OLMo2の論文中の設定と同じで、事前学習終了時の学習率から0に中間学習の間に線形に減衰させる
  2. 事前学習終了時の学習率を中間学習終了時まで維持する

その他の学習のハイパーパラメータは事前学習時の設定に合わせました。。

2.1.3. 評価指標

モデルの評価には、Swallowの評価に追従し、以下の12個のベンチマークを用いました。

  • gsm8k
  • squad2
  • triviaqa
  • hellaswag
  • openbookqa
  • xwinograd_en
  • bbh_cot
  • mmlu
  • mmlu_social_sciences
  • mmlu_humanities
  • mmlu_stem
  • mmlu_other

2.2. 評価結果

本節では、まず学習率とデータセットサイズがモデル性能に与える影響を比較し、次にModel Soup手法による性能向上効果を検証した結果を述べます。

2.2.1. 学習率、中間学習データセットサイズの比較

学習率を(線形減衰、固定)の2通り、中間学習データセットサイズを(50B、100B、300B)の3通りで変化させ、計6通りのパターンを1.3B、7.7Bのモデルに対して検証しました。検証結果は以下の図1、図2になります。

図1) 1.3Bモデルにおける中間学習データセットサイズおよび学習率戦略別の評価結果。

図2) 7.7Bモデルにおける中間学習データセットサイズおよび学習率戦略別の評価結果。

図1は1.3Bモデルの中間学習後の結果、図2は7.7Bモデルの中間学習後の結果を表します。ベンチマークの平均ランクが最も高かった中間学習後の1.3Bモデルは事前学習済みベースラインモデルを12個中10個のベンチマークで上回り、平均ランク最良の中間学習後の7.7Bモデルは12個全てのベンチマークで事前学習済みベースラインモデルを上回りました。特にgsm8kの性能向上は顕著で、1.3Bモデルは1.17%から27.90%に、7.7Bモデルは14.86%から51.25%までスコアが向上しました。

■1.3Bモデルにおける分析

データセットサイズの比較:
学習率固定時、学習率線形減衰時の双方において、データセットを拡大するにつれ、性能の向上が確認できました。gsm8kにおいては、双方の学習率設定において100Bトークンの時に50Bトークンよりもスコアが下がる現象が観測されましたが、300Bトークンの学習時においては50Bトークンよりもgsm8kのスコアが向上することが確認できました。

学習率戦略の比較:
上で述べたように、学習率固定、線形減衰のいずれの戦略においても、データセットサイズが300Bトークンの時に、ベンチマークの平均ランクが最も高くなりました。そして、300Bトークンの学習率線形減衰、300Bトークンの学習率固定の比較をすると、学習率固定時が、学習率線形減衰時よりも12個中10個のベンチマークで性能が向上することが確認されました。結果として、1.3Bモデルでは、(300Bトークン、学習率固定)の組み合わせがベンチマークの平均ランクで最良となりました。

■7.7Bモデルにおける分析

データセットサイズの比較:
学習率固定時、学習率線形減衰時の双方において、データセットサイズを拡大しても、gsm8kを除いた各ベンチマークで最大でも2%程度の性能変化に留まりました。しかし、gsm8kにおいては、50Bから300Bトークンにデータセットサイズを増やすにつれ、顕著な性能低下が観測されました。より具体的には、学習率線形減衰時は46.85%から34.04%まで低下し、学習率固定時は51.25%から17.74%にまで低下しました。これにより、ベンチマークの平均スコア(単純平均)は学習率固定時、学習率線形減衰時の双方において、50Bトークンの時が最良となりました。

学習率戦略の比較:
全ての同一サイズデータセットの比較で、学習率固定時の結果が学習率線形減衰時よりも良い平均ランクとなりました。特にgsm8kの性能向上が顕著だった50B同士の比較においては、学習率固定時が学習率線形減衰時を12個中10個のベンチマークで上回り、平均ランクで最良となりました。結果として、7.7Bモデルでは、(50Bトークン、学習率固定)の組み合わせが平均ランクで最良となりました。

2.2.2. Model Soupの検証

次に、2.2.1節で特定した最良の学習条件(1.3Bモデルでは300Bトークン・学習率固定、7.7Bモデルでは50Bトークン・学習率固定)を用いて、Model Soupによる性能向上効果を検証しました。

その結果が図3、図4になります。

図3) 学習率固定、300Bトークンのデータセット中間学習した3つの1.3B モデルのModel Soup。

図4) 学習率固定、50Bトークンのデータセット中間学習した3つの7.7B モデルのModel Soup。

図3は1.3BモデルのModel Soupを、図4は7.7BモデルのModel Soupを示します。

  • 1.3Bモデルでは、Model Soupを適用したモデルは、個別学習における最良モデル(seed best)のスコアを、12個中7個のベンチマークで上回りました。また、gsm8kでは、最良シード値が27.90%でしたが、30.17%まで向上させることができました(事前学習後は1.17%)。
  • 7.7BモデルではModel Soup後のモデルが他のseed bestを12個中10個のベンチマークで上回りました。特にgsm8kでは、最良シード値が53.53%でしたが、59.06%にまで性能が向上しました(事前学習後は14.86%)。

1.3B、7.7BともにModel Soupによって、全てのベンチマークがseed bestを上回ったわけではありませんが、平均ランクを向上させ、特にgsm8kにおいて顕著な改善が見られました。そして、性能が下がったベンチマークも1%以内のスコア低下であり、全体としてModel Soupによる性能向上が確認できました。

(おまけ)中間学習に要した時間

参考までに中間学習に要した時間を記載します。今回の検証では以下の環境でモデルを学習させました。

  • 学習プラットフォーム: ABCI3.0
  • 学習フレームワーク: llm-jp/Megatron-LM
  • ノード数: 16台 (各サーバはNVIDIA H200 GPUを8枚搭載)

学習に要した時間は表4のようになりました。本検証を行っていた期間、ABCI3.0上では各GPUの消費電力の上限が500W3に制限されていたため、GPUの消費電力上限が700Wで動作した場合、学習時間はさらに短縮されたと考えられます。

表4) 各モデルの学習に要した時間

Model SizeTraining Corpus Size (token)学習時間 (hours)
1.3B 50B2.9
100B6.1
300B18.8
7.7B50B17.6
100B36.1
300B107.3

おわりに

本記事では、Ai2が公開するOlmo2の中間学習データセットDolmino Mix 1124を事前学習済みチェックポイントに適用した時の検証結果について紹介しました。Dolmino Mix 1124を適用することで、モデルの性能向上が確認でき、特にgsm8kでは顕著な性能向上が見られました。

今回は計算資源と実験期間の制約から検証しきれませんでしたが、学習率やデータセット割合を適切に変更し、他のデータセットも組み合わせることで、さらなる性能向上の可能性があると考えられます。

今後もLLM-jpでは、オープンソースで日本語に強い大規模言語モデルの開発を進めていきます。LLM-jp の活動に興味を持たれた方はこちらのページからぜひ参加していただければと思います。


1 事前学習と事後学習の中間に位置するステップになります。
2 OLMo2の論文中の表2と同じ表になります。OLMo2の論文中の表2では、GSM8K Train splitが17.6Kのドキュメント数と記載されていますが、これはTran splitとTest splitを合わせた値であり、実際の学習に使用されたTrain Splitは、14.9Kのドキュメント数であることに注意して下さい。
3 2025年7月現在では400Wに制限されています。