Multi-stage data selection method and apparatus for post-training reinforcement learning of large language models

Abstract

A multi-stage data selection method and apparatus for post-reinforcement learning training of large language models are disclosed. The method involves selecting a warm-up subset of data from the original training set and performing rule-verifiable reward reinforcement learning post-training to obtain the trained large language model. Multiple response sampling is performed on candidate samples, and a learnability score is calculated based on the response trajectory verification results. The gradient representation of the candidate samples under the reinforcement learning training objective is calculated and mapped to low-dimensional projected gradient features. An improved gradient alignment score matrix is ​​constructed based on the learnability score and projected gradient features. Candidate samples are sorted and filtered according to the average score to obtain the target training subset. The data selection score is recalculated based on the staged model obtained during training, and the training subset is updated for training in the new stage. This invention can simultaneously consider sample representativeness and learning potential, reduce response length bias, and maintain the learning performance of post-reinforcement learning training of large language models while reducing the scale of training data.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence training optimization technology, specifically a multi-stage data selection method and apparatus for post-training reinforcement learning of large language models. Background Technology

[0002] In recent years, the performance of large language models in tasks such as mathematical reasoning, code generation, and complex decision-making has continued to improve. To further enhance the reasoning ability of these models, the industry has gradually adopted reinforcement learning post-training methods based on rule-verifiable rewards, enabling models to optimize policies around automatically verifiable correctness signals.

[0003] However, existing reinforcement learning post-training processes typically rely on large-scale training data and high computational resources. Directly post-training all original training samples is not only costly but also introduces a large number of redundant samples that offer limited benefit to the current model, resulting in low training efficiency. Existing research has shown that some capabilities of a pre-trained model can be activated by a small number of high-quality samples. Therefore, identifying the most valuable samples for post-training from the original training set has become a crucial issue for improving training efficiency.

[0004] Most existing data selection methods have been developed in supervised fine-tuning scenarios. These methods can be broadly divided into two categories: external scoring methods and self-scoring methods.

[0005] External scoring methods typically utilize more powerful external models or additional evaluators to score the complexity, diversity, quality, or task relevance of samples. For example, some methods use stronger language models to generate finer-grained labels to measure instruction complexity, some methods directly use external models to screen for high-quality instructions, and others measure sample value by comparing the difficulty of generation with and without instructions, or by using gradient similarity to select data that is closer to the target task.

[0006] Self-scoring methods do not rely on additional external models but extract signals from the behavior of the target model itself. For example, some methods combine word-level, sentence-level, and model-level uncertainty to identify high-quality samples, while others score samples based on their impact on the perplexity of the anchor set. These methods can improve the quality of training data to some extent in supervised fine-tuning scenarios, but their evaluation criteria usually revolve around likelihood fitting, perplexity changes, or static uncertainty, and the optimization objectives are not consistent with the reward improvement objectives in reinforcement learning training.

[0007] Compared to supervised fine-tuning, research on data selection for post-training reinforcement learning oriented towards rule-verifiable reward remains limited. Existing methods indicate that post-training in reinforcement learning does not necessarily require a large number of samples; even a small number of samples or a single sample can provide effective training signals. For example, some methods perform multiple rounds of post-training on the original full dataset and then evaluate the sample value based on performance changes in the training trajectory. However, these methods often require multiple expensive rounds of training on the full dataset before data selection can be performed, resulting in the selection process itself consuming significant computational resources, making it difficult to truly achieve the goal of selecting data first and then reducing costs.

[0008] Furthermore, existing methods suffer from several common limitations. First, many methods directly use perplexity, length, static difficulty, or external evaluation scores from supervised fine-tuning scenarios as selection criteria, failing to accurately reflect the policy improvement value during training after reinforcement learning. Second, even when gradient information is used to judge sample value, the original gradient norm is easily affected by response length bias, leading to over-preference for shorter response samples. Third, existing solutions often struggle to simultaneously consider both sample representativeness and learnability, i.e., they cannot simultaneously answer the questions of which samples represent the overall training direction and which samples fall within the most worthwhile difficulty range for the current model to learn.

[0009] Therefore, there is a need to provide a reinforcement learning post-training data selection scheme that can simultaneously reflect the representativeness and learnability of the samples, so as to reduce the scale of data use and computational cost while ensuring training effectiveness. Summary of the Invention

[0010] The purpose of this invention is to provide a multi-stage data selection method and apparatus for post-training reinforcement learning of large language models, which solves the technical problems of existing post-training data selection methods for reinforcement learning not matching the actual policy optimization objectives, difficulty in simultaneously taking into account sample representativeness and learning potential, and susceptibility to response length bias.

[0011] To achieve the above objectives, the present invention adopts the following technical solution:

[0012] A multi-stage data selection method for training after reinforcement learning of large language models includes the following steps:

[0013] S1. Pre-sample the original training set to obtain a pre-warmed data subset, and use the pre-warmed data subset to perform rule-verifiable reward reinforcement learning on the large language model to be trained, and then train it to obtain the trained large language model.

[0014] S2, using the trained large language model, perform multiple response trajectory sampling on the candidate samples in the original training set, and calculate the learnability score of the candidate samples based on the response trajectory verification results corresponding to the candidate samples.

[0015] S3, calculate the gradient representation of each candidate sample under the training objective after reinforcement learning, and map the gradient representation to a low-dimensional projected gradient feature;

[0016] S4, calculate the improved gradient alignment score between candidate samples based on the learnability score of candidate samples and the low-dimensional projected gradient features;

[0017] S5. The candidate samples are sorted according to the average improved gradient alignment score of each candidate sample relative to all candidate samples, and the top N candidate samples are selected to form the target training subset for training after formal reinforcement learning.

[0018] S6: During the training of the large language model, S2 to S5 are repeated in stages. The average improved gradient alignment score is recalculated based on the stage model obtained during training. The score is selected and the training subset is updated for training in the new stage.

[0019] Furthermore, the learnability score is based on the success rate of the candidate samples. The success rate was calculated to be... The learnability score is the percentage of response trajectories that pass verification within a predetermined number of response trajectory samplings for the same candidate sample. The form represents the learning potential of candidate samples.

[0020] Furthermore, the gradient is represented as a candidate sample. Policy gradient under the training objective after reinforcement learning The low-dimensional projection gradient feature is obtained by combining the policy gradient with the random projection matrix. Multiplication yields the result, preserving gradient direction similarity while reducing storage and computational overhead. The low-dimensional projected gradient features are represented as follows: .

[0021] Furthermore, before calculating the improved gradient alignment score between candidate samples, the low-dimensional projected gradient features corresponding to the candidate samples are normalized.

[0022] Furthermore, the improved gradient alignment score is calculated according to the following formula: ,in and Representing candidate samples respectively and Learnability score This represents the cosine similarity between the low-dimensional projected gradient features of two candidate samples.

[0023] Furthermore, the average improved gradient alignment score is calculated according to the following formula:

[0024]

[0025] in This represents the total number of candidate samples. Indicates candidate samples With candidate samples The improved gradient alignment score between the average improved gradient alignment scores is used to select the top N candidate samples, or the candidate samples with the highest scores in each of the multiple candidate sample clusters are selected as the target training subset.

[0026] Furthermore, when calculating the gradient representation of candidate samples under the training objective after reinforcement learning, only one verified response trajectory or a predetermined number of representative response trajectories are selected for each candidate sample to participate in the gradient calculation, so as to reduce the backpropagation overhead.

[0027] The present invention also provides a multi-stage data selection device for post-training reinforcement learning of large language models, comprising:

[0028] The pre-training module is used to pre-sample the original training set and obtain the trained large language model.

[0029] The learnability calculation module is used to perform multiple response trajectory samplings on candidate samples and calculate learnability scores;

[0030] The gradient projection module is used to calculate the gradient representation of candidate samples and generate low-dimensional projected gradient features.

[0031] The sample scoring and filtering module is used to calculate the improved gradient alignment score, the average improved gradient alignment score, and output the target training subset.

[0032] The present invention also provides a multi-stage data selection device for post-training reinforcement learning of large language models, including a preheating training module, a learnability calculation module, a gradient projection module, and a sample scoring and screening module, each module being used to perform the corresponding steps of the above methods.

[0033] Compared with the prior art, the present invention has at least the following positive effects:

[0034] 1. This invention characterizes the influence relationship between samples from the perspective of gradient alignment, and can directly select samples for policy improvement targets in reinforcement learning training, instead of using the static scoring method in supervised fine-tuning scenarios.

[0035] 2. By introducing a learnability index, this invention can characterize the learning potential of samples without relying on the original gradient norm, reduce response length bias, and improve the stability of sample scoring.

[0036] 3. This invention combines preheating training and random projection mechanisms, which can reduce the computational cost of gradient estimation and sample sorting while maintaining the gradient inner product structure, thereby improving the deployability of the method.

[0037] 4. This invention can select a high-value training subset from the original training set based on the current training model's capabilities, achieving post-training results that are close to or even better than full training while significantly reducing the size of the training samples. Attached Figure Description

[0038] Figure 1 This is an overall flowchart of the method of the present invention.

[0039] Figure 2 This is a schematic diagram of the process of calculating the learnability of candidate samples and gradient projection in this invention.

[0040] Figure 3 This is a schematic diagram of the sample scoring and screening process in this invention.

[0041] Figure 4 This is a schematic diagram of a multi-stage data selection device for post-training reinforcement learning of large models provided by the present invention.

[0042] Figure 5 This is a flowchart illustrating the application of the present invention in the scenario of a math problem-solving assistant on an online education platform. Detailed Implementation

[0043] To enable those skilled in the art to better understand the present invention, the technical solutions in the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings.

[0044] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover a non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus.

[0045] According to one aspect of the present invention, a multi-stage data selection method for post-training reinforcement learning of large language models is provided. Let the original training set be... Candidate samples are The trained large language model is The method is as follows: Figure 1 As shown, the specific steps include:

[0046] Step S1: Warm-up training.

[0047] First, starting with the original training set A subset of preheating data is formed by randomly selecting samples from the middle. ,in The sample size is significantly smaller than The sample size. Subsequently, the training objective is based on rule-verifiable rewards after reinforcement learning, in The initial large language model is pre-trained to obtain the trained large language model. This warm-up phase allows the large language model to initially adapt to the target task distribution, thereby improving the stability of subsequent sample gradient estimation. The rules for verifying rewards include at least one of the following: whether the answer is correct, whether the program passes test cases, whether the formula derivation satisfies constraints, or whether the format meets preset rules.

[0048] In one optional implementation, the post-reinforcement learning training objective is achieved using the GRPO objective function. For a given query... Model generation Response Furthermore, based on the relative advantage calculated within the group, stable training is achieved by constraining the deviation between the large language model and the reference model.

[0049] Step S2: Calculate the learnability of candidate samples.

[0050] like Figure 2 Using the trained large language model Perform for each candidate sample Response sampling is performed. For each response, automatic validation is conducted based on whether the answer is correct, whether the code passes the test cases, whether the derivation meets the rules, or whether the output format meets the requirements. Responses that pass validation are recorded as successful samples, and then further... Calculate the success rate of candidate samples ,in The success rate is an indicator function. It is obtained statistically from multiple response sampling results corresponding to candidate samples and is used to characterize the proportion of successful validation of the trained large language model on candidate samples.

[0051] Based on the success rate, a learnability score is calculated for each candidate sample. The learnability score is determined by the combination of the candidate sample's success rate and its complementary values, and is used to uniformly characterize the learning potential and difficulty of the candidate sample: when a candidate sample is too easy or too difficult, the learnability score is low; when a candidate sample is near the capability boundary of the current large language model, the learnability score is high. Specifically, the learnability score is expressed as... ,when When the value is close to 0 or 1, it indicates that the candidate samples are either too difficult or too easy for the current model, and the gain from their continued participation in training is limited; when... When the value is in the middle range, especially close to 0.5, it indicates that the candidate sample is near the current model's capability boundary and is more likely to bring effective learning signals.

[0052] Step S3: Estimate the gradient representation of the candidate sample and perform random projection. The gradient representation is the policy gradient of the candidate sample under the training objective after reinforcement learning; to reduce computational overhead, the gradient can be calculated only for one verified response corresponding to the candidate sample, or only for a predetermined number of representative responses.

[0053] like Figure 2 Based on the trained large language model Calculate the policy gradient of each candidate sample under the training objective after reinforcement learning. Considering the extremely high dimensionality of the original gradients, directly storing and comparing all gradients would be costly. Therefore, this embodiment introduces a random projection matrix. Mapping the gradient vector to a lower-dimensional space, Mapping to low-dimensional projective gradient features yields the projective gradient features. Low-dimensional projected gradient features are obtained through the mapping relationship between the random projection matrix and the target gradient vector of the candidate samples, so as to preserve the similarity of gradient directions while reducing storage and computational overhead.

[0054] In a preferred embodiment, to reduce the cost of gradient estimation, the gradient is calculated for each candidate sample by selecting only at least one validated response.

[0055] Step S4: Construct the gradient alignment and learnability score matrix. The gradient alignment and learnability score are determined based on the learnability score of the candidate samples and the directional similarity of the low-dimensional projected gradient features between candidate samples. By incorporating the learnability score, rather than the original gradient norm, into the scoring process, the impact of response length bias on sample selection can be mitigated.

[0056] like Figure 3 The projected gradient features of all candidate samples are normalized to obtain unit feature vectors representing the sample update direction. Then, for any two candidate samples... and Calculate their gradient alignment and learnability scores. and and low-dimensional projection gradient features Calculate the improved gradient alignment score:

[0057]

[0058] in, The cosine similarity between the low-dimensional projected gradient features of two candidate samples is used to measure the consistency between the two candidate samples in the policy update direction. The learnability score is used to reflect the learning potential of the two candidate samples.

[0059] Unlike directly using the gradient inner product or gradient norm, this embodiment uses an improved gradient alignment score instead of the original gradient norm in the scoring. This is because the original gradient norm is easily affected by the response length, while the improved gradient alignment score can more robustly represent the sample value, thereby alleviating the problem of over-preference for short-response samples.

[0060] Step S5: Sample sorting and target training subset generation. An average improved gradient alignment score is obtained by aggregating the improved gradient alignment scores between candidate samples and all samples. Based on this average improved gradient alignment score, the top few samples are selected, or the samples with the highest scores are selected from each clustered candidate sample cluster to form the target training subset.

[0061] like Figure 3 Based on the improved gradient alignment score, the average improved gradient alignment score is calculated for each candidate sample:

[0062]

[0063] in Indicates the total number of candidate samples. (By...) Sort the samples from highest to lowest score. Select the top N candidate samples with the highest scores to form the target training subset. Subsequently, the large language model is trained using the target training subset.

[0064] In an optional implementation, when the total number of candidate samples is large, the candidate samples can be clustered based on sample embedding or task label first, and then the candidate sample with the highest score can be selected in each cluster to balance sample representativeness and distribution coverage.

[0065] Step S6: Multi-stage selection and training.

[0066] During model training, steps S2 to S5 are repeated periodically. Based on the stage model obtained during training, the data is recalculated, scores are selected, and the training subset is updated for training in the new stage.

[0067] Through the above description of the embodiments, those skilled in the art can clearly understand that the methods according to the above embodiments can be implemented by means of software plus necessary general-purpose hardware platforms. Of course, they can also be implemented by hardware, but in many cases the former is a better implementation method. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, can be embodied in the form of a software product. This computer software product is stored in a storage medium (such as ROM / RAM, magnetic disk, optical disk) and includes several instructions to cause a terminal device (which may be a mobile phone, computer, server, or network device, etc.) to execute the methods described in the various embodiments of the present invention. For example:

[0068] One embodiment of the present invention provides a multi-stage data selection device for post-training reinforcement learning of large language models, such as... Figure 4 As shown, it includes:

[0069] The preheating training module is used to select a preheating data subset from the original training set and perform preheating training on the large language model to obtain the trained large language model.

[0070] The learnability calculation module is used to perform response sampling on candidate samples based on the trained large language model, and calculate the success rate and learnability score based on the validation results;

[0071] The gradient projection module is used to calculate the target gradient of candidate samples and generate low-dimensional projected gradient features.

[0072] The sample scoring and filtering module is used to calculate gradient alignment and learnability scores based on learnability scores and projected gradient features, and to filter out the target training subset.

[0073] For the specific implementation process of each module, please refer to the description of the method of the present invention above.

[0074] The following is combined with Figure 5 The application of the present invention in specific scenarios is described.

[0075] In today's online education platforms, particularly in scenarios like math problem-solving assistants, efficiently improving the problem-solving capabilities of large language models in mathematical reasoning tasks has become a key focus in the implementation of many artificial intelligence applications. In recent years, training large language models with reinforcement learning using large-scale mathematical data has significantly improved their reasoning performance to some extent. However, this training method still faces challenges in many practical applications. On the one hand, the candidate training questions are typically large in scale; directly sampling responses, verifying rewards, and calculating gradients from the entire dataset would incur high computational and time costs. On the other hand, the original question bank often contains a large number of overly simple, overly difficult, or redundant samples, which offer limited help in improving the current model's capabilities, making efficient training difficult. This invention relates to a multi-stage data selection method for reinforcement learning post-training of large language models, which can effectively address these problems. Specifically, a large-scale mathematical problem training set can be input into the data selection device of the present invention. First, a large language model is obtained through preheating training. Then, the success rate, learnability, and gradient alignment information of candidate problems are calculated to select the high-quality training subset that has the most training value for the current model. Subsequently, the large language model is formally trained by reinforcement learning based on the selected training subset. This can effectively improve the accuracy and practicality of the model in scenarios such as mathematical problem solving, step-by-step reasoning, and automatic question answering while reducing the scale of training data and computational costs.

[0076] The parts of the technical solutions of the present invention that are not described in detail can be implemented using existing technologies.

[0077] It should be noted that the present invention is not limited to the embodiments described above. For those skilled in the art, various modifications and substitutions can be made to the preheating strategy, learnability calculation method, random projection dimension, sorting rules, and target training algorithm without departing from the concept of the present invention, and all such modifications and substitutions should fall within the protection scope of the present invention.

Claims

1. A multi-stage data selection method for post-training reinforcement learning of large language models, characterized in that, Includes the following steps: S1. Pre-sample the original training set to obtain a pre-warmed data subset, and use the pre-warmed data subset to perform rule-verifiable reward reinforcement learning on the large language model to be trained, and then train it to obtain the trained large language model. S2, using the trained large language model, perform multiple response trajectory sampling on the candidate samples in the original training set, and calculate the learnability score of the candidate samples based on the response trajectory verification results corresponding to the candidate samples. S3, calculate the gradient representation of each candidate sample under the training objective after reinforcement learning, and map the gradient representation to a low-dimensional projected gradient feature; S4, calculate the improved gradient alignment score between candidate samples based on the learnability score of candidate samples and the low-dimensional projected gradient features; S5. The candidate samples are sorted according to the average improved gradient alignment score of each candidate sample relative to all candidate samples, and the top N candidate samples are selected to form the target training subset for training after formal reinforcement learning. S6: During the training of the large language model, S2 to S5 are repeated in stages. The average improved gradient alignment score is recalculated based on the stage model obtained during training. The score is selected and the training subset is updated for training in the new stage.

2. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, The rule can verify that the reward signal corresponding to the training after reward reinforcement learning is determined by at least one of the following automatic verification results: correctness of answer, code execution result, symbol derivation result, or output format constraint.

3. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, The learnability score is based on the success rate of candidate samples. The success rate was calculated to be... The learnability score is the percentage of response trajectories that pass verification in a predetermined number of response trajectory samplings for the same candidate sample. The form represents the learning potential of candidate samples.

4. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, The gradient is represented as a candidate sample. Policy gradient under the training objective after reinforcement learning The low-dimensional projection gradient feature is obtained by combining the policy gradient with the random projection matrix. Multiplication yields the result, preserving gradient direction similarity while reducing storage and computational overhead. The low-dimensional projected gradient features are represented as follows: .

5. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, Before calculating the improved gradient alignment score between candidate samples, the low-dimensional projected gradient features corresponding to the candidate samples are normalized.

6. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, The improved gradient alignment score is calculated according to the following formula: ,in and Representing candidate samples respectively and Learnability score This represents the cosine similarity between the low-dimensional projected gradient features of two candidate samples.

7. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, The average improved gradient alignment score is calculated according to the following formula: in This represents the total number of candidate samples. Indicates candidate samples With candidate samples The improved gradient alignment score between the average improved gradient alignment scores is used to select the top N candidate samples, or the candidate samples with the highest scores in each of the multiple candidate sample clusters are selected as the target training subset.

8. The multi-stage data selection method for post-training reinforcement learning of large language models according to claim 1, characterized in that, When calculating the gradient representation of a candidate sample under the training objective after reinforcement learning, only one validated response trajectory or a predetermined number of representative response trajectories are selected for each candidate sample to participate in the gradient calculation, so as to reduce the backpropagation overhead.

9. A multi-stage data selection device for post-training reinforcement learning of large language models, characterized in that, include: The pre-training module is used to pre-sample the original training set and obtain the trained large language model. The learnability calculation module is used to perform multiple response trajectory samplings on candidate samples and calculate learnability scores; The gradient projection module is used to calculate the gradient representation of candidate samples and generate low-dimensional projected gradient features. The sample scoring and filtering module is used to calculate the improved gradient alignment score, the average improved gradient alignment score, and output the target training subset.

Images