# VaeDiff-DocRE: End-to-end Data Augmentation Framework for Document-level Relation Extraction

arXiv: https://arxiv.org/abs/2412.13503 · HTML: https://mlsanity.com/paper/2412.13503/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 40 references that are in the map.
- First posted on arXiv 18 Dec 2024.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Deep Long-Tailed Learning: A Survey](https://mlsanity.com/paper/2110.04596/) (2021)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [DocRED: A Large-Scale Document-Level Relation Extraction Dataset](https://mlsanity.com/paper/1906.06127/) (2019)
- [Document-Level Relation Extraction with Adaptive Thresholding and Localized Context Pooling](https://mlsanity.com/paper/2010.11304/) (2020)
- [Reasoning with Latent Structure Refinement for Document-Level Relation Extraction](https://mlsanity.com/paper/2005.06312/) (2020)
- [Double Graph Based Reasoning for Document-level Relation Extraction](https://mlsanity.com/paper/2009.13752/) (2020)
- [Document-level Relation Extraction as Semantic Segmentation](https://mlsanity.com/paper/2106.03618/) (2021)
- [Revisiting DocRED -- Addressing the False Negative Problem in Relation Extraction](https://mlsanity.com/paper/2205.12696/) (2022)
- [Document-Level Relation Extraction with Reconstruction](https://mlsanity.com/paper/2012.11384/) (2020)
- [Question Answering on Freebase via Relation Extraction and Textual Evidence](https://mlsanity.com/paper/1603.00957/) (2016)
- [An End-to-end Model for Entity-level Relation Extraction using Multi-instance Learning](https://mlsanity.com/paper/2102.05980/) (2021)
- [DWIE: an entity-centric dataset for multi-task document-level information extraction](https://mlsanity.com/paper/2009.12626/) (2020)
- [Diffusion Priors In Variational Autoencoders](https://mlsanity.com/paper/2106.15671/) (2021)
- [Consistency Guided Knowledge Retrieval and Denoising in LLMs for Zero-shot Document-level Relation Triplet Extraction](https://mlsanity.com/paper/2401.13598/) (2024)
- [Document-level Relation Extraction with Relation Correlations](https://mlsanity.com/paper/2212.10171/) (2022)
- [Class-Adaptive Self-Training for Relation Extraction with Incompletely Annotated Training Data](https://mlsanity.com/paper/2306.09697/) (2023)

## Abstract (arXiv)

Document-level Relation Extraction (DocRE) aims to identify relationships between entity pairs within a document. However, most existing methods assume a uniform label distribution, resulting in suboptimal performance on real-world, imbalanced datasets. To tackle this challenge, we propose a novel data augmentation approach using generative models to enhance data from the embedding space. Our method leverages the Variational Autoencoder (VAE) architecture to capture all relation-wise distributions formed by entity pair representations and augment data for underrepresented relations. To better capture the multi-label nature of DocRE, we parameterize the VAE's latent space with a Diffusion Model. Additionally, we introduce a hierarchical training framework to integrate the proposed VAE-based augmentation module into DocRE systems. Experiments on two benchmark datasets demonstrate that our method outperforms state-of-the-art models, effectively addressing the long-tail distribution problem in DocRE.