Boundary Smoothing for Named Entity Recognition

A neural network-based NER model typically encodes the input tokens to a sequence of representations $x=x_{1},x_{2},\dots,x_{T}$ of length $T$, and then decodes these representations to task outputs, i.e., a list of entities specified by types and boundaries.

We follow Yu et al. (2020) and use the biaffine decoder. Specifically, the representations $x$ are separately affined by two feedforward networks, resulting in two representations $h^{s}\in\mathbb{R}^{T\times d}$ and $h^{e}\in\mathbb{R}^{T\times d}$, which correspond to the start and end positions of spans. For $c$ entity types (a “non-entity” type included), given a span starting at the $i$-th token and ending at the $j$-th token, a scoring vector $r_{ij}\in\mathbb{R}^{c}$ can be computed as:

where $w_{j-i}\in\mathbb{R}^{d_{w}}$ is the $(j-i)$-th width embedding from a dedicated learnable matrix; $U\in\mathbb{R}^{d\times c\times d}$, $W\in\mathbb{R}^{c\times(2d+d_{w})}$ and $b\in\mathbb{R}^{c}$ are learnable parameters. $r_{ij}$ is then fed into a softmax layer:

which yields the predicted probabilities over all entity types.

The ground truth $y_{ij}\in\mathbb{R}^{c}$ is an one-hot encoded vector, with value being 1 if the index corresponds with the annotated entity type, and 0 otherwise. Thus, the model can be optimized by the standard cross entropy loss for all candidate spans:

In the inference time, the spans predicted to be “non-entity” are first discarded, and the remaining ones are ranked by their predictive confidences. Spans with lower confidences would also be discarded if they clash with the boundaries of spans with higher confidences. Refer to Yu et al. (2020) for more details.

Figure 1(a) visualizes the ground truth $y_{ij}$ for an example sentence with two annotated entities. The valid candidate spans cover the upper triangular area of the matrix. In existing NER models, the annotated boundaries are considered to be absolutely reliable. Hence, each annotated span is assigned with the full probability to be an entity, whereas all unannotated spans are assigned with zero probability. We refer to this probability allocation as hard boundary, which is, however, probably not the best choice.

As aforementioned, the entity boundaries may be ambiguous and inconsistent, so the spans surrounding an annotated one deserve a small probability to be an entity. Figure 1(b) visualizes $\tilde{y}_{ij}$, the boundary smoothing version of $y_{ij}$. Specifically, given an annotated entity, a portion of probability $\epsilon$ is assigned to its surrounding spans, and the remaining probability $1-\epsilon$ is assigned to the originally annotated span. With smoothing size $D$, all the spans with Manhattan distance $d$ $(d\leq D)$ to the annotated entity equally share probability $\epsilon/D$. After such entity probability re-allocation, any remaining probability of a span is assigned to be “non-entity”. We refer to this as smoothed boundary.

Thus, the biaffine model can be optimized by the boundary-smoothing regularized cross entropy loss:

Empirically, the positive samples (i.e., ground-truth entities) are sparsely distributed over the candidate spans. For example, the CoNLL 2003 dataset has about 35 thousand entities, which represent only 0.93% in the 3.78 million candidate spans. By explicitly assigning probability to surrounding spans, boundary smoothing prevents the model from concentrating all probability mass on the scarce positive samples. This intuitively helps alleviate over-confidence.

In addition, hard boundary presents noticeable sharpness between the classification targets of positive spans and surrounding ones, although they share similar contextualized representations. Smoothed boundary provides more continuous targets across spans, which are conceptually more compatible with the inductive bias of neural networks that prefers continuous solutions (Hornik et al., 1989).

Table 2 presents the evaluation results on four English datasets, in which CoNLL 2003 and OntoNotes 5 are flat NER corpora, whereas ACE 2004 and ACE 2005 contains a high proportion of nested entities. Compared with previous SOTA systems, our simple baseline (RoBERTa-base + BiLSTM + Biaffine) achieves on-par or slightly inferior performance. Provided the strong baseline, our experiments show that boundary smoothing can effectively and consistently boost the $F_{1}$ score of entity recognition across different datasets. With the help of boundary smoothing, our model outperforms the best of the previous SOTA systems by a magnitude from 0.2 to 0.5 percentages.

Table 3 presents the results on four Chinese datasets, which are all flat NER corpora. Again, boundary smoothing consistently improves model performance against the baseline (BERT-base-wwm + BiLSTM + Biaffine) across all datasets. In addition, our model outperforms previous SOTA by 2.16 and 0.55 percentages on Weibo and Resume NER datasets, and achieves comparable $F_{1}$ scores on OntoNotes 4 and MSRA. Note that almost all previous systems solve these tasks within a sequence tagging framework; this work adds to the literature by introducing a span-based approach and establishing SOTA results on multiple Chinese NER benchmarks.

In five out of the above eight datasets, integrating boundary smoothing significantly increases the precision rate with a slight drop in the recall, resulting in a better overall $F_{1}$ score. This is consistent with our expectation, because boundary smoothing discourages over-confidence when recognizing entities, which implicitly leads the model to establish a more critical threshold to admit entities.

Given the use of well pretrained language models, most of the performance gains are relatively marginal. However, boundary smoothing can work effectively and consistently for different languages and datasets. In addition, it is easy to implement and integrate into any span-based neural NER models, with almost no side effects.

We perform ablation studies on CoNLL 2003, ACE 2005 and Resume NER datasets (covering flat/nested and English/Chinese datasets), to evaluate the effects of boundary smoothing parameter $\epsilon$ and $D$, as well as other components of our NER system.

The model performance (evaluated by, e.g., accuracy or $F_{1}$ score) is certainly important. However, the confidences of model predictions are also of interest in many applications. For example, when it requires the predicted entities to be highly reliable (i.e., precision is of more priority than recall), we may filter out the entities with confidences lower than a specific threshold.

However, Guo et al. (2017) have indicated that modern neural networks are poorly calibrated, and typically over-confident with their predictions. By calibration, they mean the extent to which the prediction confidences produced by a model can represent the true correctness probability. We find neural NER models also easy to become miscalibrated and over-confident. We observe that, with the standard cross entropy loss, both the development loss and $F_{1}$ score increase in the later training stage, which goes against the common perception that the loss and $F_{1}$ score should change in the opposite directions. This phenomenon is similar to the disconnect between negative likelihood and accuracy in image classification described by Guo et al. (2017). We suppose that the model becomes over-confident with its predictions, including the incorrect ones, which contributes to the increase of loss (see Appendix A for more details).

To formally investigate the over-confidence issue, we plot the reliability diagrams and calculate expected calibration error (ECE). In brief, for an NER model, we group all the predicted entities by the associated confidences into ten bins, and then calculate the precision rate for each bin. If the model is well calibrated, the precision rate should be close to the confidence level for each bin (see Appendix B for more details).

Figure 2 compares the reliability diagrams and ECEs between models with different smoothness $\epsilon$ on CoNLL 2003 and OntoNotes 5. For the baseline model ($\epsilon$ = 0), the precision rates are much lower than corresponding confidence levels, suggesting significant over-confidence. By introducing boundary smoothing and increasing the smoothness $\epsilon$, the over-confidence is gradually mitigated, and shifted to under-confidence ($\epsilon$ = 0.3). In general, the model presents best reliability diagrams when $\epsilon$ is 0.1 or 0.2. In addition, the ECEs of the baseline model are 0.072 and 0.063 on CoNLL 2003 and OntoNotes 5, respectively; with $\epsilon$ of 0.1, the ECEs are reduced to 0.013 and 0.034.

In conclusion, boundary smoothing can prevent the model from becoming over-confident with the predicted entities, and result in better calibration. In addition, as mentioned previously, spans with lower confidences are discarded if they clash with those of higher confidences when decoding. With the better calibration, the model can obtain a very marginal but consistent increase in the $F_{1}$ score.

How does boundary smoothing improve the model performance? We originally conjectured that boundary smoothing can de-noise the inconsistently annotated entity boundaries (Lukasik et al., 2020), but failed to find enough evidence – the performance improvement did not significantly increase when we injected boundary noises into the training data.⁸⁸8On the other hand, this cannot rule out the de-noising effect of boundary smoothing, because the synthesized boundary noises are distributed differently from the real noises.

As aforementioned, positive samples are very sparse among the candidate spans. Without boundary smoothing, the annotated spans are regarded to be entities with full probability, whereas all other spans are assigned with zero probability. This creates noticeable sharpness between the targets of the annotated spans and surrounding ones, although their neural representations are similar. Boundary smoothing re-allocates the entity probabilities across contiguous spans, which mitigates the sharpness and results in more continuous targets. Conceptually, such targets are more compatible with the inductive bias of neural networks that prefers continuous solutions (Hornik et al., 1989).

Li et al. (2018) have shown that residual connections and well-tuned hyperparameters (e.g., learning rate, batch size) can produce flatter minima and less chaotic loss landscapes, which account for the better generalization and trainability. Their findings provide important insights into the geometric properties of non-convex neural loss functions.

Figure 3 visualizes the loss landscapes for models with different smoothness $\epsilon$ on CoNLL 2003 and OntoNotes 5, following Li et al. (2018). In short, for a trained model, a direction of the parameters is randomly sampled, normalized and fixed, and the loss landscape is computed by sampling over this direction (refer to Appendix C for more details).

The visualization results qualitatively show that, the solutions found by the standard cross entropy are relatively sharp, whereas boundary smoothing can help arrive at flatter minima. As many theoretical studies regard the flatness as a promising predictor for model generalization (Hochreiter and Schmidhuber, 1997; Jiang et al., 2019), this result may explain why boundary smoothing can improve the model performance. In addition, boundary smoothing is associated with more smoothed landscapes – the surrounding local minima are small, shallow, and thus easy for the optimizer to escape. Intuitively, such geometric property suggests that the underlying loss functions are easier to train (Li et al., 2018).

We believe that the sharpness in the span-based NER targets is probably the reason for the sharp and chaotic loss landscape. Boundary smoothing can effectively mitigate the sharpness, and result in loss landscapes of better generalization and trainability.