Knowledge distillation (KD), introduced by Hinton et al. [4], transfers the output distribution of a teacher model to a lightweight student model through temperature-scaled soft labels. In this way, the student learns not only from ground-truth labels but also from the “dark knowledge” encoded in the relationships among classes. Beyond output-level distillation, researchers have extended knowledge transfer to intermediate features and structural relationships. FitNets introduced intermediate representations from the teacher model as hints, enabling the student to learn deeper feature representations [6]. Attention Transfer further transfers knowledge through the attention maps of the teacher and student models, allowing the student to learn the teacher’s patterns of focus on input information [14]. Relational Knowledge Distillation goes beyond matching the features of individual samples by constraining the distances and angular relationships between samples, thereby transferring the structured knowledge learned by the teacher model [15]. These studies demonstrate that the knowledge of a teacher model is not limited to its final predictions but is also encoded in intermediate representations, attention distributions, and inter-sample relationships.

Driven by progress in pre-trained language models, knowledge distillation is now commonly employed to compress and accelerate natural language processing models. During pre-training, DistilBERT combines the language modeling loss, distillation loss, and representation-level constraints to compress BERT into a smaller student model, reducing the number of parameters while preserving strong language understanding capabilities [8]. TinyBERT further performs multi-level distillation on the embedding representations, hidden states, attention matrices, and output layer of Transformer models. It also improves the performance of student models on natural language understanding tasks through the combination of general distillation and task-specific distillation [9]. MobileBERT designs a network architecture suitable for resource-constrained devices and incorporates knowledge distillation to balance model performance and deployment efficiency [16]. In addition, MiniLM distills the self-attention relations and self-attention values of the teacher model, thereby preserving important linguistic knowledge of Transformer models at a relatively low computational cost [17].

Existing knowledge distillation methods primarily focus on maximizing the capacity and information coverage of knowledge transfer, but relatively little attention has been given to two critical factors that limit distillation performance. First, student models may have limited ability to learn from difficult samples or challenging semantic segments. Second, teacher predictions inevitably contain uncertainty, which may cause erroneous or low-reliability knowledge to be propagated to the student model. To overcome these challenges, we propose a knowledge distillation framework that makes joint use of the Segment-wise Focal Loss (SFL) and the Adaptive Temperature Distillation Loss (ATL). The proposed framework jointly optimizes the distillation process from two perspectives—imbalanced sample difficulty and uncertainty in teacher predictions—thereby improving the robustness and effectiveness of knowledge transfer.

Focal loss was originally proposed by Lin et al. [18] to address the severe imbalance between positive and negative samples in object detection. In most cases, network training is guided by the cross-entropy loss. However, focal loss generally achieves superior performance to cross-entropy loss in scenarios involving imbalanced class distributions, foreground-background imbalance in the image domain [19], or imbalanced positive and negative samples in multi-label classification [20]. By modifying the cross-entropy loss, focal loss encourages neural networks to focus more on difficult-to-classify samples and assign greater attention to samples that are prone to misclassification. The conventional cross-entropy loss treats all samples equally, allowing the large number of easy negative samples to dominate the training process and consequently limiting the model’s ability to learn from difficult samples. Focal loss introduces an adjustable modulating factor that dynamically reduces the weights of easy samples while increasing the relative emphasis placed on difficult ones. Its mathematical formulation is given as follows:

here, L l c (Low-Confidence Loss) primarily focuses on the loss values of low-confidence samples. p i denotes the predicted probability of a sample i , while y i denotes its correct label. Focal loss modifies the cross-entropy loss by introducing a weighting term ( 1 − p i ) γ l c . When the predicted probability p i of a training sample approaches 1, indicating that the model has high confidence in its prediction, the loss weight assigned to that sample decreases rapidly. Consequently, compared with conventional cross-entropy loss, its loss term declines more quickly. This weighting mechanism directs the training process toward the relatively small number of low-confidence samples, thereby improving the model’s ability to learn difficult-to-classify examples.

Ridnik et al. [19] applied focal loss to multi-label classification tasks to address the imbalance between positive and negative labels. In such tasks, each image typically contains only a small number of positive labels, while most labels are negative, resulting in a much larger number of negative samples than positive samples. To address this issue, the authors proposed a variant of Focal Loss called Asymmetric Loss (ASL), which achieved better performance than the original Focal Loss in multi-label classification tasks. ASL separates the loss terms for positive and negative samples in Focal Loss and assigns an independent focusing parameter to each, thereby improving the modeling of positive-negative sample imbalance and enhancing classification performance.

(2)

here, y i denotes the ground-truth label of the sample, p i denotes the predicted probability obtained after applying the sigmoid function to the sample, γ + is the focusing parameter for positive samples, which enhances the learning of difficult positive samples, and γ − is the focusing parameter for negative samples, which suppresses the loss contribution of easily classified negative samples. Nevertheless, existing methods still cannot fully account for the characteristics of different training stages or flexibly adjust the degree of attention assigned to samples throughout the training process. Inspired by curriculum learning, this paper proposes a periodic segment-wise focal loss that divides the training process into multiple stages and dynamically adjusts the weighting assigned to the focal loss at each stage.

General Cyclical Training (GCT) [20] refers to a general machine learning training paradigm that partitions the training procedure into a series of suitably ordered stages. In the initial phase, the model prioritizes samples that are easy to classify.; in the middle stage, its attention expands to more challenging samples; and in the final stage, it returns to high-confidence samples for fine-tuning, thereby improving generalization and convergence efficiency. This strategy can be regarded as a combination of Curriculum Learning [21] and fine-tuning, while also leveraging the entire dataset for comprehensive optimization during the middle stage of training [22]. Studies have shown that the key learning processes of neural networks typically occur during the early stages of training. Therefore, properly scheduling the difficulty of samples can help improve the model’s optimization trajectory and final performance.

Within the General Cyclical Training framework, the Cyclical Learning Rate (CLR) [23] has become one of the mainstream approaches. By dynamically adjusting the learning rates for different samples, ASL can balance the importance of different samples. In addition, Learning Rate Warmup [24] and Stochastic Gradient Descent with Warm Restarts (SGDR) [25] are also regarded as variants of CLR. These methods adjust the learning rate during optimization to improve convergence speed and generalization ability. In recent years, the concept of adaptive hyperparameters has gradually been extended to weight decay [26] and dynamic batch-size adjustment strategies, further enhancing the convergence stability of models. Based on these developments, this paper further extends the principles of General Cyclical Training and explores their application to loss-function optimization. Specifically, cyclical focal loss is employed to adaptively adjust loss weights at different training stages, thereby improving the model’s ability to learn from imbalanced datasets.

The temperature parameter plays an important role in knowledge distillation. By controlling the smoothness of the soft labels output by the teacher model, it affects the learning process of the student model. Traditional knowledge distillation methods typically use a fixed temperature parameter τ , and the soft-label generation formula is given as follows:

here, z i denotes the predicted probability output for class i by the teacher model, and q i denotes the generated soft label. Although a fixed temperature parameter is simple to use, it has clear limitations in practical applications. For samples with high predictive uncertainty, an excessively high temperature value may amplify the uncertainty of the soft labels, causing the student model to inherit incorrect predictions from the teacher model. For easily classified samples, a fixed temperature may fail to fully exploit the relational information among different classes.

To address this issue, researchers have proposed various dynamic temperature scaling methods. For example, Zhang et al. [27] proposed a temperature adjustment method based on sample uncertainty. This method dynamically adjusts the temperature parameter according to the prediction entropy of each sample, as shown in the following equation:

here, H ( q ) denotes the entropy of the soft labels and is used to measure the predictive uncertainty of each sample; τ base denotes the base temperature, and α isa scaling coefficient. This method can dynamically adjust the temperature according to the uncertainty of each sample, thereby reducing the impact of incorrect predictions on the student model. However, existing methods typically focus only on predictive uncertainty and fail to fully consider the difficulty of each sample and its influence on temperature adjustment. Therefore, this paper proposes an adaptive temperature distillation loss that further extends the concept of dynamic temperature adjustment. Specifically, it assigns a customized temperature parameter to each sample and dynamically adjusts it according to both the predictive uncertainty and difficulty of the sample.

The imbalanced distribution of easy and difficult samples is prevalent in most natural language processing tasks. Conventional focal loss prioritizes difficult samples; if this bias is maintained throughout the entire training process, the model may easily overfit noisy samples, overlook the fundamental semantic features contained in easy samples, and consequently suffer a decline in generalization performance. Following the cyclical training strategy, the model should dynamically adjust the weights of easy and difficult samples at different training stages. Accordingly, this paper designs a novel loss-weighting term aimed at guiding the teacher model during fine-tuning and the student model during knowledge transfer to prioritize learning from easy samples in the early stages of training. Its specific form is given as follows:

Equation (1) represents the focal loss, which encourages the model to focus more on low-confidence samples during training. In contrast, the loss term L h c (High-Confidence Loss) places greater emphasis on high-confidence samples. Figure 2 compares the standard cross-entropy loss (CE) with the low-confidence loss L l c and high-confidence loss L h c under different parameter settings ( γ = 2 and γ = 4 ).

Figure 2. This figure compares the variations of cross-entropy loss, focal loss, and the proposed high-confidence weighting term across different ground truth probabilities.

As shown in Figure 2, different loss-weighting terms have varying effects on the loss function as the true class probability p i changes. The blue solid line represents the cross-entropy loss (CE), whose value decreases monotonically as p i increases. L l c primarily reduces the loss weight assigned to high-confidence samples ( p i → 1 ) while increasing the loss for low-confidence samples ( p i → 0 ), thereby enabling the model to focus more on difficult-to-classify samples. In contrast, L h c increases the loss for high-confidence samples while suppressing the influence of low-confidence samples. In particular, when γ h c takes a relatively large value (e.g., γ h c = 4 ), the loss weight of high-confidence samples increases significantly.

The objective of segment-wise focal loss is to dynamically adjust the model’s attention to samples at different stages of training, thereby improving learning performance. During the early training stage, SFL primarily focuses on high-confidence samples to ensure that the model’s gradient updates are more stable and that the model converges rapidly. During the middle stage, SFL gradually shifts its focus toward difficult-to-classify samples, namely those that are more likely to be misclassified. This enhances the model’s ability to learn complex patterns and improves its generalization performance.

To achieve this objective, SFL combines the two loss terms in Equation (1) and Equation (5), dynamically balancing them through a reasonable scheduling strategy. In this study, a linear scheduling strategy is adopted, allowing the loss-function weights to smoothly shift from easy-to-classify samples to difficult-to-classify samples throughout the training process. This approach not only ensures training stability but also prevents the model from becoming trapped in a local optimum, thereby further improving classification performance. Accordingly, this paper defines a parameter ξ , which varies with the training epoch as follows:

here, e i denotes the current training epoch, while e n denotes the total number of training epochs. We introduce a cycle factor f c ≥ 1 to provide flexibility for cyclic scheduling. When f c = 1 , ξ decreases linearly from 1 at the beginning of training to O at the end. When f c = 2 , the cyclic curve resembles an inverted equilateral triangle: ξ decreases from 1 to 0 during the first half of training, and then increases linearly from O to 1 during the second half. When f c = 4 , ξ decreases from 1 to 0 during the first quarter of the training epochs, and then increases linearly to 1 over the remaining three quarters.

Figure 3 is plotted according to Equation (6), illustrating the effects of different cycle factors, with the total number of training epochs set to e n = 20 . The curves correspond to f c = 1 (blue), f c = 2 (red), and f c = 4 (green), reflecting the cyclical scheduling patterns of ξ as the training epochs change. Based on Equations (1) (5), and (6), this paper defines the segment-wise focal loss as follows:

Figure 3. Hyperparameter change curves under different settings.

here, p denotes the probability predicted by the model, and y denotes the true label of the sample. By introducing the linear scheduling parameter ξ , the weighting ratio of the two focal losses is dynamically adjusted at different training stages, guiding the model to focus on samples with varying confidence levels. By introducing the scheduling parameter ξ , the model dynamically adjusts its attention to samples of varying difficulty, as well as their corresponding loss weights, at different stages of training. During the early stage, the weights assigned to easy samples are appropriately increased to facilitate the rapid establishment of stable classification representations. During the middle stage, the learning focus gradually shifts toward hard-to-classify samples, thereby enhancing the model’s ability to identify complex patterns. During the late stage, attention is again placed on easy samples to consolidate the learned feature representations and improve training stability and generalization performance.

Before transferring knowledge to the student model, the teacher model may produce uncertain or erroneous predictions, leading to inherited errors during the knowledge distillation process. Consequently, the student model may also face a high probability of uncertainty in its predictions. Some of the uncertainty in the soft labels arises from an excessively high temperature setting during distillation. Therefore, this paper introduces an adaptive-temperature distillation loss, which adaptively adjusts the temperature according to the prediction confidence of each sample, thereby mitigating the negative impact of uncertain predictions.

1) Assume that the teacher model outputs the class probabilities as p = [ p 1 , p 2 , ⋯ , p C ] , where C denotes the number of classes. To quantify the degree of uncertainty in the teacher model’s predictions for a sample, this paper defines the confusion degree w x , which is calculated as follows:

here, p y denotes the predicted probability corresponding to the true label y of the sample, and p i denotes the predicted probability for another class i ( i ≠ y ). This formula calculates the average of the absolute differences between the predicted probability of the true label and those of the other classes, thereby reflecting the degree of confusion in the model’s predictions. When w i is close to 0, the model’s prediction for the sample is relatively clear and confident; when w i is close to 1, the model’s prediction for the sample is relatively confused.

2) To further optimize the temperature adjustment mechanism in the knowledge distillation process, this paper introduces an adaptive temperature parameter based on the confusion degree w i , which is calculated as follows:

here, τ max and τ min represent the minimum and maximum temperature values, respectively. When the degree of confusion is high, the temperature approaches the lower value, thereby widening the probability gap between classes and helping the student model focus on the classes in which the teacher model has higher confidence. When the degree of confusion is low, the temperature approaches the higher value, enhancing the smoothness of the class probabilities and enabling the student model to learn more inter-class relational information.

3) The final adaptive temperature distillation loss function is defined as follows:

here, p i t and p i s denote the output probability distributions for sample i produced by the teacher model and the student model, respectively. KL ( ⋅ ∥ ⋅ ) denotes the Kullback-Leibler divergence, which is used to measure the difference between the output distributions of the teacher and student models. τ i is the adaptive temperature parameter calculated using Equation (9), which is dynamically adjusted according to the confusion degree w x of the sample.

As illustrated by the overall framework in Figure 1, the SAT training procedure introduced in this study is primarily composed of two stages, which are detailed below:

1) Stage 1: Fine-tuning of the teacher model

Before knowledge distillation, the teacher model needs to be fine-tuned according to the task-related knowledge so that it can subsequently transfer this knowledge to the student model. Previously, the teacher model used the cross-entropy loss function to adjust its parameters. However, because the task contains samples with different levels of difficulty, the teacher model should focus on samples with different confidence levels at different stages. Therefore, during the fine-tuning process, this paper adopts the segment-wise focal loss as the training objective.

Given the training sample set { ( x i , y i ) } i = 1 N , where x i denotes the input sample and y i denotes its corresponding ground-truth label, the teacher model outputs a class-probability distribution p i t = f θ T ( x i ) , where θ T denotes the parameters of the teacher model. The optimization objective of the teacher model is defined as minimizing the periodic focal loss:

here, N denotes the total number of samples. L SFL consists of two components: the low-confidence focal loss L l c and the high-confidence focal loss L h c . The periodic parameter ξ guides the model to adjust its attention to easy and difficult samples at different stages. To balance the training weights between hard and easy samples during the fine-tuning phase of the teacher model, the period factor f_c is set to 2. This allows the parameter ξ to be dynamically adjusted throughout the full training cycle: the first quarter of the epoch focuses on easy samples, the middle half focuses on hard samples, and the final quarter focuses on easy samples again.

2) Stage 2: Training the student model

After the teacher model has been fine-tuned for the specific task, its knowledge is transferred to the student model. The student model is usually smaller, and its training objective mainly consists of one part: imitating the prediction results of the teacher model.

Similar to the teacher model, given the same training sample set { ( x i , y i ) } i = 1 N , where x i denotes the input sample and y i denotes its corresponding ground-truth label, the student model outputs a class probability distribution p i s = f θ S ( x i ) , where θ S denotes the parameters of the student model. To optimize the student model, two main loss functions are designed: the periodic focal loss and the adaptive temperature confusion loss. The overall objective is defined by combining these two loss functions, enabling the student model to learn from the ground-truth labels while acquiring knowledge from the teacher model’s soft labels. The specific loss function is defined as follows:

here, α is a weighting coefficient used to balance the contributions of the two losses. During training, the student model updates its parameters θ S through backpropagation to minimize the overall loss L SAT . Based on the updated student-model parameters, the process is repeated until the student model converges. Within the knowledge distillation framework, to prevent the student model from overfitting simple samples, this mechanism directs the model to learn simple samples briefly during the early stages of training before shifting its learning focus to difficult samples, thereby enhancing the model’s generalization performance. To this end, this paper adopts a periodicity factor of f c = 4 for the periodic focal loss. This configuration enables the model to dynamically shift its learning focus throughout the training cycle: prioritizing simple samples during the first one-eighth of the cycle, focusing on hard samples during the middle five-eighths, and returning to learning simple samples during the final quarter.

All experiments are performed on a machine with an Intel Xeon E3-1231 CPU and an NVIDIA RTX 3060 GPU. We implement the proposed method and baselines in PyTorch using PyCharm.

This work validates the proposed approach on a suite of tasks from the General Language Understanding Evaluation (GLUE) benchmark [28]. The benchmark consists of single-sentence tasks including [29] and SST-2 [30], natural language inference tasks covering MNLI [31], QNLI [32], and RTE [33], as well as similarity and paraphrase tasks such as MRPC [34], STS-B [35], and QQP [36]. Multiple metrics are adopted to measure model performance: accuracy is reported for MNLI, SST-2, QNLI, QQP and RTE; F1-score for MRPC; Matthews correlation coefficient for CoLA; and Spearman rank correlation coefficient for STS-B. This combination of metrics enables a holistic assessment of our method. Detailed descriptions of these datasets are presented in Table 1.

Table 1. Details of GLUE datasets.

Furthermore, to further analyze the effect of our method on mitigating the imbalance between easy and hard samples as well as prediction uncertainty, we adopt multiple evaluation metrics for comparison. Among them, Top-1 denotes the prediction accuracy of the category with the fewest samples in the training set; a higher value indicates stronger capability of the model to learn hard samples. In addition to Top-1, we introduce the Area Under Risk-Coverage Curve (AURC, ×10³). A lower AURC implies that the model can more effectively distinguish correct predictions from incorrect ones, thereby alleviating the impacts brought by high-uncertainty samples.

Stage 1: This paper uses the pretrained BERT-base-uncased model as the teacher model. Its architecture consists of a 12-layer Transformer encoder, with a hidden-layer dimension of 768 and approximately 110M parameters. To adapt the model to different tasks, it is fine-tuned separately on six specific tasks from the GLUE benchmark. During fine-tuning, the learning rate is uniformly set to 3e−5, and the training period is set to 20 epochs to ensure that the model can learn and converge sufficiently on each task.

In addition, two new hyperparameters, γ l c and γ h c , are introduced into the periodic focal loss. According to the study by Smith et al. [20], the focal-loss parameter γ l c is not highly sensitive, and setting γ l c = 4 can achieve good performance on most tasks. Therefore, this paper sets γ l c = 4 . Furthermore, experimental results show that the model generally performs well when γ h c is set to 2 or 3.

Stage 2: This paper uses four-layer BERT models, BERT_base (4/768; 52 M) as student models. Each student model is trained on a specific task using the corresponding teacher model fine-tuned in Stage 1, and knowledge distillation is performed according to Equation (12). The learning rate is selected from {2e−5, 3e−5, 4e−5}. The hyperparameter α is selected from {0.3, 0.4, 0.5 }, while the temperature parameters are selected from τ max ∈ { 6 , 8 , 10 } and τ min ∈ { 3 , 4 , 5 } . In addition, for the two hyperparameters γ l c and γ h c in the periodic focal loss, the same settings as those used during teacher-model fine-tuning are adopted: γ l c is fixed at 4, while γ h c is selected as either 2 or 3.

To validate the effectiveness of our method, we adopt multiple well-performed knowledge distillation methods for NLP tasks as baselines.

Vanilla KD [4] employs basic knowledge distillation methods that transfer knowledge through the prediction distributions.TinyBERT [9] adopts layer-wise distillation based on hidden representations and self-attention matrices. It proposes two distillation stages—general learning and task-specific learning—and employs a data augmentation strategy.CKD [37] and MGSKD [38] extract pairwise and triplet structural relationships between the hidden representations of the teacher and student models, and then calculate the corresponding loss.MiniLMv2 [39] first performs task-agnostic distillation and then fine-tunes the model on task-specific datasets. By introducing multi-head self-attention relations, it allows greater diversity in the number of attention heads used by the student model.

Note that we re-implement several methods (MGSKD, MiniLMv2) using their original codes for fair comparisons.

This paper compresses the model size from 110 M to 52 M parameters, yielding the student model SAT_Small. Its overall accuracy on the GLUE benchmark is presented in Table 2. Compared with the teacher model BERT-base, SAT_Small reduces the model size by 52.7%, while its average performance decreases by only 3.2%, demonstrating excellent model compression effectiveness.

Detailed comparisons with other student models show that SAT_Small achieves superior performance on the majority of tasks, with an average accuracy of 80.8%. This represents improvements of 1.8% over BERT_Small, 2.1% over TinyBERT₄, 3.3% over MiniLMv2₄, 3.8% over CKD₄, and 2.8% over MGSKD₄. In particular, SAT_Small achieves the best results on the sentiment classification task SST-2, the semantic matching task MRPC, and the natural language inference task QNLI, with accuracies of 90.0%, 90.6%, and 90.4%, respectively. These results demonstrate clear performance gains over the other student models.

Across all evaluated tasks, SAT_Small achieves its most significant improvement on the MRPC dataset, attaining an F1 score of 90.6%. This is 1.0% higher than BERT_Small, 4.8% higher than TinyBERT₄, 5.3% higher than MiniLMv2₄, 8.0% higher than CKD₄, and 8.2% higher than MGSKD₄, representing a substantial improvement over the baseline methods. These results demonstrate that SAT_Small has stronger representational capability on tasks that rely on capturing semantic similarity and learning fine-grained features.

Table 2. Comparative results show that the student model obtained in this study reduces the number of model parameters from 110 M to 52 M. The optimal value of the student model is shown in bold.

Table 3 presents the performance of our proposed student model SAT_Small. The model compresses the parameter count from 110M to 52M and achieves notable improvements on hard-sample classification tasks across all datasets of the GLUE benchmark. Compared with the teacher model BERT-base, SAT_Small only suffers a 3.7% drop in Top-1 accuracy (78.6% vs. 82.3%), indicating that it can retain performance close to the teacher model even with drastic parameter compression. When compared with the undistilled student model BERT_Small (74.4%), SAT_Small obtains a 4.2% accuracy gain, which further verifies its effectiveness for model compression while preserving performance on hard samples.

For individual tasks, SAT_Small achieves an accuracy of 88.8% on QQP, outperforming BERT_Small (86.4%) and achieving comparable results to the state-of-the-art MiniLMv2₄ (90.3%). Although the positive and negative samples in QQP are relatively balanced, the promising performance of SAT_Small demonstrates its strong robustness when handling class imbalance. The model can effectively capture hard-to-distinguish semantic patterns and improve prediction performance for minority classes. On the MRPC semantic matching task, SAT_Small reaches an F1 score of 90.2%, substantially exceeding BERT_Small’s 84.6%. The MRPC dataset contains many low-frequency non-equivalent sentence pairs. This result further proves that knowledge distillation enables SAT_Small to boost prediction accuracy for hard and minority-class samples.

In summary, SAT_Small achieves strong Top-1 performance. Especially on hard-sample tasks such as QQP and MRPC, it delivers better classification ability than the undistilled BERT_Small. The experiments demonstrate that SAT_Small achieves remarkable advantages in both model compression and hard-sample classification.

Table 3. The comparison results indicate that the second student model obtained in this paper compresses the model parameters from 110 M to 52 M and achieves Top-1 prediction accuracy.

Table 4 evaluates the ability of the 4 × 768 student model to distinguish correct and incorrect predictions, i.e., the model’s capability to quantify prediction uncertainty, via AURC (Area Under the Risk-Coverage Curve). Experimental results demonstrate that the SAT_Small student model achieves outstanding performance in uncertainty estimation, with an average AURC of 50.1. This value is markedly lower than that of the undistilled BERT_Small (average AURC = 55.4) and close to the performance of the teacher model BERT-base (average AURC = 45.9).

Specifically, on the SST-2 task, SAT_Small achieves an AURC of 59.2, which is substantially lower than BERT_Small’s 64.2 and close to the teacher model’s 51.5. For the RTE task, SAT_Small obtains an AURC of 55.7, outperforming BERT_Small (61.2) and approaching the teacher model’s 55.6. These observations indicate that SAT_Small effectively learns the teacher model’s ability to quantify prediction uncertainty through knowledge distillation. While reducing parameter size, the model significantly improves its capability to identify wrong predictions. In contrast, without supervision from the teacher model, the undistilled BERT_Small yields generally higher AURC values, which reflects its weaker uncertainty estimation performance.

Overall, SAT_Small not only narrows the performance gap with the teacher model on uncertainty estimation tasks but also surpasses the undistilled student model by a notable margin. The results verify that by introducing the adaptive temperature distillation loss, the SAT method successfully alleviates the “error inheritance” problem from the teacher model during knowledge distillation.

Table 4. The ability of the student model to discriminate between correct and incorrect predictions was assessed by AURC.

To verify the language modeling capability of the small language model, we perform pre-training distillation on the Open WebText dataset. This dataset contains approximately 8 million training documents with 38 GB of text, which is suitable for large-scale natural language processing tasks. Furthermore, we further evaluate the generalization ability of the model through zero-shot task evaluation and transfer learning experiments. The main datasets used include LAMBADA [40] (for evaluating text comprehension and generation, with perplexity PPL and accuracy ACC as metrics), WikiText-103 [41] (for evaluating language modeling performance, with perplexity PPL as the metric), and The Children’s Book Test [42] (for evaluating text comprehension and generation, with accuracy ACC as the metric).

Perplexity (PPL) is a commonly used evaluation metric for language models [26], primarily used to measure a model’s ability to predict text sequences. PPL reflects the model’s average degree of uncertainty on test data; a smaller value indicates stronger predictive ability.

For a given text sequence X = ( x 1 , x 2 , ⋯ , x n ) , perplexity is calculated as:

where n denotes the total length of the text sequence, and P ( x i | x 1 , x 2 , ⋯ , x i − 1 ) denotes the model’s conditional probability prediction for the i -th word perplexity measures a language model’s ability to predict a text sequence. A smaller value indicates greater confidence when the model generates or predicts text. Generally, a lower PPL means that the text generated by the language model is closer to real-world text.

Within the framework of knowledge distillation for language modeling, the distillation process utilizes a large-scale general-purpose text corpus (specifically, the Open WebText dataset) with the aim of transferring the general language modeling knowledge acquired by the teacher model—such as implicit knowledge regarding contextual dependencies and word probability distributions—to the student model.

Stage 1: This paper uses a pretrained 12-layer GPT-2 model as the teacher model. Each hidden layer has a dimension of 768, and the model contains 125M parameters. During fine-tuning, the learning rate is fixed at 3e−5, and the training period is set to 10 epochs to ensure that the model can learn sufficiently and reach convergence on each task.

For the two hyperparameters in the periodic focal loss, the teacher model continues to use the same fine-tuning settings for natural language understanding tasks: γ l c = 4 , while γ h c is set to either 2 or 3. As with natural language understanding tasks, the period factor f c is set to 2.

Stage 2: This paper uses a 4-layer GPT-2 (4/312, 19.6 M) model as the student model and performs distillation from the teacher model after it has been fine-tuned on the specific task, following Equation (12). The learning rate is selected from {2e−5, 3e−5, 4e−5, 5e−5}, with the best configuration being chosen The hyperparameter α is selected from { 0.4 , 0.5 , 0.6 , 0.7 } , while the two temperature hyperparameters in Equation (9) are selected from τ max ∈ { 6 , 8 , 10 } and τ min ∈ { 3 , 4 , 5 } .

For the two hyperparameters in the periodic focal loss, γ l c and γ h c , the same settings as those used during teacher-model fine-tuning are adopted. The low-confidence hyperparameter is fixed at γ l c = 4 , while the high-confidence hyperparameter γ h c is selected as either 2 or 3. In contrast to the training weight allocation for natural language processing tasks—where simple and difficult samples are treated equally for the student model—the period factor f c is set to 2.

To verify the effectiveness of the proposed method, the following approaches are selected for comparison:

Vanilla KD [4] is trained by minimizing the discrepancy between the probability distributions output by the student and teacher models.Layer-wise Distillation (LWD) matches the hidden vectors at each layer and optimizes the student model based on the consistency of hidden-layer features.TED [43] uses a fully connected layer to select hidden features and calculates the knowledge discrepancy between the selected features and the corresponding teacher-model features.CI-WA [44] employs a dynamic extractor to obtain word-level correlations from each sample and transfers them as knowledge to the student model. It further enhances low-level feature representations through a top-down cross-layer integration mechanism.A2D [45] aligns attention heads to avoid heuristic selection of layer mappings.

Table 5 presents the zero-shot learning and transfer learning performance of the GPT-2₄ model on the test set, compared with the teacher model GPT-2₁₂. The experimental results show that the teacher model performs well on both zero-shot and transfer learning tasks, whereas the baseline student model (GPT-2₄) performs poorly without knowledge distillation. This is particularly evident in the zero-shot learning task on the LAMBADA dataset, where its perplexity (PPL) reaches 99.8 and its accuracy is only 16.4%.

In the zero-shot learning task, SAT achieves the best overall performance across the three datasets. This indicates that dynamically emphasizing sample difficulty at different stages and appropriately adjusting inter-class differences can help the model better adapt to unseen classes or tasks.

Through knowledge distillation (KD), layer-wise distillation (LWD), task-related distillation (TED), word-level association knowledge distillation (CI-WA), and self-adaptive temperature periodic knowledge distillation (SAT), the performance of the student model is significantly improved. For example, GPT-2₄ (SAT) achieves a transfer-learning accuracy of 80.3% on the CBT dataset, the highest among all methods.

Table 5. The zero-shot and transfer learning performance of the GPT-2₄ model on the test set.

This ablation study verifies the effectiveness of segment-wise focal loss for knowledge distillation on multiple NLU tasks. The results are presented in Table 6. The undistilled baseline Student4 underperforms Teacher12 severely, reaching only 63.2% on RTE and 89.6% on MRPC and struggling with imbalanced low-frequency samples. Vanilla KD brings moderate gains, lifting SST-2 from 89% to 90.6%.

FL + KD further improves RTE and MRPC to 62.9% and 90%. Our SFL + KD achieves the best or near-best results across all tasks, with 91.2% accuracy on SST-2 and 50.9% MCC on CoLA. Segment-wise focal loss dynamically reweights samples and boosts prediction accuracy for hard and uncertain instances.

Table 6. The impact of segment-wise focal loss on the BERT4 student model in balanced and imbalanced datasets.

To investigate the effects of segment-wise focal loss (SFL) and adaptive temperature distillation loss (ATL) on imbalanced datasets, this paper conduct experiments on QNLI, CoLA and MRPC. Following Kim et al. [46], besides Top-1 and AURC, we also report Expected Calibration Error (ECE, %). Lower ECE indicates better calibration, where model confidence aligns closely with actual accuracy. Top-1 accuracy denotes the prediction accuracy of the majority class, which contains the largest number of samples in the dataset.

1) ECE

Expected Calibration Error (ECE) is a commonly used model calibration metric, primarily used to measure the agreement between a classification model’s predicted confidence and its actual accuracy. ECE reflects whether the model tends to be overly confident or overly conservative in its predictions. It divides the model’s predicted confidence scores into M bins, and then calculates the difference between the average confidence and accuracy within each bin. The formula is:

where B m denotes the sample set in the m -th bin, | B m | denotes the number of samples in the m -th bin, n denotes the total number of samples, acc ( B m ) denotes the average accuracy of the samples in the m -th bin, and conf ( B m ) denotes the average confidence of the samples in the m -th bin. The value of ECE ranges from [0, 1]. A smaller value indicates that the model’s confidence is closer to the actual probability, suggesting better calibration.

2) AURC

AURC (Area Under the Reliability Curve) is a metric used to measure the reliability of a classification model’s predictions. It is typically used to evaluate whether the model’s confidence is consistent with its actual performance Similar to Expected Calibration Error (ECE), AURC focuses on the relationship between the model’s predicted confidence and its actual accuracy. The formula for calculating AURC is:

where acc ( conf ) is the actual accuracy at a given confidence level conf. The value of AURC generally ranges from 0 to 1. A value closer to 1 indicates a larger discrepancy between the model’s confidence and its actual accuracy, suggesting that the model’s predictions are unreliable. Conversely, an AURC value closer to 0 indicates that the model’s predicted confidence is consistent with its actual accuracy, suggesting that the model is relatively reliable.

As shown in Table 7 for QNLI, either SFL or ATL alone improves Top-1 accuracy but yields limited gains for ECE and AURC. Combining SFL and ATL (SAT) achieves the best results across all metrics: Top-1 accuracy of 80.4%, ECE of 2.39%, and AURC of 49.1. This verifies the complementarity and effectiveness of SFL and ATL in knowledge distillation.

Table 7. Comparative experiments of the student model BERT4 on the QNLI dataset.