MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actions

TL;DR

MA-RLHF introduces macro actions into reinforcement learning from human feedback, enabling higher-level language constructs to improve learning speed, stability, and performance across multiple NLP tasks without added computational costs.

Contribution

The paper presents a novel macro actions framework for RLHF that enhances credit assignment and training efficiency in large language models.

Findings

Achieves up to 30% performance improvement in text summarization and code generation.

Reaches parity with standard RLHF 1.7 to 2 times faster in training.

Demonstrates effectiveness across various NLP tasks and model sizes.

Abstract

Reinforcement learning from human feedback (RLHF) has demonstrated effectiveness in aligning large language models (LLMs) with human preferences. However, token-level RLHF suffers from the credit assignment problem over long sequences, where delayed rewards make it challenging for the model to discern which actions contributed to preferred outcomes. This hinders learning efficiency and slows convergence.In this paper, we propose MA-RLHF, a simple yet effective RLHF framework that incorporates macro actions -- sequences of tokens or higher-level language constructs -- into the learning process. By operating at higher level of abstraction, our approach reduces the temporal distance between actions and rewards, facilitating faster and more accurate credit assignment. This results in more stable policy gradient estimates and enhances learning efficiency within each episode, all without increasing computational complexity during training or inference. We validate our approach through extensive experiments across various model sizes and tasks, including text summarization, dialogue generation, question answering, and program synthesis. Our method achieves substantial performance improvements over standard RLHF, with performance gains of up to 30% in text summarization and code generation, 18% in dialogue, and 8% in question answering tasks. Notably, our approach reaches parity with vanilla RLHF 1.7 ~ 2 times faster in terms of training time and continues to outperform it with further training. We make our code and data publicly available at https://github.com/ernie-research/MA-RLHF.