A Method for Enhancing Reasoning Ability of Large Language Models Based on Reward-Adaptive Reinforcement Learning

Abstract

This invention discloses a method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration. The method includes: selecting a pre-trained large language model as a reference model; selecting a training dataset; replacing the reverse KL divergence regularization term with positive KL divergence to construct a training objective function that maximizes reward and entropy while minimizing positive KL divergence, thus achieving out-of-distribution exploration; constructing a reward-adaptive reference model distribution reweighting mechanism to dynamically reweight the reference model, obtaining a reweighted reference model to drive in-distribution adaptive exploration; combining the reweighted reference model with the training objective function to form the final objective function; training the model using the training dataset; and during application, inputting the question into the trained model to obtain the answer. This invention can solve the problem of insufficient exploration ability caused by the limitation of reverse KL divergence in existing RLVR methods.

Description

Technical Field

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

[0002] In the field of optimizing the reasoning ability of Large Language Models (LLMs), Reinforcement Learning with Verifiable Rewards (RLVR) has become a core method for improving the model's ability to solve complex problems (such as mathematical reasoning). A representative work is the DeepSeek-R1 model recently released by DeepSeek. RLVR significantly improves reasoning efficiency by using the automatic verification of the correctness of the solution as a reward signal to guide the model's policy towards high-reward regions. The reasoning ability of a model is usually measured by the Pass@k metric, which is 1 if at least one answer is correct in k sampled responses, and 0 otherwise. The higher the Pass@k, the stronger the model's reasoning ability. A commonly used low-variance estimate of the Pass@k metric is as follows:

[0003]

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

[0005] However, existing RLVR methods have 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, i.e., 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 indicates that RLVR essentially only improves the response quality under small sample conditions by adjusting the output distribution of the base model, that is, it concentrates the probability quality on the existing "familiar" inference path of the base model, but fails to endow the model with true exploration capabilities, resulting in its inference 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 employ the inverse Kullback-Leibler (KL) divergence regularization term D. KL [π θ ||π ref ]:

[0007]

[0008] Here π ref The pre-trained model, also known as the "base model" or "reference model," is represented by π. θ This represents the model to be optimized. The closer the output probability distributions of the two models are, the smaller the KL divergence. This regularization term prevents π... θ With π ref Excessive distance can cause model collapse. Current RLVR methods, based on inverse KL divergence regularization, typically maximize the following objective:

[0009]

[0010] Here, r(x,y) is the reward function. The "mode-seeking behavior" property of the inverse KL divergence forces the fine-tuned policy π to... θ Trapped in the basic model π ref Within the high-probability region. A rigorous mathematical description is employed, namely, the inverse KL divergence constraint requires strategy π. θ The support set (i.e., the region with non-zero probability density) is completely contained within the base model π. ref Within the support set. If π ref If a certain reasoning path is assigned zero probability, then π θ This path must also be assigned a zero probability (regardless of its reward), as shown in the appendix. Figure 1 As shown. Specifically, through rigorous proof, the following conclusion can be obtained:

[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 support set, such as novel reasoning steps or problem-solving approaches, may contain better 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 verified this phenomenon in multiple model families and reasoning domains, demonstrating its universality.

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

[0015]

[0016] However, theoretical analysis shows that the inherent constraint of the aforementioned reverse KL divergence remains unbroken, meaning the support set of the strategy is still strictly constrained by the base model. Specifically, through theoretical analysis, the following conclusions can be drawn:

[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 through rewards within the support set of the basic model, and cannot break through the limitation of the zero probability region.

[0020] In summary, existing RLVR methods, constrained by the inverse KL divergence, are trapped in the dilemma of dependency on the solution space of the underlying model. This results in the model failing to discover new inference paths even with high sampling budgets, leading to stagnation in inference capability improvement. Overcoming the distributional limitations of the underlying model and achieving effective exploration has become a core problem that urgently needs to be solved in current RLVR technology. Summary of the Invention

[0021] This invention provides a method for enhancing the reasoning ability of large language models based on reward-adaptive reinforcement learning exploration. This method addresses the problem of insufficient exploration ability in existing RLVR methods due to the limitation of inverse KL divergence, 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, this invention theoretically proves the following conclusion: if the reverse KL divergence is not used, but instead the forward KL divergence D is used... KL [π ref ||π θ ]:

[0023]

[0024] This allows us to break through the limitations of the basic model and explore beyond its support set. Specifically, if we replace the directional KL divergence with the positive KL divergence, then problem (2) becomes:

[0025]

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

[0027]

[0028] Through derivation, the following theorem can be obtained:

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

[0030]

[0031] Where g is composed of π ref (y i |x) and r(x,y) i And the function determined by the hyperparameters α, β, and λ. Optimal multiplier λ * By constraints The decision here is that the summation traversal is based on all the answers y obtained by sampling the input question x. i Given a fixed maximum length for the answer, the answer y i The types are limited.

[0032] In the above theorem, for π ref (y i The case where |x)=0, i.e., the answer y that the basic model cannot generate. i The trained model It is still possible to obtain a non-zero probability, and the reward is r(x,y). i The larger the ) is, A higher probability of assignment aligns with intuitive requirements. This indicates that forward KL divergence regularization breaks through the inference boundary constraints of reverse KL divergence regularization.

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

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

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

[0036] S2 replaces the reverse KL divergence regularization term with the forward KL divergence to construct a training objective function that maximizes reward and entropy while minimizing the forward KL divergence, thereby achieving out-of-distribution exploration;

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

[0038] S4. Combine the reweighted reference model from step S3 with the training objective function from step S2 to form the final objective function.

[0039] S5, based on the final objective function, combines the training dataset with the model π to be optimized. θ Train the model to obtain a trained model.

[0040] S6, During application, input the problem into the trained model. The answer was obtained.

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

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

[0043] By optimizing the objective function, the model is encouraged to explore regions outside the distribution, enabling the policy to assign high probabilities to regions with zero or low probabilities but high rewards in the base model, thereby sampling answers outside the support set of the base model.

[0044] In step S3, the reference model is dynamically reweighted to obtain the reweighted reference model, as shown in the formula:

[0045]

[0046] In the formula, The reweighted reference model is defined as follows: x represents the question, y represents the corresponding answer, 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 index is Z, where Z is the normalization factor.

[0047] The reweighting function φ(r) is a monotonically increasing function, taking values ​​between [0,1]. As the reward r(x,y) increases, φ(r) approaches 1 more closely, making... Approximate to π ref Make full use of existing reasoning abilities; the lower the reward r(x,y), the closer φ(r) is to 0, making The distribution tends to be more uniform to encourage exploration.

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

[0049]

[0050] In the formula,

[0051] P(x) represents the distribution of the problem, E represents the expected probability, and H(π) represents the expected probability. θ ) represents π θ The entropy of , where α and β represent hyperparameters, are 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. At the same time, the positive KL divergence term is calculated using the 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 regions.

[0053] The advantage value for each response is calculated using grouped relative advantage estimation, as shown in the following formula:

[0054]

[0055] In the formula, mean represents the mean, std represents the standard deviation, and A i For each answer y i The normalized dominance value; r i =r(x,y) i ), indicating the answer y i The reward; for each question x, G answers are sampled using an autoregressive approach. Then apply the reward function r to each sample pair (x, y) i ), calculate 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: By leveraging positive KL divergence, the model can explore high-reward regions outside the distribution of the basic model, effectively solving problems that the basic model cannot handle and breaking through the performance limit of the basic model.

[0058] 2. Adaptive Exploration Balance: The reward-adaptive reference model reweighting mechanism can dynamically adjust the balance between exploration and utilization based on reward signals. It actively explores in low-reward areas and makes full use of existing capabilities in high-reward areas, thereby improving the targeting and effectiveness of exploration.

[0059] 3. Significantly Improved Performance: Experimental results show that applying this method to the Qwen2.5-3B and 7B models significantly improves the model's problem-solving performance in mathematical reasoning benchmark tests such as AIME2024, AIME2025, and MATH500, especially when dealing with difficult problems. Attached Figure Description

[0060] To more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings used in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0061] Figure 1 This is the optimized policy distribution obtained by applying forward KL divergence and reverse KL divergence to the basic model, respectively.

[0062] Figure 2 A schematic diagram of a reweighted reference model based on reward adaptation.

[0063] Figure 3 This is a flowchart of a method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning, according to an embodiment of the present invention.

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

[0065] Figure 5 The accuracy of the base model and the RAPO proposed in this invention is compared on different difficulty levels of questions on the MATH500 dataset. Detailed Implementation

[0066] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0067] It should be noted that, unless otherwise specified, the features in the following embodiments and implementation methods can 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] We select a pre-trained autoregressive large language model. This large language model can be any open-source model, such as the Qwen or Llama series, or a self-pre-trained large language model. Here, we choose Qwen2.5-7B or Qwen2.5-3B as examples, serving as the pre-training reference model π. ref At the same time, the current strategy π θ Initialize to π ref .

[0071] The training dataset can be a self-constructed dataset or a publicly available dataset of mathematical problems. Here, we choose the SimpleRL-Zero dataset (《SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild》), which contains approximately 8000 mathematical problems, as an example. This dataset is a combination of preprocessed GSM8K and MATH datasets, categorized into easy, medium, and hard difficulty levels. Each data point consists of a question, an answer, and a solution. The answer includes the complete problem-solving process, while the solution is the final result, usually a numerical value. During training, supervised fine-tuning (SFT) is not performed; training begins directly from the base model.

[0072] The reward function r is based on rule matching, rather than a trained reward model. This is 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, i.e., the final answer is correct, the reward is 1; otherwise, it is 0.

[0073] Step 2: Sampling and reward calculation.

[0074] To enhance online exploration, this invention employs an on-policy training strategy. In each iteration of the training process, the training begins with the problem set... A set of problems randomly sampled in a batch Then, using the current policy π θ ,for For each question x, G answers are sampled using an autoregressive approach. Then, using the reward function r, for each sample pair (x, y) i ), calculate the reward r(x,y)i Finally, the answer y is calculated for each response by estimating the group relative advantage. i Normalized dominance value A i The formula is:

[0075]

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

[0077] Step 3: Reweight the reference model distribution.

[0078] Based on the reward r(x,y) i Using the defined reweighting function φ(r), for each sample pair (x, y) i The reference model distribution π ref Reweighting is then performed. This invention employs a reweighting function based on the tanh function. The reweighting formula is obtained as follows:

[0079]

[0080] in, Normalization is performed by summing the discrete probability outputs. When the reward is high, φ(r) approaches 1, making... Approximate to π ref Make full use of existing reasoning abilities; when the reward is low, φ(r) approaches 0, making The distribution tends towards uniformity to encourage exploration, as shown in the appendix. Figure 2 As shown.

[0081] Step 4: Calculate the forward KL divergence and update the strategy.

[0082] This invention replaces the reverse KL divergence regularization term with the forward KL divergence, constructing a training objective function that maximizes reward and entropy while minimizing the forward KL divergence, thereby achieving out-of-distribution exploration.

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

[0084] Using a low-variance estimation method, based on the above from π θ A single sample pair (x, y) i Estimate the positive KL divergence term.

[0085]

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

[0087]

[0088] Here, g(θ) represents the gradient descent algorithm used during the update process, with the policy gradient being clipped, as shown in the formula:

[0089]

[0090] Where, θ old The strategy before the update is defined, and ∈ represents the pruning parameter.

[0091] Repeat steps 2, 3, and 4 above for sampling and reward calculation, reweighting of the reference model distribution, and policy 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 benchmark tests such as AIME2024, AIME2025, and MATH500. The AIME2024 dataset contains real problems from the 2024 American Invitational Mathematics Competition (AIME), covering multiple mathematical fields such as algebra, geometry, number theory, and combinatorics, and the problems are quite difficult. The AIME2025 dataset contains real problems from the 2025 AIME, which are also quite difficult. In both datasets, problems that the basic model could not solve in n = 2048 samples are called "hard" problems, and these problems constitute the hard subset. Evaluation was performed on the hard subset and the full dataset, respectively. The MATH500 dataset consists of 500 high-quality mathematical problems, with five difficulty levels, from Level-1 to Level-5, gradually increasing in difficulty.

[0094] 5.1 Evaluation Indicators

[0095] (1) Pass@k indicator.

[0096] (2) Problem Solving Rate: The proportion of problems successfully solved by the statistical model out of the total number of problems, which directly reflects the model's problem-solving ability.

[0097] 5.2 Inference Parameter Settings

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

[0099] The top-p parameter is set to 0.95, meaning 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 Methods

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

[0103] Traditional RLVR method (GRPO-RKL): GRPO is a reinforcement learning method that uses inverse KL divergence regularization.

[0104] 5.4 Two methods of the present invention:

[0105] RAPO-light: RAPO with reweighted distribution without reference model, 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 performance of different methods on the Pass@1024 mathematical reasoning task. Here, k=1024 was chosen to test the reasoning limits of each method with a sufficient sample size n=2048. As can be seen from Table 1, the RAPO method of this invention achieves the highest inference accuracy on both models and multiple datasets, demonstrating that the method possesses superior reasoning ability after training. Specifically, Table 1 shows the following results: On the full dataset: RAPO outperforms GRPO-RKL, especially on the 7B model. For example, on the AIME2024 dataset, accuracy increased from 74.4% to 80.9% (relative gain of 8.74%); on the AIME2025 dataset, accuracy increased from 65.7% to 80.0% (relative gain of 21.77%). Furthermore, RAPO significantly outperforms the baseline model: 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 even higher, reaching 17.67% (3B) and 23.84% (7B). Regarding the Hard subset, which the baseline model cannot solve: even for problems that the baseline model cannot solve after 2048 attempts, RAPO still demonstrates a strong advantage. On AIME24 (3B model), RAPO performs comparably to GRPO-RKL, but achieves a 42.29% improvement on AIME25; on the 7B model, RAPO outperforms GRPO-RKL by 40% and 188.55% on AIME24 and AIME25, respectively. These results highlight RAPO's robustness across different model sizes and problem difficulties.

[0111] Figure 4 The performance of the 7B model at different sampling counts k is shown. The left figure shows the results for the complete AIME25 dataset: when k is small, various RLVR methods, including RAPO, outperform the base model. However, as k increases, the Pass@k score of GRPO-RKL is gradually surpassed by the base model, confirming its limitation in overcoming the capabilities of the base model. In contrast, the RAPO method of this invention maintains its stable (RAPO) or continuously improving (RAPO-light) advantage as k increases, demonstrating its effectiveness in surpassing baseline performance. The right figure focuses on the Hard problem of AIME24. Although RAPO and GRPO-RKL have similar Pass@k values ​​when k is small, the improvement trend of RAPO is steeper as the sampling count increases. At k=1024, the performance is nearly twice that of GRPO-RKL, which is consistent with the results in Table 1.

[0112] Figure 5This study presents a performance comparison between RAPO and the base model on the MATH500 dataset. RAPO outperforms the base model across all difficulty levels (from level 1 to 5), with its advantage becoming increasingly significant as the problem difficulty increases. This demonstrates that RAPO can sample high-reward regions with near-zero probability from the base model, achieving the goal of breaking through the upper limit of the base model's inference capabilities.

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

Claims

1. A method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration, characterized in that, Includes the following steps: S1. Select a pre-trained large language model as the reference model; select a training dataset, in which each data point contains a question, answer, and response. S2 replaces the reverse KL divergence regularization term with the forward KL divergence to construct a training objective function that maximizes reward and entropy while minimizing the forward KL divergence, thereby achieving out-of-distribution exploration; S3, Construct a reward-adaptive reference model distribution reweighting mechanism to dynamically reweight the reference model, obtaining a reweighted reference model, and driving adaptive exploration within the distribution; the formula is: In the formula, The reweighted reference model is defined as follows: x represents the question, y represents the corresponding answer, 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 exponent is Z, which is the normalization factor; the reweighting function φ(r) is a monotonically increasing function, with values ​​ranging from [0,1]; the higher the reward r(x,y), the closer φ(r) is to 1, making... Approximate to π ref Make full use of existing reasoning abilities; the lower the reward r(x,y), the closer φ(r) is to 0, making The distribution tends towards uniformity to encourage exploration; S4. Combine the reweighted reference model from step S3 with the training objective function from step S2 to form the final objective function: In the formula, P(x) represents the distribution of the problem, E represents the expected probability, and H(π) represents the expected probability. θ ) represents π θ The entropy, where α and β represent hyperparameters, are both greater than 0; D KL Indicates the positive KL divergence; S5, based on the final objective function, combines the training dataset with the model π to be optimized. θ Train the model to obtain a trained model. S6, During application, input the problem into the trained model. The answer was obtained.

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

3. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 1, 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. At the same time, the positive KL divergence term is calculated using the 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 regions.

4. The method for enhancing the reasoning ability of a large language model based on reward-adaptive reinforcement learning exploration according to claim 3, characterized in that, The advantage value for each response is calculated using grouped relative advantage estimation, as shown in the following formula: In the formula, mean represents the mean, std represents the standard deviation, and A i For each answer y i The normalized dominance value; r i =r(x,y) i ), indicating the answer y i The reward; for each question x, G answers are sampled using an autoregressive approach. Then apply the reward function r to each sample pair (x, y) i ), calculate reward r i =r(x,y) i ).

Images