KNOT : Knowledge Distillation Using Optimal Transport for Solving NLP Tasks (2022)

Abstract: We propose a new approach, Knowledge Distillation 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) student model by learning to minimize the optimal transport cost of its assigned probability distribution 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 models’ parameters or training data. To evaluate the quality of knowledge transfer, we introduce 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 performing on par with Entropy-based distillation on standard accuracy and F1 metrics. The implementation pertaining to this work is publicly available at https://github.com/declare-lab/KNOT .

## introduction due to recent technological advancements, more than two-thirds of the world's population use mobile phones * . a client application on these devices has access to the unprecedented amount of data obtained from user-device interactions, sensors, 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 application installed on an edge device might want to learn from user feedback (reviews) to satisfy the client's needs pertaining to distinct domains. however, directly retrieving this data comes at the cost of losing user privacy @xcite . to minimize the risk of leaking user information, federated learning (fl) is a class of algorithms that proposes an alternative learning mechanism @xcite . 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 network architectures. additionally, fl has critical limitations of being costly in terms of communication load with the increase in local model sizes @xcite @xcite . another set of algorithms, federated distillation (fd), propose to exchange only the outputs of the local model, i.e, either logits or probability measures whose dimensions are usually much smaller than the parameter size of models themselves @xcite . thus, it enables learning from an ensemble of teacher local models of dissimilar architecture types. in contrast to fl, fd trains the global student model at reduced risk of user privacy, lower communication overhead, and lesser memory space utilization. in this work, we base our problem formulation under fd setting where global model can access the local model outputs without visibility over local user data or model parameters (figure 1 ). since the kullback-leibler (kl) divergence is easy to compute, facilitates smooth backpropagation, and is widely used @xcite , it became standard practice to use it to define the objective function in most fd algorithms @xcite @xcite . however, a critical limitation of such entropy-based losses is that they ignore any metric structure in the label space. for instance, in the task of fine-grained sentiment classification of text, strongly positive sentiment is closer to positive while far from strongly negative sentiment. this information is not fully utilized by the existing distillation algorithms. contrary to entropy losses, optimal transport (ot) loss admits such inter-class relationships as demonstrated by @xcite . thus, we propose an ot-based approach, knot (knowledge distillation using optimal transport), for knowledge distillation of natural language semantics encoded in local models under fd setting. to improve the semantic knowledge transfer from local models to the global model acting as multiple teachers-single student framework, we explicitly encode the inter-label relationship in distillation loss in terms of the cost of probability mass transport. thus, the major contributions of this work are: contribution:1 (c1) for the tasks with intrinsic inter-class semantics, we propose a novel optimal transport-based knowledge distillation approach knot for distillation from an ensemble of multiple teacher networks under fd setting. the problem of bias distillation: local models are prone to possess biases that might get transferred to the global model during distillation. the bias in a local model can potentially arise from the user-specific local data that is potentially nonindependent and non-identically distributed (non-iid). one such bias is population bias, i.e., the local user may not represent the target overall population @xcite . this motivates our second contribution: contribution:2 (c2) to reduce the potential risk of bias transfer from local models to global, knot employs a weighted distillation scheme. the local model with a higher l 2 distance of predicted probabilities from its intrinsic bias (a distribution) is given higher importance during distillation. it is important to note that the application of c2 in knot does not warrant the c1 setting to be satisfied and vice versa. as shown in figure 1 , the global model m g aims to learns from the weighted sum of distance d 1 , . . . , d n with weights being w 1 , . . . , w n . where d k 's represents the divergence of the global model's prediction from k th local model's prediction. this constitutes our c1. the w k 's are obtained via c2. to validate the fitness of c1 and c2, we also derive generalization bounds of the proposed distillation mechanism. next, we introduce a metric called semantic distance (sd). it helps us evaluate the semantic closeness of the model's output from the ground truth distribution. ## related work there have been many approaches to fl such as local model parameter averaging based on local sgd updates @xcite . it warrants global and local models to have the same model architecture. another line of work is multiple-source adaptation formulations where a learner has access to source domain-specific predictor without access to the labeled data. the expected loss is the mixture of source domains @xcite . even though the formulation is close, our solution to the problem is different as we do not have access to the local or global data domain distribution. in natural language processing (nlp), ‡ the training setup is described later in section 6. @xcite al. (2021); lin et al. (2020) ; @xcite fine-tune transformer-based architecture in the federated setting. however, they do not leverage label space semantics and the analysis is restricted to small-scale datasets. closest to our work aims to improve local client training based on local data heterogeneity @xcite . knowledge distillation aims to transfer knowledge from a (large) teacher model to a (smaller) student model @xcite . given the output logit/softmaxed values of the teacher model, the student can imitate the teacher's behavior @xcite . a few works are dedicated to the distillation of the ensemble of teacher models to the student model. this includes logit averaging of teacher models @xcite or feature level knowledge extraction @xcite . to the best of our knowledge, there is no prior work that aims to leverage ot to enhance the distillation of semantic knowledge in local models under the fd paradigm. we use standard and widely used entropy-based loss (kl-divergence) as our baseline to compare with c1. we also construct two baselines for confidence score calculation from the prior works, i.e., logit averaging and weighting scheme based on local model dataset size @xcite . this is to compare the contribution of c2. ## optimal transport traditional divergences, such as kl, ignore metric structure in the label space y. in other words, they do not allow the incorporation of inter-class relationships in the loss function. contrary to this, optimal transport metrics can be extremely useful in defining inter-class semantic relationships in the label space. entropic regularisation of ot convexifies the loss function and thus is a computational advantage in computing gradients @xcite @xcite . as ε → 0 + , we retrieve the unregularized wasserstein distance. ## sinkhorn-based distillation as an entropy-based loss, we adopt kl divergence. as discussed in section 3, we employ sinkhorn distance to implement ot-based loss which is the proposed knot algorithm. ## experiments baselines. we setup the following baselines for a thorough comparison between sinkhorn || and entropy-based losses. let [method] be the placeholder for sinkhorn and entropy. tasks. we set up the three natural language understanding tasks that possess inter-class semantics: 1) fine-grained sentiment analysis (sa), 2) emotion recognition in conversation (erc) @xcite , and 3) natural language inference (nli). nli is the task of determining the inference relation between two texts. the relation can be entailment, contradiction, or neutral @xcite . for a given transcript of a conversation, the erc task aims to identify the emotion of each utterance from the set of pre-defined emotions @xcite . for our experiments, we choose the five most common emotions that are sadness, anger, surprise, happiness, and no emotion. datasets. for the sa task, we use four largescale datasets: 1) toys: toys and games; 2) cloths: architecture. we set up a compact transformerbased model used by both global and local models (figure 4 ), although, the federation does not restrict both the local and model architectures to be the same. the input is fed to the pretrained bertbased classifier @xcite . thus, we obtain probabilities with support in the space of output labels, i.e., y. we keep all the parameters trainable, hence, bert will learn its embeddings specific to the classification task. for the nli task, we append premise and hypothesis at input sepa-rated by special token [sep] token, followed by a standard classification setup. ## conclusion this work proposed knot, i.e., a novel optimal transport (ot)-based natural language semantic knowledge distillation. for the tasks with intrinsic label similarities, the ot distance between the predicted probability of the central (global) model and user-specific (local) models is minimized. to reduce the potential hazard of bias transfer from local model distillation, we introduced a weighting scheme based on the l 2 distance between the local model's prediction and probability bias. our experiments on three language understanding tasks -fine-grained sentiment analysis, emotion recognition in conversation, and natural language inference-show consistent semantic distance improvements while performing as good as the entropybased baselines on the accuracy and f1 metrics. a what have we kept for the appendix? we include proofs of theorem 5.2 and an analysis of the l 2 -based confidence score. there are a few experiments that we consider to be important and may help compare the ot loss-based learning with kullback-leibler (kl). these results build a firm base to choose sinkhorn-based (ot) losses on the task of federated distillation of sentiments. for the experiments, we work on the global model m g that has acquired knowledge from local models in the learning without forgetting the paradigm.