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: This paper proposed Grad2Task, a model for few-shot text classification tasks. Grad2Task first computes gradients when fine-tuning the base model on a support set; then the task embedding network decodes task representations using these gradients; finally the task representations go through MLP layers and output scaling and shifting parameters, which are plugged into the base model for task-specific adaptation.   The paper provides comprehensive evaluation on 10 tasks. Grad2Task outperforms strong baselines (Leopard, ProtoNet-BERT) in most settings. Analysis with a toy dataset shows the task representations are informative. Ablation study justified several design choices in Grad2Task.