Exploring RL-based LLM Training for Formal Language Tasks with Programmed Rewards

TL;DR

This paper investigates the feasibility of using PPO for direct reinforcement learning (RL) from explicitly programmed reward signals, as opposed to indirect learning from human feedback via an intermediary reward model, and proposes a novel batch-entropy regularization term to aid exploration.

Abstract

Proximal Policy Optimization (PPO) is commonly used in Reinforcement Learning from Human Feedback to align large language models (LLMs) with downstream tasks. This paper investigates the feasibility of using PPO for direct reinforcement learning (RL) from explicitly programmed reward signals, as opposed to indirect learning from human feedback via an intermediary reward model. We focus on tasks expressed through formal languages, such as mathematics and programming, where explicit reward functions can be programmed to automatically assess the quality of generated outputs. We apply this approach to a sentiment alignment task, a simple arithmetic task, and a more complex game synthesis task. The sentiment alignment task replicates prior research and serves to validate our experimental setup. Our results show that pure RL-based training for the two formal language tasks is challenging, with success being limited even for the simple arithmetic task. We propose a novel batch-entropy regularization term to aid exploration, although training is not yet entirely stable. Our findings suggest that direct RL training of LLMs may be more suitable for relatively minor changes, such as alignment, than for learning new tasks altogether, even if an informative reward signal can be expressed programmatically.

Figures (9)

• Figure 1: Comparison between three different base models, being trained with PPO for the sentiment alignment task, using (top) DistillBERT as a trained reward model, or (bottom) VADER as programmatic reward function.
• Figure 2: Comparison between three different values for the entropy regularization coefficient $\beta_{ENT}$ on the sentiment alignment task, using DistillBERT as a reward model.
• Figure 3: Comparison between training with PPO using three different values for $\beta_{BENT}$ on the sentiment alignment task, using DistillBERT as a reward model.
• Figure 4: Comparison between training using PPO with and without the KL penalty on the sentiment alignment task, using DistillBERT as a reward model.
• Figure 5: Validation Loss (left) and Mean Reward (right) during supervised MLM training on the arithmetic task.