Grad2Task: Improved Few-shot Text Classification Using Gradients for Task Representation

1 Introduction . Transformer-based pretrained large-scale language models ( LMs ) have achieved tremendous success on many NLP tasks [ 15 , 34 ] , but require a large number of in-domain labeled examples for finetuning [ 54 ] . One approach to alleviate that issue is to fine-tune the model on an intermediate ( source ) task that is related to the target task . While previous work has focused on this setting for a single source task [ 33 , 50 ] , the problem of transferring from a set of source tasks has not been thoroughly considered . Wang et al . [ 51 ] showed that a naïve combination of multiple source tasks may negatively impact target task performance . Additionally , transfer learning methods typically consider the setting where more than a medium ( N > 100 ) number of training examples are available for both the source and target tasks . In this work we develop a method to improve a large pre-trained LM for few-shot text classification problems by transferring from multiple source tasks . While it is possible to directly adapt recent advances in meta-learning developed for few-shot image classification , a unique challenge in the NLP setting is that source tasks do no share the same task structure . Namely , on standard few-shot image classification benchmarks , the training tasks are sampled from a “ single ” larger dataset , and the label space contains the same task structure for all tasks . In contrast , in text classification tasks , the set of source tasks available during training can range from sentiment analysis to grammatical acceptability judgment . Different source tasks could not only be different in terms of input domain , but also their task structure ( i.e . label semantics , and number of output labels ) . This challenging problem requires resistance to overfitting due to its few-shot nature and more task-specific adaptation due to the distinct nature among tasks . Fine-tuning based approaches are known to suffer from the overfitting issue and approaches like MAML [ 17 ] and prototypical networks 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ( ProtoNet ) [ 41 ] were not designed to meta-learn from such a diverse collection of tasks . CNAP [ 37 ] , a conditional neural process [ 18 ] , explicitly designed for learning from heterogeneous task distributions by adding a task conditional adaptation mechanism is better suited for solving tasks with a shifted task distribution . We design our approach based on the CNAP framework . Compared with CNAP , we use pretrained transformers for text sequence classification instead of convolutional networks for image classification . We insert adapter modules [ 21 ] into a pretrained LM to ensure efficiency and maximal parameter-sharing , and learn to modulate those adapter parameters conditioned on tasks , while CNAP learns to adapt the whole feature extractor . Additionally , we use gradient information for task representation as opposed to average input encoding in CNAP , since gradients can capture information from both input and ouput space . In summary , our main contributions and findings are : • We propose the use of gradients as features to represent tasks under a model-based metalearning framework . The gradient features are used to modulate a base learner for rapid generalization on diverse few-shot text classification tasks . • We use pre-trained BERT with adapter modules [ 21 ] as the feature extractor and show that it works better than using only the BERT model on few-shot tasks , while also being more efficient to train as it has much fewer parameters to learn . • We compare our approach with traditional fine-tuning , sequential transfer learning , and stateof-the-art meta-learning approaches on a collection of diverse few-shot text classification tasks used in [ 5 ] and three newly added tasks , and show that our approach achieves the best performance . Our codes are publicly available1 . 2 Related Work . 2.1 Few-shot text classification . There is a wide range of text classification tasks , such as user intent classification , sentiment analysis , among others . Few-shot text classification tasks refer to those that contain novel classes unseen in training tasks and where only a few labeled examples are given for each class [ 55 ] . Meta-learning approaches have been proposed for this problem [ 7 , 19 , 5 ] . To increase the amount and diversity of training episodes , previous work trained few-shot learners by semi-supervised learning [ 40 ] and selfsupervised learning [ 6 ] . We compare different meta-learning approaches following the experiment procedures of [ 5 ] . 2.2 Meta-learning . Metric-based meta-learning approaches [ 49 , 41 , 45 ] try to learn embedding models on which simple classifiers can achieve good performance . Optimization-based approaches [ 17 , 38 , 30 ] are designed to learn good parameter initialization that can quickly adapt to new tasks within a few gradient descent steps . “ Black box ” or model-based meta-learning approaches use neural networks to embed task information and predict test examples conditioned on the task information [ 39 , 36 , 28 , 31 , 18 , 37 ] . Our approach falls into this third category and is mostly related to CNAP [ 37 ] . Compared with CNAP we use a different feature extractor and different mechanisms for modulation , as well as focusing on different problems . An important difference with CNAP is that we use gradient information as features to represent tasks , instead of using average input encoding . Gradient information has also been used to generate task-dependent attenuation of parameter initialization under the MAML framework [ 4 ] . 2.3 Transfer learning . Our work is also related to transfer learning . Instead of fine-tuning a pretrained language model directly on a target task , Phang et al . [ 33 ] and Talmor and Berant [ 46 ] showed that intermediate training on data-rich supervised tasks is beneficial for downstream performance . Wang et al . [ 51 ] found that there is no guarantee that multi-task learning or intermediate training yields better performance than direct fine-tuning on target tasks . Transferability estimation is the task of predicting whether it is 1https : //github.com/jixuan-wang/Grad2Task beneficial to transfer from some task to another one . There have been works estimating transferability of NLP tasks based on handcrafted features [ 8 , 23 ] , while data-driven approaches have been studied mainly for computer vision tasks [ 56 , 47 , 43 , 29 ] . TASK2VEC is a task embedding approach based on the Fisher Information Matrix ( FIM ) [ 2 ] . Transferability between tasks can be easily evaluated by the distance between their task embeddings according to some simple metric . In the NLP domain , given a target task and a predefined set of source tasks , Vu et al . [ 50 ] proposed to use task embedding based on FIM to measure task similarity and predict transferability according to the distance between task embeddings . We draw inspiration from this work to use gradients as features for task representation . Instead of one-to-one transfer , our work can be seen as transferring from a set of source tasks to new tasks through meta-learning . 3 Problem Definition . Following the terminologies in meta learning , each task ( or episode ) , t = ( S , Q ) , is specified by a support set ( training set ) S and a query set ( test set ) Q , where S = { ( xi , yi ) } |S|i=1 , Q = { ( xi , yi ) } |Q| i=1 , xi is a text sequence and yi is a discrete value corresponding to a class . Different tasks could have different input domains , different numbers of classes , different label space , and so on . Our goal is to train a text classification model on a set of labeled datasets { Dsi } Ni=1 of N source tasks . The model is expected to achieve good performance on new target tasks after training . For each target task , a small labeled dataset , e.g. , five shots for each class , is given for adaptation or fine-tuning and the full test dataset is used for evaluation . Note that when learning on target tasks , we assume no access to the source tasks nor other target tasks . This is different with multi-task learning where target tasks are learned together with source tasks . 4 Model Design . To handle diverse tasks with different structures a single model may not be flexible enough . Thus , for every task , we utilize a task embedding network to capture the task nature and adapt a base model conditioned on the current task . Task-specific adaptation is done by generating shifting and scaling parameters , named adaptation parameters , that are applied on the hidden representations inside the base model . Figure 1 shows an overview of our model architecture , mainly consisting of three parts : a base model denoted by f , a task embedding network denoted by d , and an adaptation network denoted by a . To make a prediction , our model takes the following steps : 1 ) given a support set , it uses the base model ( Section 4.1 ) to compute gradients w.r.t . a subset of its parameters to use as input to the task embedding network ( Section 4.2 ) , 2 ) the task embedding network maps these gradients to task representations , 3 ) layer-wise adaptation networks ( Section 4.3 ) take as input the task representations and output adaptation parameters , and lastly 4 ) the adaptation parameters are applied to the base model before it predicts on the query set . The training of our model consists of two stages : 1 ) we first train the base model episodically as a prototypical network ; 2 ) we then freeze the base model and train the task embedding network and the adaptation network using the same loss function of the first stage ( Section 4.4 ) . Next we introduce our model in details . The list of notations used throughout the paper can be found in Appendix A . 4.1 Base Model : BERT & Bottleneck Adapters . Our base model is built upon transformer-based pretrained LMs . We use the pretrained BERTBASE model throughout this paper , although other transformer-based pretrained LMs are also applicable [ 35 , 25 ] . Following Houlsby et al . [ 21 ] , we insert two adapter modules containing bottleneck layers into each transformer layer of BERT . Each bottleneck adapter is denoted as ↵l , l = 1 ... 2L where L is the total number of transformer layers in BERT . We use the bottleneck adapters because , first , BERT with adapters are more efficient to train and less vulnerable to overfitting – a desirable property in the few-shot setting . Additionally , this simplifies our goal of using gradients for task representation . It is infeasible to use the gradients of the whole BERT model because the dimensionality is too large . Instead , we compute gradients w.r.t . parameters in these bottleneck layers ( < 0.5 % of the number of parameters in BERT ) . Since text classification is a sequence classification task , we need a pooling method to represent each sequence by a single embedding before the embedding is fed into a classifier . Following Devlin et al . [ 15 ] , this single embedding is obtained by applying a linear transformation on top of the contextual embedding of the special ‘ [ CLS ] ’ token inserted at the beginning of every sequence . We further apply a fully connected layer on top of the BERT output . Together , we name the pretrained BERT with adapters and the last linear layer as the ‘ base model ’ , denoted by f . We denote the parameters of the pretrained BERT as ✓ , bottleneck adapters as { ↵l } 2Ll=1 , and linear layer as ! . For the classifier , we use the ProtoNet classifier [ 41 ] . Given a support set S , a query example x and our base model f , a nearest cluster classifier can be described as follows : pbase ( y = c |x ) = softmax ( euc ( f ( x ; ✓ , ↵ , ! ) , µc ) ) , ( 1 ) µc = 1 |Sc| X xi2Sc f ( xi ; ✓ , ↵ , ! ) where Sc = { ( xi , yi ) | ( xi , yi ) 2 S , yi = c } , µc is the cluster center in the embedding space for support examples with class c , and euc refers to the Euclidean distance .

Summary: The paper presents a novel method for few-shot text classification tasks. The method is built on top of the CNAP  The key contribution is to effectively encode the gradients of networks as the task representation so that the meta-learning stage can use such task representations to improve the few-shot generalization ability. The two key technical contributions are a task embedding network (an RNN-based layer-by-layer gradient encoder via Fisher Information Matrix) and an adaptation network for auto-regressively adapting parameters from the bottom to the top layer.  The implementation is based on BERT-like LMs with Adapter bottleneck modules, which are parameter efficient to compute.   The authors compared the proposed method with a few meta-learning methods on a diverse set of tasks, and show that their method outperforms other methods for most of the time. The empirical analysis and ablations also provide insights for understanding the proposed method.