A knowledge distillation method based on inference step disassembly and differentiated supervision

Abstract

The present application relates to a kind of knowledge distillation method based on inference step deconstruction and differentiating supervision, belong to artificial intelligence field.The structured thinking chain generated by acquiring teacher model is separated according to line, each inference step is deconstructed into independent unit;Each step is type labeled, according to the five-level classification system of basic calculation, basic fact, operation execution, logical reasoning, strategy planning, differentiating weight is distributed, and weight superposition is used to composite step;Weighted loss function is constructed, step weight is introduced into cross-entropy loss, and student model is distilled training.The present application focuses on key inference link in the learning process by step-level type perception and differentiating supervision, overcomes the problem of detail loss and global understanding damage caused by flat sequence supervision in traditional distillation.Experimental results show that, on mathematical reasoning task, compared with equal-weight step distillation method, accuracy is improved by 4.6 percentage points, and efficient migration of large model reasoning ability to lightweight model is realized.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, specifically to a knowledge distillation method based on reasoning step deconstruction and differential supervision. Background Technology

[0002] Large language models exhibit superior performance in complex reasoning tasks, but their parameter counts often reach hundreds of billions or even trillions, leading to high inference latency and high computational resource consumption, severely restricting their deployment in resource-constrained scenarios (such as mobile terminals and edge devices). Knowledge distillation, as an effective model compression technique, significantly reduces model size while maintaining performance by having the student model mimic the output of the teacher model, becoming an important technical path for lightweighting large models. However, existing knowledge distillation methods have significant limitations when facing complex reasoning tasks: traditional output distillation only requires the student model to fit the final answer of the teacher model, ignoring the inherent logical thinking in the reasoning process itself. This results in the student model only learning "what" but failing to grasp "why," leading to severe performance degradation in tasks requiring multi-step reasoning. Some improved methods attempt to introduce the teacher model's thought chain as a supervision signal, but still treat the entire thought chain as a flat sequence of tokens for overall modeling, failing to distinguish the differences in semantic type and importance of different reasoning steps. This makes it difficult for the student model to simultaneously capture detailed information and global structure in the reasoning process, resulting in the dual dilemma of "loss of details" and "impaired global understanding."

[0003] To address the aforementioned issues, existing technologies can be broadly categorized into three evolutionary directions. The first is standard output distillation, which uses only the teacher model's final answer as the supervision target. Its advantage lies in its simple training, but its drawback is that it completely fails to transfer reasoning ability, leaving student models lacking an understanding of problem-solving paths. The second is sequence distillation, which uses the complete thought chain generated by the teacher model as a sequence generation task for supervision. While this method can transfer some reasoning processes, it treats different reasoning steps with equal weight, failing to guide student models to focus on key steps such as logical reasoning and strategy planning. Furthermore, it is prone to overfitting or underfitting during training due to differences in ability. The third is step-by-step distillation, such as the Distilling Step-by-Step method. Its advancement lies in explicitly treating reasoning steps as additional supervision signals, but it still treats each step as an equally important flat unit, failing to differentiate steps based on their semantic features and cognitive complexity.

[0004] In summary, existing technologies have not solved the core technical problems of reasoning step type perception and differentiated supervision. How to achieve fine-grained differentiation and targeted enhancement of different reasoning steps during the distillation process remains a technical bottleneck that urgently needs to be overcome in this field. Summary of the Invention

[0005] To address the problems of existing technologies, this invention provides a knowledge distillation method based on reasoning step deconstruction and differential supervision.

[0006] To solve the above-mentioned technical problems, the present invention is achieved through the following technical solution: Firstly, a knowledge distillation method based on reasoning step deconstruction and differentiated supervision, comprising the following steps: Step S1: Obtain the reasoning process generated by the teacher model for the input question. The reasoning process is deconstructed into multiple reasoning steps separated by rows, and each reasoning step corresponds to an independent reasoning unit. Step S2: Label each reasoning step with its type, determine its reasoning type, and assign a preset weight value to each reasoning type; when a reasoning step is identified as multiple reasoning types at the same time, calculate the comprehensive weight of the step by weight superposition. Step S3: Construct training data, which includes the input question, the inference step sequence of the teacher model, and their corresponding comprehensive weights; Step S4: Distill the student model using the training data. During the training process, the loss function is weighted according to the comprehensive weight of each inference step to obtain the weighted loss. Step S5: Update the student model parameters according to the weighted loss until the model converges, and obtain the distilled student model.

[0007] In one specific implementation of the first aspect, the type labeling in step S2 is automatically completed by the teacher model. The labeling is based on a preset reasoning type classification system, and the reasoning types include: basic calculation, basic facts, operation execution, logical reasoning, and strategy planning.

[0008] In one specific implementation of the first aspect, the identification of the reasoning type is achieved through keyword matching or semantic judgment; when a reasoning step belongs to multiple reasoning types at the same time, the weight superposition method is to add the weight values ​​corresponding to each type, and the upper limit of the superimposed comprehensive weight is 1.0.

[0009] In one specific implementation of the first aspect, the inference type and its weight are set as follows: Basic calculations: weight 0.3, corresponding to arithmetic operation steps; Basic facts: weighted at 0.4, corresponding to the fact retrieval or information extraction steps; Operation execution: Weight is 0.6, corresponding to the application method, formula or rule steps; Logical reasoning: weighted at 0.8, corresponding to conditional judgment, causal analysis, or logical deduction steps; Strategy planning: with a weight of 1.0, corresponding to the steps of problem-solving plan formulation, method selection, or problem decomposition.

[0010] In one specific implementation of the first aspect, the weighted loss in step S4 includes a weighted thought chain loss and a final answer loss, and the weighted loss is calculated as follows: in, This represents the total number of tokens in the MindChain. For the first The overall weight of the reasoning step to which each token belongs. Let cross-entropy be the loss function. and These represent the output probability distributions of the teacher model and the student model on the final answer, respectively. and These are the preset weighting coefficients.

[0011] In one specific implementation of the first aspect, the The value is 0.7. The value is 0.3.

[0012] In one specific implementation of the first aspect, the teacher model is a large-scale language model with more than or equal to 30 billion parameters, and the student model is a medium-sized language model with less than or equal to 10 billion parameters.

[0013] In one specific embodiment of the first aspect, during the distillation training process in step S4, multiple baseline methods are used for comparative experiments, the baseline methods including: Baseline Method 1: Standard output distillation supervised only for the final answer; Baseline Method 2: Supervised sequence distillation that treats the entire thought chain as a flat sequence; Baseline Method 3: Treat inference steps as supervised step distillation with equal weights; The superiority of this method is verified by comparing the accuracy of answers on the test set.

[0014] Secondly, a knowledge distillation system based on reasoning step deconstruction and differential supervision includes: The reasoning step deconstruction module is used to obtain the reasoning process generated by the teacher model and divide it into multiple reasoning steps by row. The type labeling module is used to identify the type of each inference step and assign corresponding weights; when the same inference step corresponds to multiple types, the weight is calculated by weight superposition. The training data construction module is used to generate training samples that include input questions, inference step sequences, and comprehensive weights. The weighted distillation training module is used to weight the loss according to the comprehensive weight of the inference steps and update the model parameters during the training of the student model. The output module is used to output the trained student model.

[0015] In one specific implementation of the second aspect, the type labeling module has a built-in reasoning type classification system, wherein the reasoning types include basic calculation, basic facts, operation execution, logical reasoning, and strategy planning, and each corresponds to a different weight.

[0016] The beneficial effects of this invention are as follows: 1. This invention achieves structured modeling and differentiated knowledge transfer of the teacher model's reasoning process by constructing a closed-loop technical architecture of step deconstruction, type labeling, and weighted supervision. Specifically, firstly, the thought chain generated by the teacher model is divided into independent reasoning units by row, solving the problem of treating the reasoning process as a black box in traditional methods. Based on this, a five-level reasoning type system is established, including basic calculation, basic facts, operation execution, logical reasoning, and strategy planning. Differentiated weights are preset for each type. Step-level type labeling is completed through a combination of keyword matching and semantic judgment. For compound reasoning steps, weighted additive superposition with an upper limit is used to accurately quantify the semantic complexity of the reasoning steps. Furthermore, a weighted loss function is constructed in the distillation training, introducing the differentiated weights of each reasoning step into the cross-entropy loss, applying stronger supervision signals to key reasoning steps while retaining the final answer loss to balance the learning of the reasoning process and the result. This technical solution enables student models to learn the reasoning abilities of teacher models in a hierarchical and focused manner, effectively solving the problems of "loss of details" and "impaired global understanding" caused by flat sequence supervision in existing methods, and significantly improving the fineness and interpretability of knowledge distillation.

[0017] Based on the aforementioned technical architecture, this invention demonstrates quantifiable technical advantages in practical deployment. Experimental results show that, in mathematical reasoning tasks, compared to standard output distillation that only supervises the answer, flat sequence distillation that treats the thought chain as a whole, and equal-weighted step distillation, the method of this invention improves the accuracy on the test set by 14.8%, 9.0%, and 4.6 percentage points, respectively, verifying the effectiveness of the differentiated supervision mechanism in transferring reasoning ability. From a system implementation perspective, this invention, through modular decoupling, modularizes functions such as reasoning step deconstruction, type labeling, and weighted training. Each module has clear inputs and outputs and adjustable parameter ranges. The architecture design, with no less than 30 billion parameters for the teacher model and no more than 10 billion parameters for the student model, balances reasoning ability and deployment efficiency. Key hyperparameters such as learning rate, batch size, and weight coefficients all have clear optimal ranges and adjustment space, ensuring the reproducibility and engineering applicability of the technical solution. In summary, this invention achieves a technological leap from "result imitation" to "process understanding" in the field of model compression, providing a complete and efficient technical path for lightweight deployment of large models in complex reasoning tasks. Attached Figure Description

[0018] Figure 1 This is a schematic diagram of the overall process of the method of the present invention.

[0019] Figure 2 This is a schematic diagram of the system architecture of the present invention. Detailed Implementation

[0020] The technical solutions of the present invention will be clearly and completely described below with reference to the embodiments of the present invention. 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 of ordinary skill in the art without creative effort are within the scope of protection of the present invention.

[0021] like Figures 1 to 2 This paper presents a knowledge distillation method based on reasoning step deconstruction and differential supervision.

[0022] I. Overall Methodology Figure 1 This is an overall flowchart of the method of the present invention. (Refer to...) Figure 1 The method includes the following steps S1 to S5, and there are clear temporal and data dependencies between the steps: Step S1: Obtain the teacher model reasoning process This step corresponds to Figure 1 Step S1 in the module. Specifically, input the problem. Sent into the teacher model The teacher model is a large-scale language model, typically with no fewer than 30 billion parameters (e.g., Qwen3-32B, LLaMA-3-70B, etc.). During inference, the teacher model is configured in a "step-by-step deduction" mode, requiring each inference step to be separated by a newline character, forming a structured chain of thought. ,in Indicates the first One reasoning step, This represents the total number of steps. Simultaneously, the teacher model outputs the final answer. .

[0023] The output of this step is a triplet. ,in This is a sequence of reasoning steps. The final answer flows to step S2.

[0024] Step S2: Step type labeling and weight allocation This step corresponds to Figure 1 The S2 module in step S1. Each inference step generated in step S1... Input is sent to the type labeling module. This module can use the teacher model itself for labeling, specifically by constructing prompt words and requiring the teacher model to classify each step based on a preset reasoning type system. The prompt words can include keyword matching rules as an aid, such as "if the step contains numbers and operators, it tends to label it as basic calculation."

[0025] The reasoning type system and its preset weights are shown in the table below: When a reasoning step matches multiple types simultaneously, the overall weight is calculated using a weighted summation method. : That is, the weights corresponding to each type are added together, with a maximum of 1.0. For example, if a step contains both "operation execution" (0.6) and "logical reasoning" (0.8), then the total weight is 0.6 + 0.8 = 1.4, with a maximum of 1.0.

[0026] The output of this step is a set of type labels. and comprehensive weight It flows to step S3 along with the original reasoning step sequence.

[0027] Step S3: Construct training data This step corresponds to Figure 1 Step S3 in the module. The original problem. The sequence of reasoning steps in the teacher model The overall weight corresponding to each step And the final answer Combine them into a single training sample. Repeat steps S1-S2 to construct a large-scale training dataset. .

[0028] The output of this step is a structured training dataset. The flow proceeds to step S4.

[0029] Step S4: Weighted distillation training This step corresponds to Figure 1 Step S4 in the module. The training dataset... Used for training the student model S. The student model is a small to medium-sized language model, typically with no more than 10 billion parameters (e.g., Qwen3-8B, LLaMA-3-8B, etc.). The training objective is to minimize the weighted loss function.

[0030] The loss function consists of two parts: Weighted Mind Chain Loss: The cross-entropy loss of each token in the mind chain is weighted according to the overall weight of the step to which the token belongs.

[0031] Final answer loss: Cross-entropy loss on the final answer.

[0032] Total loss is defined as: in: Total number of tokens in the MindChain; : No. The overall weight of the reasoning step to which each token belongs; Cross-entropy loss function; Teacher model and student model in the first Output probability distribution on each token; The output probability distributions of the teacher model and the student model on the final answer; , Weighting coefficient, with a value range of: Preferably, =0.7, =0.3.

[0033] During training, mini-batch gradient descent is used to update the student model parameters. The batch size can be set to 8, 16, or 32, and the learning rate can be set to... to The AdamW optimizer is used, and the number of training epochs is set to 3 to 10 epochs depending on the size of the dataset. An early stopping mechanism is used to prevent overfitting.

[0034] The output of this step is the updated student model parameters, which flow to step S5.

[0035] Step S5: Model Convergence and Output This step corresponds to Figure 1 Step S5 in the process. Repeat step S4 until the loss on the validation set no longer decreases or the preset number of training rounds is reached, and save the best student model as the output of the distilled model.

[0036] The output of this step is the distilled student model. It can be used for downstream task deployment.

[0037] II. System Architecture Figure 2 This is a schematic diagram of the module structure of the system of the present invention. (Refer to...) Figure 2 The system comprises the following five modules, which are connected and interact with each other via data flow: 1. Reasoning Step Deconstruction Module The input to this module is the raw text output generated by the teacher model. Internally, the module segments the text based on newline characters, identifies each independent reasoning step, filters out blank lines and formatting characters, and outputs a structured sequence of reasoning steps. The output of this module serves as the input to the type annotation module.

[0038] 2. Type labeling module This module receives the sequence of steps output by the reasoning step deconstruction module. The module incorporates a built-in reasoning type classification system to identify the type of each reasoning step. Labeling can be performed using a teacher model based on prompts or a separate lightweight classifier. When the same step corresponds to multiple types, an additive weighting method is used to calculate the overall weight, with an upper limit of 1.0. The module's output is a set of type labels and a sequence of overall weights for each step. It flows along with the original step sequence to the training data construction module.

[0039] 3. Training Data Construction Module This module receives a sequence of steps from the reasoning step deconstruction module. Comprehensive weight sequence from the type annotation module and the problem of raw input and the final answer The module assembles this information into a standard training sample format, with each sample containing a question field, a step sequence field, a weight sequence field, and an answer field. The output of this module is a structured training dataset. The flow is directed to the weighted distillation training module.

[0040] 4. Weighted Distillation Training Module This module receives the training data and constructs the dataset output by the module. The module provides the initial parameters of the student model. Internally, it implements the weighted loss function calculation logic, including: summing the cross-entropy loss of each token in the thought chain according to the weight of its corresponding step, and combining this sum with the final answer loss. The module uses gradient descent to update the student model parameters and supports mechanisms such as batch training and early stopping. The output of this module is the updated student model parameters during training, and finally, the converged optimal model is passed to the output module.

[0041] 5. Output Module This module receives the student model parameters output by the weighted distillation training module, saves the model in a specified format (such as a PyTorch .pt file, a HuggingFace bin file, etc.) after training, and outputs the distilled student model. It is used for deployment of downstream tasks.

[0042] III. Examples The following uses a specific mathematical reasoning task as an example to illustrate the implementation process of this invention in detail.

[0043] Example Task: Solving Elementary School Math Word Problems Teacher Model: Qwen3-32B Student model: Qwen3-8B Training dataset: Contains 5000 elementary school math word problems and their manually annotated solutions. Test dataset: Contains 500 untrained math word problems Step 1: Teacher Model Reasoning Input problem: Xiaoming had 15 candies. He gave 3 to Xiaohong and 5 to Xiaohua. How many candies did Xiaoming have left? The teacher model (Qwen3-32B) generates the following structured reasoning process (each step is a newline): Step 1: Xiaoming originally had 15 candies.

[0044] Step 2: 3 stones were given to Xiaohong, so 15 - 3 = 12 stones are left.

[0045] Step 3: Xiaohua was given 5 more, so there are 12 - 5 = 7 left.

[0046] Step 4: Therefore, Xiaoming has 7 candies left.

[0047] Final answer: 7.

[0048] Step 2: Step Type Labeling and Weight Allocation Using the teacher model itself for type labeling, construct prompt words: "Please determine which category each of the following reasoning steps belongs to: basic calculation, basic facts, operation execution, logical reasoning, strategy planning. If there are multiple categories, please list all categories." Annotation results: Step 3: Constructing Training Data The question, the sequence of steps, the comprehensive weights, and the final answer are assembled into a training sample: question: Xiaoming had 15 candies. He gave 3 to Xiaohong and 5 to Xiaohua. How many candies did Xiaoming have left? Mind chain: [Step 1, Step 2, Step 3, Step 4]; Weight sequence: ; Final answer: .

[0049] Step 4: Weighted distillation training The above samples (and the remaining 4999 samples) were input into the student model Qwen3-8B for training.

[0050] Taking this sample as an example, calculate the loss: Assuming the student model predicts a probability of 0.85 for token "12" in step 2, and the teacher model predicts 1.0, then the token-level cross-entropy loss is: .

[0051] The weighted loss of the thought chain is obtained by summing the weighted losses of all tokens and then averaging them.

[0052] Final answer section: The student model predicts the probability of "7" as 0.92, and the teacher model as 1.0. Therefore, the answer loss is: .

[0053] Total loss: ,but .

[0054] Training uses a batch size of 16 and a learning rate of [missing information]. The AdamW optimizer was trained for 5 rounds, and the validation set loss converged in the 4th round.

[0055] Step 5: Model Output After training, the distilled student model Qwen3-8B-Distilled was obtained. The accuracy comparison results on 500 test sets are as follows: Experimental results show that the method of the present invention is significantly better than the existing distillation method in mathematical reasoning tasks, with an accuracy improvement of about 4.6 percentage points relative to baseline 3 and about 14.8 percentage points relative to baseline 1, verifying the effectiveness of the differentiated supervision mechanism.

[0056] Although embodiments of the invention have been shown and described, it will be understood by those skilled in the art that various changes, modifications, substitutions and alterations can be made to these embodiments without departing from the principles and spirit of the invention, the scope of which is defined by the appended claims and their equivalents.

Claims

1. A knowledge distillation method based on reasoning step deconstruction and differentiated supervision, characterized in that, Includes the following steps: Step S1: Obtain the reasoning process generated by the teacher model for the input question. The reasoning process is deconstructed into multiple reasoning steps separated by rows, and each reasoning step corresponds to an independent reasoning unit. Step S2: Label each reasoning step with its type, determine its reasoning type, and assign a preset weight value to each reasoning type; when a reasoning step is identified as multiple reasoning types at the same time, calculate the comprehensive weight of the step by weight superposition. Step S3: Construct training data, which includes the input question, the inference step sequence of the teacher model, and their corresponding comprehensive weights; Step S4: Distill the student model using the training data. During the training process, the loss function is weighted according to the comprehensive weight of each inference step to obtain the weighted loss. Step S5: Update the student model parameters according to the weighted loss until the model converges, and obtain the distilled student model.

2. The method according to claim 1, characterized in that, The type labeling in step S2 is automatically completed by the teacher model. The labeling is based on a preset reasoning type classification system, which includes: basic calculation, basic facts, operation execution, logical reasoning, and strategy planning.

3. The method according to claim 2, characterized in that, The identification of the reasoning type is achieved through keyword matching or semantic judgment; when a reasoning step belongs to multiple reasoning types at the same time, the weight superposition method is to add the weight values ​​corresponding to each type, and the upper limit of the comprehensive weight after superposition is 1.

0.

4. The method according to claim 2, characterized in that, The inference types and their weights are set as follows: Basic calculations: weight 0.3, corresponding to arithmetic operation steps; Basic facts: weighted at 0.4, corresponding to the fact retrieval or information extraction steps; Operation execution: Weight is 0.6, corresponding to the application method, formula or rule steps; Logical reasoning: weighted at 0.8, corresponding to conditional judgment, causal analysis, or logical deduction steps; Strategy planning: with a weight of 1.0, corresponding to the steps of problem-solving plan formulation, method selection, or problem decomposition.

5. The method according to claim 1, characterized in that, The weighted loss mentioned in step S4 includes weighted thought chain loss and final answer loss. The weighted loss is calculated as follows: in, This represents the total number of tokens in the MindChain. For the first The overall weight of the reasoning step to which each token belongs. Let cross-entropy be the loss function. and These represent the output probability distributions of the teacher model and the student model on the final answer, respectively. and These are the preset weighting coefficients.

6. The knowledge distillation method based on reasoning step deconstruction and differential supervision according to claim 5, characterized in that: The The value is 0.

7. The value is 0.

3.

7. The method according to claim 1, characterized in that: The teacher model is a large-scale language model with 30 billion or more parameters, and the student model is a medium-sized language model with 10 billion or less parameters.

8. The method according to claim 1, characterized in that: In the distillation training process described in step S4, multiple baseline methods are used for comparative experiments. These baseline methods include: Baseline Method 1: Standard output distillation supervised only for the final answer; Baseline Method 2: Supervised sequence distillation that treats the entire thought chain as a flat sequence; Baseline Method 3: Treat inference steps as supervised step distillation with equal weights; The superiority of this method is verified by comparing the accuracy of answers on the test set.

9. A knowledge distillation system based on reasoning step deconstruction and differentiated supervision, characterized in that, include: The reasoning step deconstruction module is used to obtain the reasoning process generated by the teacher model and divide it into multiple reasoning steps by row. The type labeling module is used to identify the type of each inference step and assign corresponding weights. When the same reasoning step corresponds to multiple types, the comprehensive weight is calculated by weight superposition. The training data construction module is used to generate training samples that include input questions, inference step sequences, and comprehensive weights. The weighted distillation training module is used to weight the loss according to the comprehensive weight of the inference steps and update the model parameters during the training of the student model. The output module is used to output the trained student model.

10. The system according to claim 9, wherein the type labeling module has a built-in reasoning type classification system, wherein the reasoning type includes basic calculation, basic facts, operation execution, logical reasoning, and strategy planning, and each corresponds to a different weight.

Images