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 on 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.

heading: 1 Introduction
text: Owing to unsupervised pre-training on large-scale text corpora, large language models (LLMs) have shown remarkable performance on various text generation tasks (Zhao et al., 2023a; Google, 2023), such as question answering, summarization and translation (OpenAI, 2023). To further improve the task solving capacity, researchers (Touvron et al., 2023; Bai et al., 2023) 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 (Ouyang et al., 2022; Longpre et al., 2023) incorporate annotated input-output
∗ Equal contribution. † Corresponding author.
Red and green background colors denote negative and
positive rewards, respectively.
pairs (e.g., question and solution, instruction and response) to train the LLM for learning the sequenceto-sequence pattern; RL methods (Schulman et al., 2017; Christiano et al., 2017) 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 (Luo et al., 2023; Wang et al., 2023b)). 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-
5694
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 (Yuan et al., 2023) that determine the correctness. We argue that instance-level RL approaches (Ouyang et al., 2022; Christiano et al., 2017; Zheng et al., 2023b) 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) (Schulman et al., 2017). 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.
Table 1 presents the major differences between our method and previous work. To evaluate the effectiveness of our methods, we conduct the experiment on two types of complex reasoning tasks, i.e., question answering (Aggarwal et al., 2021; Mihaylov et al., 2018a) and mathematical reasoning (Cobbe et al., 2021; Hendrycks et al., 2021c). In these evaluation tasks, our RLMEC mostly outperforms other competitive SFT and RL methods, based on 7B and 13B LLMs. Moreover, our analysis experiments also show that our method is able to stabilize the RL training process and reduce the erroneous steps in the sampled outputs of LLMs.

heading: 2 Related Work
text: Reinforcement Learning for LLMs. With the development of the LLMs, reinforcement learning (RL) (Christiano et al., 2017; Ziegler et al., 2019) is widely utilized to further improve the ability of LLMs. Proximal Policy Optimization (PPO) (Schulman et al., 2017) is the traditional algorithm to employ RL. To provide fine-grained supervision signals, previous work (Mnih et al., 2016; Zheng et al., 2023b) utilizes the critic model to calculate the reward of the current stage. Because of the instability of the training procedure of reinforcement learning, recent work (Rafailov et al., 2023; Liu et al., 2023a; Lu et al., 2022; Zhao et al., 2023c) has utilized supervised-finetuning (SFT) to simulate the RL procedure. These methods fuse the quality of the responses into the supervision signals. Moreover, existing work (Uesato et al., 2022; Luo et al., 2023; Wang et al., 2023b,a; Yang et al., 2023) has found that process-supervision signals can better guide the training process of LLMs. Besides, other methods (Swamy et al., 2024; Chen et al., 2024) improve the ability of LLMs during self-play procedure. In this work, we proposed a new RL framework with generative reward model to directly provide the fine-grained supervisions, which enable to focus on few key tokens.
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) (Wei et al., 2022; Kojima et al., 2022) 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 (Dua et al., 2022), uti-
lizes the external tools to help LLMs (Gao et al., 2022; Yao et al., 2022; Schick et al., 2023; Chen et al., 2023), designs the specific agents to perform reasoning (Yin et al., 2023; Du et al., 2023), or post-process the generated response (Madaan et al., 2023; Wang et al., 2022). Besides, existing work also guides LLMs to perform reasoning in the specific structure, e.g., tree (Yao et al., 2023; Ding et al., 2023) or graph (Besta et al., 2023). For the training methods, previous work (Lewkowycz et al., 2022; Zhao et al., 2022) 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 (Yue et al., 2023; Yu et al., 2023; Gou et al., 2023; Zhao et al., 2023b; Zhou et al., 2024). In this work, we aim to train the LLMs via fine-grained RL to improve their reasoning ability.

heading: 3 Preliminary
text: 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 (Wei et al., 2022; Kojima et al., 2022)) 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 (Bang et al., 2023; Zhang et al., 2023). 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 asD = {⟨𝑞𝑖 , 𝑠𝑖⟩}𝑛𝑖=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 (Schulman et al., 2017) 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:
J (𝜃) = 𝑛∑︁ 𝑖=1 𝑟 (𝑞𝑖 , 𝑠𝑖) × 𝑅𝑠𝑖 , 𝑟 (𝑞𝑖 , 𝑠𝑖) = 𝑃𝜃 (𝑠𝑖 |𝑞𝑖) 𝑃𝜃 ′ (𝑠𝑖 |𝑞𝑖) , (1)
where 𝑟 (𝑞𝑖 , 𝑠𝑖) is the coefficient of importance sampling, 𝜃 and 𝜃′ are the parameters of policy model and reference model, respectively.

heading: 4 Approach
text: 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.

heading: 4.1 Generative Reward Model Training
text: To provide fine-grained supervision for RL, we train a generative model based on the sequence-tosequence loss as the reward model. For a given task, the reward model aims to offer estimations for all the output tokens about their correctness. To achieve this, we design an erroneous solution rewriting task with the constraint of minimum editing distance to train the reward model, enabling it to focus on the key tokens that lead to the final wrong answer for punishing.
Erroneous Solution Rewriting. This task aims to correct the error tokens in the LLM generated solutions with minimum edits. Formally, given the question 𝑞, ground-truth solution 𝑠, and the generated solution 𝑠, we rewrite 𝑠 into a correct solution 𝑠. Specifically, we decompose it into two sub-tasks, i.e., error locating and solution rewriting. For error locating, the model requires to locate the first erroneous reasoning step in 𝑠, which would mislead the following steps into erroneous ones. Concretely, we split 𝑠 into a sequence of reasoning steps according to the full stop or question mark: 𝑠 = {𝑟0, 𝑟1, . . . , 𝑟𝑛}. Then, the reward model needs to find the first undesired reasoning step 𝑟𝑡 based on the given question and ground-truth solution:
𝑅𝑀 (𝑝𝐿 , 𝑞, 𝑠, 𝑠) → 𝑟𝑡 , (2)
where 𝑝𝐿 is the prompt to guide the model. Then, for solution rewriting, we leverage another prompt 𝑝𝑅 to guide the reward model that rewrites the erroneous steps after 𝑟𝑡 in 𝑠 into the correct 𝑠:
𝑅𝑀 (𝑝𝑅 , 𝑞, 𝑠, 𝑠, 𝑟𝑡 ) → 𝑠. (3)

heading: Gen Solution �
text: By training on the two tasks, the generative reward model would be able to rewrite erroneous solutions with the minimum editing constraint.

heading: Distillation with Minimum Editing Constraint.
text: To train the reward model for fulfilling the above two subtasks, we collect the data from a powerful teacher LLM (i.e., Claude 2 (Anthropic, 2023)) to distill the task knowledge for our reward model, while other models (e.g., GPT-4) or human annotators can be also applied. Concretely, we first sample the generated solutions from our LLM, and select the wrong ones to compose the erroneous solution set {𝑠}. Then, we feed the given question 𝑞, ground-truth solution 𝑠, and the generated erroneous solution 𝑠 into the teacher LLM, and add several annotated exemplars into the prompt, to guide the generation of the first error step 𝑟𝑡 and the correct rewritten solution 𝑠. Here, in-context exemplars are human-crafted high-quality instances, and the ones for solution rewriting strictly satisfy the minimum editing constraint with only very few revised tokens. Therefore, we can obtain highquality synthetic distilled data for the two subtasks. Finally, following Eq. (2) and Eq. (3), we prepare the inputs and outputs for the two subtasks, and merge them for training our reward model.

heading: 4.2 RL with Fine-grained Supervision
text: After training the generative reward model, we can leverage it to produce fine-grained supervision for the RL training of the policy model (i.e., our LLM). We obtain the token-level rewards based on the generated probabilities from the reward model, and design the token-level RL objective with the imitationbased regularization for training our LLM.
Token-level Reward Generation. After distillation, the generative reward model can rewrite the original solution to provide the correct one. Owing to the minimum editing constraint, the error tokens would receive lower probabilities because they should be replaced by other tokens, and the correct tokens would obtain higher probabilities. Therefore, we can utilize the token probabilities from the reward model to assign the token-level rewards. This is quite distinct from the conventional reward model (Ouyang et al., 2022) which only produces instance-level reward scores. Concretely, given the prompt 𝑝𝑅, question 𝑞, ground-truth solution 𝑠, and the sampled solution 𝑠 from our LLM, the token rewritten probabilities from the generative reward model are used as the reward scores for the tokens in 𝑠. To better indicate the token quality, we normalize the reward scores by subtracting
them from the median value of the probability (i.e., 0.5) and then clip the extreme values as:
𝑅𝑠,𝑡 𝑗 = CLIP(𝑃𝑅𝑀 ( 𝑡 𝑗 |𝑝𝑅 , 𝑞, 𝑠, 𝑠, 𝑡< 𝑗 ) − 0.5, 𝛼, 𝛽 ) , (4)
where 𝑃𝑅𝑀 (𝑡 𝑗 |𝑝𝑅, 𝑞, 𝑠, 𝑠, 𝑡< 𝑗) is the predicted probability of the correct token from the reward model for the 𝑗-th token in 𝑠, and 𝛼 and 𝛽 denote the minimum and maximum thresholds for the reward value. For implementation, we employ 𝛼 = −0.1 and 𝛽 = 0 for the negative samples while adopt 𝛼 = 0 and 𝛽 = 0.5 for the positive samples. In this way, for negative samples, the upper threshold 𝛽 = 0 would lead to zero reward scores for all the non-error tokens, making the policy model only focus on punishing the error tokens. Otherwise, for positive samples, the lower threshold 𝛼 = 0 would assign zero reward score to error tokens, enabling the policy to focus on learning the correct tokens.
Token-level RL Objective. Given the token-level reward scores, we perform RL training on the policy model to correct its behaviors to avoid making errors. As mentioned in Section 3, we incorporate the PPO framework for RL, and revise its loss function to incorporate token-level reward scores. Concretely, we aim to maximize the expectation that generates the desired correct tokens in the solution. Thus, the gradients to optimize the policy model is given as:
∇J𝑅𝐿 (𝜃) = 𝑛∑︁ 𝑖=1 ∑︁ 𝑡 𝑗 ∈𝑠𝑖 𝑟 (𝑞𝑖 , 𝑡 𝑗 )×𝑅𝑠𝑖 ,𝑡 𝑗 ×∇ log 𝑃𝜃 (𝑡 𝑗 |𝑞𝑖 , 𝑡< 𝑗 ),
(5)
where 𝜃 is the parameters of the policy model, 𝑃𝜃 (𝑡 𝑗 |𝑞𝑖 , 𝑡< 𝑗) is the predicted probability of the 𝑗-th token by the policy model, and 𝑟 (𝑞𝑖 , 𝑡 𝑗) is the coefficient of the importance sampling in PPO as:
𝑟 (𝑞𝑖 , 𝑡 𝑗 ) = 𝑃𝜃 (𝑡 𝑗 |𝑞𝑖 , 𝑡< 𝑗 ) 𝑃𝜃 ′ (𝑡 𝑗 |𝑞𝑖 , 𝑡< 𝑗 ) . (6)
Moreover, inspired by existing work (Schulman et al., 2017; Chen et al., 2019) that clips the gradients of RL, we design a simplified way that clips the coefficient of the gradient to reduce the variance of the reward and prevent the large difference between the policy and reference model: min ( 𝑟 (𝑞𝑖 , 𝑡 𝑗 ) × 𝑅𝑠𝑖 ,𝑡 𝑗 , CLIP(𝑟 (𝑞𝑖 , 𝑡 𝑗 ), 1 − 𝜀, 1 + 𝜀) × 𝑅𝑠𝑖 , 𝑗 ) ,
(7)
where 𝜀 is a hyperparameter that controls the upper and lower bounds for positive and negative reward scores, respectively.

heading: Methods NS RL TLS RM

heading: RLMEC ✔ ✔ ✔ GEN
text: Imitation-based Regularization. As the RL training process is prone to be unstable, we further design a regularization loss based on imitation learning. The policy model is trained to imitate the generation of the rewritten solution 𝑠𝑖, only based on the question 𝑞𝑖. To compute the regularization term, we sample the generated wrong outputs 𝑠 from the policy model, and utilize our generative reward model to rewrite it into a correct one 𝑠 for learning. As discussed before, the original solution 𝑠𝑖 may contain only few error tokens that lead to the wrong solution. Therefore, we consider focusing on these error tokens in 𝑠, and identify them for targeted learning. Specifically, we leverage the Levenshtein Distance algorithm (Levenshtein, 1965), an effective method to find the revised tokens in 𝑠, and employ the token-level weights to emphasize them. The Levenshtein Distance algorithm utilizes dynamic programming (DP) to calculate the edit distance between 𝑠 and 𝑠, and the replaced and added tokens are selected into the error token set T . Then, the token-level weight is computed as:
𝑤 𝑗 = { 𝛾, 𝑡 𝑗 ∈ T 𝜙 × 𝛾, 𝑡 𝑗 ∉ T , (8)
where 𝛾 denotes the weight for emphasized tokens in T , and 𝜙 is the penalty coefficient for unimportant tokens. By incorporating term-level weights, the gradients of the imitation regularization are:
∇L𝐼𝑅 (𝜃) = − 𝑛∑︁ 𝑖=1 ∑︁ 𝑡 𝑗 ∈𝑠𝑖 ∇ log 𝑃𝜃 (𝑡 𝑗 |𝑞𝑖 , 𝑡< 𝑗 ) × 𝑤 𝑗 . (9)
Finally, the policy model is optimized by both the RL objective and imitation-based regularization.

heading: 4.3 Summary and Discussion
text: Here, we present the summary of our approach and discuss its difference with existing methods.
Summary. We present the pseudo-code of RLMEC in Algorithm 1 to better demonstrate our approach. The procedure of RLMEC can be divided into two parts, i.e., generative reward model training and reinforcement learning with fine-grained supervision. For generative reward model training, we leverage a teacher model (i.e., Claude 2) to synthesize the examples for the error locating and solution rewriting subtasks, to compose the dataset for distilling our generative reward model the capability of erroneous solution rewriting. Then, for RL training, we first generate the rewards for all the tokens in the sampled solutions from the policy model using Eq. 4, where we set suitable thresholds 𝛼 and 𝛽 to control our model to focus on important tokens in the generated solutions. Based on the token-level reward, we perform RL training using the PPO framework with the optimization function Eq. (5), and we design the reward clip strategy using Eq. (7) to prevent extreme rewards and stabilize the training process. Besides, we also add the imitation-based regularization using Eq. (9), to further help our policy model focus on learning key tokens.
Discussion. In Tabel 1, we present the difference between RLMEC and the existing work. Previous work mostly adopts the instance-level reward model, and only FIGA employs the tokenlevel supervision but does not utilize RL. Besides, there are several methods (e.g., WizardMath, MathShepherd) that leverage step-level reward to perform RL. As a comparison, our proposed RLMEC enables token-level supervision in the RL framework, and thus can benefit from more fine-grained supervision and focus on punishing error tokens during training procedure. A major novelty of our implementation is that we design the generative reward model trained by the erroneous solution rewriting task, to replace the conventional discriminative reward model, which can produce a rewritten probability of each token that can be naturally used as token-level supervision. Besides, by comparing with supervised fine-tuning methods (e.g., SFT and RFT), our approach can utilize the negative samples that will not be used by them, which extends the understanding of failed examples and fully utilizes the data.

heading: Task Train/Test Dataset Num. Data

heading: 5 Experiment

heading: 5.1 Experimental Settings
text: We simply introduce the experimental settings in this part. More details are shown in Appendix A.2.
Datasets. We employ mathematical tasks and question-answering tasks for evaluation. The specifics of each dataset are delineated in Table 2. Mathematical tasks include GSM8k (Cobbe et al., 2021), MATH (Hendrycks et al., 2021c), SVAMP (Patel et al., 2021) and the mathematical problems in MMLU (MM) (Hendrycks et al., 2021b,a). We adopt MathInstruct (Yue et al., 2023) as the training set and eliminate the code samples. Question-answering tasks contain ECQA (Aggarwal et al., 2021), QASC (Khot et al., 2020), OpenbookQA (Mihaylov et al., 2018b) and ARCEasy (Clark et al., 2018). We merge the training set of ECQA and QASC, and adopt the mixture as the training set in the experiment.
Baselines. For a more comprehensive assessment, we incorporate three categories of methods as baseline approaches. We conduct the SFT (Ouyang et al., 2022) and the Rejection sampling FineTuning (RFT) (Liu et al., 2023b; Yuan et al., 2023) as the baseline methods of supervised fine-tuning. To conduct more persuasive experiment, we also evaluate the variants of RFT, including adding the ground truth solution, rewritten solution from the teacher model, and rewritten solution from the generative reward model, named RFT w/ GT, RFT w/ TD, and RFT w/ RD, respectively. Besides, the representative methods of alignment without reinforcement learning, e.g., DPO (Rafailov et al., 2023), CoH (Liu et al., 2023a), and FIGA (Guo
et al., 2023) are conducted as the baseline. Moreover, We conduct the vanilla PPO (Schulman et al., 2017) and Actor-Critic version of PPO (PPO A2C) (Zheng et al., 2023b) as the baseline of RL methods. Additionally, we also report the performance of base LLMs, including LLaMA 2 (Touvron et al., 2023), Vicuna (Zheng et al., 2023a), and WizardLM (Xu et al., 2023).

heading: 5.2 Main Results
text: The evaluation results of RLMEC and the baseline methods are presented in Table 3.
First, RLMEC outperforms other baselines on the average accuracy of both scenarios. RLMEC demonstrates a strong capacity to further enhance the specific ability (e.g., reasoning ability) of LLMs. With the limited training data, compared with the previous methods (e.g., RFT, PPO), RLMEC leverages both positive and negative samples to provide fine-grained supervision signals, guiding LLMs to focus on the mistakes and correct them.
Second, RLMEC can prevent overfitting during domain adaption. Previous methods (e.g., SFT) utilize the data from the training set or generated by LLMs to fine-tune the LLMs which might cause overfitting. We can observe that the performance decreases after SFT on the unseen tasks (e.g., OBQA and SVAMP) of the 7B LLM. In contrast, the performance of LLMs on all of the unseen tasks is improved after RLMEC. The reason is that RLMEC makes LLMs focus on mistakes rather than correct components and utilize the clip mechanism to avoid overfitting.
Third, RLMEC can better leverage the generated response containing undesired components than other methods. Comparing the performance of RLMEC and DPO, we can observe that RLMEC enhance the reasoning ability of LLMs in both scenarios, but DPO only works on question-answer tasks. That is because RLMEC utilizes soft rewards to indicate positive or negative responses, while DPO
collects the positive-negative response pairs to train LLMs which can be regarded as utilizing the hard labels to identify the quality of generated responses. Given the quality of generated responses is difficult to assess, it is hard to collect response pairs in the challenge tasks (e.g., mathematical tasks). On mathematical tasks, the performance of DPO is even worse than the backbone LLM because of the low quality of the training data.
Finally, token-level supervision signals can further improve the performance of the policy model. The results of vanilla PPO, PPO A2C, and RLMEC present the importance of fine-grained supervision signals. Vanilla PPO utilizes instance-level signals to train the LLMs, which do not conform to reality because the generated response might contain both desired and undesired components. PPO A2C trains the critic model to provide fine-grained supervision signals which will increase the requirement of the computation resources. In RLMEC, the generative reward model is competent to implement the functionality of the reward model and the critic model in the PPO A2C at the same time.

heading: 5.3 Detailed Analysis
text: To further verify the effectiveness of RLMEC, we conduct the ablation study and analyze the model performance during the training process. Besides, we analyze the scaling of the generative reward model and present the case study of supervision signals and the model outputs in Appendix B and C.
Ablation Study. We evaluate the effectiveness of token-level supervision, reinforcement learning, and imitation-based regularization. Results are presented in Table 4. Given the results of the QA tasks (i.e., ECQA and ARC), we can observe that removing any of the modules will hurt the performance of the LLMs. In the mathematical tasks, without token-level supervision and reinforcement learning, LLMs overfit the training set, which brings
RLMEC RFT DPO
Training Steps
A cc
ur ac
y on
T ra
in in
g Se
t ( %
)
Figure 3: The performance of 7B LLMs on questionanswering tasks during different training strategies. To better present the difference, we smooth out the lines.
the improvement on the seen task (i.e., GSM8k) and hurts the performance on the unseen task (i.e., MM). The evaluation results demonstrate the ability of RLMEC to prevent overfitting and achieve the balance between seen tasks and unseen tasks. Besides, imitation-based regularization is also an important module in RLMEC. Without regularization, LLMs learn to generate correct responses only through token-level rewards. Because of the large search space, it is very difficult for LLMs to find the correct behavior in the challenge tasks. In the setting of removing imitation-based regularization, the decreasing performance on all of the tasks can verify our analysis.
Performance During Training Process. To comprehensively assess the performance of RLMEC, we conduct experiments on the accuracy of the training set during the training process. In Figure 3, we can observe that RLMEC can fit the training set more effectively and rapidly than other methods (i.e., RFT and DPO). Around 120 training steps, the policy model almost fits the training set through RLMEC. That is because our methods focus on the mistakes in the generated response and guide LLMs to correct these errors, which is more efficient. In contrast, RFT optimizes the whole tokens in the correct solution which might include many unimportant tokens, and DPO is overemphasized about the negative samples. These futures will decrease the speed of optimization and hurt the performance.

heading: 6 Conclusion
text: 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.

heading: Limitations
text: In this section, we discuss the limitations of our work. First, in this work, we focus on the complex reasoning tasks and only conduct experiments on the QA tasks and mathematical tasks. However, RLMEC can also be employed in other scenarios, e.g., human alignment and reducing hallucination, which has not been verified in this work. We leave it as the future work. Second, due to the limitation of computing resources, we only assess the performance of RLMEC on 7B and 13B LLMs, without the experiments on larger LLMs. Actually, by comparing the performance of baseline methods and RLMEC on 7B and 13B LLMs, we can observe the effectiveness of RLMEC. Third, our approach mainly focuses on enhancing LLMs on complex reasoning tasks, and does not consider the possible bias and ethic risks when using LLMs. It is also a promising direction that our RLMEC can be applied to, and we will investigate it in the future.

heading: Acknowledgement
text: This work was partially supported by National Natural Science Foundation of China under Grant No. 62222215, Beijing Natural Science Foundation under Grant No. L233008 and 4222027. Xin Zhao is the corresponding author.

heading: A Details for RLMEC

heading: A.1 Prompts for Generative Reward Model Training
text: We present the template of the prompt for teacher model distillation, and generative reward model training and inference in Table 5 and Table 6, respectively. In practice, the information (i.e., question 𝑞, Ground-Truth Solution 𝑠 and Generated Erroneous Solution 𝑠) should be filled into the corresponding curly brackets. For error locating task, to better guide teacher model and generative reward model to figure out the first undesired step, we utilize the index to format the ground-truth solution. The formatted solution is as follows,
[0] The First Reasoning Step 𝑟0 [1] The Second Reasoning Step 𝑟1 · · · [𝑛] The Last Reasoning Step 𝑟𝑛 For the generative reward model, the training instruction and inference prompt are similar. The target output of the training procedure (i.e., the bold sentence in the table) will be removed during inference.
A.2 Implementation Details for Experiments
Datasets. We employ mathematical tasks and question-answering tasks for evaluation. Success-
fully solving these tasks necessitates LLMs to possess domain-specific knowledge and engage in systematic, step-by-step reasoning to reach the ultimate answer. The specifics of each dataset are delineated in Table 2. • Mathematical tasks include GSM8k (Cobbe et al., 2021), MATH (Hendrycks et al., 2021c), SVAMP (Patel et al., 2021) and the mathematical problems in MMLU (MM) (Hendrycks et al., 2021b,a). We adopt MathInstruct (Yue et al., 2023) as the training set and eliminate the code samples. Given MathInstruct contains the training set of GSM8k and MATH, they are seen tasks for LLMs, while SVAMP and MM are unseen tasks. • Question-answering tasks contain ECQA (Aggarwal et al., 2021), QASC (Khot et al., 2020), OpenbookQA (Mihaylov et al., 2018b) and ARCEasy (Clark et al., 2018). We merge the training set of ECQA and QASC, and adopt the mixture as the training set in the experiment, Therefore, ECQA and QASC are seen tasks for LLMs, while OpenbookQA and ARC are unseen tasks for LLMs.
Baselines. For a more comprehensive assessment, we incorporate three categories of methods as baseline approaches. including supervised fine-tuning, alignment without reinforcement learning, and reinforcement learning. • Supervised Fine-tuning trains LLMs to imitate the human desired behavior. We conduct the SFT (Ouyang et al., 2022) and the Rejection sampling Fine-Tuning (RFT) (Liu et al., 2023b; Yuan et al., 2023) as the baseline methods.
• Alignment without Reinforcement Learning is the method to align LLMs to human preference and prevent instability in reinforcement learning. Representative methods, e.g., DPO (Rafailov et al., 2023), CoH (Liu et al., 2023a), and FIGA (Guo et al., 2023) are conducted as the baseline. • Reinforcement Learning is the traditional method to guide LLMs to explore the world and learn from external feedback. PPO (Schulman et al., 2017) is the classical algorithm to employ reinforcement learning. We conduct the vanilla PPO (Schulman et al., 2017) and Actor-Critic version of PPO (PPO A2C) (Zheng et al., 2023b) in the experiment.
Moreover, we also report the performance of base LLMs, including LLaMA 2 (Touvron et al., 2023), Vicuna (Zheng et al., 2023a), and WizardLM (Xu et al., 2023).
Hyper-Parameters Setting. In the experiment, we adopt Claude 2 (Anthropic, 2023) as the teacher model. For backbone LLMs, we utilize the mixture dataset of ECQA and QASC to fine-tune LLaMA 2 (Touvron et al., 2023) to obtain the domainadapted SFT backbone model in QA tasks, and adopt MAmmoTH (Yue et al., 2023) as the backbone model for mathematical tasks. The backbone LLMs of the policy model and the generative reward model are the same SFT LLMs. In the training procedure, we employ 5 × 10−6 as the learning rate for all tasks and train LLMs for 1 epoch. Besides, we set 128 and 768 as the batch size for
QA tasks and mathematical tasks. For the value of 𝜀, we leverage 0.3 and 0.4 for 7B model and 13B model, respectively. Because the LLMs have adapted to the corresponding domain after training, we adopt the 0-shot setting during evaluation.

heading: B Performance Analysis of RLMEC

heading: B.1 Analysis of Generative Reward Model.
text: The effectiveness of the generative reward model will influence the quality of the token-level rewards and the refined response. Thus, we present the comparison of the teacher model and the generative reward model on QA tasks in Figure 4. We can observe that both the teacher model and the generative reward model can significantly reduce the edit distance and even perform slightly better than the teacher model. That is because we utilize the two-stage prompting strategy to distillate knowledge from the teacher model and conduct the high-quality data to fine-tune the generative reward model. Through fine-tuning, it can adapt to the erroneous solution rewriting task well. Moreover, the teacher model and the generative reward model have shown similar performance on the accuracy of the refined responses, which verifies that the rewriting task can be easily learned by the LLMs with smaller parameters. Besides, given the performance of RFT w/ TD and RFT w/ RD in Table 3, we can observe that the higher accuracy of the refined responses will lead to higher performance in downstream tasks through simply supervised fine-
tuning.

heading: B.2 Scaling Analysis of Reward Model.
text: To explore the influence of the scale of the generative reward model, we conduct the experiment and present the results in Table 7. For both 7B and 13B LLMs, the rewriting model trained from the same backbone LLMs with the policy model performs better. The potential reason might be that the policy model and the rewriting model with the same backbone model will have a similar distribution. In this situation, the rewriting model can provide appropriate supervision signals and better guide the training process.

heading: B.3 Position of the First Error
text: We conduct experiments about the position of the first error in the generated response after training.
The results are shown in Figure 5, respectively. Rhe experiment on the position of the first error can verify the effectiveness of RLMEC. Compared with the backbone LLMs, the first error appears later after RLMEC. For example, after RLMEC, the number of problems where the first error occurs before the final answer 7 steps (i.e., the third column on the right) has increased, while the number of problems where the first error occurs more than 7 steps has decreased. The reason is that LLMs focus on the mistakes and learn to correct the early errors during RLMEC. In the ideal situation, all of the mistakes will be corrected through further training. In contrast, after training through other methods, the position of the first error is irregular, which means that these methods do not consider the mistakes in the generated response and guide LLMs to learn to generate the correct solution without purposiveness.

heading: C Case Study

heading: C.1 Analysis of the Supervision Signals
text: We present the case study about the reward from different methods in Table 8. To better express the difference, we do not employ the clip mechanism in the case study. From the results, we can observe that the reasoning step of the generated solution is
Algorithm 1: The RLMEC algorithm. Input :Training set D = {⟨𝑞𝑖 , 𝑠𝑖⟩}𝑛𝑖=1, the teacher
model (Claude 2), and the parameters of SFT model 𝜃𝑆𝐹𝑇 .
Output :A well trained policy model. Initialize the parameters of the generative reward
model: 𝜃𝐺𝑅𝑀 ← 𝜃𝑆𝐹𝑇 ; Initialize the parameters of the policy model: 𝜃 ← 𝜃𝑆𝐹𝑇 ;
// Generative Reward Model Training for each instance ⟨𝑞𝑖 , 𝑠𝑖⟩ in D do
The policy model generates 𝑠𝑖 based on 𝑞𝑖 ; if the data is sampled then
The teacher model locals the first error at 𝑟𝑖 using Eq. 2; The teacher model rewrites 𝑠𝑖 to obtain 𝑠𝑖 using Eq. 3;
Use 𝑞𝑖 , 𝑠𝑖 , 𝑠𝑖 , 𝑟𝑖 , and 𝑠𝑖 to construct D′; Leverage D′ to supervised-finetune the generative
reward model through Seq2Seq training paradigm;
// RL with Fine-grained Supervision for each instance ⟨𝑞𝑖 , 𝑠𝑖 , 𝑠𝑖 , 𝑟𝑖 , 𝑠𝑖⟩ in D′ do
Generate the token-level rewards using Eq. 4; Compute the reward J𝑅𝐿 (𝜃) using Eq. 5; Use Levenshtein Distance algorithm to compute
the token-level weight; Compute the loss of imitation-based
regularization L𝐼𝑅 (𝜃); Update 𝜃 through J𝑅𝐿 (𝜃) and L𝐼𝑅 (𝜃);
correct but the final answer is error. In PPO A2C, the reward will be calculated by the reward model and the critic model. The tokens generated earlier will receive a lower reward, which is contradictory to reality. That is because PPO A2C has assumed that the previous token will influence the last token. In this case, once the generated solution contains the wrong answer, the rewards of the previous tokens are likely lower than the last tokens. In contrast, we leverage the generative reward model to generate the reward in RLMEC. The reward of the current token is calculated based on the previous tokens. Therefore, the rewriting model in RLMEC can better indicate whether the token is correct and provide high-quality token-level supervision signals. Besides, for the outcome-supervised method (i.e., Vanilla PPO), the reward of each token is equivalent and is based on whether the generated solution is correct. This method cannot describe the correctness of the tokens in the generated response.

heading: C.2 Analysis of the Generated Responses
text: To further demonstrate the effectiveness of RLMEC, we present the case study about the performance of the LLMs trained by different methods in Table 9 and Table 10. In both cases, our proposed
RLMEC can help the LLMs to focus on the previous errors and correct the errors in the next time generation. Concretely, in the question-answering tasks, the keywords of the problem are “even if they get it”. After being trained through RLMEC, the LLMs can understand the meaning of the problem, figure the key point, and reach the correct answer. However, through other methods, the LLM is still unable to grasp the key works in the problem and generate the answer about the emotion of losing the job. Moreover, for mathematical problem, the LLM have made the mistake in calculating “12−15”. The LLM trained by baseline methods still make similar mistakes. This case has shown that it is difficult for the previous methods to generate the supervised signals which can directly indicate the mistakes in the generated content and guide the LLMs to correct the errors. In contrast, RLMEC leverages the generative reward model to provide the token-level supervised signals and guide the LLMs to focus on the mistakes. Therefore, through RLMEC, the LLMs can correct the previous errors and obtain the correct answer.
Problem Look at this series: 1.25, 7.5, ____, 270, 1620, ... What number should fill the blank?Answer Choices: (A) 50 (B) 65 (C) 45 (D) 55 (E) 40
Ground Truth Solution Let’s think about the multi-choice question. This is a simple multiplication series. Each number is 6 times greater than the previous number. The answer is C
Generated Solution Let’s think about the multi-choice question step by step. This is a simple multiplication series. Each number is 6 times greater than the previous number. The answer is B

heading: Future Work
text: We leave it as the future work. Second, due to the limitation of computing resources, we only assess the performance of RLMEC on 7B and 13B LLMs, without the experiments on larger LLMs. Actually, by comparing the performance of baseline methods and RLMEC on 7B and 13B LLMs, we can observe the effectiveness of RLMEC. Third, our approach mainly focuses on enhancing LLMs on complex reasoning tasks, and does not consider the possible bias and ethic risks when using LLMs. It is also a promising direction that our RLMEC can be applied to, and we will investigate it in the future