Federated Distillation of Natural Language Understanding with Confident Sinkhorns

Rishabh Bhardwaj, Tushar Vaidya, Soujanya Poria

Keywords: Sinkhorn, NLP, Wasserstein, Random probability skew, Federated Distillation

Abstract: Enhancing the user experience is an essential task for application service providers. For instance, two users living wide apart may have different tastes of food. A food recommender mobile application installed on an edge device might want to learn from user feedback (reviews) to satisfy the client's needs pertaining to distinct domains. Retrieving user data comes at the cost of privacy while asking for model parameters trained on a user device becomes space inefficient at a large scale. In this work, we propose an approach to learn a central (global) model from the federation of (local) models which are trained on user-devices, without disclosing the local data or model parameters to the server. We propose a federation mechanism for the problems with natural similarity metric between the labels which commonly appear in natural language understanding (NLU) tasks. To learn the global model, the objective is to minimize the optimal transport cost of the global model's predictions from the confident sum of soft-targets assigned by local models. The confidence (a model weighting scheme) score of a model is defined as the L2 distance of a model's prediction from its probability bias. The method improves the global model's performance over the baseline designed on three NLU tasks with intrinsic label space semantics, i.e., fine-grained sentiment analysis, emotion recognition in conversation, and natural language inference.

One-sentence Summary: The study of confident Sinkhorn-based federation of natural language understanding, its theoretical properties, and empirical evidences.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2110.02432/code)