Large language model inference ability enhancement method based on reward adaptive reinforcement learning exploration

Abstract

The invention discloses a large language model inference capability enhancement method based on reward adaptive reinforcement learning exploration. The method comprises the following steps: selecting a pre-trained large language model as a reference model; selecting a training data set; a forward KL divergence is adopted to replace a reverse KL divergence regularization item, a training objective function containing maximized rewards and entropies and minimizing the forward KL divergence is constructed, and out-of-distribution exploration is achieved; constructing a reward self-adaptive reference model distribution reweighting mechanism, performing dynamic reweighting on the reference model to obtain a reweighted reference model, and promoting self-adaptive exploration in the distribution; combining the re-weighted reference model with the training objective function to form a final objective function; training the model in combination with the training data set; in the application process, questions are input into the trained model, and answers are obtained. According to the method, the problem of insufficient exploration capability caused by reverse KL divergence limitation of an existing RLVR method can be solved.

Description

Technical Field

[0001] The present invention belongs to the field of large language model reasoning optimization, and in particular relates to a method for enhancing the reasoning capability of large language models based on reward-adaptive reinforcement learning exploration. Background Art

[0002] In the field of reasoning optimization of large language models (LLMs), reinforcement learning with verifiable rewards (RLVR) technology has become a core method for improving models' ability to solve complex problems (such as mathematical reasoning). Representative work is the DeepSeek-R1 model recently released by DeepSeek. RLVR uses the results of automatic verification of the correctness of the solution as a reward signal to guide the model strategy to converge to the high reward area, thereby significantly improving reasoning efficiency. The reasoning ability of a model is usually measured by the Pass@k indicator, that is, if at least one of the k sampled answers is correct, it is 1; otherwise, it is 0. The higher the Pass@k, the stronger the model's reasoning ability. A commonly used low-variance estimate of the Pass@k indicator is as follows:

[0003]

[0004] Here, P(x) represents the set of questions, n represents the number of sampled answers (n ≥ k), and c represents the number of correct answers. As demonstrated by models such as DeepSeek-R1, when the number of samples k is low, the RLVR-trained model is able to densely sample within the high-probability reasoning paths of the base model, rapidly improving the sampling efficiency of correct answers. Consequently, it significantly outperforms the pre-trained base model in metrics such as Pass@k.

[0005] However, existing RLVR methods suffer from a fundamental flaw: as the sampling budget increases, the advantages of RLVR models gradually weaken or even disappear. Specifically, when the model is allowed to generate more candidate solutions, that is, when the value of k increases, the performance of the RLVR-trained model on the Pass@k metric is surpassed by the base model. This phenomenon suggests that RLVR essentially improves response quality under small sample sizes only by adjusting the output distribution of the base model. This means that the probability mass is concentrated on the "familiar" reasoning paths of the base model, but it fails to give the model true exploration capabilities, resulting in its reasoning strategy being restricted to the inherent solution space of the base model.

[0006] In-depth analysis shows that this limitation stems from the fact that existing RLVR methods generally use the reverse Kullback-Leibler (KL) divergence regularization term D KL [π θ ||π ref ]:

[0007]

[0008] Here π ref Represents the pre-trained model, also known as the "base model" or "reference model", π θ Represents the model to be optimized. The closer the output probability distributions of the two models are, the smaller the KL divergence is. This regularization term prevents π θ and π ref The distance is too far and the model collapses. Based on the reverse KL divergence regularization, the current RLVR method usually maximizes the following objectives:

[0009]

[0010] Here r(x,y) is the reward function. The "mode-seeking behavior" of the reverse KL divergence will force the fine-tuned policy π θ Trapped in the base model π ref Using rigorous mathematical description, the reverse KL divergence constraint requires strategy π θ The support set (i.e. the area with non-zero probability density) is completely contained in the basic model π ref If π ref Assign zero probability to a certain reasoning path, then π θ This path must also be assigned a probability of zero (no matter how high its reward is), as shown in the following Figure 1 Specifically, through rigorous proof, we can get the following conclusions:

[0011] Lemma 1: The optimal solution to problem (1) has the following closed form:

[0012]

[0013] This limitation prevents the model from exploring the underlying model π ref High-reward regions outside the distribution's support set, such as novel reasoning steps or problem-solving ideas, occur even when these regions may contain superior solutions. Existing research, such as "Assessing Diversity Collapse in Reasoning" and "Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMS Beyond the Base Model?", has experimentally validated this phenomenon across multiple model families and reasoning domains, demonstrating its universality.

[0014] Although existing methods attempt to add a maximum entropy term H(π θ) to enhance exploration, that is, to maximize the following objective function:

[0015]

[0016] However, theoretical analysis shows that the inherent constraint of the reverse KL divergence mentioned above has not been broken, that is, the support set of the strategy is still strictly limited to the basic model. Specifically, through theoretical analysis, we can draw the following conclusions:

[0017] Lemma 2: The optimal solution to problem (2) has the following closed form:

[0018]

[0019] Therefore, the optimal strategy after introducing the entropy term can only reweight the probability by rewarding within the support set of the basic model, and cannot break through the limitations of the zero probability area.

[0020] In summary, existing RLVR methods, constrained by the inverse KL divergence, are trapped by a dependency on the underlying model's solution space. This results in the model being unable to discover new inference paths even with a high sampling budget, leading to stagnant improvements in inference capabilities. Breaking through the distributional limitations of the underlying model and achieving effective exploration has become a core issue currently underway in RLVR technology. Summary of the Invention

[0021] The present invention provides a method for enhancing the reasoning capability of large language models based on reward-adaptive reinforcement learning exploration, which solves the problem of insufficient exploration capability of existing RLVR methods due to the reverse KL divergence limitation, promotes breakthroughs in the distribution limitations of basic models, effectively solves complex problems that were originally difficult to handle, and improves the accuracy of answer generation.

[0022] First, the present invention theoretically proves the following conclusion: if the reverse KL divergence is not used and the forward KL divergence D is used instead KL [π ref ||π θ ]:

[0023]

[0024] It is possible to break through the limitations of the basic model and explore beyond the support set of the basic model. Specifically, if the forward KL divergence is used instead of the directional KL divergence, then problem (2) is transformed into:

[0025]

[0026] Under the constraint ∫π θ When (y|x)dy=1, the Lagrange form of problem (3) is

[0027]

[0028] By deduction, we can obtain the following theorem:

[0029] Theorem 1: The optimal solution to problem (4) It has the following closed form:

[0030]

[0031] Where g is given by π ref (y i |x) and r(x,y i ) and the function determined by the hyperparameters α, β, and λ. The optimal multiplier λ * By constraints Decision, where the sum traverses all the answers y sampled from the input question x i , when the maximum length of the answer is fixed, answer y i The types are limited.

[0032] In the above theorem, for π ref (y i |x)=0, that is, the answer y that the basic model cannot generate i , the trained model Still get a non-zero probability and reward r(x,y i ) is larger, The higher the probability of the assignment, the better the result, which is in line with the intuitive requirement. This shows that the forward KL divergence regularization breaks through the inference boundary constraint of the reverse KL divergence regularization.

[0033] Based on the above theoretical analysis, the present invention proposes to use forward KL divergence instead of reverse KL divergence for training.

[0034] A method for enhancing the reasoning capability of a large language model based on reward-adaptive reinforcement learning exploration includes the following steps:

[0035] S1: Select a pre-trained large language model as a reference model; select a training dataset, where each data in the training dataset contains a question, a response, and an answer;

[0036] S2 uses forward KL divergence instead of reverse KL divergence regularization term to construct a training objective function that maximizes reward and entropy while minimizing forward KL divergence to achieve out-of-distribution exploration;

[0037] S3: Build a reward-adaptive reference model distribution reweighting mechanism to dynamically reweight the reference model, obtain the reweighted reference model, and promote adaptive exploration within the distribution;

[0038] S4, combining the reweighted reference model in step S3 with the training objective function in step S2 to form a final objective function;

[0039] S5, based on the final objective function, combined with the training data set to optimize the model π θ Perform training to obtain a trained model

[0040] S6, during the application process, the problem is input into the trained model Get the answer.

[0041] In step S2, the forward KL divergence D is used KL [π ref ||π θ ]Replace the reverse KL divergence regularization term D KL [π θ ||π ref ];

[0042] Among them, π ref represents the reference model, π θ Represents the model to be optimized; the closer the output probability distributions of the two models are, the smaller the KL divergence is.

[0043] By optimizing the objective function, the model is prompted to explore areas outside the distribution, so that the strategy can assign high probability to zero-probability or low-probability but high-reward areas in the basic model, thereby sampling answers outside the support set of the basic model.

[0044] In step S3, the reference model is dynamically reweighted to obtain a reweighted reference model, and the formula is:

[0045]

[0046] Where, is the reweighted reference model, x represents the question, y represents the answer to the question, and r(x,y) is the reward; φ(r) is the reweighting function, which adjusts the reference model π according to the reward r(x,y). ref The exponential of , Z is the normalization factor.

[0047] The reweighting function φ(r) is a monotonically increasing function with a value range between [0,1]. When the reward r(x,y) is higher, φ(r) is closer to 1, making Close to π ref , making full use of existing reasoning capabilities; when the reward r(x,y) is lower, φ(r) is closer to 0, making Tends to be evenly distributed to encourage exploration.

[0048] In step S4, the final objective function is:

[0049]

[0050] Where,

[0051] P(x) represents the distribution of the problem, E represents the expected probability of calculation, and H(π θ ) represents π θ The entropy of , α and β represent hyper parameters, both greater than 0.

[0052] In step S5, an on-policy training strategy is adopted. During the training process, the advantage value of each answer is calculated by grouping relative advantage estimation, and the forward KL divergence term is calculated using a low-variance estimation method to achieve effective reward adaptive exploration and improve the model's reasoning ability in both in-distribution and out-of-distribution areas.

[0053] The odds value of each response is calculated by group relative odds estimation, as follows:

[0054]

[0055] In the formula, mean represents the mean, std represents the standard deviation, and A i For each answer y i Normalized advantage value of r i =r(x,y i ), indicating the answer is y i Reward; for each question x, sample G answers through autoregression Then use the reward function r for each sample pair (x,y i ), calculate the reward r i =r(x,y i ).

[0056] Compared with the prior art, the present invention has the following beneficial effects:

[0057] 1. Breaking through the limitations of the basic model: With the help of positive KL divergence, the model can explore high-reward areas outside the distribution of the basic model, effectively solving problems that the basic model is difficult to handle, and breaking through the performance ceiling of the basic model.

[0058] 2. Adaptive exploration balance: The reward-adaptive reference model reweighting mechanism can dynamically adjust the balance between exploration and exploitation based on the reward signal, actively exploring in low-reward areas and fully utilizing existing capabilities in high-reward areas, thereby improving the targetedness and effectiveness of exploration.

[0059] 3. Significant performance improvement: Experimental results show that by applying this method to the Qwen2.5-3B and 7B models, the problem-solving performance of the model is significantly improved in mathematical reasoning benchmarks such as AIME2024, AIME2025, and MATH500, especially when dealing with difficult problems. BRIEF DESCRIPTION OF THE DRAWINGS

[0060] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the following briefly introduces the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0061] Figure 1 The strategy distribution obtained by optimizing the basic model after applying forward KL divergence and reverse KL divergence respectively.

[0062] Figure 2 Schematic diagram of reference model reweighting based on reward adaptation.

[0063] Figure 3 This is a flow chart of a method for enhancing the reasoning capability of a large language model based on reward-adaptive reinforcement learning exploration in an embodiment of the present invention.

[0064] Figure 4 Comparison of the performance of the basic model, the comparative method GRPO-RKL and the RAPO proposed in this invention on the AIME25 full dataset (left figure) and the AIME24 difficult problem subset (right figure).

[0065] Figure 5 The accuracy of the basic model and the RAPO proposed in this invention are compared on questions of different difficulty levels on the MATH500 dataset. DETAILED DESCRIPTION

[0066] The following will clearly and completely describe the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of the present invention.

[0067] It should be noted that, unless there is any conflict, the features in the following embodiments and implementations may be combined with each other.

[0068] like Figure 3 As shown, a method for enhancing the reasoning ability of a large language model based on reward adaptive reinforcement learning exploration includes the following steps:

[0069] Step 1: Model initialization, dataset and reward function preparation.

[0070] Select a pre-trained autoregressive large language model. The large language model here can be any open source model, such as Qwen, Llama series, etc., or a self-pre-trained large language model. Here we choose Qwen2.5-7B or Qwen2.5-3B model as an example, as the pre-trained reference model π ref At the same time, the current strategy π θ Initialized to π ref .

[0071] The training dataset can be a self-constructed dataset or a public mathematical problem dataset. Here, the SimpleRL-Zero dataset (SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild), which contains approximately 8,000 mathematical problems, is selected as an example. This dataset is a combination of the GSM8K and MATH datasets after preprocessing and is divided into easy, medium, and hard difficulty levels. Each data item consists of a question, an answer, and an answer. The answer includes the complete conclusion process, while the answer is the final result of the answer, usually a numerical value. During the training process, no supervised fine-tuning (SFT) is performed, and training starts directly from the base model.

[0072] The reward function r is based on rule matching rather than a trained reward model to prevent the model from learning shortcuts that maximize the reward function (reward hacking). For each question-answer pair (x, y), the reward is r(x, y): if the answer is correct, that is, the final answer is correct, then the reward is 1, otherwise it is 0.

[0073] Step 2: Sampling and reward calculation.

[0074] In order to increase online exploration, the present invention adopts the On-Policy training strategy. In each iteration of the training process, the problem set A set of problems where a batch is randomly sampled from Then, using the current policy π θ ,for For each question x in the , we sample G answers by autoregression Then use the reward function r to calculate the reward for each sample pair (x,y i ), calculate the reward r(x,yi Finally, the Group Relative Advantage (GRO) estimate is used to calculate the value of each answer y. i The normalized advantage value A i , the formula is:

[0075]

[0076] Among them, r i =r(x,y i ), mean represents the mean, and std represents the standard deviation.

[0077] Step 3: Reference model distribution is reweighted.

[0078] According to the reward r(x,y i ), using the set reweighting function φ(r), for each sample pair (x,y i ) of the reference model distribution π ref The present invention adopts a re-weighting function based on the tanh function. The reweighted formula is:

[0079]

[0080] in, Normalization is performed by summing the discrete probability outputs. When the reward is high, φ(r) approaches 1, making Close to π ref , making full use of existing reasoning capabilities; when the reward is low, φ(r) approaches 0, making Tends to be evenly distributed to encourage exploration, as shown in the following Figure 2 shown.

[0081] Step 4: Forward KL divergence calculation and policy update.

[0082] The present invention adopts forward KL divergence to replace the reverse KL divergence regularization term, constructs a training objective function that maximizes reward and entropy while minimizing forward KL divergence, and realizes out-of-distribution exploration.

[0083] like Figure 1 As shown in FIG, the forward KL divergence regularization proposed in the present invention can expand the support set range of the probability density function of the basic model, while the reverse KL divergence does not have such an advantage.

[0084] Using the low variance estimation method, based on the above θ The sampled single sample pair (x,y i ), estimate the forward KL divergence term

[0085]

[0086] Then, by maximizing the following objective function J of RAPO RAPO (θ), update policy π θ :

[0087]

[0088] Here g(θ) indicates that the gradient descent algorithm is used in the update process and the policy gradient is clipped. The formula is:

[0089]

[0090] Among them, θ old is the strategy before updating, and ∈ is the clipping parameter.

[0091] Repeat the above steps 2, 3, and 4 of sampling and reward calculation, reference model distribution reweighting, and strategy update until the preset number of training steps is reached.

[0092] Step 5: Model evaluation.

[0093] After training, the model was evaluated on mathematical reasoning benchmarks such as AIME2024, AIME2025, and MATH500. The AIME2024 dataset contains real questions from the 2024 American Invitational Mathematics Exam (AIME), covering multiple mathematical fields such as algebra, geometry, number theory, and combinatorics, and the questions are quite difficult. The AIME2025 dataset contains real questions from the 2025 AIME, which are also quite difficult. In these two datasets, those questions that the basic model cannot solve even with n=2048 samples are called "hard" problems, and these problems constitute the hard problem subset. Evaluation was performed on the hard problem subset and the full dataset. The MATH500 dataset consists of 500 high-quality mathematical problems, divided into five levels of difficulty, gradually increasing in difficulty from Level-1 to Level-5.

[0094] 5.1 Evaluation Metrics

[0095] (1) Pass@k indicator.

[0096] (2) Problem solving rate: The ratio of the number of problems successfully solved by the statistical model to the total number of problems, which directly reflects the problem-solving ability of the model.

[0097] 5.2 Inference Parameter Setting

[0098] Temperature parameter: set to 0.6 to control the randomness of the answers generated by the model.

[0099] The top-p parameter is set to 0.95, which means that only tokens with a cumulative probability of 95% are considered.

[0100] The maximum input length is set to 1024 tokens. The maximum output length is set to 8196 tokens.

[0101] 5.3 Experimental comparison method

[0102] Base Model: Directly use the Qwen2.5-7B or Qwen2.5-3B pre-trained model without any fine-tuning or reinforcement learning training as the baseline.

[0103] Traditional RLVR method (GRPO-RKL): reinforcement learning method GRPO with reverse KL divergence regularization.

[0104] 5.4 Two methods of the present invention:

[0105] RAPO-light: RAPO with no reference model distribution reweighting, i.e. the reweighting function φ(r) is always 1.

[0106] RAPO: The full version of RAPO.

[0107] Table 1

[0108]

[0109] 5.5 Experimental Results and Analysis

[0110] Table 1 reports the Pass@1024 performance of different methods on mathematical reasoning tasks. Here k=1024 is selected to test the reasoning ability limit of each method under sufficient sample size n=2048. As can be seen from Table 1, the RAPO method of the present invention achieves the highest reasoning accuracy on both models and multiple datasets, proving that the method has better reasoning ability after training. Specifically, Table 1 shows the following results: On the full dataset (Fulldataset): RAPO outperforms GRPO-RKL, especially on the 7B model. For example, on the AIME2024 dataset, the accuracy increased from 74.4% to 80.9% (relative gain 8.74%); on the AIME2025 dataset, the accuracy increased from 65.7% to 80.0% (relative gain 21.77%). Furthermore, RAPO significantly outperforms the baseline models: on the AIME2024 dataset, the relative gains for the 3B and 7B models are 2.32% and 4.12%, respectively; on the AIME25 dataset, the gains are as high as 17.67% (3B) and 23.84% (7B). On the difficult problems (Hard subset) that the baseline models cannot solve, RAPO still demonstrates a strong advantage, even on problems that the baseline models cannot solve after 2048 attempts. On AIME24 (3B model), RAPO performs on par with GRPO-RKL, but improves by 42.29% on AIME25. On the 7B model, RAPO improves on GRPO-RKL by 40% on AIME24 and 188.55% on AIME25. These results highlight RAPO's robustness across different model sizes and problem difficulties.

[0111] Figure 4 The Pass@k performance of the 7B model under different sampling times k is shown. The left figure shows the results of the complete AIME25 dataset: when k is small, various RLVR methods including RAPO are better than the basic model, but as k increases, the Pass@k score of GRPO-RKL is gradually surpassed by the basic model, confirming its limitation of being unable to break through the capabilities of the basic model. In contrast, the RAPO method of the present invention can still maintain stability (RAPO) or continue to improve (RAPO-light) when k increases, proving its effectiveness in surpassing the baseline performance. The right figure focuses on the difficult (Hard) problem of AIME24. Although RAPO and GRPO-RKL have similar Pass@k when k values ​​are small, RAPO's improvement trend is steeper as the number of sampling increases. When k = 1024, the performance is close to twice that of GRPO-RKL, which is consistent with the results in Table 1.

[0112] Figure 5The performance comparison of RAPO and the base model on the MATH500 dataset is presented. RAPO outperforms the base model on problems of varying difficulty levels (from 1 to 5), with the advantage becoming increasingly significant as the problem difficulty increases. This demonstrates that RAPO can sample from high-reward regions with near-zero probability in the base model, achieving its goal of breaking the upper limit of the base model's reasoning capabilities.

[0113] The embodiments described above provide a detailed description of the technical solutions and beneficial effects of the present invention. It should be understood that the above are only specific embodiments of the present invention and are not intended to limit the present invention. Any modifications, supplements and equivalent substitutions made within the scope of the principles of the present invention should be included in the scope of protection of the present invention.

Claims

1. A method for enhancing the reasoning ability of large language models based on reward-adaptive reinforcement learning exploration, characterized in that: The following steps are involved: S1: Select a pre-trained large language model as a reference model; select a training dataset, where each data in the training dataset contains a question, a response, and an answer; S2 uses forward KL divergence instead of reverse KL divergence regularization term to construct a training objective function that maximizes reward and entropy while minimizing forward KL divergence to achieve out-of-distribution exploration; S3: Build a reward-adaptive reference model distribution reweighting mechanism to dynamically reweight the reference model, obtain the reweighted reference model, and promote adaptive exploration within the distribution; S4, combining the reweighted reference model in step S3 with the training objective function in step S2 to form a final objective function; S5, based on the final objective function, combined with the training data set to optimize the model π θ Perform training to obtain a trained model S6, during the application process, the problem is input into the trained model Get the answer.

2. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 1 is characterized in that: In step S2, forward KL divergence is used D KL [π ref ||π θ ]Replace the reverse KL divergence regularization term D KL [π θ ||π ref ]; Among them, π ref represents the reference model, π θ Represents the model to be optimized; the closer the output probability distributions of the two models are, the smaller the KL divergence is.

3. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 2 is characterized in that: In step S3, the reference model is dynamically reweighted to obtain a reweighted reference model, and the formula is: Where, is the reweighted reference model, x represents the question, y represents the answer to the question, and r(x,y) is the reward; φ(r) is the reweighting function, which adjusts the reference model π according to the reward r(x,y). ref The exponential of , Z is the normalization factor.

4. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 3 is characterized in that: The reweighting function φ(r) is a monotonically increasing function with a value range between [0,1].

5. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 3 is characterized in that: When the reward r(x,y) is higher, φ(r) is closer to 1, making Close to π ref , making full use of existing reasoning capabilities; when the reward r(x,y) is lower, φ(r) is closer to 0, making Tends to be evenly distributed to encourage exploration.

6. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 5 is characterized in that: In step S4, the final objective function is: In the formula, P(x) represents the distribution of the problem, E represents the expected probability of calculation, and H(π θ ) represents π θ The entropy of , α and β represent hyper parameters, both greater than 0.

7. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 1 is characterized in that: In step S5, an on-policy training strategy is adopted. During the training process, the advantage value of each answer is calculated by grouping relative advantage estimation, and the forward KL divergence term is calculated using a low-variance estimation method to achieve effective reward adaptive exploration and improve the model's reasoning ability in both in-distribution and out-of-distribution areas.

8. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 7 is characterized in that: The odds value of each response is calculated by group relative odds estimation, the formula is as follows: In the formula, mean represents the mean, std represents the standard deviation, and A i For each answer y i Normalized advantage value of r i =r(x,y i ), indicating the answer is y i Reward; for each question x, sample G answers through autoregression Then use the reward function r for each sample pair (x,y i ), calculate the reward r i =r(x,y i ).

Images