Improving Coreference Resolution by Learning Entity-Level Distributed Representations

Introduction

Coreference resolution, the task of identifying which mentions in a text refer to the same real-world entity, is fundamentally a clustering problem. However, many recent state-of-the-art coreference systems operate solely by linking pairs of mentions together [Durrett and Klein, 2013, Martschat and Strube, 2015, Wiseman et al., 2015].

An alternative approach is to use agglomerative clustering, treating each mention as a singleton cluster at the outset and then repeatedly merging clusters of mentions deemed to be referring to the same entity. Such systems can take advantage of entity-level information, i.e., features between clusters of mentions instead of between just two mentions. As an example for why this is useful, it is clear that the clusters {Bill Clinton} and {Clinton, she} are not referring to the same entity, but it is ambiguous whether the pair of mentions Bill Clinton and Clinton are coreferent.

Previous work has incorporated entity-level information through features that capture hard constraints like having gender or number agreement between clusters [Raghunathan et al., 2010, Durrett et al., 2013]. In this work, we instead train a deep neural network to build distributed representations of pairs of coreference clusters. This captures entity-level information with a large number of learned, continuous features instead of a small number of hand-crafted categorical ones.

Using the cluster-pair representations, our network learns when combining two coreference clusters is desirable. At test time it builds up coreference clusters incrementally, starting with each mention in its own cluster and then merging a pair of clusters each step. It makes these decisions with a novel easy-first cluster-ranking procedure that combines the strengths of cluster-ranking [Rahman and Ng, 2011] and easy-first [Stoyanov and Eisner, 2012] coreference algorithms.

Training incremental coreference systems is challenging because the coreference decisions facing a model depend on previous decisions it has already made. We address this by using a learning-to-search algorithm inspired by SEARN [Daumé III et al., 2009] to train our neural network. This approach allows the model to learn which action (a cluster merge) available from the current state (a partially completed coreference clustering) will eventually lead to a high-scoring coreference partition.

Our system uses little manual feature engineering, which means it is easily extended to multiple languages. We evaluate our system on the English and Chinese portions of the CoNLL 2012 Shared Task dataset. The cluster-ranking model significantly outperforms a mention-ranking model that does not use entity-level information. We also show that using an easy-first strategy improves the performance of the cluster-ranking model. Our final system achieves CoNLL F1 scores of 65.29 for English and 63.66 for Chinese, substantially outperforming other state-of-the-art systems.Code and trained models are available at https://github.com/clarkkev/deep-coref.

System Architecture

Our cluster-ranking model is a single neural network that learns which coreference cluster merges are desirable. However, it is helpful to think of the network as being composed of distinct sub-networks. The mention-pair encoder produces distributed representations for pairs of mentions by passing relevant features through a feedforward neural network. The cluster-pair encoder produces distributed representations for pairs of clusters by applying a pooling operation over the representations of relevant mention pairs, i.e., pairs where one mention is in each cluster. The cluster-ranking model then scores pairs of clusters by passing their representations through a single neural network layer.

We also train a mention-ranking model that scores pairs of mentions by passing their representations through a single neural network layer. Its parameters are used to initialize the cluster-ranking model, and the scores it produces are used to prune which candidate cluster merges the cluster-ranking model considers, allowing the cluster-ranking model to run much faster. The system architecture is summarized in Figure 1.

Building Representations

In this section, we describe the neural networks producing distributed representations of pairs of mentions and pairs of coreference clusters. We assume that a set of mentions has already been extracted from each document using a method such as the one in Raghunathan et al. (2010).

Embedding Features: Word embeddings of the head word, dependency parent, first word, last word, two preceding words, and two following words of the mention. Averaged word embeddings of the five preceding words, five following words, all words in the mention, all words in the mention’s sentence, and all words in the mention’s document.

Additional Mention Features: The type of the mention (pronoun, nominal, proper, or list), the mention’s position (index of the mention divided by the number of mentions in the document), whether the mentions is contained in another mention, and the length of the mention in words.

Document Genre: The genre of the mention’s document (broadcast news, newswire, web data, etc.).

Distance Features: The distance between the mentions in sentences, the distance between the mentions in intervening mentions, and whether the mentions overlap.

Speaker Features: Whether the mentions have the same speaker and whether one mention is the other mention’s speaker as determined by string matching rules from Raghunathan et al. (2010).

String Matching Features: Head match, exact string match, and partial string match.

The vectors for all of these features are concatenated to produce an $I$-dimensional vector $\bm{h}_{0}$, the input to the neural network. If $a=\text{\sc na}$, the features defined over mention pairs are not included. For this case, we train a separate network with an identical architecture to the pair network except for the input layer to produce anaphoricity scores.

Our set of hand-engineered features is much smaller than the dozens of complex features typically used in coreference systems. However, we found these features were crucial for getting good model performance. See Section 6.1 for a feature ablation study.

Hidden Layers. The input gets passed through three hidden layers of rectified linear (ReLU) units [Nair and Hinton, 2010]. Each unit in a hidden layer is fully connected to the previous layer:

where $\bm{W}_{1}$ is a $M_{1}\times I$ weight matrix, $\bm{W}_{2}$ is a $M_{2}\times M_{1}$ matrix, and $\bm{W}_{3}$ is a $d\times M_{2}$ matrix.

The output of the last hidden layer is the vector representation for the mention pair: $\bm{r}_{m}(a,m)=\bm{h}_{3}(a,m)$.

2 Cluster-Pair Encoder

The cluster-pair encoder first combines the information contained in the matrix of mention-pair representations $\bm{R}_{m}(c_{i},c_{j})=[\bm{r}_{m}(m^{i}_{1},m^{j}_{1}),\bm{r}_{m}(m^{i}_{1},m^{j}_{2}),...,\bm{r}_{m}(m^{i}_{|c_{i}|},m^{j}_{|c_{j}|})]$ to produce $\bm{r}_{c}(c_{i},c_{j})$. This is done by applying a pooling operation. In particular it concatenates the results of max-pooling and average-pooling, which we found to be slightly more effective than using either one alone:

Mention-Ranking Model

Rather than training a cluster-ranking model from scratch, we first train a mention-ranking model that assigns each mention its highest scoring candidate antecedent. There are two key advantages of doing this. First, it serves as pretraining for the cluster-ranking model; in particular the mention-ranking model learns effective weights for the mention-pair encoder. Second, the scores produced by the mention-ranking model are used to provide a measure of which coreference decisions are easy (allowing for an easy-first clustering strategy) and which decisions are clearly wrong (these decisions can be pruned away, significantly reducing the search space of the cluster-ranking model).

The mention-ranking model assigns a score $s_{m}(a,m)$ to a mention $m$ and candidate antecedent $a$ representing their compatibility for coreference. This is produced by applying a single fully connected layer of size one to the representation $\bm{r}_{m}(a,m)$ produced by the mention-pair encoder:

where $\bm{W}_{m}$ is a $1\times d$ weight matrix. At test time, the mention-ranking model links each mention with its highest scoring candidate antecedent.

Training Objective. We train the mention-ranking model with the slack-rescaled max-margin training objective from Wiseman et al. (2015), which encourages separation between the highest scoring true and false antecedents of the current mention. Suppose the training set consists of $N$ mentions $m_{1},m_{2},...,m_{N}$. Let $\mathcal{A}(m_{i})$ denote the set of candidate antecedents of a mention $m_{i}$ (i.e., mentions preceding $m_{i}$ and na), and $\mathcal{T}(m_{i})$ denote the set of true antecedents of $m_{i}$ (i.e., mentions preceding $m_{i}$ that are coreferent with it or $\{\textsc{na}\}$ if $m_{i}$ has no antecedent). Let $\hat{t}_{i}$ be the highest scoring true antecedent of mention $m_{i}$:

where $\Delta(a,m_{i})$ is the mistake-specific cost function

for “false new,” “false anaphoric,” “wrong link,” and correct coreference decisions. The different error penalties allow the system to be tuned for coreference evaluation metrics by biasing it towards making more or fewer coreference links.

Finding Effective Error Penalties. We fix $\alpha_{\textsc{wl}}=1.0$ and search for $\alpha_{\textsc{fa}}$ and $\alpha_{\textsc{fn}}$ out of $\{0.1,0.2,...,1.5\}$ with a variant of grid search. Each new trial uses the unexplored set of hyperparameters that has the closest Manhattan distance to the best setting found so far on the dev set. We stopped the search when all immediate neighbors (within 0.1 distance) of the best setting had been explored. We found $(\alpha_{\textsc{fn}},\alpha_{\textsc{fa}},\alpha_{\textsc{wl}})=(0.8,0.4,1.0)$ to be best for English and $(\alpha_{\textsc{fn}},\alpha_{\textsc{fa}},\alpha_{\textsc{wl}})=(0.7,0.4,1.0)$ to be best for Chinese on the CoNLL 2012 data. We attribute our smaller false new cost from the one used by Wiseman et al. (they set $\alpha_{\textsc{fn}}=1.2$) to using more precise mention detection, which results in fewer links to na.

Training Details. We initialized our word embeddings with 50 dimensional ones produced by word2vec [Mikolov et al., 2013] on the Gigaword corpus for English and 64 dimensional ones provided by Polyglot [Al-Rfou et al., 2013] for Chinese. Averaged word embeddings were held fixed during training while the embeddings used for single words were updated. We set our hidden layer sizes to $M_{1}=1000,M_{2}=d$ = 500 and minimized the training objective using RMSProp [Hinton and Tieleman, 2012]. To regularize the network, we applied L2 regularization to the model weights and dropout [Hinton et al., 2012] with a rate of 0.5 on the word embeddings and the output of each hidden layer.

Pretraining. As in Wiseman et al. (2015), we found that pretraining is crucial for the mention-ranking model’s success. We pretrained the network in two stages, minimizing the following objectives from Clark and Manning (2015): All-Pairs Classification

Where $\mathcal{F}(m_{i})$ is the set of false antecedents for $m_{i}$ and $p(a,m_{i})=\text{\it sigmoid}(s(a,m_{i}))$. The top-pairs objective is a middle ground between the all-pairs classification and mention ranking objectives: it only processes high-scoring mentions, but is probabilistic rather than max-margin. We first pretrained the network with all-pairs classification for 150 epochs and then with top-pairs classification for 50 epochs. See Section 6.1 for experiments on the two-stage pretraining.

Cluster-Ranking Model

Although a strong coreference system on its own, the mention-ranking model has the disadvantage of only considering local information between pairs of mentions, so it cannot consolidate information at the entity-level. We address this problem by training a cluster-ranking model that scores pairs of clusters instead of pairs of mentions.

Given two clusters of mentions $c_{i}$ and $c_{j}$, the cluster-ranking model produces a score $s_{c}(c_{i},c_{j})$ representing their compatibility for coreference. This is produced by applying a single fully connected layer of size one to the representation $\bm{r}_{c}(c_{i},c_{j})$ produced by the cluster-pair encoder:

where $\bm{W}_{c}$ is a $1\times 2d$ weight matrix. Our cluster-ranking approach also uses a measure of anaphoricity, or how likely it is for a mention $m$ to have an antecedent. This is defined as

where $\bm{W}_{\textsc{na}}$ is a $1\times d$ matrix.

At test time, the cluster ranker iterates through every mention in the document, merging the current mention’s cluster with a preceding one or performing no action. We view this procedure as a sequential decision process where at each step the algorithm observes the current state $x$ and performs some action $u$.

Specifically, we define a state $x=(C,m)$ to consist of $C=\{c_{1},c_{2},...\}$, the set of existing coreference clusters, and $m$, the current mention being considered. At a start state, each cluster in $C$ contains a single mention. Let $c_{m}\in C$ be the cluster containing $m$ and $\mathcal{A}(m)$ be a set of candidate antecedents for $m$: mentions occurring previously in the document. Then the available actions $U(x)$ from $x$ are

Merge$[c_{m},c]$, where $c$ is a cluster containing a mention in $\mathcal{A}(m)$. This combines $c_{m}$ and $c$ into a single coreference cluster.

Pass. This leaves the clustering unchanged.

After determining the new clustering $C^{\prime}$ based on the existing clustering $C$ and action $u$, we consider another mention $m^{\prime}$ to get the next state $x^{\prime}=(C^{\prime},m^{\prime})$.

Using the scoring functions $s_{c}$ and $s_{\textsc{na}}$, we define a policy network $\pi$ that assigns a probability distribution over $U(x)$ as follows:

During inference, $\pi$ is executed by taking the highest-scoring (most probable) action at each step.

2 Easy-First Cluster Ranking

The last detail needed is the ordering in which to consider mentions. Cluster-ranking models in prior work order the mentions according to their positions in the document, processing them left-to-right [Rahman and Ng, 2011, Ma et al., 2014]. However, we instead sort the mentions in descending order by their highest scoring candidate coreference link according to the mention-ranking model. This causes inference to occur in an easy-first fashion where hard decisions are delayed until more information is available. Easy-first orderings have been shown to improve the performance of other incremental coreference strategies [Raghunathan et al., 2010, Stoyanov and Eisner, 2012] because they reduce the problem of errors compounding as the algorithm runs.

We also find it beneficial to prune the set of candidate antecedents $\mathcal{A}(m)$ for each mention $m$. Rather than using all previously occurring mentions as candidate antecedents, we only include high-scoring ones, which greatly reduces the size of the search space. This allows for much faster learning and inference; we are able to remove over 95% of candidate actions with no decrease in the model’s performance. For both of these two preprocessing steps, we use $s(a,m)-s(\textsc{na},m)$ as the score of a coreference link between $a$ and $m$.

3 Deep Learning to Search

We face a sequential prediction problem where future observations (visited states) depend on previous actions. This is challenging because it violates the common i.i.d. assumption made in machine learning. Learning-to-search algorithms are effective for this sort of problem, and have been applied successfully to coreference resolution [Daumé III and Marcu, 2005, Clark and Manning, 2015] as well as other structured prediction tasks in natural language processing [Daumé III et al., 2014, Chang et al., 2015a].

We train the cluster-ranking model using a learning-to-search algorithm inspired by SEARN [Daumé III et al., 2009], which is described in Algorithm 1. The algorithm takes as input a dataset $\mathcal{D}$ of start states $x$ (in our case documents with each mention in its own singleton coreference cluster) and structured labels $y$ (in our case gold coreference clusters). Its goal is to train the policy $\pi$ so when it executes from $x$, reaching a final state $e$, the resulting loss $\mathcal{L}(e,y)$ is small. We use the negative of the B3 coreference metric for this loss [Bagga and Baldwin, 1998]. Although our system evaluation also includes the MUC [Vilain et al., 1995] and CEAF${}_{\phi_{4}}$ [Luo, 2005] metrics, we do not incorporate them into the loss because MUC has the flaw of treating all errors equally and CEAF${}_{\phi_{4}}$ is slow to compute.

For each example $(x,y)\in\mathcal{D}$, the algorithm obtains a trajectory of states $x_{1},x_{2},...,x_{n}$ visited by the current policy by running it to completion (i.e., repeatedly taking the highest scoring action until reaching an end state) from the start state $x$. This exposes the model to states at train time similar to the ones it will face at test time, allowing it to learn how to cope with mistakes.

Given a state $x$ in a trajectory, the algorithm then assigns a cost $l(u)$ to each action $u\in U(x)$ by executing the action, “rolling out” from the resulting state with a reference policy $\pi^{\text{ref}}$ until reaching an end state $e$, and computing the resulting loss $\mathcal{L}(e,y)$. This rolling out procedure allows the model to learn how a local action will affect the final score, which cannot be otherwise computed because coreference evaluation metrics do not decompose over cluster merges. The policy network is then trained to minimize the risk associated with taking each action: $\sum_{u\in U(x)}\pi(u|x)l(u)$.

Reference policies typically refer to the gold labels to find actions that are likely to be beneficial. Our reference policy $\pi^{\text{ref}}$ takes the action that increases the B3 score the most each step, breaking ties randomly. It is generally recommended to use a stochastic mixture of the reference policy and the current learned policy during rollouts when the reference policy is not optimal [Chang et al., 2015b]. However, we find only using the reference policy (which is close to optimal) to be much more efficient because it does not require neural network computations and is deterministic, which means the costs of actions can be cached.

Training details. We update $\pi$ using RMSProp and apply dropout with a rate of 0.5 to the input layer. For most experiments, we initialize the mention-pair encoder component of the cluster-ranking model with the learned weights from the mention-ranking model, which we find to greatly improve performance (see Section 6.2).

Runtime. The full cluster-ranking system runs end-to-end in slightly under 1 second per document on the English test set when using a GPU (including scoring all pairs of mentions with the mention-ranking model for search-space pruning). This means the bottleneck for the overall system is the syntactic parsing required for mention detection (about 4 seconds per document on the English test set).

Experiments and Results

Experimental Setup. We run experiments on the English and Chinese portions of the CoNLL 2012 Shared Task data [Pradhan et al., 2012]. The models are evaluated using three of the most popular coreference metrics: MUC, B3, and Entity-based CEAF (CEAF${}_{\phi_{4}}$). We generally report the average F1 score (CoNLL F1) of the three, which is common practice in coreference evaluation. We used the most recent version of the CoNLL scorer (version 8.01), which implements the original definitions of the metrics.

Mention Detection. Our experiments were run using system-produced predicted mentions. We used the rule-based mention detection algorithm from Raghunathan et al. (2010), which first extracts pronouns and maximal NP projections as candidate mentions and then filters this set with rules that remove spurious mentions such as numeric entities and pleonastic it pronouns.

Feature Ablations. We performed a feature ablation study to determine the importance of the hand-engineered features included in our model. The results are shown in Table 1. We find the small number of non-embedding features substantially improves model performance, especially the distance and string matching features. This is unsurprising, as the additional features are not easily captured by word embeddings and historically such features have been very important in coreference resolvers [Bengtson and Roth, 2008].

The Importance of Pretraining. We evaluate the benefit of the two-step pretraining for the mention-ranking model and report results in Table 2. Consistent with Wiseman et al. (2015), we find pretraining to greatly improve the model’s accuracy. We note in particular that the model benefits from using both pretraining steps from Section 4, which more smoothly transitions the model from a mention-pair classification objective that is easy to optimize to a max-margin objective better suited for a ranking task.

2 Cluster-Ranking Model Experiments

We evaluate the importance of three key details of the cluster ranker: initializing it with the mention-ranking model’s weights, using an easy-first ordering of mentions, and using learning to search. The results are shown in Table 3.

Pretrained Weights. We compare initializing the cluster-ranking model randomly with initializing it with the weights learned by the mention-ranking model. Using pretrained weights greatly improves performance. We believe the cluster-ranking model has difficulty learning effective weights from scratch due to noise in the signal coming from cluster-level decisions (an overall bad cluster merge may still involve a few correct pairwise links) and the smaller amount of data used to train the cluster-ranking model (many possible actions are pruned away during preprocessing). We believe the score would be even lower without search-space pruning, which stops the model from considering many bad actions.

Easy-First Cluster Ranking. We compare the effectiveness of easy-first cluster-ranking with the commonly used left-to-right approach. Using a left-to-right strategy simply requires changing the preprocessing step ordering the mentions so mentions are sorted by their position in the document instead of their highest scoring coreference link according to the mention-ranking model. We find the easy-first approach slightly outperforms using a left-to-right ordering of mentions. We believe this is because delaying hard decisions until later reduces the problem of early mistakes causing later decisions to be made incorrectly.

Learning to Search. We also compare learning to search with the simpler approach of training the model on a trajectory of gold coreference decisions (i.e., training on a fixed cost-sensitive classification dataset). Using this approach significantly decreases performance. We attribute this to the model not learning how to deal with mistakes when it only sees correct decisions during training.

3 Capturing Semantic Similarity

Using semantic information to improve coreference accuracy has had mixed in results in previous research, and has been called an “uphill battle” in coreference resolution [Durrett and Klein, 2013]. However, word embeddings are well known for being effective at capturing semantic relatedness, and we show here that neural network coreference models can take advantage of this.

Perhaps the case where semantic similarity is most important is in linking nominals with no head match (e.g., “the nation” and “the country”). We compare the performance of our neural network model with our earlier statistical system [Clark and Manning, 2015] at classifying mention pairs of this type as being coreferent or not. The neural network shows substantial improvement (18.9 F1 vs. 10.7 F1) on this task compared to the more modest improvement it gets at classifying any pair of mentions as coreferent (68.7 F1 vs. 66.1 F1). Some example wins are shown in Table 4. These types of coreference links are quite rare in the CoNLL data (about 1.2% of the positive coreference links in the test set), so the improvement does not significantly contribute to the final system’s score, but it does suggest progress on this difficult type of coreference problem.

4 Final System Performance

In Table 5 we compare the results of our system with state-of-the-art approaches for English and Chinese. Our mention-ranking model surpasses all previous systems. We attribute its improvement over the neural mention ranker from Wiseman et al. (2015) to our model using a deeper neural network, pretrained word embeddings, and more sophisticated pretraining.

The cluster-ranking model improves results further across both languages and all evaluation metrics, demonstrating the utility of incorporating entity-level information. The improvement is largest in CEAF${}_{\phi_{4}}$, which is encouraging because CEAF${}_{\phi_{4}}$ is the most recently proposed metric, designed to correct flaws in the other two [Luo, 2005]. We believe entity-level information is particularly useful for preventing bad merges between large clusters (see Figure 4 for an example). However, it is worth noting that in practice the much more complicated cluster-ranking model brings only fairly modest gains in performance.

Related Work

There has been extensive work on machine learning approaches to coreference resolution [Soon et al., 2001, Ng and Cardie, 2002], with mention-ranking models being particularly popular [Denis and Baldridge, 2007, Durrett and Klein, 2013, Martschat and Strube, 2015].

We train a neural mention-ranking model inspired by Wiseman et al. (2015) as a starting point, but then use it to pretrain a cluster-ranking model that benefits from entity-level information. Wiseman et al. (2016) extend their mention-ranking model by incorporating entity-level information produced by a recurrent neural network running over the candidate antecedent-cluster. However, this is an augmentation to a mention-ranking model, and not fundamentally a clustering model as our cluster ranker is.

Entity-level information has also been incorporated in coreference systems using joint inference [McCallum and Wellner, 2003, Poon and Domingos, 2008, Haghighi and Klein, 2010] and systems that build up coreference clusters incrementally [Luo et al., 2004, Yang et al., 2008, Raghunathan et al., 2010]. We take the latter approach, and in particular combine the cluster-ranking [Rahman and Ng, 2011, Ma et al., 2014] and easy-first [Stoyanov and Eisner, 2012, Clark and Manning, 2015] clustering strategies. These prior systems all express entity-level information in the form of hand-engineered features and constraints instead of entity-level distributed representations that are learned from data.

We train our system using a learning-to-search algorithm similar to SEARN [Daumé III et al., 2009]. Learning-to-search style algorithms have been employed to train coreference resolvers on trajectories of decisions similar to those that would be seen at test-time by Daumé et al. (2005), Ma et al. (2014), and Clark and Manning (2015). Other works use structured perceptron models for the same purpose [Stoyanov and Eisner, 2012, Fernandes et al., 2012, Björkelund and Kuhn, 2014].

Conclusion

We have presented a coreference system that captures entity-level information with distributed representations of coreference cluster pairs. These learned, dense, high-dimensional feature vectors provide our cluster-ranking coreference model with a strong ability to distinguish beneficial cluster merges from harmful ones. The model is trained with a learning-to-search algorithm that allows it to learn how local decisions will affect the final coreference score. We evaluate our system on the English and Chinese portions of the CoNLL 2012 Shared Task and report a substantial improvement over the current state-of-the-art.

Acknowledgments

We thank Will Hamilton, Jon Gauthier, and the anonymous reviewers for their thoughtful comments and suggestions. This work was supported by NSF Award IIS-1514268.

References