はじめに
LLMを自作•改造する際に、RAGやファインチューニングだけでは思い通りの出力が得られないことはありませんか？
強化学習をうまく使えば、LLMの望ましい出力と望ましくない出力を学習させることができ、LLMの性能向上につながります。
今回はLLMの強化学習を行う方法をプログラム付きで説明します。
本記事で取り扱う内容
- 強化学習(PPO)のプログラム
- 強化学習に必要な報酬モデルの作成
前提:強化学習とは
強化学習は、エージェントが環境と相互作用し、報酬を最大化するように学習するプロセスです。エージェントは、どの行動が最も報酬をもたらすかを学習します。
今回のLLMの強化学習の場合、エージェント&行動&環境を以下の用に対応させることができます。
| 項目 | 説明 |
|---|---|
| エージェント | 学習するモデル(LLM) |
| 行動 | LLMの出力 |
| 環境 | 報酬モデル（学習するモデルの出力を評価して、スコア=報酬を算出する） |
学習用モデルと参照用のモデルの準備
まず強化学習させたいモデル(LLM)を定義します。
今回はTRLTrainerを使って強化学習を行っていくので、AutoModelForCausalLMWithValueHead.from_pretrained()を使ってモデルを定義します。
from trl import AutoModelForCausalLMWithValueHead
from transformers import pipeline, AutoTokenizer
model_name = 'モデル名'
model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
ここで学習するモデルmodelとは別に参照用のモデルref_modelを定義します。
ref_modelは重みの更新は行わず、強化学習時のKLダイバージェンスを計算するためにだけ用います。（後述します。）
参照用のモデルは、前述した報酬モデルとは異なるモデルです。
報酬モデルの準備
学習モデルの出力を評価し、スコアを出力するモデルである報酬モデルを準備します。
まず報酬モデルを作成するためのデータセットを準備します。
ここでは簡単のため、望ましいテキストを1、望ましくないテキストを0のラベルをつけてモデルを学習させます。
報酬モデルのデータセットは以下のように準備します。
X = ['望ましいテキスト','望ましくないテキスト',...]
y = [1,0,...]
次に報酬モデルの学習を行います。報酬モデルのベースはナイーブベイズを採用します。（報酬モデルのベースは、BERTやLLMなどの別モデルでも構いません、また予算が潤沢な方はGPT-4のAPIを報酬モデルとして扱うことも可能です。)
ナイーブベイズには、テキストをそのまま入力できないので、TfidfVectorizerで特徴ベクトルを抽出して学習させます。
ここではsklearn.pipeline.make_pipelineを使ってプログラムをまとめています。
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
reward_base_model = make_pipeline(TfidfVectorizer(), MultinomialNB())
reward_base_model.fit(X, y)
最後に報酬を正負の値で出力するようにします。
具体的には、望ましくないテキストが出力される場合は負の値に、望ましい出力の場合は正の値を出力するような報酬モデルreward_modelを作ります。
class RewardModel:
def __init__(self,reward_base_model,range_value=1):
self.reward_base_model = reward_base_model
self.range_value = range_value * 2 #報酬の値域を決める変数.-range_value<報酬の値<range_valueとなる。
def __call__(self,X):
output = self.reward_base_model.predict_proba(X)[:,1]
return (output-0.5) * self.range_value
reward_model = RewardModel(reward_base_model)
self.reward_base_model.predict_proba(X)[:,1]の[:,1]は、ラベルが1の時の値(つまり望ましい出力の確率)を取得するという意味です。
これで報酬モデルの準備完了です!
報酬モデルは強化学習にとって重要です。
今回は簡単に報酬モデルを作っていますが、強化学習で高い性能を出すためには報酬モデルの工夫が必要となります。
強化学習
強化学習はPPOというアルゴリズムを用いて行います。PPOは学習過程を安定させ、効率を向上させるためにポリシーの更新を小さく保つ(大きく変わりすぎない)というアイデアに基づくアルゴリズムです。
具体的にはモデルの準備で用意した比較対象のモデル(ref_model)と学習対象のモデル(model)の出力のKLダイバージェンスを制限します。
PPOのTrainer
まずPPO用のTrainerを準備します。Trainerはtrl.PPOTrainerを用います。
from trl import PPOTrainer, PPOConfig
config = PPOConfig(batch_size=1) #今回はメモリサイズ短縮のためバッチサイズ1で行います。
ppo_trainer = PPOTrainer(
config=config,
model=model,
ref_model=ref_model,
)
強化学習のデータセット
次に強化学習のためのデータセットを準備します。
データセットはdatasets.Datasetで定義します。
特徴量featuresにはトークン化されたinput_idsがあれば良いです。
サンプルコードはこちらです。
from datasets import Dataset
file_name = "xxx.txt" #改行区切りのテキストデータを用意
with open(file_name, "r", encoding="utf-8") as file:
text_content = file.read()
text_data = text_content.split("\n")
#トークン化
tokenized_text = tokenizer(text_data, padding=True, truncation=True)
#学習のためDataset化
dataset = Dataset.from_dict(tokenized_text)
#データセットの形式
Dataset({
features: ['input_ids', 'attention_mask'],
num_rows: 10,000
})
強化学習の実行
最後に強化学習を行います。
バッチサイズを2以上に設定する場合は、適宜プログラムを修正してください。
#modelのテキスト生成時のオプション
generation_kwargs = {
"top_k": 500,
"top_p": 0.95,
"min_length":50,
"max_length": 100,
"do_sample": True,
"pad_token_id": tokenizer.eos_token_id
}
iteration = 0
for batch in dataset:
query_tensors = [torch.tensor(batch['input_ids'])]
#テキストの生成
response = ppo_trainer.generate(query_tensors,**generation_kwargs)
response_text = tokenizer.decode(response[0])
#報酬の計算
rewards = [torch.tensor(reward_model([response_text]))]
#強化学習•重みの更新
stats = ppo_trainer.step(query_tensors, response, rewards)
iteration += 1
学習には大量のメモリを消費します。モデルを量子化するなどして、モデルのサイズを小さくしてから実行すると良いです。
statsにはPPOのイテレーションでの学習結果が格納されています。statsに含まれている情報一覧は以下です。(stats.keys()の結果です。)
dict_keys(['objective/kl', 'objective/kl_dist', 'objective/logprobs', 'objective/ref_logprobs', 'objective/kl_coef', 'objective/entropy', 'ppo/mean_non_score_reward', 'ppo/mean_scores', 'ppo/std_scores', 'tokens/queries_len_mean', 'tokens/queries_len_std', 'tokens/queries_dist', 'tokens/responses_len_mean', 'tokens/responses_len_std', 'tokens/responses_dist', 'ppo/loss/policy', 'ppo/loss/value', 'ppo/loss/total', 'ppo/policy/entropy', 'ppo/policy/approxkl', 'ppo/policy/policykl', 'ppo/policy/clipfrac', 'ppo/policy/advantages', 'ppo/policy/advantages_mean', 'ppo/policy/ratio', 'ppo/returns/mean', 'ppo/returns/var', 'ppo/val/vpred', 'ppo/val/error', 'ppo/val/clipfrac', 'ppo/val/mean', 'ppo/val/var', 'ppo/val/var_explained', 'ppo/learning_rate', 'time/ppo/forward_pass', 'time/ppo/compute_rewards', 'time/ppo/compute_advantages', 'time/ppo/optimize_step', 'time/ppo/calc_stats', 'time/ppo/total'])
おわりに
以上でLLMの強化学習の解説は以上になります。
LLM×強化学習の記事が少なかったので、少しでも参考になれば幸いです！！！