On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting

2022/06/01 by Tomasz Korbak, Korbak, Tomasz, Hady Elsahar +5 · 20 citations
Computer Science · Mathematics · #Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Multimodal Machine Learning Applications #Natural Language Processing Techniques #Topic Modeling #cs.CL #cs.LG #stat.ML

paper · pdf · doi:10.48550/arxiv.2206.00761

openalex publication_date 2022/06/01 · openalex created_date 2022/06/13 · arxiv created 2022/11/14 · arxiv updated 2022/11/15 · openalex updated_date 2026/07/28

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.

Cited by

• Whatever Remains Must Be True: Filtering Drives Reasoning in LLMs, Shaping Diversity
• KL-Regularized Reinforcement Learning is Designed to Mode Collapse
• Ranking-based Preference Optimization for Diffusion Models from Implicit User Feedback
• Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting
• ThinkBrake: Mitigating Overthinking in Tool Reasoning
• Reinforcement Learning Meets Large Language Models: A Survey of Advancements and Applications Across the LLM Lifecycle
• LATTS: Locally Adaptive Test-Time Scaling
• Reward-Shifted Speculative Sampling Is An Efficient Test-Time Weak-to-Strong Aligner
• MaPPO: Maximum a Posteriori Preference Optimization with Prior Knowledge
• BifrostRAG: Bridging Dual Knowledge Graphs for Multi-Hop Question Answering in Construction Safety
• Concept-TRAK: Understanding how diffusion models learn concepts through concept-level attribution
• Curriculum-Adapted Robust Reinforcement Learning for UAV Deconfliction in Adversarial Environments
• SPECS: Faster Test-Time Scaling through Speculative Drafts
• Differential Information Distribution: A Bayesian Perspective on Direct Preference Optimization
• Frictional Agent Alignment Framework: Slow Down and Don't Break Things
• Optimal Policy Minimum Bayesian Risk
• BLEUBERI: BLEU is a surprisingly effective reward for instruction following
• Efficient Reinforcement Finetuning via Adaptive Curriculum Learning
• Sample, Don't Search: Rethinking Test-Time Alignment for Language Models
• Online Difficulty Filtering for Reasoning Oriented Reinforcement Learning

Related

• Distributionally Robust Reinforcement Learning with Human Feedback
• Beyond Sparse Rewards: Enhancing Reinforcement Learning with Language Model Critique in Text Generation
• Fine-Tuning Language Models from Human Preferences
• RM-Bench: Benchmarking Reward Models of Language Models with Subtlety and Style
• Generalist Reward Models: Found Inside Large Language Models