Confidence Preservation Property in Knowledge Distillation Abstractions (2401.11365v1)

Abstract: Social media platforms prevent malicious activities by detecting harmful content of posts and comments. To that end, they employ large-scale deep neural network LLMs for sentiment analysis and content understanding. Some models, like BERT, are complex, and have numerous parameters, which makes them expensive to operate and maintain. To overcome these deficiencies, industry experts employ a knowledge distillation compression technique, where a distilled model is trained to reproduce the classification behavior of the original model. The distillation processes terminates when the distillation loss function reaches the stopping criteria. This function is mainly designed to ensure that the original and the distilled models exhibit alike classification behaviors. However, besides classification accuracy, there are additional properties of the original model that the distilled model should preserve to be considered as an appropriate abstraction. In this work, we explore whether distilled TinyBERT models preserve confidence values of the original BERT models, and investigate how this confidence preservation property could guide tuning hyperparameters of the distillation process.

• The paper reveals that even when distilled models maintain classification accuracy, they may not consistently preserve the teacher model's confidence scores.
• The study employs black-box equivalence checking and a composite loss function that integrates cross-entropy with confidence alignment to evaluate model behavior.
• Experimental results on GLUE benchmark tasks indicate that fine-tuning distillation hyperparameters can help retain confidence levels without compromising accuracy.

Introduction

Knowledge distillation has emerged as a compelling technique for the compression of large deep learning models, especially within the field of natural language processing. BERT, one of the seminal works in LLM architecture, presents challenges such as computational intensity and memory requirements which hinder its deployment in limited-resource environments. Knowledge distillation addresses these constraints by enabling the training of a smaller, more efficient model—termed a student—that approximates the performance of the larger, original model—the teacher.

Background and Significance

Until recently, the primary focus of knowledge distillation has been to maintain the student model's accuracy relative to the teacher's. This paper, however, poses the question of whether additional properties beyond mere classification accuracy, specifically confidence values, are preserved in the distilled versions of BERT, specifically TinyBERT. The exploration begins with the backdrop of the crucial role that BERT-based models play in content moderation on social media platforms, where the precision of sentiment analysis directly impacts the detection and handling of harmful content. The practical application demands that models not only accurately classify but also maintain reliable confidence scores.

Methodology

To attain the distillation objectives, the loss function typically involves a composite of the student's cross-entropy loss and a term that penalizes prediction differences between student and teacher logits. This research introduces black-box equivalence checking to investigate the preservation of the confidence property during knowledge distillation. The paper further explores how this newly outlined confidence preservation criterion could influence the fine-tuning of hyperparameters in the distillation training process. The key insight is that even when knowledge distillation succeeds in maintaining decision accuracy, it may not necessarily uphold confidence values—an essential consideration for use cases such as AI-assisted diagnostics and safety-critical systems like autonomous driving.

Experimental Results and Analysis

The results from six different tasks within the General Language Understanding Evaluation (GLUE) benchmark reveal that while some distilled TinyBERT models maintain the confidence levels of their BERT predecessors, others do not. The successful instances suggest that when knowledge distillation is executed with specific parameters, confidence values can be preserved. This research further demonstrates that adjusting distillation hyperparameters, particularly at the prediction layer, can lead to preservation of confidence values without compromising on accuracy.

Impact and Future Directions

This work's most significant contributions lie in viewing knowledge distillation as implicit abstraction while establishing the criterion for confidence property preservation. Interestingly, criteria like the proposed pairwise confidence difference become instrumental in identifying whether confidence is preserved post-distillation—not just in terms of aggregate distributions, but in the analysis of per-input confidence alignments. Ultimately, these insights could guide adjustments in distillation processes to ensure integral model properties are retained. Therefore, while the findings open new avenues in model compression techniques, the implications of this research beckon further exploration, potentially extending to a formal verification framework for the identification and tuning of hyperparameters critical for the conservation of other model properties during knowledge distillation.