lucidrains/PaLM-rlhf-pytorch

What it solves

このプロジェクトは、PaLM アーキテクチャを使用した Reinforcement Learning from Human Feedback (RLHF) の PyTorch 実装を提供します。人間が作成したフィードバックに基づいた報酬モデル (reward model) を学習させ、強化学習を用いてモデルの出力を最適化することで、開発者が事前学習済み言語モデルを人間の好みに合わせることを可能にします。

How it works

プロセスは主に 3 つの段階に分けられます：

1. Base Model Training: PaLM transformer を標準的な自己回帰モデルとして学習させます。

2. Reward Model Training: 与えられたシーケンスに対する人間の評価を予測する RewardModel を作成します (オプションで効率的な微調整のために LoRA を使用可能)。 n

3. RLHF Optimization: RLHFTrainer は、事前学習済み PaLM モデルと報酬モデルの両方を使用して、強化学習のエピソードを通じて言語モデルの生成能力を洗練させます。

Who it’s for

ChatGPT などのモデルで使用されている RLHF パイプラインを、PaLM アーキテクチャのオープンソース実装を使用して再現したい AI 研究者や開発者向けに設計されています。

Highlights

• LoRA Support: 過学習を防ぐために、Low-Rank Adaptation (LoRA) を使用して報酬モデルを微調整するオプションが含まれています。
• Flash Attention: より高速でメモリ効率の高いアテンションメカニズムを実現するために、統合されています。
• Complete Pipeline: ベースとなる transformer 訓練から報酬モデル構築、最終的な RLHF 学習まで、構造化されたワークフローを提供します。