TeaMs-RL: Teaching LLMs to Generate Better Instruction Datasets via Reinforcement Learning

TL;DR

TeaMs-RL introduces a reinforcement learning driven approach to generate the foundational instruction data used for fine tuning LLMs, eliminating the need for multi-stage RLHF. An instructor LLM is trained with an MDP and TRPO to produce diverse, complex prompts, with rewards given by a reviewer LLM. The generated instruction-response dataset is then used to fine-tune a pre-aligned LLM in a single SFT step, achieving competitive results on ARC and HellaSwag with much lower data and external queries, and enabling stronger privacy protection. Across benchmarks and tasks, TeaMs-RL demonstrates data efficiency, transferability across base models, and privacy benefits, challenging the necessity of traditional human-in-the-loop alignment pipelines.

Abstract

The development of Large Language Models (LLMs) often confronts challenges stemming from the heavy reliance on human annotators in the reinforcement learning with human feedback (RLHF) framework, or the frequent and costly external queries tied to the self-instruct paradigm. In this work, we pivot to Reinforcement Learning (RL) -- but with a twist. Diverging from the typical RLHF, which refines LLMs following instruction data training, we use RL to directly generate the foundational instruction dataset that alone suffices for fine-tuning. Our method, TeaMs-RL, uses a suite of textual operations and rules, prioritizing the diversification of training datasets. It facilitates the generation of high-quality data without excessive reliance on external advanced models, paving the way for a single fine-tuning step and negating the need for subsequent RLHF stages. Our findings highlight key advantages of our approach: reduced need for human involvement and fewer model queries (only 5.73% of the strong baseline's total), along with enhanced capabilities of LLMs in crafting and comprehending complex instructions compared to strong baselines, and substantially improved model privacy protection. Code is available at the link: https://github.com/SafeRL-Lab/TeaMs-RL

Figures (6)

• Figure 1: Comparative overview of LLM alignment techniques. Current methods (red shaded region) typically involve a two-phase process, starting with Supervised Fine-Tuning (SFT) of a pre-aligned LLM using a dataset of human-crafted instructions and corresponding responses (often sourced from an expert LLM like ChatGPT), leading to a post-SFT LLM. This is then fine-tuned using RLHF, where human feedback on preferences is incorporated, resulting in a post-RLHF LLM. In contrast, our TeaMs-RL method (green shaded region) employs a single-phase SFT approach, initially utilizing RL for teaching expert LLMs to generate high-quality instructions. We train an RL policy (instructor LLM) to create diverse instructions (with the diversity evaluated by a reviewer LLM as a reward signal). Once trained, the instructor LLM produces a set of actions to teach an expert LLM to generate high-quality instructions, and the instructions are leveraged to query the expert LLM to form the SFT instruction dataset. This approach capitalizes on the strengths of RL to enhance the complexity of instructions and consequently the value of responses from the expert LLM. Note that the expert LLMs are not involved in the training of the RL policy; we only use the instructor and the reviewer LLM for training the RL policy. The expert LLM is used to generate instructions and corresponding responses under the guidance of the trained RL policy.
• Figure 2: RL policy search for LLM instruction generation. A denotes actions (prompts), S denotes states (instructions), and the green line indicates the optimal instruction generation achieved during the policy search.
• Figure 3: Comparison of our method with WizardLM in terms of data diversity.
• Figure 4: Comparison of our method with WizardLM 7B on LM-Eval Benchmarks (the higher the value, the better the method's performance).
• Figure 5: Comparison of our method with WizardLM 7B on dataset size used for training LLMs (a) and querying count of advanced LLMs (b) (the lower the value, the better the method's performance).