KARAKURI LMの解説
日本語事前学習済み言語モデルであるKARAKURI LM 70B v0.1と、それを会話形式にファインチューニングしたKARAKURI LM 70B Chat v0.1をご紹介します。
KARAKURI LMは、Llama 2を基に開発した事前学習済み言語モデルです。
日本語の語彙を追加し、日本語と多言語コーパスを混ぜて追加の事前学習を行うことで、Llama 2の日本語能力を強化しています。
KARAKURI LM Chatは、KARAKURI LMをファインチューニングしたモデルです。
公開されている会話データセットと独自で開発した非公開の会話データセットを混ぜて学習させています。
ファインチューニングで用いた会話データセットには、日本語トークンがわずか2.5%しか含まれていないにもかかわらず、公開時点で、MT-Bench-jpベンチマークにおいて日本語のオープンモデルの中で最高性能を記録しています。
両モデルはともにHuggingFace Hubにて公開されています。
本記事ではKARAKURI LMの技術的な詳細について解説します。
継続事前学習
訓練データの作成に際しては、公開されているmC4やRedPajamaといったコーパスに加え、独自に収集した日本語コーパスを使用しました。
これらを組み合わせて、合計で約100Bトークンのデータセットを構築しました。
このデータセットには、全体の約20%を占める日本語以外のトークンも含まれています。
モデルの初期重みとしては、Llama 2 70Bを使用し、日本語の語彙を追加しました。
ただし、追加する語彙には、常用漢字以外のマイナーな漢字や絵文字などが含まれないようにフィルタリングを行いました。
このモデルには、16Bトークンのデータを学習させました。
ハイパーパラメータについては、Llama 2の論文で述べられていた値をそのまま用いました。
学習はAmazon EC2のtrn1.32xlargeインスタンスを32ノードで利用し、約3日間かけて行いました。
分散学習のライブラリとしてはneuronx-nemo-megatronを使用しました。
ファインチューニング
訓練データとしては、OASST2と独自に構築した会話データセットを使用しました。
独自に構築したデータセットには約1000件の会話が含まれています。
OASST2に含まれる日本語以外の会話は日本語に翻訳せずにそのまま用いました。
これは機械翻訳したデータを学習させると、機械翻訳由来のノイズをそのまま学習してしまい、性能劣化につながると考えたためです。
訓練データに含まれるトークン数は約36Mで、そのうち日本語のトークンは約2.5%です。
日本語のトークン数が少ないため、そのまま学習させると破滅的忘却が発生し、日本語の性能が大幅に低下する可能性があります。
これを防ぐために、継続学習のアプローチを採用しました。
具体的には、継続事前学習の際に用意したデータセットの中からまだ学習させていないものを使用し、事前学習タスクも同時に解かせるマルチタスク学習を行いました。
これにより、継続事前学習によって得られたパラメータから大きく逸脱しないように正則化する効果が期待できます。
事前学習タスクを混ぜる割合は全体の約20%としました。
ファインチューニングには、SFTやRLHFではなく、SteerLMを用いました。
SteerLMは、NVIDIAによって提案されたRLHFよりもシンプルなアラインメントの手法です。
プロンプトに属性を埋め込み、その属性も含めて条件付けして言語モデルを学習させます。
生成時には、この属性の値を変化させることで、言語モデルの挙動をコントロールできます。
教師あり学習の枠組みで学習できるため、強化学習よりも実装が容易で学習が安定します。
学習率は、事前学習時よりも一桁小さい1e-5を使用しました。
また、バッチサイズは少し大きめの256を用いました。
これは、バッチ内に日本語と英語のトークンがともに十分な量含まれるようにするためです。
両言語のトークンが十分に含まれた状態で逆伝播することで、両言語の情報を併せ持った勾配が得られ、クロスリンガル転移が起きやすくなることを期待しています。
SteerLMの属性予測モデルにはSwallow 13Bを使用し、OASST2とHelpSteerを学習させました。
学習はtrn1.32xlargeインスタンスを2ノードで利用し、約10時間かけて行いました。
分散学習のライブラリには、継続事前学習と同様に、neuronx-nemo-megatronを使用しました。
評価
モデルの性能は、MT-BenchとMT-Bench-jpを用いて評価しました。
MT-BenchはLMSYSによって開発された、マルチターンの会話能力を評価するベンチマークです。
モデルの指示に従う能力とマルチターンにおける一貫性を評価することができます。
MT-Bench-jpはMT-Benchの日本語版です。
Stability AIによってタスクの日本語化が行われ、W&Bによってリーダーボードがホストされています。
MT-Bench-jp
KARAKURI LM Chat v0.1は、モデルの公開時点で、MT-Bench-jpベンチマークにおいて日本語のオープンモデルの中で最高性能を記録しています。
KARAKURI LM Chat v0.1のスコアは自社で実験した値で、その他のモデルはNejumi LLMリーダーボード Neoからの引用です。
MT-Bench
オリジナルの英語のMT-Benchにおいても、Llama 2 70B Chatと同水準の性能を記録しています。
KARAKURI LM Chat v0.1のスコアは自社で実験した値で、Llama 2 70B Chatのスコアは、LMSYSの公開している実験結果から計算しています。
謝辞
このプロジェクトはAWS ジャパンによるAWS LLM 開発支援プログラムの支援を受けています。
