Graph Refinement for Coreference Resolution (2022)

Abstract: The state-of-the-art models for coreference resolution are based on independent mention pair-wise decisions. We propose a modelling approach that learns coreference at the document-level and takes global decisions. For this purpose, we model coreference links in a graph structure where the nodes are tokens in the text, and the edges represent the relationship between them. Our model predicts the graph in a non-autoregressive manner, then iteratively refines it based on previous predictions, allowing global dependencies between decisions. The experimental results show improvements over various baselines, reinforcing the hypothesis that document-level information improves conference resolution.

## introduction current state-of-the-art (sota) solutions for coreference resolution such as @xcite @xcite formulate the problem in an end-to-end manner where the models jointly learn to detect mentions and link coreferent mentions. the objective is to predict the antecedent of each mention-span in a document, so the model performs pair-wise decisions of all mentions. after having the model predictions, related mentions are grouped into clusters. under this scenario, each decision (i.e., whether two mentions are related to the same entity or not) is independent. @xcite proposed an iterative method to update the representation of a mention with information of its probable antecedents. however, the final decisions are still made locally. we propose a modeling approach that learns coreference at the document-level and takes global decisions. we propose to model mentions and coreference links in a graph structure where the nodes are tokens in the text, and the edges represent the relationships between them. figures 1 and 2 * work done as a phd student at epfl/idiap the connection types are encoded as, 0: no links, 1: mention links, 2: coreference links. we follow a similar approach to the graph-to-graph transformer (g2gt) proposed in @xcite for syntactic parsing, but instead of encoding sentences, we en-code documents. our model predicts the graph in a non-autoregressive manner, then iteratively refines it based on previous predictions. this recursive process introduces global dependencies between decisions. unlike @xcite , we define different structures for input and output graphs, to reflect the different roles of these graphs. to ensure that locality in the input graph reflects all the relevant relationships, the input graph encodes relations for all mention tokens. this makes the encoding process easier. to provide a unique specification of the target graph, the output only encodes a minimal set of connections. this facilitates prediction. we initialize the transformer with pre-trained language models, either bert @xcite , or spanbert @xcite . another difference with @xcite is that our model predicts two levels of representation. while they predict the whole graph at each iteration, during the first iteration our model only predicts edges that identify mention-spans. this is because mention detection is a sentence-level phenomenon whose outputs are required as inputs to coreference resolution, which is a discourse-level phenomenon. but we do not organise these two tasks in a pipeline. starting at the second iteration, the model predicts the complete graph. this allows the model to refine mention decisions given coreference decisions, and vice versa. in this way, we propose to use iterative graph refinement as an alternative to pipeline architectures for multi-level deep learning models. the iterative process finishes when there are no more changes in the graph or when a maximum number of iterations is reached. ideally, the whole document should be encoded at once, but in practice there is a limit on the maximum length. in order to deal with this issue, we propose two strategies: overlapping windows and reduced document. in the first strategy, we split documents into overlapping windows of the maximum allowed size k. the segments overlap for a length k/2. at decoding time, segments are input in order, and we construct the final graph by joining all graphs from different segments. in the second strategy, we use two networks. the mention-span network is the previously described overlapping model, and we use it for predicting the first graph. for the second network, we reduce the document by including only the tokens of candidate mention-spans, separated by a special token. this network refines the initial graph for the following iterations. the experiments show improvements over the relevant baselines and state-of-the-art. they also indicate that the models reach the best solution in a maximum of three iterations. given that we predict the graph at once for each iteration, our model's complexity is lower than the baselines. our contributions are the following: the rest of the paper is organized as follows. section 2 presents a summary of coreference resolution approaches related to this paper. section 3 briefly describes the fundamentals of state-of-theart approaches. in section 4, we define entity mentions and their coreference links as a graph, and fomulate the task as a sequence-to-graph problem. in section 5, we present our iterative refinement solution to global modelling of the coreference graph, and in section 6, we present two proposed architectures to address the resulting computational issues. sections 7, 8 and 9 contain the experimental setup, results and discussion, respectively. finally, section 10 draws the conclusions of this paper. ## related work the first approaches to coreference resolution (cr) were rule-based systems @xcite , but eventually, they were outperformed by machine learning approaches @xcite @xcite due to annotated corpora's creation. in genral, there are three coreference approaches : mention-pair, entity-mention, and ranking models. mention-pair models set coreference as a binary classification problem. the initial stage is the mention detection, where the input is raw text, and the output is the locations of each entity mention in the text. mention detection is done as an independent task in a pipeline model @xcite or as part of an end-to-end model @xcite . the next stage is the classification of mention pairs. at first, the best classifiers were decision trees @xcite @xcite , but later, neural networks became the sota. the final stage is reconciling the pair-wise decisions to create entity chains, usually by utilizing greedy algorithms or clustering approaches. entity-mention models focus on maintaining single underlying entity representation for each cluster, contrasting the independent pair-wise decisions of mention-pair approaches @xcite . ranking models aim at ranking the possibles antecedent of each mention instead of making binary decisions @xcite ). an alternative modeling approach is to perform clustering instead of classification @xcite . sota models for cr are mostly based on @xcite . they introduced the first end-to-end model that jointly optimizes mention detection and coreference resolution tasks. these neural networkbased models also simplify the mention input representation to be word embedding vectors, instead of the traditional pipeline of different linguistic feature extraction tools such as part-of-speech (pos) tagging and dependency parsing. the following models proposed improvements over this work. @xcite improved the previous model by introducing higher order inference so the entity's mention representation will get iteratively updated with the weighted average of antecedent representations, where the weights are the predictions from the model at the previous iteration. this contrasts with our approach in that we iterate over the whole coreference link graph and we perform discrete decisions at each iteration. @xcite use reinforcement learning to directly optimize the model on the evaluation metrics. @xcite uses bert embeddings @xcite as input. @xcite introduced a new span-bert embedding model, which is shown to outperform bert for the cr task. @xcite showed that higher order inference has low impact on strong models such as spanbert. @xcite proposed a bounded memory model trained to manage limited memory by learning to forget entities. finally, @xcite formulated the problem of coreference resolution as questionanswering and trained a model for span prediction. this model has the advantage of being pretrained with larger data-sets from the question-answering task. ## baseline: neural coreference resolution neural coreference resolution, as formulated in @xcite , is a mention-pair approach. it uses an exhaustive method defining mentions as any text span of any size in a document. there, a document d represents a sequence of tokens of size n . the objective is to assign an antecedent y i to each of the m text spans m i in d. the set of possible antecedents of the span m i is denoted as y(i). this set contains all text spans with index less than i, plus a null antecedent ϵ, the null antecedent is assigned when: (a) the span is not an entity mention, (b) the span is the first mention of an entity in the document. the final mention clusters are constructed greedily by grouping connected spans based on the model predictions during decoding time. the model is trained to learn a conditional probability distribution over documents p(y 1 , ..., y n |d), assuming independence among each decision of antecedent assignment y i , as follows: in @xcite , the probability distribution p(y i |d) is inferred over t iterations of the model over the same input document. at each iteration t, the span representations are updated with the weighted average of all possible antecedents at time t -1 where the weights are given by the probability distribution of the model at time t -1. they called this model high-order coreference resolution since each mention representation considers information from its probable antecedents. the training optimization is done using crossentropy. given that a mention-span m i can have more than one true antecedent, the loss considers the sum of probabilities of all true antecedents in the annotated data: ## graph modeling we propose to model the set of coreference links of a document in a graph structure where the nodes are tokens foot_0 and the edges are links of different types. given a document @xmath0, ..., x n ] of size n , the coreference graph is defined as the matrix g ⊂ n n ×n of links between tokens. here, the relation type between two tokens, x i and x j , is encoded with integers and is denoted as g i,@xmath1, 1, 2}. we define three relation types: (0) no link, (1) mention link, and (2) coreference link, as illustrated in figure 2 . mention links this type of link serves to identify mentions. we define mention links in two different manners depending on whether the graph is an input or output of the model, for functional reasons. when the graph is an input g in , there is a directed link from each mention's token to the mention head, including the head to itself. when the graph is the model's output g out , there is only one directed link from the last token of the mention-span to the first token. both encoding methods define a mentionspan uniquely, even when having nested mentions; every mention has a unique start-end combination and a unique head. the model utilizes the output for prediction, so it is simpler to predict one single link, whereas, in the input, the model uses links to all tokens to provide a more direct representation of the role of every token in the mention. ## iterative refinement the strong independence assumption made in equation (3) does not reflect the real scenario and could lead to poor performance. therefore, we use an iterative refinement approach to model interdependencies between relations, similar to g2gt (mohammadshahi and henderson, 2021). under this approach, the model makes t iterations over the same document d. at each iteration t, the predicted coreference graph g t is conditioned on the previously predicted one g t-1 . the model's conditional probability distribution is now defined as follows: this means that the graph should be input to the transformer model @xcite . following @xcite , the graph is encoded by inputting an embedding for the type of each relation into the self-attention function of the transformer : to illustrate the iterative refinement of a graph, figure 3 shows an example of two iterations of the model. the mention links are indicated with solid line arrows and the coreference links with dotted arrows. the initial graph matrix g in 0 is full of zeros, so no connections are drawn. the first predicted graph g out 1 only has mention-links because initially there were no mention heads to be connected. this graph is transformed to serve as input g in 1 for the next iteration. finally, during the second iteration, the model predicts the coreference graph g out 2 . the model can continue iterating for a maximum of t times. ## architectures there exists in practice a maximum length for encoding a document due to limited hardware memory. in this section, we describe two strategies to manage this issue: overlapping windows and reduced document. in the experiments we also report results for a naive strategy of truncating the documents at the maximum segment length of k for both training and testing. ## results analysis this section describes the results of various baselines and our models. first, we analyze the optimum number of refinement iterations, and then we show results using the best models. as expected, pre-training with spanbert results in better scores than with bert, and bertlarge is better than bert-base. not surprisingly, 'g2gt bert-base truncated' and 'g2gt bertlarge truncated' perform poorly in comparison to the baseline because their information is incomplete. for bert-base, both the 'overlap' and 'reduce' models have better scores than the comparable baseline. for bert-large and spanbert, the 'overlap' model has similar scores to the baseline, but the 'reduce' model consistently improves over the baseline. preliminary experiments with g2gt 'overlap' in a pipeline approach, where mention detection is performed before coreference, showed that it is not better than in a joint approach showed here. overall, our g2gt 'reduce' method consistently shows the highest scores across all the models for each pre-trained model. our models do not surpass sota @xcite (shown in grey), but as mentioned before, this sota model is also trained on the much more abundant data from the questionanswering task, and so it is not directly comparable to our model. we leave the issue of incorporating additional data into the training of our model to future work. ## discussion these results support our claim that coreference resolution benefits from making global coreference decisions using document-level information. first, refinement of coreference decisions using global information about other coreference decisions clearly improves accuracy, as indicated by the improved scores for models with more than one coreference iteration in one issue with our method is the necessity to iteratively pass the input through an expensive encoder model more than once. however, the number of iterations needed is small, and results in significant improvement. the length management methods would not be necessary if we had more efficient pre-trained transformer models or larger-memory gpu hardware which could handle longer sequences. however, the computational cost of very large transformers will always be an issue, so in general there is a need to address the issue of how to reduce the number of inputs when modelling phenomena which require large contexts, such as coreference resolution. this paper contributes towards addressing this general issue. ## conclusion we proposed a g2gt model with iterative refinement for coreference resolution. for this purpose, we define a graph structure to encode coreference links contained in a document. that enables our model to predict the complete coreference graph at once. the graph is then refined in a recursive manner, iterating the model conditioned on the document and the graph prediction from the previous step. this allows global modelling of all coreference decisions using all document-level information, but it introduces computational issues for longer documents. we experimented with two methods to manage long documents and maintain computational efficiency. the first method encodes the document in overlapping segments. the second method reduces the set of tokens which are input. the evaluation shows that both methods can outperform a comparable baseline, and that the second method has better performance than the first one and than all other comparable models. this experiment shows that global decisions and documentlevel information are useful to improve coreference and thus should not be ignored. it also shows that the models can benefit from increasingly powerful pre-trained language models, bert-base @xcite , bert-large @xcite , and spanbert @xcite . by empirically showing the benefits of making global decisions and using document-level information in coreference resolution, this work motivates further work on this topic. in addition, the model designs developed in this work provide a viable approach to addressing the related issues. addressing the computational issues with modelling large documents in transformers is an area of active research, and our proposed methods could be improved in future work.