Automated Concatenation of Embeddings for Structured Prediction

1 INTRODUCTION . Recent developments on pretrained contextualized embeddings have significantly improved the performance of structured prediction tasks in natural language processing . Approaches based on contextualized embeddings , such as ELMo ( Peters et al. , 2018 ) , Flair ( Akbik et al. , 2018 ) , BERT ( Devlin et al. , 2019 ) , and XLM-R ( Conneau et al. , 2020 ) , have been consistently raising the state-of-the-art for various structured prediction tasks . Concurrently , research has also showed that word representations based on the concatenation of multiple pretrained contextualized embeddings and traditional non-contextualized embeddings ( such as word2vec ( Mikolov et al. , 2013 ) and character embeddings ( Santos & Zadrozny , 2014 ) ) can further improve performance ( Peters et al. , 2018 ; Akbik et al. , 2018 ; Straková et al. , 2019 ; He & Choi , 2020 ) . Given the ever-increasing number of embedding learning methods that operate on different granularities ( e.g. , word , subword , or character level ) and with different model architectures , choosing the best embeddings to concatenate for a specific task becomes non-trivial , and exploring all possible concatenations can be prohibitively demanding in computing resources . Neural architecture search ( NAS ) is an active area of research in deep learning to automatically search for better model architectures , and has achieved state-of-the-art performance on various tasks in computer vision , such as image classification ( Real et al. , 2019 ) , semantic segmentation ( Liu et al. , 2019a ) , and object detection ( Ghiasi et al. , 2019 ) . In natural language processing , NAS has been successfully applied to find better RNN structures ( Zoph & Le , 2017 ; Pham et al. , 2018b ) and recently better transformer structures ( So et al. , 2019 ; Zhu et al. , 2020 ) . In this paper , we propose the Automated Concatenation of Embeddings ( ACE ) approach to automate the process of finding better concatenations of embeddings for structured prediction tasks , formulated as an NAS problem . In this approach , an iterative search process is guided by a controller based on its belief that models the effectiveness of individual embedding candidates in consideration for a specific task . At each step , the controller samples a concatenation of embeddings according to the belief model and feeds the concatenated word representations as inputs to a task model , which in turn is trained on the task dataset and returns the model accuracy as a reward signal to update the belief model . We use the policy gradient algorithm ( Williams , 1992 ) in reinforcement learning ( Sutton & Barto , 1992 ) to solve the optimization problem . In order to improve the efficiency of the search process , we also design a special reward function by accumulating all the rewards based on the transformation between the current concatenation and all previously sampled concatenations . Our approach is different from previous work on NAS in the following aspects : 1 . Unlike most previous work , we focus on searching for better word representations rather than better model architectures . 2 . We design a unique search space for the embedding concatenation search . Instead of using RNN as in previous work of Zoph & Le ( 2017 ) , we design a more straightforward controller to generate the embedding concatenation . We design a novel reward function in the objective of optimization to better evaluate the effectiveness of each concatenated embeddings . 3 . Our approach is efficient and practical . ACE can find a strong word representation on a single GPU with only a few GPU-hours for structured prediction tasks , while a lot of the NAS approaches require dozens of or even thousands of GPU-hours to search for good neural architecture . 4 . The task model from ACE achieves high accuracy without the need for retraining , while in previous work of NAS the resulting neural network usually requires retraining from scratch . Empirical results show that ACE outperforms strong baselines . Furthermore , we show that when ACE is applied to concatenate pretrained contextualized embeddings which are already fine-tuned on specific tasks , we can achieve state-of-the-art or competitive accuracy on 6 structured prediction tasks including Named Entity Recognition ( Sundheim , 1995 ) , Part-Of-Speech tagging ( DeRose , 1988 ) , chunking ( Tjong Kim Sang & Buchholz , 2000 ) , aspect extraction ( Hu & Liu , 2004 ) , syntactic dependency parsing ( Tesnière , 1959 ) and semantic dependency parsing ( Oepen et al. , 2014 ) over 21 datasets . Besides , we also analyze the advantage of ACE and reward function design over the baselines and show the advantage of ACE over ensemble models . 2 RELATED WORK . 2.1 EMBEDDINGS . Non-contextualized embeddings , such as word2vec ( Mikolov et al. , 2013 ) , GloVe ( Pennington et al. , 2014 ) , and fastText ( Bojanowski et al. , 2017 ) , help lots of NLP tasks . Character embeddings ( Santos & Zadrozny , 2014 ) are trained together with the task and applied in many structured prediction tasks ( Ma & Hovy , 2016 ; Lample et al. , 2016 ; Dozat & Manning , 2018 ) . For pretrained contextualized embeddings , ELMo ( Peters et al. , 2018 ) , a pretrained contextualized word embedding generated with multiple Bidirectional LSTM layers , significantly outperforms previous state-of-the-art approaches on several NLP tasks . Following this idea , Akbik et al . ( 2018 ) proposed Flair embeddings , which is a kind of contextualized character embeddings and achieved strong performance in sequence labeling tasks . Recently , Devlin et al . ( 2019 ) proposed BERT , which encodes contextualized sub-word information by Transformers and significantly improves the performance on a lot of NLP tasks . Much research such as RoBERTa ( Liu et al. , 2019c ) has focused on improving BERT model ’ s performance through stronger masking strategies . Moreover , multilingual contextualized embeddings become popular . Pires et al . ( 2019 ) and Wu & Dredze ( 2019 ) showed that Multilingual BERT ( M-BERT ) could learn a good multilingual representation effectively with strong crosslingual zero-shot transfer performance in various tasks . Conneau et al . ( 2020 ) proposed XLM-R , which is trained on a larger multilingual corpus and significantly outperforms M-BERT on various multilingual tasks . 2.2 NEURAL ARCHITECTURE SEARCH . Recent progress on deep learning has shown that network architecture design is crucial to the model performance . However , designing a strong neural architecture for each task requires enormous efforts , high level of knowledge , and experiences over the task domain . Therefore , automatic design of neural architecture is desired . A crucial part of NAS is search space design , which defines the discoverable NAS space . Previous work ( Baker et al. , 2017 ; Zoph & Le , 2017 ; Xie & Yuille , 2017 ) designs a global search space ( Elsken et al. , 2019 ) which incorporates structures from handcrafted architectures . For example , Zoph & Le ( 2017 ) designed a chained-structured search space with skip connections . The global search space usually has a considerable degree of freedom . As an example , the approach of Zoph & Le ( 2017 ) takes 22,400 GPU-hours to search on CIFAR-10 dataset . Based on the observation that existing hand-crafted architectures contain repeated structures ( Szegedy et al. , 2016 ; He et al. , 2016 ; Huang et al. , 2017 ) , Zoph et al . ( 2018 ) explored cell-based search space which can reduce the search time to 2,000 GPU-hours . In recent NAS research , reinforcement learning and evolutionary algorithms are the most usual approaches . In reinforcement learning , the agent ’ s actions are the generation of neural architectures and the action space is identical to the search space . Previous work usually applies an RNN layer ( Zoph & Le , 2017 ; Zhong et al. , 2018 ; Zoph et al. , 2018 ) or use Markov Decision Process ( Baker et al. , 2017 ) to decide the hyper-parameter of each structure and decide the input order of each structure . Evolutionary algorithms have been applied to architecture search for many decades ( Miller et al. , 1989 ; Angeline et al. , 1994 ; Stanley & Miikkulainen , 2002 ; Floreano et al. , 2008 ; Jozefowicz et al. , 2015 ) . The algorithm repeatedly generates new populations through recombination and mutation operations and selects survivors through competing among the population . Recent work with evolutionary algorithms differ in the method on parent/survivor selection and population generation . For example , Real et al . ( 2017 ) , Liu et al . ( 2018a ) , Wistuba ( 2018 ) and Real et al . ( 2019 ) applied tournament selection ( Goldberg & Deb , 1991 ) for the parent selection while Xie & Yuille ( 2017 ) keeps all parents . Suganuma et al . ( 2017 ) and Elsken et al . ( 2018 ) chose the best model while Real et al . ( 2019 ) chose several latest models as survivors . 3 AUTOMATED CONCATENATION OF EMBEDDINGS . In ACE , a task model and a controller interact with each other repeatedly . The task model predicts the task output , while the controller searches for better embedding concatenation as the word representation for the task model to achieve higher accuracy . Given an embedding concatenation generated from the controller , the task model is trained over the task data and returns a reward to the controller . The controller receives the reward to update its parameter and samples a new embedding concatenation for the task model . Figure 1 shows the general architecture of our approach . 3.1 TASK MODEL . For the tasks model , we emphasis on sequence-structured and graph-structured outputs . Given a structured prediction task with input sentence x and structured output y , we can calculate the probability distribution P ( y|x ) by : P ( y|x ) = exp ( Score ( x , y ) ) ∑ y′∈Y ( x ) exp ( Score ( x , y′ ) ) where Y ( x ) represents all possible output structures given the input sentence x . Depending on different structured prediction tasks , the output structure y can be label sequences , trees , graphs or other structures . In this paper , we use sequence-structured and graph-structured outputs as two exemplar structured prediction tasks . We use BiLSTM-CRF model ( Ma & Hovy , 2016 ; Lample et al. , 2016 ) for sequence-structured outputs and use BiLSTM-Biaffine model ( Dozat & Manning , 2017 ) for graph-structured outputs : P seq ( y|x ) = BiLSTM-CRF ( V , y ) ; P graph ( y|x ) = BiLSTM-Biaffine ( V , y ) where V = [ v1 ; · · · ; vn ] , V ∈ Rd×n is a matrix of the word representations for the input sentence xwith nwords , d is the hidden size of the concatenation of all embeddings . The word representation vi of i-th word is a concatenation of L types of word embeddings : vli = embed l i ( x ) ; vi = [ v 1 i ; v 2 i ; . . . ; v L i ] where embedl is the model of l-th embeddings , vi ∈ Rd , vli ∈ Rd l . dl is the hidden size of embedl . 3.2 SEARCH SPACE DESIGN . The neural architecture search space can be represented as a set of neural networks ( Elsken et al. , 2019 ) . A neural network can be represented as a directed acyclic graph with a set of nodes and directed edges . Each node represents an operation , while each edge represents the inputs and outputs between these nodes . In ACE , we represent each embedding candidate as a node . The input to the nodes is the input sentence x , and the outputs are the embeddings vl . Since we concatenate the embeddings as the word representation of the task model , there is no connection between nodes in our search space . Without considering the connections between nodes , the search space can be significantly reduced . For each node , there are a lot of options to extract word features . Taking BERT embeddings as an example , Devlin et al . ( 2019 ) concatenated the last four layers as word features while Kondratyuk & Straka ( 2019 ) applied a weighted sum of all twelve layers . However , the empirical results ( Devlin et al. , 2019 ) do not show a significant difference in accuracy . We follow the typical usage for each embedding to further reduce the search space . As a result , each embedding only has a fixed operation and the resulting search space contains 2L−1 possible combinations of nodes . In NAS , weight sharing ( Pham et al. , 2018a ) shares the weight of structures in training different neural architectures to reduce the training cost . In comparison , we fixed the weight of pretrained embedding candidates in ACE except for the character embeddings . Instead of sharing the parameters of the embeddings , we share the parameters of the task models at each step of search . However , the hidden size of word representation varies over the concatenations , making the weight sharing of structured prediction models difficult . Instead of deciding whether each node exists in the graph , we keep all nodes in the search space and add an additional operation for each node to indicate whether the embedding is masked out . To represent the selected concatenation , we use a binary vector a = [ a1 , · · · , al , · · · , aL ] as an mask to mask out the embeddings which are not selected : vi = [ v 1 i a1 ; . . . ; v l ial ; . . . ; v L i aL ] ( 1 ) where al is a binary variable . Since the input V is applied to a linear layer in the BiLSTM layer , multiplying the mask with the embeddings is equivalent to directly concatenating the selected embeddings : W = [ W1 ; W2 ; . . . ; WL ] ; W > vi = L∑ l=1 W > l v l ial ; W ∈ Rd×h andWl ∈ Rd l×h ( 2 ) Therefore , the model weights can be shared after applying the embedding mask to all embedding candidates ’ concatenation . Another benefit of our search space design is that we can remove the unused embedding candidates and the corresponding weights inW for a lighter task model after the best concatenation is found by ACE .

Summary: This paper introduced an interesting application of reinforcement learning in the selection of concatenation of contextual/non-contextual word embeddings.  It is clever to limit the search space on the selection of embedding sources rather than search the whole network structure, as the current strategy is much easier in the training step. The author(s) conducted many experiments that compared many other models (including SOTA models, and ablation study). Those results are pretty good and impressive.