Improving Large Language Models via Fine-grained Reinforcement Learning with Minimum Editing Constraint (2024)

Abstract: Reinforcement learning (RL) has been widely used in training large language models (LLMs) for preventing unexpected outputs, e.g., reducing harmfulness and errors. However, existing RL methods mainly adopt instance-level reward, which cannot provide fine-grained supervision for complex reasoning tasks. As a result, the RL training cannot be fully aware of the specific part or step that actually leads to the incorrectness in model response. To address it, we propose a new RL method named RLMEC that incorporates a generative model as the reward model, which is trained by the erroneous solution rewriting task under the minimum editing constraint, which can produce token-level supervision for RL training. Based 0on the generative reward model, we design the token-level RL objective for training and an imitation-based regularization for stabilizing RL process. And these two objectives focus on the revision of the key tokens for the erroneous solution, reducing the effect of other unimportant tokens. Experiment results on 8 tasks have demonstrated the effectiveness of our approach. Our code and data will be publicly released.

## introduction owing to unsupervised pre-training on large-scale text corpora, large language models (llms) have shown remarkable performance on various text generation tasks @xcite , such as question answering, summarization and translation @xcite . to further improve the task solving capacity, researchers @xcite have proposed supervised fine-tuning (sft) and reinforcement learning (rl) methods, which can better adapt llms to specific domains or downstream tasks after pre-training. typically, sft methods @xcite incorporate annotated input-output pairs (e.g., question and solution, instruction and response) to train the llm for learning the sequenceto-sequence pattern; rl methods @xcite adopt a reward model to measure the quality of the generated outputs from the llm, and then guide its training for maximizing and minimizing the expectation of generating high-quality and low-quality ones, respectively. as rl methods are capable of directly reducing the probability of llms for producing unexpected outputs, they have been widely used in optimizing llms towards better human alignment (e.g., reducing harmfulness) and stronger ability (e.g., reducing errors @xcite ). generally, rl methods first train a discrimination model for distinguishing desirable and undesirable outputs. then, the model is used to produce the reward scores for the sampled outputs from the llm, and the llm would be trained by encouraging and punishing the generation of high-score and low-score ones accordingly. despite the success, as existing rl methods mostly utilize instance-level reward for each sam-pled output, it is often difficult to provide accurate fine-grained supervision on complex reasoning tasks (e.g., mathematical reasoning). concretely, given a complex task, the sampled outputs from the llm tend to be highly similar in surface expression, only with key differences in few specific words or steps @xcite that determine the correctness. we argue that instance-level rl approaches @xcite @xcite have two major limitations. first, as the unimportant parts would often occupy a large amount of supervision signals, instance-level rewards can not accurately emphasize the more important evidence related to correctness, leading to inefficient or redundant supervision. second, the paired correct and incorrect outputs may share the overlapping content but receive opposite optimization goals, which may lead to the optimization conflict issue on such overlapped content, making it still infeasible to provide accurate fine-grained supervisions. to address these issues, in this paper, we propose a novel method, reinforcement learning with minimum editing constraint (rlmec), to improve the training of llms by fine-grained supervision signals. our approach is inspired by the homework correction process of professional teachers, in which she/he first identifies the incorrect parts and then provides necessary revisions or comments accordingly. following such an idea, we train a generative reward model by an erroneous solution rewriting task under the constraint of minimum editing distance. the reward model plays a similar role to teachers by producing fine-grained supervision, i.e., token-level quality assessment scores. instead of using a new demonstration as positive, our reward model tries to correct the output with minimum edits. specially, we utilize the specially trained reward model to produce the token probabilities for computing the token-level rewards, and optimize the llm using the proximal policy optimization method (ppo) @xcite . in this way, by contrasting the original and corrected outputs, the llm would be instructed more informatively, thus becoming aware of the correct way to generate the response. figure 1 illustrates the comparison of the vanilla ppo and our proposed rlmec approach. the major novelty of this paper lies in the incorporation of a generative reward models with minimum editing constraint for rl training of llms. llms for reasoning. previous work utilizes two types of methods (i.e., prompting and training) to enhance the reasoning ability of llms. for the prompting methods, chain-of-thought (cot) @xcite guides llms to generate the intermediate reasoning steps before generating the final answer. based on cot, previous work decomposes the problem into several simple sub-problems @xcite , uti-lizes the external tools to help llms @xcite @xcite , designs the specific agents to perform reasoning @xcite , or post-process the generated response @xcite . besides, existing work also guides llms to perform reasoning in the specific structure, e.g., tree @xcite or graph @xcite . for the training methods, previous work @xcite has leveraged domainspecific data to fine-tune the llms. because of the limitation of the training data, the data generated by teacher model (e.g., gpt-4) is utilized to augment the training data @xcite @xcite @xcite . in this work, we aim to train the llms via fine-grained rl to improve their reasoning ability. ## preliminary in this work, we focus on improving the performance of llms on complex reasoning tasks with reinforcement learning (rl) algorithm. typically, complex reasoning tasks require llms to perform step-by-step reasoning (e.g., chain-of-thought @xcite ) for each question, where llms progressively generate the solution for reaching the answer. in this process, llms are prone to make mistakes at the intermediate steps, which likely lead to totally wrong answer @xcite . our goal is to optimize a pre-trained llm using rl algorithm, to reduce its errors and improve the task performance. formally, we are given a collection of questionsolution pairs, denoted as @xmath0, 𝑠 𝑖 ⟩} 𝑛 𝑖=1 , where each question and solution are both composed by a sequence of tokens, denoted as {𝑡 0 , • • • , 𝑡 𝑚 }. then, we follow the proximal policy optimization (ppo) framework @xcite for rl, and make improvements about reward model and training loss. in ppo, the llm to be optimized is the policy model, and its original parameters would be copied to compose the reference model. during training, the reference model outputs the sampled solutions for the given question, denoted as ŝ, and then the policy model would learn from the feedback from a reward model, which produces the reward 𝑅 ŝ for the sampled output ŝ. based on it, the parameters of the policy model will be optimized to maximize the reward expectation of all the sampled outputs, and the target function is: ## approach in this section, we present our proposed rlmec, a new rl approach for improving llms on complex reasoning tasks. in rlmec, we train a generative reward model to produce token-level reward scores for the sampled outputs from the policy model (i.e., the llm), then optimize the policy model via rl based on the fine-grained rewards. figure 2 illustrates the overall framework of our rlmec. ## conclusion in this paper, we proposed rlmec, a new reinforcement learning framework with minimum editing constraint, to leverage fine-grained supervision signals to further improve the ability of llms. in our rlmec, we first trained the generative reward model via the erroneous solution rewriting task under the minimum editing constraint, with the help of a teacher llm. then, we leveraged it to produce token-level rewards, and devised the token-level rl objective and an imitation-based regularization for training our llm, which both focus on the revision of the key tokens leading to errors in the solution. experimental results on mathematical tasks and question-answering tasks have demonstrated the effectiveness of rlmec. as future work, we will consider implementing our rl method on more advanced llms to further improve their performance on complex reasoning tasks. besides, we will also evaluate the capacity of our approach on enhancing human alignment and reducing hallucination.