ChatGPT自作入門 - 七枝のブログ

本記事は、そこそこ自力で ChatGPTを自作するための解説＆実装例となっています。ただ、現時点では ChatGPTの詳細は明かされていないようなので (2023-06-01)、実際には姉妹モデルである InstructGPTモデルを実装していくことになります。また、 InstructGPTモデルを作成するには、複数のモデルを訓練する必要があり、多量の計算資源や手動でのラベル付けが必要となります。なので、高性能なモデルを訓練するのは少し難しいですが、少なくとも実装だけはすることができます。

間違いなどを見つけた場合は、優しく教えていただけると嬉しいです。

ChatGPTの概要
ベースとなる一般 GPTモデルを訓練しよう
Supervised Fine-Tuningモデルを訓練しよう
- 概要
- 訓練
SFTモデルを更にチューニングしよう
最後に

ChatGPTの概要

ChatGPTの生みの親である OpenAIのホームページには、 ChatGPTがどのように作られたのか、ということが軽く書かれています。しかし、 ChatGPTが InstructGPTモデルの姉妹モデルであるという以上のことは書かれていません (2023-06-01)。この InstructGPTモデルは、質問に対して答えを生成する NNモデルです。作成方法はこちらの論文に書かれており、以下の手順を辿ることで作成されます。

(1) ベースとなる一般 GPTモデルを訓練する。
(2) 1のモデルをプロンプト応答用にチューニングする。
(3) 2のモデルを強化学習を使い、更にチューニングする。

ということで、ここからはこれらの工程を１つずつ辿っていくことになります。

ベースとなる一般 GPTモデルを訓練しよう

GPTモデルの概要とアーキテクチャ

GPTモデルは文章を渡された時に、次の文字を予測するモデルです。例えば、 吾輩はという文字列が渡された時に、 猫などと出力します。また、出力された 猫を付け加えた 吾輩は猫を次の入力としてモデルに渡すと、 でなどと出力します。このステップを繰り返すことにより、 吾輩は猫であるといった文章が生成されます。

では、そんな GPTモデルの NNアーキテクチャはどんな感じかというと、本記事で使用するものは以下のようになっています。

各コンポーネントについては、各々の実装時に詳細な説明をするので、ここではマクロ的な説明をします。 GPTモデルは受け取った文章の次の文字を予測するモデルです。しかし、上の図では、ほぼ全ての文字を受け取って、全ての文字を出力しています。少し不思議に感じますが、これは GPTモデルが実際に全ての文字を予測しているためです。 GPTモデルはアーキテクチャの構造的に、 i番目の文字を予測する時には i番目以降の情報が使われないようになっています。例えば、下の図のように、 猫を予測する時はそれよりも左下の情報しか使われていません。

このように未来の文字の情報を受け取らないようにすることで、全ての文字を同時に予測することが可能になっています。

ちなみに GPTモデルの論文では、以下のような図がアーキテクチャとして書かれています。いくつかの違いを除けば (最終層、 Transformer層の数)、上の図と下の図のアーキテクチャはほぼ同じものとなっています。

データを用意する

GPTモデルは、文章の次の文字を予測するので、日本語の文章さえあれば訓練することができます。ということで、まずは日本語の文章を用意していこうと思います。今回は、こちらのAhmedさんのデータセットをお借りしようと思います。そこも自分で用意したいという方は Wikipediaのダンプなどから抽出してもいいのかもしれないです。

~ > head -n 5 wiki-sentences.txt
テネシー大学、デューク大学、フロリダ大学などからのオファーもある中、彼が選んだのはノートルダム大学であった。
9月5日、シカゴ・ベアーズとプラクティス・スクワッドとして契約を結んだ。
12月19日、クリス・コンテが故障者リスト入りするのと入れ違いにアクティブロースター入りした。
クロード・ドビュッシーの曲を原曲にした楽曲をリリースした。
又、のちに山口裕加里が同曲をカバーした。

データを読み込む

では、データを Pythonで読み込みます。ただ、 NNモデルで文字を直接扱うのは少し厳しいので、ここではシンプルに各文字を数字にマッピングしておこうと思います。つまり、 {吾 → 0, 輩 → 1, は →2, …}といった感じのマッピングを作成します。

GPT-2や GPT-3では、 {吾輩 → 0, は → 1, 猫 → 2, である → 3, …}のような、もう少し纏まったレベルでのマッピングを行います。ただ、初めにトークナイザーばかり作り込むのも何とも言えないので、本格的な実装は後に回そうと思います。

コード
- 実装例
- 実行例

事前知識: PyTorch

本記事では、主に PyTorchを使います。ただ、 PyTorchの基礎的な機能しか使わないので、この公式チュートリアルがわかるくらいの知識があれば、問題ないと思います。

実装: 入力と出力だけ

一度に GPTモデルを全て実装するのはキツイので、全体の流れや入出力の確認も兼ねて、必要最低限の部分だけを実装しようと思います。ということで、まずは下の図のような Word Embeddingと全結合層だけからなるモデルを実装します。

コード
- 実装例
- 実行例

実装: Transformer

モデルの基礎部分が実装できたので、次に Transformerを実装していこうと思います。これは GPTモデルの論文内にある図の、水色の部分に相当します。

つまり、以下の一連のブロックを12回繰り返したものになります。

(1) Masked Multi Self Attention
(2) Layer Norm
(3) 全結合
(4) Layer Norm

一番目の Masked Multi Self Attentionというのは、 Multi-Head Self Attentionをベースにした NNブロックです。そこに GPTモデル用の変更を加え、 i文字目では i-1文字目までの情報しか使わないように抑制 (Mask)をかけています。ということで、まず初めに Multi-Head Self Attentionが何かという話なのですが、それに関する解説記事は巷に数多くあるので、申し訳ないのですが、今回はそれらを参考にしていただけるとありがたいです。個人的には、このMishaさんのツイートがとてもわかりやすいと思います。

では、次にどうやってマスクをかけるのかという話なのですが、 GPTモデルは i文字目を予測する時に i-1文字目までの情報しか使ってはいけません。そのため、 Attention機構中で i文字目以降を参照されると困ります。なので、 Attentionの計算中に、 weightの該当部分に -infをセットすることにより、 i文字目以降の情報を取ってくることを抑制 (Mask)することになります。詳細は、以下の実装例を参考にしていただけるとありがたいです。

コード
- 実装例
- 実行例

実装: Positional Embedding / Positional Encoding

Attention機構は語順を理解できず、 Cats like dogsと Dogs like catsの見分けがつきません。なので、 Transformerを提唱した Attention is All You Needという論文では、文字位置に関する情報を入力に追加して、モデルに渡していました。私が勘違いしていない限り、 GPT用のマスキングを施した Attentionは少し事情が違い、 2層目以降の Transformerでは、文字位置を把握することも可能な気がします。ただ、性能の面からか、 GPTモデルの論文でも文字位置を入力に追加しています。

文字位置を伝える方法は色々とあり、GPTモデルの論文で採用されているのは、学習可能なバイアスを Word Embeddingに足し合わせて、トレーニングを通じて最適な方法を探す方法です。これは Positional Embeddingと呼ばれています。非常にシンプルなので、詳細については下にある実装例を確認した方が速いかもしれないです。

また、別の方法として、 Attention is All You Needの中では、位置情報を伝えるためのベクターを予め人力でデザインしておいて、 Word Embeddingに足し合わせています。これは Positional Encodingと呼ばれています。今回は GPTモデルの実装ということで詳細は省きますが、このkazemnejadさんのブログが直感的に理解しやすく、非常にわかりやすいと思いました。

コード
- 実装例
- 実行例

実装: Byte Pair Encoding

ここまでは 1文字ずつモデルに渡して、 吾輩などの複数文字からなる単語を学習して認識してくれることを願っていました。しかし、頻出単語は単語レベルで纏めて渡してあげた方が性能が上がりそう、と言われればそんな気もします。

GPT-3の論文によると GPT-2や GPT-3では Byte Pair Encoding (BPE)というトークナイザーを使っているようです。 BPEでは、文字のマージ規則に従って文字を纏めていくことになります。例えば、マージ規則が [(吾, 輩) → 吾輩, (で, あ) → であ, (であ, る) → である]であるとします。その場合、 吾輩は猫であるという文章は

(1) 文字レベルに分解され 吾 / 輩 / は / 猫 / で / あ / るとなる。
(2) マージ規則を適用して 吾輩 / は / 猫 / であると纏められる。
(3) 各トークンは、対応する数字の IDに変換されてモデルに渡される。

という手順を辿ります。

また、辞書の作成方法もかなりシンプルで、

(1) 全ての文字を辞書に加える。
(2) 出現頻度の多い、隣接する文字のペアを規則に加えていく。

というロジックになっています。以下に説明用の簡単なコードを置いておきます。

コード (説明用なので実装しなくても大丈夫です)
- 実装例
- 実行例

以上が BPEの実装となります。上の実行例からもわかる通り、 BPEでは単語レベルではなく、 ableのようなサブワードと呼ばれる単位で区切られることになります (able自体は単語でもありますが...)。これにより、 backpropagationableのような未知語も、 back / propagation / ableのようなサブワードに分割し対応できるようになる、と言われています。

では、ここで BPEを日本語に適用したい訳ですが、日本語の文章をそのまま BPEにかけると、 [(猫, は) → 猫は, (猫, に) → 猫に, (猫, を) → 猫を]のように、助詞や助動詞とのマージが頻出するようになります。私が参考にした論文 (『日本語 Tokenizerの違いは下流タスク性能に影響を与えるか?』)では、既存の形態素解析器で単語にパースした後に、 BPEをかけているので、本記事でもそのように実装しようと思います。解析機は MeCabの Pythonラッパーを使用します。

以下の実装例では、マージ規則を生成しファイルに保存しています。トークナイズされた文章を保存することもできたのですが、後で他のデータもトークナイズする必要があるので、マージ規則を保存しています。

コード
- 実装例
- 実行例

では、最後にトークナイザーを実装します。また、トークナイズされた文章を使ってモデルを訓練します。これで一般 GPTモデル完成となります。

コード
- 実装例
- 実行例

Supervised Fine-Tuningモデルを訓練しよう

概要

ここでは、先ほど作成した GPTモデルを会話用にチューニングしていきます。今までは、 吾輩は猫であるのような文章を学習させていましたが、ここからは 吾輩は誰ですか？:猫のような 質問:答えもしくは クエリ:応答形式の文章を渡して、モデルを fine-tuningしていきます。このモデルを Supervised Fine-Tuning (SFT)モデルと呼びます。

この章から InstructGPTモデルの論文に書かれている内容に入るので、 InstructGPTモデルの概要についても書いておこうと思います。

GPTモデルのパラメーター数を増やして精度を良くしたからといって、ユーザーの意図した答えを返せるようになるとは言い切れません。なので、人間からのフィードバックを利用することにより、ユーザーの意図した答えを返せるようにしたい、というのが InstructGPTの考えらしいです。実際にここからは、初めに作った GPTモデルを fine-tuningしていくことにより、人間が見て良いと感じる返答を返せるように頑張ることになります。この章で訓練する SFTモデルはその第一歩となります。ちなみに、ここからは会話用にモデルをチューニングするだけなので、一般的な NLPタスクにおける性能は徐々に下がっていきます。

訓練

ということで、 SFTモデルを訓練するのですが、訓練に使用するデータを自力で用意するのは大変です。なので、こちらのMasaさんのgithubからデータをお借りして、以前訓練した GPTモデルを fine-tuningしていこうと思います。

コード
- 実装例
- 実行例

SFTモデルを更にチューニングしよう

概要

ここからは強化学習を使い、前回作った SFTモデルを更にチューニングしていくことになります。簡単に言うと、 SFTの生成した文章を採点するモデルを使い、 SFTモデルがより得点の高い答えを生成するように訓練します。そのために、以下のような工程が必要になります。

(1) SFTモデルに質問 (例: 吾輩は誰ですか？:)を投げて、答えをサンプリングする
(2) サンプリングした答えを人間が採点する
(3) 採点された答えを使って、採点をするNNモデルを作る (Rewardモデル)
(4) Rewardモデルを使い、 SFTモデルをさらにチューニングする (InstructGPTモデル)

最後に InstructGPTモデルと書いてある通り、これが最終章となります。では、一つ目のステップから見ていこうと思います。

SFTの答えをサンプリングする

答えのサンプリングをしていきます。ここでサンプリングされた答えは、次のステップで採点されることになります。ここで言う採点とは、実際に点数を付ける訳ではなく、同一の質問に対して複数の答えをサンプリングして、それらを順位付けするという形で行われます。なので、ここでは１つの質問に対して複数の答えをサンプリングします。

サンプリングするために使用する質問は、 SFTモデルを訓練した時に使ったものを再利用しようと思います。訓練時に使ったデータを再利用するのは良くない気もしますが、実装の簡略化ということで許していただけるとありがたいです。生成された答えはファイルに書き出されます。

コード
- 実装例
- 実行例
- ファイルへの出力例

サンプリングした答えを採点する

次に、モデルの答えを手動で採点します。ここで言う採点とは、同一の質問に対する複数の答えを順位付けするという形で行われます。なので、前のステップでファイルに出力した答えを、望ましい順に上から下に並べ替えることになります。

コード
- 採点例

Rewardモデルを訓練する

では、 SFTモデルの生成した答えを採点するモデル (Rewardモデル)を作っていこうと思います。 Rewardモデルは、アーキテクチャ的にはほぼ SFTモデルと同じです。唯一の違いは、最終層がスカラーを出力するということだけです。なので、 SFTモデルの最終層を取り替えた後に fine-tuningすることになります (実装の簡略化のため、今回の実装例では最終層を取り除かず追加だけしています) 。

ということで、 Rewardモデルを訓練するために以下の手順を繰り返します。

(1) 1つの質問と、それに対応する順位付けされた答え達を取ってくる。
(2) Rewardモデルでそれぞれの答えを採点する。
(3) 採点された答えを 2つ取ってきて、 $\displaystyle -log(sigmoid(score_i - score_j))$ を計算する (ここで $\displaystyle i$ 番目の答えは $\displaystyle j$ 番目の答えよりも手動採点の順位が高いものとします)。
(4) 3の値を全ての答えのペアに対して計算し、その平均を損失としてモデルを訓練する。

では、以下のコードでモデルを訓練していこうと思います。

コード
- 実装例
- 実行例

InstructGPTモデルを訓練する。

では、最後に強化学習 (RL)を使って InstructGPTモデルを訓練します。 InstructGPTモデルは、 SFTモデルを以下の目的関数に対して fine-tuningすることにより作成されます (論文)。

$\displaystyle E_{(x, y) \sim D_{\pi^{RL}}}[ reward(x, y) - \beta log \left( \dfrac{\pi^{RL}(y|x)}{\pi^{SFT}(y|x)} \right) ]+\gamma E_{x \sim D_{pretrain}}[log(\pi^{RL}(x))]$

- $\displaystyle x$ : 質問
- $\displaystyle y$ : 答え
- $\displaystyle reward(x, y)$ : 質問 $\displaystyle x$ と答え $\displaystyle y$ に対して、 Rewardモデルの出力する点数
- $\displaystyle \pi^{RL}(y|x)$ : InstructGPTモデルが質問 $\displaystyle x$ に対して、答え $\displaystyle y$ を生成する確率
- $\displaystyle \pi^{SFT}(y|x)$ : SFTモデルが質問 $\displaystyle x$ に対して、答え $\displaystyle y$ を生成する確率
- $\displaystyle \beta$ : 重み付けのハイパーパラメーター
- $\displaystyle \gamma$ : 重み付けのハイパーパラメーター
- $\displaystyle E_{x \sim D_{pretrain}}$ [f(x)]: $\displaystyle x$ を一般 GPTモデル訓練用のデータから取ってきた時のf(x)の平均値

$\displaystyle [ reward(x, y) - \beta log (\pi^{RL}(y|x) / \pi^{SFT}(y|x))$ ] についてですが、これは Proximal Policy Optimization (PPO)というテクニックから来ています。軽く説明をすると、本来であれば 1項目の Rewardモデルの項だけで訓練したいです。ただ、それだけだとモデルのパラメーターが急激に変化しすぎてしまい、訓練が不安定になるかもしれません。なので、 2項目の $\displaystyle log (\pi^{RL}(y|x) / \pi^{SFT}(y|x))$ という KL-divergenceを加えることで、元の SFTモデルから距離が離れすぎないようにしています。また、実装例にも含まれていますが、目的関数だけではなく、トレーニングの方でも Clipped Surrogate Objectiveと呼ばれる手法が使われています。申し訳ないのですが、 PPOについての詳しい説明は、元論文やめんだこさんのブログなどを参照してくださると助かります。

最終項についてですが、これは通常の GPTモデル用の目的関数となっています。 InstructGPTモデルは元の GPTモデルに比べ、一般的な NLPタスクに対しての性能が悪くなってしまいます。なので、ある程度の汎化性能を残そうということらしいです。今回は実装の簡略化のため、 $\displaystyle \gamma=0$ を採用して、論文中でいうところの PPOモデルを実装しようと思います。 InstructGPTと呼ばれるモデルは $\displaystyle \gamma \neq 0$ を採用しています。