KNOT: Knowledge Distillation Using Optimal Transport for Solving NLP Tasks

Proceedings of the 29th International Conference on Computational Linguistics, pages 4801–4820 October 12–17, 2022. 4801 KNOT: Knowledge Distillation using Optimal Transport for Solving NLP Tasks Rishabh Bhardwaj DeCLaRe Tushar Vaidya Soujanya Poria DeCLaRe DeCLaReSingapore University of Technology and Design, Singapore Nanyang Technological University, Singapore rishabh_bhardwaj@mymail.sutd.edu.sg, tushar.vaidya@ntu.edu.sg sporia@sutd.edu.sg Abstract We propose a new approach, Knowledge Dis- tillation using Optimal Transport (KNOT), to distill the natural language semantic knowledge from multiple teacher networks to a student network. KNOT aims to train a (global) stu- dent model by learning to minimize the optimal transport cost of its assigned probability distri- bution over the labels to the weighted sum of probabilities predicted by the (local) teacher models, under the constraints that the student model does not have access to teacher mod- els’ parameters or training data. To evaluate the quality of knowledge transfer, we intro- duce a new metric, Semantic Distance (SD), that measures semantic closeness between the predicted and ground truth label distributions. The proposed method shows improvements in the global model’s SD performance over the baseline across three NLP tasks while perform- ing on par with Entropy-based distillation on standard accuracy and F1 metrics. The im- plementation pertaining to this work is pub- licly available at https://github.com/ declare-lab/KNOT. 1 Introduction Due to recent technological advancements, more than two-thirds of the world’s population use mo- bile phones*. A client application on these de- vices has access to the unprecedented amount of data obtained from user-device interactions, sen- sors, etc. Learning algorithms can employ this data to provide an enhanced experience to its users. For instance, two users living wide apart may have different tastes in food. A food recommender ap- plication installed on an edge device might want to learn from user feedback (reviews) to satisfy the client’s needs pertaining to distinct domains. How- ever, directly retrieving this data comes at the cost of losing user privacy (Jeong et al., 2018). *https://datareportal.com/ global-digital-overview Figure 1: KNOT framework: Local models acting as multiple teachers trained locally on user data while the global model acts as a student. The global model can only observe predictions of local models. To minimize the risk of leaking user informa- tion, Federated Learning (FL) is a class of algo- rithms that proposes an alternative learning mecha- nism (Koneˇcn`y et al., 2016; McMahan et al., 2017). The parameters from the teacher networks, i.e., user (domain)-specific local models are retrieved to train a student network, i.e., a user (domain)- generic global model. The classic FL algorithms such as federated averaging and its successors are based on averaging of local model parameters or local gradient updates, and thus only applied when the global and local models possess similar net- work architectures. Additionally, FL has critical limitations of being costly in terms of communi- cation load with the increase in local model sizes (Mohri et al., 2019; Li et al., 2019; Jeong et al., 2018; Lin et al., 2020). Another set of algorithms, Federated Distillation (FD), propose to exchange only the outputs of …

Content truncated to 500 words