[論文レビュー] On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting

Tomasz Korbak, Hady Elsahar+2 more|arXiv (Cornell University)|Jun 1, 2022

Topic Modeling被引用数 6

ひとこと要約

この論文は、災害的忘却を引き起こさずに言語モデルのファインチューニングを行うために、報酬最大化（RM）と分布マッチング（DM）を統合する。RMにおけるKL制御を再解釈し、DMの一種とみなすことにより、DMにベースラインを導入することで、感情、トピック、性別制御などの制御的テキスト生成タスクにおけるサンプル効率、安定性、制約満足度が向上する。

ABSTRACT

The availability of large pre-trained models is changing the landscape of Machine Learning research and practice, moving from a training-from-scratch to a fine-tuning paradigm. While in some applications the goal is to "nudge" the pre-trained distribution towards preferred outputs, in others it is to steer it towards a different distribution over the sample space. Two main paradigms have emerged to tackle this challenge: Reward Maximization (RM) and, more recently, Distribution Matching (DM). RM applies standard Reinforcement Learning (RL) techniques, such as Policy Gradients, to gradually increase the reward signal. DM prescribes to first make explicit the target distribution that the model is fine-tuned to approximate. Here we explore the theoretical connections between the two paradigms, and show that methods such as KL-control developed for RM can also be construed as belonging to DM. We further observe that while DM differs from RM, it can suffer from similar training difficulties, such as high gradient variance. We leverage connections between the two paradigms to import the concept of baseline into DM methods. We empirically validate the benefits of adding a baseline on an array of controllable language generation tasks such as constraining topic, sentiment, and gender distributions in texts sampled from a language model. We observe superior performance in terms of constraint satisfaction, stability and sample efficiency.

研究の動機と目的

• 新しい好みに適応する際、元の能力を失うという災害的忘却を回避する言語モデルファインチューニングの課題に対処する。
• 報酬最大化（RM）と分布マッチング（DM）という2つの主要なパラダイムを統合し、それぞれの長所を活かす。
• 強化学習のベースライン概念を導入することで、DM手法の高い勾配分散を克服し、学習の安定性を向上させる。
• シーケンスレベルの報酬では実現不可能な分布制約（例：性別やトピックのバランス）を可能にする。
• 複数の制御的生成タスクにおいて、サンプル効率、制約遵守、学習安定性の観点で優れた性能を示す。

提案手法

• RMにおけるKL制御を再解釈し、前向きKLダイバージェンスの最小化という観点から分布マッチング（DM）の一種と示し、理論的同等性を示す。
• 強化学習のポリシー勾配法にインspiredし、勾配分散を低減するため、分布マッチングアルゴリズムにベースラインを導入する。
• 感情、トピック、性別バランスなどの望ましい好みをエンコードするエネルギーに基づくモデル（EBM）を用いて、ターゲット分布を定義する。
• 分布ポリシー勾配（DPG）を用い、ターゲット分布とモデル出力分布の間の前向きKLダイバージェンスを最小化する。
• インジケータ関数とターゲット平均を用いて、複数の特徴（例：女性の言及が50％、スポーツ用語の含む）を組み合わせたハイブリッド制約を定義する。
• ベースライン減算を伴うポリシー勾配更新を組み合わせて使用し、DMに基づく最適化における学習の安定化を実現する。

実験結果

リサーチクエスチョン

• RQ1 報酬最大化におけるKL制御は、分布マッチングの一種と解釈可能か？
• RQ2 言語モデルファインチューニングの文脈において、RMとDMの理論的・実用的関係は何か？
• RQ3 DM手法にベースラインを導入することで、学習の安定性とサンプル効率が向上するか？
• RQ4 DMに基づく手法は、テキスト生成における複雑な分布制約（例：性別やトピックのバランス）を効果的に実装できるか？
• RQ5 提案されたベースライン強化DM手法は、標準的なRMおよびDMベースラインと比較して、制約遵守性と安定性の点で優れているか？

主な発見

• RMにおけるKL制御は、事前学習モデルとターゲット分布の間の前向きKLダイバージェンスを最小化するという観点から理論的にDMフレームワーク内に位置づけられる。
• 提案されたDMにおけるベースラインは、勾配分散を顕著に低減し、標準的なDM手法と比較してより安定した学習ダイナミクスを実現する。
• 本手法は、生成テキストにおける女性の言及を50％に維持するなど、優れた制約遵守性を達成するが、同時に文の自然さと多様性も維持する。
• 性別バランス＋トピックの含むなどのハイブリッド制約を含むタスクにおいて、ベースライン強化DMのモデルは、標準的なDMおよびRMベースラインを上回るサンプル効率と安定性を示す。
• 実験結果から、本手法は標準的なRLファインチューニングよりも元のモデルの分布的特性をより良く保持しており、災害的忘却を回避している。
• 本手法は、感情制御、トピック制御、性別バランスの生成といった多様な制御的生成タスクに一般化でき、すべての指標で一貫した改善を示す。