Document-Level Relation Extraction with Adaptive Thresholding and Localized Context Pooling

Wenxuan Zhou; Kevin Huang; Tengyu Ma; Jing Huang

Source-linked AI summary

Document-Level Relation Extraction with Adaptive Thresholding and Localized Context Pooling

Wenxuan Zhou, Kevin Huang, Tengyu Ma, Jing Huang

arXiv:2010.11304v3Oct 21, 2020cs.CL

TL;DR

Document-level relation extraction must classify multiple entity pairs and multiple possible relations while identifying relevant context within a document. ATLOP addresses these challenges with adaptive thresholding and localized context pooling, achieving 63.40% F1 and new state-of-the-art results on DocRED and other benchmarks.

Problem
Document-level relation extraction involves multiple entity pairs and possible relations per document, requiring models to identify relevant pair-specific context.
Method
ATLOP combines a learnable entity-pair-dependent threshold with localized pooling of pre-trained attention to classify relations using relevant context.
Results
63.40% F1 on DocRED is a new state-of-the-art result, while experiments report state-of-the-art performance across three document-level relation extraction benchmarks.
Takeaways & Limitations
The experiments support adaptive thresholding and localized context pooling as effective techniques for document-level relation extraction.

Abstract

from arXiv · show

Document-level relation extraction (RE) poses new challenges compared to its sentence-level counterpart. One document commonly contains multiple entity pairs, and one entity pair occurs multiple times in the document associated with multiple possible relations. In this paper, we propose two novel techniques, adaptive thresholding and localized context pooling, to solve the multi-label and multi-entity problems. The adaptive thresholding replaces the global threshold for multi-label classification in the prior work with a learnable entities-dependent threshold. The localized context pooling directly transfers attention from pre-trained language models to locate relevant context that is useful to decide the relation. We experiment on three document-level RE benchmark datasets: DocRED, a recently released large-scale RE dataset, and two datasets CDRand GDA in the biomedical domain. Our ATLOP (Adaptive Thresholding and Localized cOntext Pooling) model achieves an F1 score of 63.4, and also significantly outperforms existing models on both CDR and GDA.

Introduction

Document-level relation extraction must classify multiple entity pairs and labels while focusing on pair-relevant context amid irrelevant document information. ATLOP addresses these challenges with adaptive thresholding and localized context pooling, achieving strong results across three benchmark datasets.

Challenges: Document-level RE classifies multiple entity pairs simultaneously and must identify relevant context despite irrelevant information.This differs from sentence-level RE, where a sentence contains only one entity pair to classify.
Method: Localized context pooling transfers attention heads from pretrained language models to add pair-relevant context to entity representations.The technique avoids introducing graph structures and addresses the problem of using one entity embedding for all pairs.
Challenges: A global probability threshold for multilabel prediction requires heuristic tuning and can cause instance-specific decision errors.The tuned development-set threshold may not be optimal for all instances.
Method: Adaptive thresholding learns an entity-pair-dependent threshold using a rank-based loss that separates positive and negative class logits.At test time, labels above the threshold are returned, or NA when no class exceeds it.
ATLOP: ATLOP combines adaptive thresholding and localized context pooling to address multilabel and multientity document-level RE.Experiments cover DocRED, CDR, and GDA, and the model significantly outperforms state-of-the-art methods.

Problem Formulation

Document-level relation extraction predicts relations for every ordered pair of entities in a document, accounting for multiple mentions of each entity and the possibility of no relation. A relation holds when any pair of entity mentions expresses it.

Problem Formulation: The task predicts a subset of relations from R ∪ {NA} for each ordered entity pair (e_s, e_o), where R is the predefined relation set.The subject and object entities are denoted e_s and e_o, respectively.
Problem Formulation: Each entity e_i may occur multiple times through its entity mentions {m_ij}.The formulation represents repeated appearances of an entity as a set of mentions.
Problem Formulation: A relation exists between entities when any pair of their mentions expresses it; otherwise, the entity pair is labeled NA.At test time, the model predicts labels for all entity pairs in the document.

Enhanced BERT Baseline

The enhanced BERT baseline encodes documents once, pools mention representations into entity embeddings, and classifies entity-pair relations with a parameter-efficient bilinear model. It uses a tuned global threshold during inference and significantly outperforms existing BERT baselines.

The model builds on existing BERT baselines and integrates additional techniques to improve document-level relation extraction.
Entity mentions are marked with special symbols, and a pretrained language model produces contextual embeddings used for all entity-pair classifications.The document is encoded once, so all entity pairs share the same contextual representation.
Logsumexp pooling aggregates mention signals into entity embeddings and performs better than mean pooling in experiments.
Group bilinear classification splits embedding dimensions into k equal groups, reducing parameters from d^2 to d^2/k.Relation probabilities are computed with sigmoid activation after nonlinear entity transformations.
During inference, a global threshold θ tuned on the development set determines predicted relations or NA when no relation exceeds it.
The enhanced base model achieves near state-of-the-art performance and significantly outperforms existing BERT baselines.

Adaptive Thresholding

Adaptive thresholding replaces a single global decision threshold with a learned threshold class for each entity pair, separating positive from negative relation classes. Its specialized loss trains these comparisons and yields a large performance gain over global thresholding in experiments.

Adaptive thresholding: A learned threshold class TH separates positive and negative relation classes for each entity pair.At test time, classes with logits higher than TH are returned as positive relations.
Adaptive thresholding: The adaptive-thresholding loss decomposes into two parts, L = L1 + L2.L1 handles positive classes with TH, while L2 handles negative classes with TH.
Adaptive thresholding: L1 sums categorical cross-entropy losses over positive classes to push their logits above the TH class.L1 is not used when an entity pair has no positive label.
Adaptive thresholding: The proposed adaptive-thresholding loss obtains a large performance gain over the global threshold in experiments.A global threshold may not reflect different confidence levels across entity pairs or relation classes.

Localized Context Pooling

Localized context pooling augments globally pooled entity representations with pair-specific context identified from transformer attention. It emphasizes tokens important to both entities, producing representations that differ across entity pairs.

Method: For an entity pair, multiplying the two entities’ attention identifies context important to both and yields the localized context embedding c(s,o).The localized context is formed as a weighted combination of contextual token embeddings.
Motivation: Global pooling aggregates all entity mentions across the document, but some mention contexts may be irrelevant to a particular entity pair.The paper motivates localization with an example where a mention of John Stanistreet is irrelevant to the pair (John Stanistreet, Bendigo).
Method: Localized context pooling enhances each entity pair’s embedding with an additional local context embedding related to both entities.It directly transfers token-level dependencies learned by pre-trained transformer self-attention.
Method: Entity-level attention is obtained by averaging mention-level attention across mentions of the same entity.Mention-level attention comes from attention to the “*” symbol in the pre-trained multi-head attention matrix.
Implementation: The localized embedding is fused into globally pooled entity embeddings, producing pair-specific representations; experiments use the last transformer layer’s attention matrix.The fusion modifies the original linear layers with parameters Wc1 and Wc2.

Experiments

ATLOP is evaluated against sequence-, graph-, and transformer-based models on DocRED and against prior biomedical relation-extraction models on CDR and GDA. Ablation and targeted analyses show that adaptive thresholding and localized context pooling improve performance, with each contributing substantially to ATLOP.

Biomedical results: ATLOP with SciBERT outperforms current state-of-the-art results on the biomedical CDR and GDA datasets.CDR predicts Chemical–Disease interactions, while GDA predicts Gene–Disease interactions.
DocRED results: ATLOP is compared with sequence-based, graph-based, and transformer-based models on DocRED using F1 and Ign F1.Ign F1 excludes relational facts shared by the training and development/test sets.
DocRED results: 58.52% F1 is achieved by the enhanced baseline BERT-EBASE, close to the state-of-the-art BERT-LSRBASE.The re-implemented BERT baseline also outperforms BiLSTM-LSR by 1.2%.
Ablation studies: 0.89% and 0.97% dev F1 drops occur when adaptive thresholding and localized context pooling are separately removed from ATLOP.The adaptive-thresholding benefit requires the adaptive-thresholding loss; using binary cross entropy instead yields 41.74% dev F1.
Ablation studies: About 1% dev F1 increases result from group bilinear and logsumexp pooling in BERT-EBASE, while entity markers provide a 0.24% increase.Entity markers are retained because they simplify deriving mention embeddings and mention-level attention.
Targeted analyses: Localized context pooling consistently outperforms its removal as document entity counts increase, except in one group, while adaptive thresholding consistently outperforms other thresholding strategies on the DocRED test set.The localized pooling advantage is reported when documents contain more than five entities and grows as entity counts increase.

Related Work

Relation extraction research progressed from sentence-level methods to document-level approaches as relationships spanning multiple sentences became increasingly important. Existing document-level models use document graphs or pretrained language models, but ATLOP specifically addresses multi-entity and multi-label problems with adaptive thresholding and localized context pooling.

Sentence-level relation extraction: Early relation extraction methods predicted relationships between two entities within a sentence using sequence-based, graph-based, transformer-based, and pre-training approaches.These approaches are represented by works including Zeng et al. (2014), Miwa and Bansal (2016), Alt, H¨ubner, and Hennig (2019), and Soares et al. (2019).
Document-level relation extraction: As many relationships are expressed across multiple sentences, research began exploring document-level relation extraction.This shift responds to the limitations of restricting relation prediction to individual sentences.
Document graphs: Most document-level models construct document graphs with words as nodes and intra- and inter-sentential dependencies as edges to extract entity-pair features.Later work extended this document-graph idea by improving the underlying neural architectures.
Pretrained language models: Other models apply pretrained language models directly, while Wang et al. (2019a) first predicts relationship existence and then predicts the specific relation.These models avoid explicit document graphs because pretrained language models can learn structures such as dependencies and coreferences.
Research gap and contribution: ATLOP addresses document-level RE’s multi-entity and multi-label problems with adaptive thresholding and localized context pooling, significantly outperforming existing models.The paper identifies these problems as key differences from sentence-level relation extraction.

Conclusion

The paper proposes ATLOP for document-level relation extraction, combining adaptive thresholding with localized context pooling. These techniques learn entity-pair-specific classification thresholds and use pre-trained attention heads to identify relevant context.

Conclusion: ATLOP is a document-level relation extraction model built around adaptive thresholding and localized context pooling.The paper presents both as novel techniques.
Conclusion: Adaptive thresholding replaces the global multi-label classification threshold with a learnable threshold that selects the best value for each entity pair.
Conclusion: Localized context pooling uses pre-trained attention heads to locate relevant context for entity pairs.The technique is intended to help alleviate document-level relation extraction difficulties.