Source-linked AI summary
Joint Entity and Relation Extraction with Set Prediction Networks
Dianbo Sui, Yubo Chen, Kang Liu, Jun Zhao, Xiangrong Zeng, Shengping Liu
TL;DR
Joint extraction models must handle unordered relational triples without imposing artificial generation order. This paper uses transformer-based non-autoregressive set prediction with bipartite matching loss and reports state-of-the-art performance across two benchmark datasets.
Problem
Seq2seq joint extraction models must sort unordered relational triples into sequences, while cross-entropy penalizes differences in their positions.
Method
The proposed set prediction networks use a BERT sentence encoder, a transformer-based non-autoregressive decoder, and bipartite matching loss.
Results
The proposed networks outperform state-of-the-art baselines across different scenarios on two benchmark datasets.
Takeaways & Limitations
Set prediction removes the burden of predicting an extraction order and evaluates predictions in a permutation-invariant way.
Takeaways & Limitations
The method assumes that the non-autoregressive decoder generates three triples for each sentence, with m set to 3.
Abstract
from arXiv · showhide
The joint entity and relation extraction task aims to extract all relational triples from a sentence. In essence, the relational triples contained in a sentence are unordered. However, previous seq2seq based models require to convert the set of triples into a sequence in the training phase. To break this bottleneck, we treat joint entity and relation extraction as a direct set prediction problem, so that the extraction model can get rid of the burden of predicting the order of multiple triples. To solve this set prediction problem, we propose networks featured by transformers with non-autoregressive parallel decoding. Unlike autoregressive approaches that generate triples one by one in a certain order, the proposed networks directly output the final set of triples in one shot. Furthermore, we also design a set-based loss that forces unique predictions via bipartite matching. Compared with cross-entropy loss that highly penalizes small shifts in triple order, the proposed bipartite matching loss is invariant to any permutation of predictions; thus, it can provide the proposed networks with a more accurate training signal by ignoring triple order and focusing on relation types and entities. Experiments on two benchmark datasets show that our proposed model significantly outperforms current state-of-the-art methods. Training code and trained models will be available at http://github.com/DianboWork/SPN4RE.
Introduction
Joint entity and relation extraction is recast as unordered set prediction to avoid sequence-order burdens, using parallel decoding and permutation-invariant matching. The proposed method reports state-of-the-art results on two benchmark datasets.
- Motivation: Relational triples are unordered, but seq2seq models sort them into sequences and cross-entropy penalizes predictions that shift triple positions.These models must learn both triple extraction and an artificial extraction order.
- Proposed approach: The paper formulates joint entity and relation extraction as a set prediction problem, avoiding the need to consider multiple-triple order.This directly targets the ordering burden in prior seq2seq approaches.
- Proposed approach: SPN combines a sentence encoder, a set generator, and a set-based loss built around non-autoregressive parallel decoding and bipartite matching.The encoder represents the sentence, while the generator predicts the triple set in parallel.
- Proposed approach: Bipartite matching loss is invariant to prediction permutations, unlike cross-entropy, so it evaluates predicted and ground-truth sets without relying on triple order.The loss assigns predicted triples to unique ground-truth triples.
- Results: The proposed method achieves state-of-the-art results on two benchmark datasets and is evaluated through various experiments.The datasets are New York Times and WebNLG.
Related Work
Prior joint extraction methods include pipeline, table-filling, sequence-labeling, and seq2seq paradigms. Non-autoregressive models offer parallel generation and are applied here to generate triple sets in one shot.
- Prior paradigms: Pipeline methods recognize entities first and predict relations afterward, which can ignore cross-task relevance and suffer error propagation.Joint extraction was proposed to model entity-relation dependencies explicitly.
- Prior paradigms: Joint extraction research spans table filling, sequence labeling, and seq2seq generation, with different trade-offs in representation and annotation requirements.Sequence labeling requires a complex tagging schema, whereas seq2seq models flexibly handle overlapping triples.
- Non-autoregressive models: Non-autoregressive models generate target tokens in parallel and can speed up inference in language and speech processing.The paper applies this modeling approach to information extraction.
- Non-autoregressive models: This work uses a non-autoregressive model to generate a set of relational triples in one shot.The approach avoids converting the triple set into an ordered sequence.
Method
SPN formulates joint entity and relation extraction as set prediction, using BERT, a non-autoregressive decoder, and bipartite matching to generate and train unordered triple sets.
- Set Prediction Formulation: SPN models the target as a set of relational triples rather than an ordered sequence.The formulation avoids requiring an intrinsic order among multiple triples.
- Sentence Encoder: BERT encodes each sentence into context-aware token representations for downstream triple prediction.The encoder output represents tokens using their bidirectional context.
- Non-Autoregressive Decoder: A transformer-based non-autoregressive decoder generates all triples in one pass without learning an extraction order.It uses bidirectional information rather than only left-to-right decoding.
- Non-Autoregressive Decoder: The decoder receives m shared learnable triple queries, where m exceeds the typical number of triples per sentence.This fixed-size prediction set simplifies modeling the target-set size.
- Triple Decoding: Each decoder output embedding is independently decoded into a relation type and subject and object span boundaries.Feed-forward networks predict relation types and four entity-span indices for each output.
- Bipartite Matching Loss: Bipartite matching loss first finds a minimum-cost assignment between padded ground-truth triples and predictions, then computes loss over matched pairs.The assignment considers relation-type and entity-span predictions and is computed with the Hungarian algorithm in O(m^3) time.
Experiments
The experiments evaluate SPN on NYT and WebNLG using multiple research questions covering overall performance, design choices, triple counts, and overlap patterns.
- Research Questions: The experiments address overall performance, component effectiveness, performance by triple count, and adaptation to overlapping triple patterns.These correspond to research questions RQ1 through RQ4.
- Datasets: SPN is evaluated on the widely used NYT and WebNLG joint entity and relation extraction datasets.NYT contains 24 predefined relation types, while WebNLG contains 246.
- Datasets: NYT is derived from distantly supervised New York Times articles aligned with Freebase and is treated as supervised after noise is ignored.The experiments use the preprocessed version from Zeng et al. (2018).
- Datasets: WebNLG instances contain sets of triples paired with human-written sentences that express all triples in each instance.The dataset was originally created for natural language generation.
- Overlap Patterns: The study categorizes sentences as Normal, Entity Pair Overlap, or Single Entity Overlap according to how triples share entities.These categories support analysis of performance under different overlap patterns.
Evaluation Metrics
The evaluation uses standard micro-F1 and compares SPN with established joint extraction baselines, including tagging, seq2seq, graph, cascade, and attention-based methods.
- Metrics: Performance is measured with standard micro-F1, and a triple is correct only when its relation type and both entities are correct.Entity correctness can be judged using partial or exact matching.
- Metrics: Partial matching treats an extracted entity as correct when its predicted subject and object match the ground-truth entity head words.The passage notes that entity head words are often not annotated.
- Baselines: The comparison includes NovelTagging, CopyRE, GraphRel, CopyRRL, PNDec, Attention as Relation, and CasRel.These baselines represent sequence labeling, seq2seq, graph-convolutional, pointer, attention, and cascade approaches.
- Element Evaluation: Table 4 evaluates extraction of relational-triple elements separately from the aggregate triple-level metrics.The table is specifically described as reporting results on extracting elements of relational triples.
- Baselines: CasRel is identified as the current state-of-the-art baseline and uses subject-first cascade tagging to identify relations and objects.Its first stage identifies possible subjects, followed by relation-specific object tagging.
Main Results
On NYT and WebNLG, the proposed model outperforms baselines, with gains over the current state-of-the-art and balanced extraction performance across components.
- The proposed model significantly outperforms baselines on NYT and WebNLG.
- 2.9% and 3.0% F1 improvements over Wei et al. (2020) are achieved on NYT under partial and exact matching, respectively.
- Entity pair extraction is the main bottleneck of joint extraction in NYT, while WebNLG results show balanced precision and recall.
- 93.4% partial matching F1 on WebNLG exceeds Wei et al. (2020)'s 91.8% by 1.6%.
- Table 5 reports ablation-study results on the NYT dataset.
Ablation Studies
Ablations support both core design choices: deeper non-autoregressive decoding improves results, and bipartite matching substantially outperforms cross-entropy loss.
- The ablation study evaluates decoder depth and replaces bipartite matching loss with cross-entropy loss.
- 91.4%, 92.0% and 92.15% are the best results with 1, 2 and 3 decoder layers, respectively.
- Deeper non-autoregressive decoders may better model relationships between triple queries and integrate sentence information.The passage attributes this to additional self-attention and inter-attention modules.
- Table 5 reports ablation-study results on the NYT dataset.
- 8.9% improvement is obtained when bipartite matching loss replaces cross-entropy loss.
Detailed Results on Sentences with Different Number of Triples
The proposed model performs best across sentence complexities and overlap patterns, with especially large gains for sentences containing at least five triples.
- Different Numbers of Triples: The proposed model achieves the best results across sentences containing 1, 2, 3, 4 or ≥5 triples.
- Different Numbers of Triples: 6.9% is the NYT F1 improvement over Wei et al. (2020) for sentences with ≥5 triples.
- Overlapping Patterns: The proposed model outperforms all baselines in Normal, SEO and EPO overlap classes.
- Overlapping Patterns: 3.5%, 2.6% and 2.1% improvements over Wei et al. (2020) occur in Normal, SEO and EPO classes, respectively.
- Overlapping Patterns: Figure 2 reports F1 scores for relational-triple extraction under different overlapping patterns in NYT.
- Overlapping Patterns: Seq2seq models decline across harder overlap patterns, whereas the proposed model remains consistently strong across all three.
Conclusion and Future Work
The paper introduces set prediction networks that remove triple-order prediction through parallel decoding and bipartite matching, achieving strong results across scenarios while leaving long-tailed relations for future work.
- Set prediction networks formulate joint entity and relation extraction as set prediction rather than sequence prediction.
- Non-autoregressive parallel decoding and bipartite matching loss avoid predicting triple order and validate unique predictions.
- Experiments on two widely used datasets show performance exceeding state-of-the-art baselines across different scenarios.
- Relation types have an imbalanced or long-tailed distribution in both NYT and WebNLG.
- Future work will combine cost-sensitive learning with the proposed set prediction networks.
Hungarian Algorithm
The Hungarian method solves the assignment problem using a square cost matrix. It iteratively transforms the matrix until zeros support an optimal assignment.
- The algorithm takes an m × m square cost matrix as input and solves the assignment problem in polynomial time.
- Steps 1–2: Matrix reduction: Step 1 subtracts each row’s minimum element from every element in that row.
- Steps 1–2: Matrix reduction: Step 2 subtracts each column’s minimum element from every element in that column.
- Steps 3–4: Iterative optimization: Step 3 covers all zeros with the minimum number of horizontal and vertical lines.
- Steps 3–4: Iterative optimization: If m lines cover the zeros, an optimal assignment exists; otherwise, Step 4 creates additional zeros and the process continues.
An Example of Bipartite Matching Loss
The example demonstrates bipartite matching loss for three predicted triples by constructing a cost matrix, finding its minimum-cost assignment, and evaluating the matched loss.
- Example setup: The example uses three relation types plus a special ∅ relation, with a non-autoregressive decoder generating m = 3 triples per sentence.
- Example setup: The sentence describes Aarhus Airport serving Aarhus, which is led by Jacob Bundsgaard, and the ground truth includes a padded no-triple entry.
- Example setup: Entity spans and relation types are converted into integer representations for the ground-truth sets Y0, Y1, and Y2.
- Loss computation: Bipartite matching loss first finds an optimal permutation between ground-truth and predicted triples, then computes the loss from that assignment.
- Loss computation: The example computes individual matching costs, including Cmatch(Y0, ˆY0) = −0.4, Cmatch(Y0, ˆY1) = −2.6, and Cmatch(Y1, ˆY0) = −3.3.
- Loss computation: The Hungarian algorithm selects the assignment with minimum total cost, reported as −5.9 = (−2.6) + (−3.3) + 0.
- Loss computation: Using the optimal assignment, the final loss is calculated as 7.52.