Exploring All-In-One Knowledge Distillation Framework for Neural Machine Translation (2023)

Abstract: Conventional knowledge distillation(KD) approaches are commonly employed to compress neural machine translation(NMT) models. However, they only obtain one lightweight student each time. Consequently, we have to conduct KD multiple times when different students are required at the same time, which could be resource-intensive. Additionally, these students are individually optimized, and thus lack interactions with each other, leading to their potential not being fully exerted. In this work, we propose a novel All-In-One Knowledge Distillation(AIO-KD) framework for NMT, which generates multiple satisfactory students at once. Under AIO-KD, we first randomly extract fewer-layer subnetworks from the teacher as the sample students. Then, we jointly optimize the teacher and these students, where the students simultaneously learn the knowledge from the teacher and interact with other students via mutual learning. When utilized, we re-extract the candidate students, satisfying the specifications of various devices. Particularly, we adopt carefully-designed strategies for AIO-KD: 1) we dynamically detach gradients to prevent poorly-performed students from negatively affecting the teacher during the knowledge transfer, which could subsequently impact other students; 2) we design a two-stage mutual learning strategy, which alleviates the negative impacts of poorly-performed students on the early-stage student interactions. Extensive experiments and in-depth analyses on three benchmarks demonstrate the effectiveness and eco-friendliness of AIO-KD. Our source code is available at https://github.com/DeepLearnXMU/AIO-KD.

## introduction in recent years, transformer @xcite has become the dominant architecture in neural machine translation (nmt). to further improve model performance, there have been many efforts in exploring wider or deeper transformers @xcite @xcite . however, deploying such models with billions of parameters on edge devices (e.g., mobile phones) remains to be a challenge. to solve this problem, various methods for model compression have been proposed @xcite . among them, knowledge distillation (kd) @xcite has been widely adopted due to its effectiveness and simplicity. in this regard, researchers explore many kd approaches for nmt, such as word-kd @xcite and selective-kd @xcite . generally, in these conventional kd approaches, the knowledge of a large teacher model is transferred to only a compact student model, which usually adopts the same model architecture as the teacher but with fewer parameters @xcite @xcite . nevertheless, due to hardware differences, we are often required to deploy models of varying sizes on different devices @xcite @xcite . in this scenario, conventional kd approaches have two defects: 1) they have to be conducted multiple times for different students, which leads to substantial costs; 2) students are individually optimized, and thus they are unable to interact with each other. note that in the human learning process, communication between students will benefit their learning @xcite . in light of this, we believe that these students can be further improved through collaborative interactions. in this paper, we propose a novel all-in-one knowledge distillation (aio-kd) framework for nmt, which constructs students from the teacher and jointly optimizes both the teacher and students from scratch. employing aio-kd, we regard the fewer-layer subnetworks extracted from the teacher as the candidate students. during training, we first randomly select from all candidate students to obtain the sample ones at each training step. then, we jointly optimize the teacher and the sample students. during this process, the students simultaneously learn the knowledge from the teacher and interact with other students via mutual learning. when utilized, we re-extract the candidate students from the teacher, satisfying the specifications of various devices. to better accommodate aio-kd, we carefully design the following strategies: empirical experiments and in-depth analyses on three translation benchmarks demonstrate that aio-kd is superior to conventional kd approaches in terms of translation quality and training costs. as a bonus, the teacher in aio-kd is significantly enhanced through knowledge transfer with the students. ## related work our related works mainly include the following three lines of studies: transformer with variable depths. to reduce computation costs, plenty of researchers investigate the variable-depth transformer architecture @xcite @xcite @xcite @xcite @xcite . however, the majority of these works focus on the transformer encoder, while paying less attention to the overall transformer architecture. our work focuses on the latter, i.e., the overall transformer with variable depths for nmt. ## all-in-one knowledge distillation in this section, we first give a brief overview of aio-kd (section 3.1), and then describe the training objective (section 3.2). finally, we detail carefully-designed strategies (section 3.3). ## conclusion in this work, we present aio-kd, a novel kd framework for nmt that constructs various candidate students from the teacher itself. with aio-kd, we jointly optimize the teacher and the sample students from scratch. during this process, the students learn from the teacher and interact with other students via mutual learning, resulting in multiple satisfactory students with efficient training. carefully-designed strategies are also introduced to accommodate aio-kd. extensive experiments and in-depth analyses on three benchmarks demonstrate the superiority of our aio-kd. in the future, we plan to explore more compact subnetworks of teacher as students using parameter pruning methods. additionally, we aim to extend aio-kd to large language models (llms) to validate its generalizability.