A Flexible Multi-Task Model for BERT Serving (2022)

Abstract: We present an efficient BERT-based multi-task (MT) framework that is particularly suitable for iterative and incremental development of the tasks. The proposed framework is based on the idea of partial fine-tuning, i.e. only fine-tune some top layers of BERT while keep the other layers frozen. For each task, we train independently a single-task (ST) model using partial fine-tuning. Then we compress the task-specific layers in each ST model using knowledge distillation. Those compressed ST models are finally merged into one MT model so that the frozen layers of the former are shared across the tasks. We exemplify our approach on eight GLUE tasks, demonstrating that it is able to achieve 99.6% of the performance of the full fine-tuning method, while reducing up to two thirds of its overhead.

## introduction in this work we explore the strategies of bert @xcite serving for multiple tasks under the following two constraints: 1) memory and computational resources are limited. on edge devices such as mobile phones, this is usually a hard constraint. on local gpu stations and cloud-based servers, this constraint is not as hard but it is still desirable to reduce the computation overhead to cut the serving cost. 2) the tasks are expected to be modular and are subject to frequent updates. when one task is updated, the system should to be able to quickly adapt to the task modification such that the other tasks are not affected. this is a typical situation for applications (e.g. ai assistant) under iterative and incremental development. in principle, there are two strategies of bert serving: single-task serving and multi-task serving. in single-task serving, one independent single-task model is trained and deployed for each task. typically, those models are obtained by fine-tuning a copy of the pre-trained bert and are completely different from each other. single-task serving has the advantage of being flexible and modular as there is no dependency between the task models. the downside is its inefficiency in terms of both memory usage and computation, as neither parameters nor computation are shared or reused across the tasks. in multi-task serving, one single multi-task model is trained and deployed for all tasks. this model is typically trained with multi-task learning (mtl) @xcite . compared to its single-task counterpart, multi-task serving is much more computationally efficient and incurs much less memory usage thanks to its sharing mechanism. however, it has the disadvantage in that any modification made to one task usually affect the other tasks. the main contribution of this work is the proposition of a framework for bert serving that simultaneously achieves the flexibility of single-task serving and the efficiency of multi-task serving. our method is based on the idea of partial finetuning, i.e. only fine-tuning some topmost layers of bert depending on the task and keeping the remaining bottom layers frozen. the fine-tuned layers are task-specific, which can be updated on a per-task basis. the frozen layers at the bottom, which plays the role of a feature extractor, can be shared across the tasks. ## related work the standard practice of using bert is fine-tuning, i.e. the entirety of the model parameters is adjusted on the training corpus of the downstream task, so that the model is adapted to that specific task @xcite . there is also an alternative feature-based approach, used by elmo @xcite . in the latter approach, the pre-trained model is regarded as a feature extractor with frozen parameters. during the learning of a downstream task, one feeds a fixed or learnable combination of the model's intermediate representations as input to the task-specific module, and only the parameters of the latter will be updated. it has been shown that the fine-tuning approach is generally superior to the feature-based approach for bert in terms of task performance @xcite . a natural middle ground between these two approaches is partial fine-tuning, i.e. only fine-tuning some topmost layers of bert while keeping the remaining bottom layers frozen. this approach has been studied in @xcite , where the authors observed that finetuning only the top layers can almost achieve the performance of full fine-tuning on several glue tasks. the approach of partial fine-tuning essentially regards the bottom layers of bert as a feature extractor. freezing weights from bottom layers is a sensible idea as previous studies show that the mid layer representations produced by bert are most transferrable, whereas the top layers representations are more task-oriented @xcite @xcite . notably, @xcite showed that fine-tuning primarily affects weights from the top layers while weights from bottom layers do not alter much. @xcite showed that it is possible to achieve state-of-the-art results on a number of probing tasks with linear models trained on frozen mid layer representations of bert. ## method in what follows, we denote by t the set of all target tasks. we always use the 12-layer uncased version of bert as the pre-trained language model foot_0 . the proposed framework features a pipeline (fig. 1 ) that consists of three steps: 1) single task partial fine-tuning; 2) single task knowledge distillation; 3) model merging. we give details of these steps below. ## experiments in this section, we compare the performance and efficiency of our model with various baselines on eight glue tasks @xcite . more details on these tasks can be found in appendix a. ## conclusion we have presented our framework that is designed to provide efficient and flexible bert-based multitask serving. we have demonstrated on eight glue datasets that the proposed method achieves both strong performance and efficiency. we release our code foot_2 and hope that it can facilitate bert serving in cost-sensitive applications.