Constrained Decoding for Computationally Efficient Named Entity Recognition Taggers (2020)

Abstract: Current state-of-the-art models for named entity recognition (NER) are neural models with a conditional random field (CRF) as the final layer. Entities are represented as per-token labels with a special structure in order to decode them into spans. Current work eschews prior knowledge of how the span encoding scheme works and relies on the CRF learning which transitions are illegal and which are not to facilitate global coherence. We find that by constraining the output to suppress illegal transitions we can train a tagger with a cross-entropy loss twice as fast as a CRF with differences in F1 that are statistically insignificant, effectively eliminating the need for a CRF. We analyze the dynamics of tag co-occurrence to explain when these constraints are most effective and provide open source implementations of our tagger in both PyTorch and TensorFlow.

## introduction named entity recognition (ner) is the task of finding phrases of interest in text that map to real world entities such as organizations ("org") or locations ("loc"). this is normally cast as a sequence labeling problem where each token is assigned a label that represents its entity type. multi-token entities are handled by having special "beginning" and "inside" indicators that specify which tokens start, continue, or change the type of an entity. @xcite show that the iobes tagging scheme, where entity spans must begin with a "b" token, end with an "e" token and where single token entities are labeled with an "s", performs better than the traditional bio scheme. the iobes tagging scheme dictates that some token sequences are illegal. for example, one cannot start an entity with an "e" tag (such as a transition from an "o", meaning it is outside of an entity, to "e-org") nor can they change types in the middle of an entity-for example, transitioning from "i-org" to "i-loc". most approaches to ner rely on the model learning which transitions are legal from the training data rather than injecting prior knowledge of how the encoding scheme works. it is conventional wisdom that, for ner, models with a linear-chain conditional random field (crf) @xcite layer perform better than those without, yielding relative performance increases between 2 and 3 percent in f1 @xcite . a crf with viterbi decoding promotes, but does not guarantee, global coherence while simple greedy decoding does not @xcite . therefore, in a bidirectional lstm (bilstm) model with a crf layer, illegal transitions are rare compared to models that select the best scoring tag for each token. due to the high variance observed in the performance of ner models @xcite it is important to have fast training times to allow for multiple runs of these models. however, as the crf forward algorithm is o(n t 2 ), where n is the length of the sentence and t is the number of possible tags, it slows down the training significantly. moreover, substantial effort is required to build an optimized, correct implementation of this layer. alternately, training with a cross-entropy loss runs in o(n ) for sparse labels and popular deep learning toolkits provide an easy to use, parallel version of this loss which brings the runtime down to o(log n ). we believe that, due to the strong contextualized local features with infinite context created by today's neural models, global features used in the crf do little more than enforce the rules of an encoding scheme. instead of traditional crf training, we propose training with a cross-entropy loss and using viterbi decoding @xcite with heuristically determined transition probabilities that prohibit illegal transitions. we call this constrained decoding and find that it allows us to train models in half the time while yielding f1 scores comparable to crfs. ## method training a tagger with a crf is normally done by minimizing the negative log likelihood of the sequence of gold tags given the input, parameterized by the model, where the probability of the sequence is given by by creating a feature function, f j , that is spanencoding-scheme-aware, we can introduce constraints that penalize any sequence that includes an illegal transition by returning a large negative value. note the summation over all possible tag sequences. while efficient dynamic programs exist to make this sum tractable for linear-chain crfs with markov assumptions, this is still a costly normalization factor to compute. in neural models, these feature functions are represented as a transition matrix that represents the score of moving from one tag y at index i to another at i+1. we implement a mask that effectively eliminates invalid iobes transitions by setting those scores to large negative values. by applying this mask to the transition matrix we can simulate feature functions that down-weigh illegal transitions. contrast the crf loss with the token-level crossentropy loss where y is the correct labels and ŷ is the model's predictions. here we can see that the loss for each element in the input i can be computed independently due to the lack of a global normalization factor. this lack of a global view is potentially harmful, as we lose the ability to condition on the previous label decision to avoid making illegal transitions. we hypothesize that, using our illegal transition heuristics, we can create feature functions that do not have to be trained, but can be applied at test time and allow for contextual coherence while using a cross-entropy loss. we can use the mask directly as the transition matrix to calculate the maximum probability sequence while avoiding illegal transitions for models that were not trained with a crf. using these transitions scores in conjunction with cross-entropy trained models, we can achieve comparable models that train more quickly. we call this method constrained decoding. constrained decoding is relatively easy to implement, given a working crf implementation, all one needs to do is apply the transition mask to the crf transition parameters to create a constrained crf. replacing the transition parameters with the mask yields our constrained decoding model. starting from scratch, one only needs to implement viterbi decoding, using the mask as transition parameters, to implement the constrained decoding model-avoiding the need for the crf forward algorithm and the crf loss. for constrained decoding, we leverage the iobes tagging scheme rather than bio tagging, allowing us to inject more structure into the decoding mask. early experiments with bio tagging failed to show the large gains we realized using iobes tagging for the reasons mentioned in section 4. ## experiments & results to test if we can replace the crf with constrained decoding we use two sequential prediction tasks: ner (conll 2003 @xcite , wnut-17 @xcite , and ontonotes @xcite ) and slot-filling (snips @xcite ). for each (task, dataset) pair we use common embeddings and hyperparameters from the literature. the baseline models are bilstm-crfs with character compositional features based on convolutional neural networks (dos santos and zadrozny, 2014) and our models are identical except we train with a crossentropy loss and use the encoding scheme constraints as transition probabilities instead of learning them with a crf. our hyper-parameters mostly follow @xcite , except we use multiple pre-trained word embeddings concatenated together @xcite . for ontonotes we follow @xcite . see section a.7 or the configuration files in our implementation for more details. the only dataset with a statistically significant difference in performance. we explore this discrepancy in section 4. similarly, the models were trained using mead-baseline @xcite , an open-source framework for creating, training, evaluating and deploying models for nlp. the constrained decoding tagger performs much faster at training time. even when compared to the optimized, batched crf provided by mead-baseline, it trained in 51.2% of the time as the crf. in addition to faster training times, training our constrained models produces only 65% of the co 2 emissions that the crf does. while gpu computations for the constrained model draw 1.3 times more power-due to the greater degree of possible parallelism in the cross-entropy loss functionthan the crf, the reduction in training time results in smaller carbon emissions as calculated in @xcite . constrained decoding can also be applied to a crf. the crf does not always learn the rules of a transition scheme, especially in early training iterations. applying the constraints to the crf can improve both f1 and convergence speed. we establish this by training bilstm-crf models with and without constraints on conll 2003. we find that the constraint mask yields a small (albeit statistically insignificant) boost in f1 as shown in our experiments suggest that injecting prior knowledge of the transition scheme helps the model to focus on learning the features for sequence tagging tasks (and not the transition rules themselves) and train faster. ## analysis the relatively poor performance of constrained decoding on ontonotes suggests that there are several classes of transition that it cannot model. for example, the transition distribution between entity types, or the prior distribution of entities. we analyzed the datasets to identify the characteristics that cause constrained decoding to fail. one such presumably obvious characteristic is the number of entity types. however, our experiments suggest that number of entity types does not affect performance: snips has more entity types than ontonotes yet constrained decoding works better for snips. we define an ambiguous token as a token whose type has multiple tag values in the dataset. for example the token "chicago" could be "i-loc" or "i-org" in the phrases "the chicago river" and "the chicago bears" respectively. such ambiguous tokens are the ones for which we expect global features to be particularly useful. a "strictly dominated token" is defined as a token that can only take on a single value due to the legality of the transition from the previous tag. in the above example given that "the" was a "b-loc" then "chicago" is strictly dominated and forced to be an "i-loc'. contrast this with a non-strictly dominated token that can still have multiple possible tag values when conditioned on the previous tag. as constrained decoding eliminates illegal transitions we would expect that it would perform well on datasets where a large proportion of ambiguous tokens are strictly dominated. this tends to hold true-only 15.9% of ontonotes' ambiguous tokens are strictly dominated while 70.7% of conll's tokens are and for wnut-17 73.6% are. we believe that the ambiguity of the first and last token of an entity also plays a role. once we start an entity, constrained decoding vastly narrows the scope of decisions that need to be made. instead of making a decision over the entire set of tags, we only decide if we should continue the entity with an "i-" or end it with an "e-". therefore, we expect constrained decoding to work well with datasets that have fairly unambiguous entity starts and ends. we quantify this by finding the proportion of entities that begin (or end) with an unambiguous type, that is, the first token of an entity only has a single label throughout the dataset, for example, "kuwait" is only labeled with "s-loc" in the conll dataset. we call these metrics "easy first" and "easy last" respectively and find that datasets with higher constrained decoding performance also have a higher percentages of entities with an easy first or last token. a summary of these characteristics for each dataset is found in this also explains why constrained decoding doesn't work as well for bio-encoded conll as it does for iobes. when using the iobes format, more tokens are strictly dominated. the other stark difference is the proportion of "easy last" entities. without the "e-" token, much less structure can be injected into the model, resulting in decreased performance of constrained decoding. these trends also hold true in internal datasets, where the automotive dataset had the fewest incidences of each of these phenomena. while not perfect predictors for the performance of constrained decoding, the metrics chosen are good proxies and can be used as a prescriptive measure for new datasets. ## previous work our approach is similar in spirit to previous work in nlp where constraints are introduced during training and inference time @xcite to lighten the computational load, and to @xcite where prior knowledge is injected into the model by manual manipulation. in our approach, however, we focus specifically on manipulating the model weights themselves rather than model features. there have been attempts to eliminate the crf layer, notably, @xcite found that an additional lstm greedy decoder layer is competitive with the crf layer, though their baseline is much weaker than the models found in other work. additionally, their decoder has an auto-regressive relationship that is difficult to parallelize and, in practice, there is still significant overhead at training time. chiu and nichols (2016) mention good results with a similar technique but don't provide in-depth analysis, metrics, or test its generality. ## conclusion for sequence tagging tasks, a crf layer introduces substantial computational cost. we propose replacing it with a lightweight technique, constrained the broad theme of the work is to find simple and computationally efficient modifications of current networks and suggest possible failure cases. while larger models have shown significant improvements, we believe there is still relevance in investigating small, targeted changes. in the future, we want to explore similar techniques in other common nlp tasks.