こんにちは、CCCMKホールディングス TECH LAB三浦です。
なんだか急に涼しくなってきました。過ごしやすくなってありがたいのですが、急な気温の変化に体が付いていけていないです・・・。こういう時期はちゃんと睡眠をとらないと、と意識するようになりました。
今回は以前から気になっていた、Reinforcement Learning from Human Feedback(RLHF)という強化学習の手法について調べてみました。
LMがより好ましいテキストを生成出来るようにする
大量のテキストデータによって自然なテキストを生成できることが出来るようになった言語モデル(Language model: LM)は、そのままでは人間にとって好ましいテキストを生成するとは限りません。たとえばインターネットにはポジティブな内容のテキストだけでなく、ネガティブな内容のテキストも存在します。
LMがトークンを選択しながらテキストを生成していく際、より人の評価が高いテキストを生成するようにチューニング出来れば解決しそうですが、これは人の"評価"を報酬と考え、トークンの選択を"行動"と捉えた強化学習の仕組みに落とし込むことが出来ます。
一方あまりにも人の評価を重視しすぎてしまうとテキストの文脈が崩壊してしまう可能性も考えられます。そのため、チューニング前のLMの自然なテキストを生成できる能力を残しつつ、その内容をより人の評価が高いものになるようにする方法が必要になります。
これを実現するのがReinforcement Learning from Human Feedback(RLHF)です。
RLHFの実行ステップ
RLHFは最終的に行われる強化学習のステップでは学習対象のモデルを含めると合計3つのモデルを動かすことになります。このモデルたちを用意するステップも含め、RLHFに必要になるステップについて簡単にまとめていきます。
まとめるにあたり、HuggingFaceのこちらのブログを参考にしました。
Pretraining language modelの用意
大規模テキストデータセットで事前学習済みのLMです。RLHFでチューニング対象になるものです。大規模データセットでの事前学習後、自前の小規模なデータで追加でチューニングを行ったものでも構いません。
Reward modelの用意
LMに入力したqueryと、それに対しLMが生成したanswerを結合して入力すると、それに対する評価をスカラー(数値)で返すモデルです。
Reward modelを学習させるためのデータセットをどうやって集めるのかが重要になります。方法は色々あるようですが、例えば同じqueryに対して複数のLMでanswerを生成させ、それを人が好ましいと思う順番に順位を振っていく。最終的にその値を正規化して評価値を計算し、Reward model学習用のデータセットを作る、といった手順があります。
スクラッチでReward modelを学習させる方法以外にも、感情分類用途で学習済みのモデルをそのまま使用する、といった方法も考えられます。たとえばテキストの内容がポジティブかネガティブかを判定するモデルがすでにあった場合、ポジティブであるスコアを評価値と見なせば、このモデルをポジティブな内容について高い評価を与えるReward modelとして使用することが出来ます。
強化学習ステップ
Pretraining language modelをReward modelの評価が高いテキストを生成出来るようにチューニングをしていくステップです。ただしReward modelの評価値のみを指標にすると、生成されるテキストそのものの質が落ちてしまう可能性があります。そこでReward modelの評価値を高めつつ、元のPretraining language modelのトークン生成分布とかけ離れないようなペナルティを設けて学習を行う、とった方法で解決します。
具体的には強化学習実行前にPretraining language modelをコピーしてInitial language modelとTuned language modelの2つのLMを作成します。学習中はInitial language modelのパラメータは固定し、Tuned language modelの方のパラメータを学習対象にします。学習ステップでは学習データセットからテキストを選択し、ランダムなテキスト長で末尾を切り取ります。このテキストをqueryとしてTuned language modelとInitial language modelに与えます。2つのモデルから同じqueryに対するトークンの分布が得られるので、分布の近さを表すKullback–Leibler (KL) divergenceを求めることが出来ます。このKL divergenceを用いて2つのモデルが"離れすぎない"ことを制約として与えることが出来ます。
まとめるとTuned language modelが生成したanswerをqueryと結合したテキストをReward modelに与えることで、そのテキストの評価値を得ることが出来、この評価値からInitial language modelとTuned language modelの間のKL divergenceに重みを付けた値を引き、分布が離れすぎるとペナルティになるようにした値を最終的な報酬としてRLHFでは強化学習を行います。
強化学習のテクニックとして、パラメータ更新後と更新前で出力される結果が大きく変動しないようにするProximal Policy Optimization (PPO)が使用されます。
HuggingFaceのライブラリ(TRL)で実装してみました。
RLHFの流れの全体像は見えてきたので、実際にRLHFを試してみたいと思いました。HuggingFaceではRLHFをサポートするライブラリTRL(Transformer Reinforcement Learning)があり、今回はこちらを使ってみました。
まずRLHFで使用するPretraining language modelとReward modelをどうやって用意するか考える必要があります。今回は以下の様にしました。
Pretraining language model
HuggingFaceで公開されている事前学習済みのGPT-2モデル"rinna/japanese-gpt2-small"を使用しました。
Reward model
Reward modelは事前学習済みのモデルをファインチューニングして用意することにしました。 まず事前学習済みのモデルとしてHuggingFaceで公開されている"xlm-roberta-base"を使用しました。
このモデルを感情分析用のデータセット"tyqiangz/multilingual-sentiments"の日本語データセットでファインチューニングしました。このデータセットは各テキストに対してポジティブ(0)、中立(1)、ネガティブ(2)のラベルが付与されています。
Reward modelのファインチューニング
ここからはコードを交えて流れをご紹介していきます。まずReward modelのファインチューニングです。ここはHuggingFace Transformersのチュートリアルの"Fine-tune a pretrained model"を参考にしました。
データセットをロードします。
from datasets import load_dataset dataset = load_dataset( "tyqiangz/multilingual-sentiments","japanese")
事前学習済みのモデルとトークナイザをロードします。
from transformers import AutoModelForSequenceClassification, AutoTokenizer model_name = "xlm-roberta-base" model = AutoModelForSequenceClassification.from_pretrained( model_name,num_labels=3) tokenizer = AutoTokenizer.from_pretrained(model_name)
事前処理としてトークン化の処理をデータセットに施します。
max_length = 128 def tokenizer_function(example): return tokenizer( example["text"],max_length=max_length,padding="max_length",truncation=True) train_ds = dataset["train"].shuffle().map(tokenizer_function,remove_columns=["text","source"]) valid_ds = dataset["validation"].shuffle().map(tokenizer_function,remove_columns=["text","source"])
accuracy計算用の関数を定義します。
import numpy as np import evaluate metric = evaluate.load("accuracy") def evaluate_function(eval_pred): logits, labels = eval_pred predictions = np.argmax(logits, axis=-1) return metric.compute(predictions=predictions, references=labels)
学習用の設定と学習の実行です。
from transformers import TrainingArguments, Trainer from transformers import EarlyStoppingCallback model_path = "/path/to/model_dir/" training_args = TrainingArguments( output_dir=model_path, evaluation_strategy="epoch", weight_decay=0.01, per_device_train_batch_size=64, per_device_eval_batch_size=64, save_strategy="epoch", load_best_model_at_end=True) trainer = Trainer( model=model, args=training_args, train_dataset=train_ds, eval_dataset=valid_ds, compute_metrics=evaluate_function, callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]) trainer.train()
学習後、検証用データで正解率84.8%程度になりました。
TRLによる強化学習の実装
では強化学習の部分の実装に入ります。実装は以下のTRLの実装例からたどることが出来る、"gpt2-sentiment.ipynb"を参考にしています。
最初に強化学習(PPO)の設定です。使用するPretraining language modelの指定も行っています。
from trl import PPOConfig ppo_config = PPOConfig( model_name="rinna/japanese-gpt2-small", learning_rate=1.41e-5, ppo_epochs=2)
データセットをロードします。
from datasets import load_dataset dataset = load_dataset( "tyqiangz/multilingual-sentiments","japanese")
データセットを学習用に加工します。LengthSamplerを用いて指定した最小値と最大値の範囲でランダムな長さにテキストを左から切り取って、RLHFにおけるqueryに加工しています。
from transformers import AutoTokenizer from trl.core import LengthSampler tokenizer = AutoTokenizer.from_pretrained(ppo_config.model_name) tokenizer.pad_token = tokenizer.eos_token input_min_text_length=10 input_max_text_length=25 def tokenizer_function(example): input_size = LengthSampler(input_min_text_length,input_max_text_length) input_ids = tokenizer.encode(example["text"])[:input_size()] query = tokenizer.decode(input_ids) return {"input_ids":input_ids, "query":query} train_ds = dataset["train"].shuffle().map(tokenizer_function) train_ds.set_format(type="torch")
ミニバッチデータをdict型に変換するCollatorを定義します。
def collator(data): return dict((key,[d[key] for d in data]) for key in data[0])
RLHFにおけるTuned language model(model)とInitial language model(ref_model)を用意します。
from trl import AutoModelForCausalLMWithValueHead model = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_config.model_name) ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_config.model_name)
PPOTrainerを用意します・・・が、一点注意点があります。ppo_trainer.is_distributed = Falseの箇所です。この箇所はTRLのexampleのnotebookには書いていなかったのですが、これが無いと私の環境では学習ステップ中にRuntimeError: Tensors must be CUDA and denseのエラーが発生し、学習を進めることが出来ませんでした。
trlのコードの中を見ていたところ、今実行している環境が分散学習の環境ではないのに何故か分散学習の環境と見なされてしまい、ppo_trainerのis_distributedがTrueにセットされていることで発生しているエラーのようでした。そこで手動でis_distributedの値をFalseに設定する、という方法を取ったところ、エラーの発生は回避できました。その場しのぎの対応ですし、正しい対応なのかもう少し調査が必要です。
from trl import PPOTrainer ppo_trainer = PPOTrainer( ppo_config, model, ref_model, tokenizer, dataset=train_ds, data_collator=collator) ppo_trainer.is_distributed = False #RuntimeError: Tensors must be CUDA and denseの回避のため
次にReward modelをロードします。先ほど学習したモデルはローカルパスに保存したため、そこから読み込んでpipeline化して使用します。
import torch from transformers import pipeline device = ppo_trainer.accelerator.device if ppo_trainer.accelerator.num_processes == 1: device = 0 if torch.cuda.is_available() else "cpu" # to avoid a `pipeline` bug sent_kwargs = {"return_all_scores":True,"function_to_apply":"none","batch_size":32} reward_pipeline = pipeline(task="text-classification",model="./reward_model",device=device)
テキストを生成する際の設定値です。ここは色々試してみたいのですが、また別の機会に・・・。
gen_kwargs = {"min_length":-1,"top_k":0.0,"do_sample":True,"top_p":1.0,"pad_token_id":tokenizer.eos_token_id}
output_min_text_length = 10
output_max_text_length = 25
output_length_sampler = LengthSampler(output_min_text_length,output_max_text_length)
最後が学習ループの実行です。テキストを生成し、Reward modelでrewardを計算し、学習ステップを実行します。
with mlflow.start_run(experiment_id=experiment_id) as run: for epoch, batch in enumerate(ppo_trainer.dataloader): query_tensor = batch["input_ids"] response_tensors = [] # テキストを生成する for query in query_tensor: gen_len = output_length_sampler() gen_kwargs["max_new_tokens"] = gen_len response = ppo_trainer.generate(query, **gen_kwargs) response_tensors.append(response.squeeze()[-gen_len:]) batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors] # queryと生成されたテキストを結合し、Reward modelに入力し、rewardを受け取る texts = [p + r for p, r in zip(batch["query"],batch["response"])] pipe_outputs = reward_pipeline(texts, **sent_kwargs) rewards = [torch.tensor(output[0]["score"]).cuda(0) for output in pipe_outputs] # 学習ステップを実行し、metricを取得 stats = ppo_trainer.step(query_tensor, response_tensors, rewards) metrics = { "ppo/mean_scores":stats["ppo/mean_scores"], "ppo/loss/total":stats["ppo/loss/total"], "objective/kl":stats["objective/kl"] } mlflow.log_metrics(metrics)
処理時間はかなりかかり、私の実行したGPU(V100)1基の環境では13時間ほどかかりました。
テキストを生成してみる
実は各種精度が上手く記録出来ておらず、モデルの収束状況が確認出来なかったのですが、ひとまずチューニング後のモデルとチューニング前のモデルで生成されるテキストに違いが出るのか試してみました。
まず、"今日の天気は"に続くテキストを生成させた場合ですが、
- チューニング前→今日の天気は今のところは曇りだから涼しい日だ
- チューニング後→今日の天気は本当に暖かくて気持ち良いお話しになりました
チューニング後の方が少しポジティブになってかも・・・という感じがします。
次に"この商品は"に続くテキストを生成してみました。
- チューニング前→この商品は積み荷車からトレーラーまである
- チューニング後→この商品は、皆様により安心してお使いいただけるよう誠心誠意努めております
少し雰囲気が丁寧な感じになったようにも思います。とりあえずこれでRLHFの一連の流れを手を動かして試すことが出来ました。(まだまだ未調査の箇所もたくさんありますが・・・)
まとめ
今回はReinforcement Learning from Human Feedback(RLHF)という、LMが生成するテキストをより人が好ましいものにチューニングするテクニックについて、調べてまとめてみました。また、理解をより深めるため、HuggingFaceのライブラリTRLを用いて実際に実装まで試してみました。
RLHFの全体像は掴めたものの、実装の部分はまだまだ調査不足だな・・・と感じています。今度はもう少し大きなサイズのLMを使い、詳細を調整しながら試してみたいと思います！
