Alleviate Dataset Shift Problem in Fine-grained Entity Typing with Virtual Adversarial Training

1 Introduction

Fine-grained entity typing (FET) aims at assigning types to mentions of entities based on thier conext. It is an important task as the type information provided by it is useful for many downstream tasks, such as entity linking and event extraction [Yang et al., 2019; Ding et al., 2015]. State-of-the-art FET systems usually utilize distant supervision (DS) to fetch abundant training data by first linking a target mention to an existing entity in a knowledge base (KB) and then assigning all possible types of the entity to the target mention [Onoe and Durrett, 2019; Chen et al., 2019]. Despite of the convenience in reducing human annotation, DS inevitably introduces noises to the heuristically generated labels. For example, the mention Shelton in the sentence "The telephone number for the charity in Shelton, Conn., has been disconnected" is assigned with types {person, state province} by DS, while only {state province} is the correct label for manual labeling. Such inconsistency brings a significant distribution gap, namely dataset shift [Moreno-Torres et al., 2012], between the distantly labeled training set and the manually curated test set. Figure 1 is a hypothetical demonstration of the dataset shift problem, where ∆s denotes the dataset shift. As shown in Figure 1, due to the existence of dataset shift, even though we achieve a local minimum of loss function on training set, the loss on test set can still be large.

∗Corresponding Author

With aware of the dataset shift problem, prior arts attempt to alleviate this problem from the label perspective can be briefly summaried into two categories: (1) the first kind of

"denoising" models aim to directly model the label noise and denoise the training labels. For instance, [Gillick et al., 2014] uses heuristic rules to filter out noisy labels, [Ren et al., 2016b] uses the prior knowledge of KB and the partial loss based label embedding to remove noisy labels, [Onoe and Durrett, 2019] relabels and filters out noisy data with a neural network. (2) the other kind of methods try to infer the correct labels dynamically and use the inferred labels to direct the model during the training phase. For example, [Xu and Barbosa, 2018] proposes a variant of partial label loss to handle noisy labels and [Chen et al., 2019] proposes to infer the correct labels with a graph-based algorithm. Actually, all the above denoising methods are dedicated to reduce the gap of joint distribution P(y, x) between the training set and test set, that is, the dataset shift ∆s in Figure 1. However, since it is impossible to completely eliminate the datataset shift in distantly-supervised FET, the test loss can still be large even we achieve the minimum loss on the training set. Specifically, the denoising methods try to get low test loss by solving argminθ f(∆s|θ) ∗ f(Ltrain|θ). As is shown in Figure 1, even though the model achieves the minimum loss on the training set at θs and the dataset shift is reduced to a small value ∆s, the test loss can still be large.

In this work, complementary to the label perspective, we reconsider this problem from the model perspective: Can we learn a more robust typing model with the existence of dataset shift? As is shown in Figure 1, when dataset shift ∆s ̸= 0 is fixed, test loss not only depends on training loss f(Ltrain|θ) but also depends on the gap ∆L between test loss and training loss. Therefore, we propose to learn a robust typing model by solving argminθ f(∆L|∆s, θ) ∗ f(∆s|θ) ∗ f(Ltrain|θ), where f(∆L|∆s, θ) is a measure of the generalization ability and robustness to dataset shift of models. As is shown in Figure 1 and pointed out in [Keskar et al., 2016], the vulnerability of the models to dataset shift (i.e., large f(∆L|∆s, θ)) is due to the fact that models tend to converge to sharp minimizers. Based on this observation, we propose to improve the robustness of typing models to dataset shift with a variant of VAT. Specifically, we first use a self-paced sample selection function to select suitable samples for VAT, then construct virtual adversarial perturbations masked by the noisy labels for selected smaples, finally regularize the model to be robust to such perturbations. Since the construct perturbations can also be seen as a series of dataset shifts, the model is regularized to be robust to dataset shifts and the generalization ability of the model is finally improved. Since this process is orthogonal to denoising process, both our model and denoising models can serve as plug-and-play modules to enhance existing neural typing systems in a complementary way. Experiments on two benchmarks shows the effectiveness of the proposed method, with average 3.8%, 2.5% and 3.2% improvements in accuracy, macro F1 and Micro F1 respectively compared to next best method.

The major contributions of this paper are summarized as follows:

1. This work provides a new perspective on the problems caused by distant supervision for fine-grained entity typing and explores a new way to improve FET systems by
2. regularizing the model to be robust to dataset shift.
3. A novel regularization module for FET is proposed, where a variant of VAT and a sample selection function is proposed to control the robustness of typing model.
4. Extensive experiments on standard benchmarks with two different base models demonstrate that our method brings stable and significant improvement over the base models. Finally, the proposed method consistently outperforms several state-of-the-art (SOTA) FET systems by a significant margin.

2 RelatedWork

Fine-grained entity typing was first proposed by [Ling and Weld, 2012] who used distant supervision to induce a relatively large training corpus for FET. Due to the contextagnostic nature of the distant supervision, the training data labeled in this way is inevitably noisy by assigning mentions with the types that cannot be inferred from the context of the mentions. Some works ignore such noise: [Yogatama et al., 2015] proposed to jointly learn feature and type representations utilizing embedding techniques and [Lin and Ji, 2019] proposed to use a hybrid type classifier to capture latent type interdependency.

However, the major challenge of FET remains the problems posed by the noisy labels of the training corpus. With respect to the label noise, [Gillick et al., 2014] proposed context dependent FET and cleaned the noisy labels with heuristics, which suffers from losing training data. To mine the information in noisy labels, [Ren et al., 2016a] proposed partial label loss (PLL) to distinguish the impact of clean data and noisy data. [Xu and Barbosa, 2018] proposed a variant of PLL. From data perspective, [Ren et al., 2016b] proposed to reduce label noise using partial-label embedding, which brings a significant improvement to FET and can be viewed as a milestone for FET. More recently, [Onoe and Durrett, 2019] further proposed to reduce the label noise of the Ontonotes dataset augmented by [Choi et al., 2018] via relabeling and filtering out noisy data, which is the SOTA FET system on the Ontonotes dataset before this work. From data consistency perspective, [Chen et al., 2019] proposed Compact Latent Space Clustering (CLSC) method to regularize the representation of mentions with the same types to form compact clusters, which achieves a great success in FET. Instead of just trying to reduce the noise in the dataset, this work is devoted to train a robust model with the noisy dataset and denoise the training data jointly. In this work, we propose a regularization based framework to enhance the robustness of the model to dataset shift utilizing a variant of VAT.

3 Task Formulation

The task of FET is to uncover the entity type information for entity mentions (i.e., a sequence of tokens representing the appearance of an entity) in natural language sentences [Ren et al., 2016a]. The task takes a corpus D labeled with a predefined entity type hierarchy Y as input and predicts the most suitable type-path in Y for each entity mention from the test set Dt based on the mention's context.

Because manually labeling a large training set for FET is too expensive and error-prone, current entity typing systems annotate the training corpus utilizing DS. Specifically, the entity typing systems first detect mentions mi and link them to one or more entity ei in a KB, and then assign the entity types Yi of ei in KB to mi as candidate types. Formally, a labeled training corpus can be represented as a set of triples D = {(mi, ci, Yi)}, where mi is the i-th detected mention, ci is the context of mi and Yi is the set of candidate types of mi. We denote the all terminal types for each type path in Yi as target type set Y t i (e.g. for Yi = {person, doctor, artist}, Y t i = {doctor, artist}). This setting is also adopted by [Xu and Barbosa, 2018; Chen et al., 2019].

Note that candidate types in Yi can form multiple type paths, which may cause out-of-context noise (i.e. Yi contains type paths that are irrelevant to the mention mi in context ci). We treat samples labeled with single type path (i.e. triples (mi, ci, Yi) in D whose corresponding |Y t i | = 1) as clean data and others as noisy data. The major challenge of DS entity typing systems is to utilize both the clean data and the noisy data to produce a high-performance typing classifier.

4 Methodology

4.1 Overview

Our method is mainly based on the following assumption: the model should predict smoothly around the mention points in the feature space because mention points close to each other in feature space have similar context. For this assumption, we apply masked VAT to construct gradient-based local perturbation in the feature space for mentions and regularize the classifier to be robust to such kind of perturbation, so that the model's prediction is smoothed.

As demonstrated in Figure 2, our model mainly consists of three parts: (1) a feature extractor to project mentions with their contexts to the feature space, (2) a self-paced sample selection function to select some noisy samples for VAT, (3) and a classifier regularized by the VAT loss. All the clean data is used for adversarial training, but some noisy data will be select for VAT. Specifically, for the i-th sample (ei, ci, Y t i ) in the training set, the mention with its context (mi, ci) is transformed into a vector zi in the feature space using a feature extractor F(z|(mi, ci), θf) parameterized by θf . Then the posterior p(y | zi, θC) will be given by a classifier C parameterized by θC. After that a self-paced filter function Γ is applied to decide whether the sample will be used for VAT.

4.2 Feature Extractors and Classifier

We have explored two different feature extractors, namely NFETC [Xu and Barbosa, 2018] and BERT [Devlin et al., 2019].

NFETC. For fair comparison, We first adopt NFETC as the feature extractor, which is the same feature extractor as in [Xu and Barbosa, 2018] and [Chen et al., 2019]. In shot, NFETC can be summarized as follows: (1) the word embedding of NFETC is the concatenation of Glove [Pennington et al., 2014] embedding and word position embedding; (2) the mention representation umi of NFETC is the concatenation of the average of the embedding sequence of mi and the last hidden state of a LSTM over mi; (3) the context representation uci is the result of a word-level attention over the hidden states of a bidirectional LSTM on the context; (4) the final representation of the mention with the context is zi = FN([umi , uci ]), where FN is a feedforward neural network of n layers.

BERT. To further explore the potential of our method, we have also performed experiments using the popular pretrained language model BERT as a new baseline for finegrained entity typing. Given a mention mi = (wl, ...,wr) with its context ci = (w1, ...,wL), we simply feed the sequence ([CLS], ci, [SEP],mi, [SEP]) to BERT encoder and use the output of [CLS] token as the representation of the mention with its context zi.

Classifier. With the representation zi of a mention with its context, we employ a softmax classifier parameterized by θC = [WC, bC] to get the posterior: P(y|zi) = softmax(Wczi + bC), where WC ∈ RK×dz can be treated as the type embeddings, bc ∈ Rdz is the type bias, where K is the number of types. The predicted type ˆ y is the type with maximum posterior probability: ˆ y = arg maxy P (y|zi).

4.3 Masked Virtual Adversarial Training for FET

The basic idea of VAT is to find the adversarial perturbation radv (i.e. the perturbation on the input that lead to the most different posterior from the original posterior) and regularize the model to be robust to such perturbation, so that the prediction of the model is smoothed. In this work, we propose to improve distantly-supervised FET which is a partially labeled problem with a masked VAT. The process of applying masked VAT to FET can be decomposed into the following three parts: (1) sample selection; (2) masked adversarial perturbation generation; (3) local distributional smoothness (LDS) regularization. The details of the process are introduced in following sections.

Sample Selection

In the early stage of training, the model is prone to give wrong prediction for a mention, in which case the use of VAT will make the classifier further fit the wrong result. To alleviate this problem, we design a self-paced sample selection function to select some samples for VAT, on which the model is not confident in its predictions. Specifically, given a sample (mi, ci, Y t i ), we first find all the types Yi in the type path terminating at Y t i , then convert Yi to an type index vector ˜ Yi ∈ RK , where ˜ Yij = 1 if the j-th type is in Yi, otherwise ˜ Yij = 0. After that, we could get the posterior P(y|(mi, ci)) using the feature extractor and the classifier. With the type index vector ˜ Yi and the posterior P(y|(mi, ci)), we define the self-paced filter function Γ as following:

where 1 is an indicator function, τi = sigmoid(β(| ˜ Yi| - η)) is a real value indicating how noisy the i-th sample is, β and η are two hyper-parameters. Only when a sample's Γ(P(y|(mi, ci), ˜ Yi) = 1, the sample will be used for VAT.

Perturbation Generation and Regularization

Adversarial training. Since our method is closely related to adversarial training (AT), we first introduce the process of AT. Given the representation zi of a mention mi with its context ci, and a classifier parameterized by θC, AT adds the following regularization term to the objective function:

where radv = arg maxr,||r||<=ϵ D(q(y|zi), p(y|zi + r, θC)), q(y|zi) is the ground truth, radv is the perturbation on the representation zi, ϵ is the maximum perturbation step size which is a hyper-parameter, D(q, p) is a non-negative function measuring the divergence between two distributions q and p, and in practice we use KL divergence. Since exact minimization with respect to r is intractable for many neural networks, [Goodfellow et al., 2014] proposed to approximate radv using the Fast Gradient Signed Method (FGSM):

where sign() is the sign function.

Masked virtual adversarial training. Since some type information is unreachable in semi-supervised learning, [Miyato et al., 2018] proposed using the posterior p(y|zi; θC) given by classifier as an estimation of the ground truth q(y|zi) and defined the local distributional smoothness (LDS) to be a measure of the local smoothness of the current classifier at each data point zi:

θC denotes the current parameters of the classifier. For the task of FET which is a partially labeled problem, we propose to use masked VAT to further utilize the noisy labels. Given the representation zi of a sample (mi, ci, Y t i ), we first compute the type index vector ˜ Yi, then apply the modified virtual adversarial loss:

θC), p(y|zi + rv-adv; θ)) (4) where rv-adv = argmaxr,||r||<=ϵ D(p(y|zi; ˆ θC), p(y|zi + r; ˆ θ)), ˆ

where ˜ p(y|∗; ˆ θC) = softmax(log( ˜ Yi ⊗ ( ˆWC ∗ +ˆbC))), rv-adv = argmaxr,||r||<=ϵ D(˜ p(y|zi; ˆ θC), ˜ p(y|zi + r; ˆ θ)), ˜ p(y|zi; ˆ θC) is the posterior masked by all candidate types Yi, and ⊗ denotes element-wise product. The intuition of masked VAT is to prevent the model from strengthening totally wrong predictions and to encourage the model to distinguish between the noisy candidate types.

Solving virtual adversarial perturbations. The evaluation of rv-adv cannot be performed using Eq.(3) because the gradient of D(˜ p(y|zi; ˆ θC), ˜ p(y|zi + r; ˆ θ)) with respect to r is always 0 at r = 0. To solve this problem, [Miyato et al., 2018] proposed to approximate rv-adv using the second order Taylor expansion of D and solve the rv-adv via the power iteration method. Specifically, we can approximate rv-adv by repeatedly applying the following update nt times (nt is a hyper-parameter):

where the sign v means the unit vector of v.

4.4 Overall Objective

The final loss function consists of two parts: the supervision loss Lsup and the VAT loss Lvat. Given a batch of samples {(mi, ci, Y t i )}Bi=1 , we use the modified version of partial label loss for softmax classifier to calculate Lsup: Lsup = 1B ∑Bi=1 ∑Kk=1 yiklog(P(yi|zi; θC))k, where K is the number of target types and B is the batch size. Besides, the corresponding VAT loss is Lvat = 1B ∑Bi=1 Γ(P(y|(mi, ci), ˜ Yi) ∗ Lvat,i. So the final loss function is Lfinal = Lsup + λvat × Lvat, where λvat is a hyper-parameter controlling the influence of VAT.

5 Experiments

5.1 Dataset

We evaluate out method on two benchmarks: Ontonotes and BBN. Statistics of the datasets are shown in Table 1.

Ontonotes. The Ontonotes dataset is derived from the Newswire part of Ontonotes corpus and annotated by [Gillick et al., 2014]. The training set of Ontonotes is annotated utilizing DBpedia spotlight, while the test set is manually labeled.

BBN. The BBN dataset is derived from 2,311 Wall Street Journal articles. We use the version processed by [Ren et al., 2016a].

In this work, we use the preprocessed dataset provided by [Chen et al., 2019; Onoe and Durrett, 2019].

5.2 ComparedMethods

We compare the proposed method with several SOTA FET systems: NFGEC+LME [Xin et al., 2018] measures the compatibility between context sentences and labels utilizing a language model; NFETC [Xu and Barbosa, 2018] models type hierarchy with the hierarchical loss function; NFETCCLSC [Chen et al., 2019] compresses the clusters in the latent space with a manifold regularization loss; UFET [Choi et al., 2018] augments the Ontonotes training set with new sources of distant supervision; LABELGCN [Xiong et al., 2019] introduces a label-relational inductive bias with a graph propagation layer; FET [Lin and Ji, 2019] models type interdependency with a hybrid type classifier; DenoiseET [Onoe and Durrett, 2019] augments the Ontonotes training set and reduces the label noise with a filter and a relabel function; NFETC-VAT and BERT-VAT are the proposed models using VAT to regularize the classifiers. The methods with prefix "D" indicates the model is trained on Ontonotes dataset denoised by DenoiseET, which can be seen as the combination of our method with DenoiseET.

5.3 Evaluation Settings

For evaluation metrics, we evaluate the performance by strict accuracy, loose macro F-score and loose micro F-score, which is the most widely used evaluation setting for FET systems [Ling and Weld, 2012].

For BERT feature extractor, we use the pretrained BERTBase, cased model with a step size of 2e-5 and batch size 32. For NFETC feature extractor, we follow the setting of [Xu and Barbosa, 2018; Chen et al., 2019] using the 300 dimensional pretrained GloVe [Pennington et al., 2014] word vectors.

5.4 Performance Comparison and Analysis

Table 3 shows the performance comparison between the two VAT improved base models and several SOTA FET systems. Note that the proposed method can be further improved when combining with DenoiseET by training the model on the augmented and denoised Ontonotes dataset. On both benchmarks, the proposed method consistently outperforms other methods by a significant margin on all three metrics.

To evaluate the influence of individual components of our method, an ablation study is conducted as shown in Table 2. To further evaluate the proposed method, we conduct another ablation study, where several variants of our method are evaluated, and the result is shown in Table 4. Our analysis of the results in Table 4 is as following.

What if we use adversarial training. By comparing the base models with AT improved models, we can see that AT improves the performance of FET systems. Although AT only regularize model to be robust to perturbations around clean data, it can still improve the generalization ability of models. Besides, with the observation that VAT improved models consistently outperform AT improved models, we cloud draw a conclusion that masked VAT makes full use of noisy data to improve FET systems.

Why not generate perturbation in input space. We performed the experiment of generating perturbations on word embedding (NFETC-VAT*), and there is an improvement over the base model NFETC. But performing VAT in such way is inefficient, because it requires much more computational resources than constructing perturbations in the feature space. Besides, the performance gap between NFETC-VAT* and NFETC-VAT maybe because of the accumulation of a series of perturbations on word embedding may cause a large drift in semantics of the sentences.

Can masked VAT be effectively combined with different denoising methods. We combined masked VAT with two different denoising methods, namely NFETC-CLSC and DenoiseET. Experiment results show that the combined methods are superior to the base methods to a great extent. This is because mask VAT and denoising focus on handling dataset shift problem from different perspectives, namely the robustness to dataset shift and dataset shift reduction.

5.5 Case Study and Visualization

Does masked VAT really regularize the model to predict smoothly. As illustrated in Figure 3, the model's average LDS on the sample points of test set has been significantly improved, although we do not touch the samples of the test set during the training phase. It demonstrates that the prediction of our model is effectively smoothed, and the final model is robust to the dataset shifts constructed by virtual adversarial perturbation. Further more, the robustness of the model to dataset shifts indicates that the parameters of our model have converged to the flat minimum θ∗ in Figure 1.

Does masked VAT have effect on mention representation. We measure the quality of the mention representation by the quality of the k-NN (k=20) ranking of mentions. We use three commonly used metrics to evalute the quality of the model's k-nearest neighbors, namely mean accuracy (MA), mean average precision (MAP) and mean Normalized Discounted Cumulative Gain (mean NDCG). As shown in Figure 4, NFETCVAT consistently outperforms NFETC by a significant margin. Masked VAT constrains that the type of a mention cannot be inverted via a small perturbation on representation, so the mentions with similar initial representations but different types become separated.

6 Conclusion

In this work, we propose to alleviate dataset shift problem in FET by combining the proposed masked VAT with denoising methods. Experiment results demonstrate the proposed method consistently outperforms SOTA models by a significant margin. The proposed method is general and can be used in other domains. As a part of future work, we plan to explore the proposed method on other tasks with sever dataset shift problem, such as relation extraction.

Acknowledgements

This work has been supported in part by National Key Research and Development Program of China (2018AAA010010), NSFC (U1611461, 61572425, 61751209), Key Technologies and Systems of Humanoid Intelligence Based on Big Data (Phase II), Zhejiang University iFLYTEK Joint Research Center, University-Tongdun Technology Joint Laboratory of Artificial Intelligence, Chinese Knowledge Center of Engineering Science and Technology (CKCEST), Engineering Research Center of Digital Library, Ministry of Education, the Fundamental Research Funds for the Central Universities. Xiang Ren's research was funded by the Defense Advanced Research Projects Agency with award W911NF-19-20271.