TTM-RE: Memory-Augmented Document-Level Relation Extraction (2024)

Abstract: Document-level relation extraction aims to categorize the association between any two entities within a document. We find that previous methods for document-level relation extraction are ineffective in exploiting the full potential of large amounts of training data with varied noise levels. For example, in the ReDocRED benchmark dataset, state-of-the-art methods trained on the large-scale, lower-quality, distantly supervised training data generally do not perform better than those trained solely on the smaller, high-quality, human-annotated training data. To unlock the full potential of large-scale noisy training data for documentlevel relation extraction, we propose TTM-RE, a novel approach that integrates a trainable memory module, known as the Token Turing Machine, with a noisy-robust loss function that accounts for the positive-unlabeled setting. Extensive experiments on ReDocRED, a benchmark dataset for document-level relation extraction, reveal that TTM-RE achieves state-ofthe-art performance (with an absolute F1 score improvement of over 3%). Ablation studies further illustrate the superiority of TTM-RE in other domains (the ChemDisGene dataset in the biomedical domain) and under highly unlabeled settings.

heading: 1 Introduction
text: Relation extraction aims to classify the relationships between two specified entities into predefined categories. This task is pivotal in natural language processing, as it involves identifying and categorizing the connections between two entities (for example, "Pacific Fair" and "Queensland" in Figure 1) into predefined classes (for example, "located in" in Figure 1). Its importance spans across various downstream applications, encompassing question answering (Veena et al., 2017), knowledge graph
†Equal Senior Contribution
construction (Distiawan et al., 2019), and the extraction of general patterns (Han et al., 2020).
Previous work for relation extraction mainly focuses on sentence-level relations (Alt et al., 2020). For example, Sainz et al. (2021) characterized each relation class using a label verbalizer and addressed the relation extraction task through a textual entailment model, as well as other models such as DeepStruct (Wang et al., 2022a). However, techniques that are primarily designed and evaluated for extracting relationships at the sentence level, face challenges that limit their suitability for datasets focused on document-level relation extraction such as ReDocRED (Tan et al., 2022c).
These challenges include a large label imbalance and a large number of possible combinations be-
443
tween possible head and tail entity pairs in each document, which is quadratic. Previous work has generally addressed DocRE’s label imbalance with custom loss functions and its quadratic entity computation by minimizing document processing. Usually, a common processing step for the document inspired by Zhou et al. (2021a) is used. Initially, a pre-trained encoder processes the entire document as a whole. Subsequently, indexing over the entity mentions can be employed to retrieve the head and tail entities for classification. Recent work has mainly focused on novel loss functions (Tan et al., 2022a; Wang et al., 2022b) or additional inputs such as evidence (Ma et al., 2023).
However, less effort has been made in effectively leveraging the large amount of the distantly labeled data in ReDocRED and DocRED (Yao et al., 2019; Tan et al., 2022c). Most work (Tan et al., 2022b; Ma et al., 2023) use the distantly labeled dataset for knowledge distillation–that is, a model is first trained on the fully human annotated train data, and then used to obtain output logits on the distantly supervised data. These logits, along with the distant relation labels, are then used as guidance in training a secondary student model. However, previous work investigating fine-tuning on the distantly supervised dataset has failed to significantly boost performance (as we show later in our results). Previously, this could be explained by the lower quality and the lack of human annotation on the distantly supervised dataset. However, we assert that this is due to an architectural limitation on the prevailing framework.
Recently, numerous studies have highlighted the advantages of incorporating memory in both computer vision and NLP to acquire pertinent representations from past data points that facilitate improved classification performance. For instance, Barraco et al. (2023) showcased the enhanced performance of an encoder-decoder model for image captioning by integrating memory of past observations into the attention mechanism. Additionally, the integration of memory has been demonstrated to enhance performance in knowledge-intensive tasks such as long-form question answering and dialogue (Wu et al., 2022). More recently, Google’s Token Turing Machine (TTM) (Ryoo et al., 2023) has showcased state-of-the-art performance in realworld long-sequential visual understanding using an autoregressive Transformer model equipped with memory.
Drawing from recent advancements in memory-
Pacific Fair1 is a major shopping centre in Broadbeach Water on the Gold Coast, Queensland1, Australia. It was Queensland2’s largest regional shopping centre until 2006. Pacific Fair2 was developed by Hooker Retail Developments and opened in 1977 on what was swampland with 96 specialty stores and two anchor tenants.
Figure 2: Sample document relation extraction document from DocRED (Yao et al., 2019). Here, the head entity is related to the tail entity by "P131: located in the administrative territorial entity".
augmented models, we introduce TTM-RE, the inaugural memory-augmented architecture designed specifically for document-level relation extraction. Through empirical evidence, we demonstrate that this architecture enables notably enhanced finetuning on extensive distantly labeled data from empirical observations. Specifically, we show that adding memory tokens from TTM (Ryoo et al., 2023) empirically enhances downstream relation classification by allowing reprocessing of head and tail entities while jointly considering the learned memory tokens in mind (see Figure 1).
1. We propose TTM-RE, the first memoryaugmented document-level relation extraction model. By incorporating pseudo entities, it significantly enhances downstream relation classification performance on datasets such as ReDocRED (+3 F1 score) and ChemDisGene (+5 F1 score).
2. We show that without any human-labeled data, TTM-RE achieves impressive relation extraction performance on unseen data (+9 F1 score). Furthermore, in an extremely unsupervised scenario (19% of training labels), TTM-RE outperforms the previous SOTA by an impressive margin (+12 F1 score).
3. We perform a thorough analysis of ablations examining the performance of TTM-RE across less/more frequent relation classes, assessing the impact of memory size, layer size, and the utilization of different base models.

heading: 2 Related Work

heading: 2.1 Document-level Relation Extraction
text: Document-level relation extraction stands as a critical endeavor in natural language processing, given that over 40.7% of relations necessitate the extraction of information spanning multiple sentences and multiple entity mentions (Yao et al., 2019; Tan et al., 2022c; Zhang et al., 2022a). In Figure 2, we illustrate an instance of document-level relation extraction. Here, the objective is to discern the relationship between a pair of entities ("Pacific Fair" and "Queensland") within the provided document. Each entity is referenced twice in the text (indicated by superscripts).
Thus, a document with n entities will have n(n− 1) possible relation predictions. Zhou et al. (2021a) proposed ATLOP, which uses a single pass to encode the document using Roberta-large and indexing relevant entities in the token embeddings. Many works (Tan et al., 2022a; Wang et al., 2022b; Tan et al., 2022b) have proposed new loss functions on top of a single encoder pass (usually added on top of ATLOP (Zhou et al., 2021a)) to alleviate the combinatorial bottleneck and propose loss functions to tackle the class-imbalance problem in document-level relation extraction. Zhou et al. (2021a) introduced the ATLOP, which uses Roberta-large as the encoder and adaptive thresholding for the multilabel relation classification. Tan et al. (2022a) introduced KD-DocRE, which combines axial attention over the entity mentions with adaptive focal loss and knowledge distillation. Wang et al. (2022b) introduced a loss that considered a positive-unlabeled nature of the label distribution. Ma et al. (2023) proposed using the evidence labels along with the distantly supervised labels to achieve the current SOTA. Like other works, TTM-RE also utilizes this single-pass encoder approach, as it is elegant and efficient. However, we introduce a new model component described in the next section.

heading: 2.2 Memory-based Models in NLP
text: Memory-based models have begun to see rising usage in the CV and image captioning areas. However, their usage in NLP has been surprisingly limited. Still, there are some interesting and relevant work to our application. De Jong et al. (2021) utilized ’mention memory’ to represent knowledge—a table of dense vector representations of each entity mention in the corpus, and demon-
strated good performance on multiple open-domain tasks including claim verification, and entity-based QA. Zhong et al. (2022) introduced a training objective that directly utilizes memory sets from local, long-term, and external and showed reduced perplexity on WIKITEXT-103. Chen et al. (2022) introduced a question-answer augmented encoderdecoder model and pretraining strategy, demonstrating improved performance on single-hop and multi-hop QA datasets. Wu et al. (2022) learns keys and values that represent questions and corresponding answers respectively; at inference time, the model would retrieve information from the memory using maximum inner product search.
Inspired by Token Turing Machines, TTM-RE’s memory mechanism differs from these approaches in that it does not necessitate learning relevant portions of real text. It simply learns memory tokens (dense embeddings). This is uniquely suited to our application, as retrieving and reprocessing text requires additional LLM encoder calls for each entity-entity pair, which is quadratic in nature.

heading: 3 Methodology
text: We propose TTM-RE, a memory-augmented, document-level relation extraction method. An illustration of the overall framework of TTM-RE is shown in Figure 3.

heading: 3.1 Problem Definition
text: In the task formulation, we examine a document D that comprises M sentences (s1, s2, ..., sM ), N entities (e1, e2, ..., eN ), and R relation classes. Given this document D and a specified pair of entities (eh, et), the goal is to forecast a set of positive relations (r̂1, r̂2, ..., r̂p) between the entity pair based on the information derived from the document. It’s important to note that each entity may appear multiple times within the document D and that each possible entity-entity pair needs to be considered (i.e. if n entities, we need to consider R×N × (N − 1) possible relations).

heading: 3.2 Token Turing Machines
text: The memory, denoted as M ∈ Rm×d, comprises a collection of m tokens each with a dimensionality of d. The input, consisting of n tokens represented by I ∈ Rn×d, is combined with the memory M . This concatenated input is then further processed to generate an output denoted as O ∈ Rr×d, where r represents the desired number of retrieved tokens. The outputs from this process, in conjunction
with the preceding inputs and the current memory, constitute the output of the TTM. In our case, I ∈ R2×d, O ∈ R2×d for the head and tail entities.
Token Turing Machines add support for external memory in the form of tokens (Figure 3 Memory Module). In Token Turing Machines (TTMs), the interface between the processing unit and memory are done purely in terms of “read” and “write” operations. Note that in the original paper, the output from the processing unit is “written” to the memory, but in our case, since we are not applying this model sequentially, we can ignore this step and focus solely on the reading portion.
Initializing Memory Tokens: We follow the original implementation of TTM and initialize memory tokens from scratch, with one major difference. While the original code initialized tokens from zeros, we found that this led to a lack of gradient updates. Therefore, we initialize from a normal distribution to allow for improved learning. Note that we cannot simply use entity text embeddings since the memory layer is after all processing steps and only before the final classification layer. We found this setup to work the best empirically, but further research is needed on the placement of the memory mechanism.
Reading from Memory: While the memory is intended to encapsulate condensed information deemed significant by the model, not all of this data may be relevant. Additionally, redundancies
in the input, denoted as I , can arise due to information already stored in our memory, M , or inherent within the data itself. Selective reading where only a smaller subset of tokens is considered should encourage the model to create a memory repository containing relevant information over the entire relation classification task.
We summarize a token set I ∈ Rn×d by deriving an importance weight vector, wi ∈ Rn, which we utilize to compute a weighted aggregation across the n tokens. Notably, each output token, indexed as i ∈ 1, . . . , k, possesses its corresponding weight wi, computed using a learnable function that takes the input I itself, denoted as αi(I) : Rd → R. This importance weighting function is implemented through a Multi-Layer Perceptron (MLP) determined as wi = αi(I) = softmax(MLP(I)).
Subsequently, these weights facilitate a weighted summation of the inputs. Let V be a generic list of input tokens we wish to analyze. For our specific case, it is the concatenation of both the memory tokens M and the input tokens I . I.e. V = [M∥I]. Let us obtain encoded token zi = wi · V = αi(V ) · V , where each token zi effectively condenses all tokens from the complete set V , guided by the dynamic weighting wi = αi(V ). As the model learns to summarize p tokens into r tokens, it generates a matrix W = [w1, · · · , wr] comprising importance weights relative to the memory tokens. Also, to allow the model to read from a location, take advantage of memory token posi-
tion, and distinguish them from input tokens, there is a learnable positional embedding (Dosovitskiy et al., 2020) before each read module. All of this is captured in a memory reading function defined as: Z = Read(M , I) = Sr([M ||I]). Where [M ||I] is the concatenated memory M and input I , and Sr(·) : R(|I|+|M |)×d → Rr×d. In our application, we set r = 2 to learn memory augmented head eh′ ∈ Rd and tail et′ ∈ Rd entities for the head and tail entity relation classification problem.

heading: 3.3 Processing of Head and Tail Entities
text: After retrieving our memory-augmented head S1 and tail S2 entities, we use the group bilinear approach as specified in Zhou et al. (2021b); Tang et al. (2019) to reduce the number of parameters to enable more efficient learning. Each entity is split into k sections of dimension d/k, eh′ = [e 1 h|| . . . ||ekh′ ], et′ = [e1t′ || . . . ||ekt′ ]
p(r|eh′ , et′)s = σ ( k∑
i=1
eih′B ieit′
)
Bi ∈ Rd/k×d/k denotes bilinear layers, and the sum of the products represents the grouped bilinear layer. This reduces parameters from d2 → d2/k and enables much better performance empirically.
Furthermore, the final output is a prediction vector of dimension R+1 (number of all relations + 1 to learn the threshold value), as we adopt the adaptive thresholding approach implemented by Zhou et al. (2021b), which most other recent Document RE work as done as well (Ma et al., 2023; Tan et al., 2022a).

heading: 3.4 Noise-Robust Loss Function (SSR-PU)
text: There exists a large number of false negatives in the labeled relation triples. Gao et al. (2023) demonstrated difficulty in learning to ignore the false negatives for zero-shot prompting, revealing the difficulty of prompting LLMs for document relation extraction. To address this problem, we adopt Positive Unlabeled (PU) learning with the prior shift for each class as in Wang et al. (2022b) (Plessis et al., 2015; du Plessis et al., 2014).
Ordinary PU learning assumes that the overall distribution is the same as the distribution of the unlabeled data, which may not be true in our case. To address this problem, PU learning under the prior shift of training data needs to be considered (Charoenphakdee and Sugiyama, 2019). For each class, assume that the original prior πi =
p(yi = +1). Let πlabeled,i = p(si = +1) and (1−πlabeled,i) = (1− p(si = +1)) = p(si = −1) where si = +1 or si = −1 mean that the i-th class is labeled or unlabeled, respectively.
The conditional probability of a positive sample under unlabeled data is:
πu,i = p(yi = 1 | si = −1)
= πi − πlabeled,i 1− πlabeled,i
The non-negative risk estimator under class prior shift of training data is obtained as follows (Kiryo et al., 2017; Wang et al., 2022b):
R̂S−PU(f) = K∑
i=1
( πi nPi
nPi∑
j=1
ℓ(fi(x Pi j ),+1)
+ max(0, [ 1
nUi 1− πi 1− πu,i
nUi∑
j=1
ℓ(fi(x Ui j ),−1)
− 1 nPi πu,i − πu,iπi 1− πu,i
nPi∑
j=1
ℓ(fi(x Pi j ),−1)]))
where πi = p(yi = +1) denotes probability of positive prior for relation class i. nPi are the number of positive and nUi are the unlabelled samples of class i, respectively. ℓ is a convex loss function, and fi(·) is a score function that predicts class i. xPij and x Ui j denotes that the j-th sample of class i is positive and unlabeled as class i respectively. Please see Appendix F for more details.

heading: 4 Experimental Settings

heading: 4.1 Datasets
text: ReDocRED To evaluate our methodology, we primarily use ReDocRED (Tan et al., 2022c), an open-access, document-level relation extraction dataset that improves upon the popular DocRED
dataset (Yao et al., 2019) by resolving incompleteness, addressing logical inconsistencies, and correcting coreferential errors. Table 1 shows the amount of training data available for all data splits as well as the average number of entities. Note that we primarily use the Dev set of ReDocRED for our experiments for computational practicality.
ChemDisGene ChemDisGene (Zhang et al., 2022a) is a biomedical multi-label document RE dataset. Entity mentions were obtained using PubTator Central (Wei et al., 2019), and the relationships are based on the Comparative Toxicogenomics Database (Davis et al., 2021). Table 2 shows the stats for the data. It comprises 523 abstracts meticulously curated by a team of biologists. Our training uses the larger distantly supervised training set, while evaluation is conducted using the fully expert-labeled test set. The average number of relations per document in the test set across both datasets significantly exceeds the average number of relations per document in the training set. This indicates the incomplete labeling phenomenon in the training set with a large number of false negatives, much like DocRED, before the updated ReDocRED.

heading: 4.2 Baselines
text: For ReDocRED, we compared baselines ranging from fully supervised to distantly supervised. We will compare three settings, all evaluated on the same human-annotated test and development set. All models were chosen using the best scores on the dev set.
1. Human Annotated Only: Denotes training only on the 3053 training dataset
2. Distantly Supervised Only: Denotes training only on the 101,873 distant dataset per DocRED, as ReDocRED does not revise this dataset.
3. Human Annotated + Distantly Supervised: Combines these two datasets.
Baseline models do this in a variety of different ways, with some using knowledge distillation (i.e. teacher training on human-annotated, student training from teacher output on distantly supervised). In TTM-RE, we fine-tune on distantly supervised training data via the regular loss, freeze the memory tokens, and then fine-tune on the training set.
For fully supervised (Human Annotation Only setting), we compare against ATLOP (Zhou et al., 2021b), DREEAM (Ma et al., 2023), KD-DocRE (Tan et al., 2022b), and SSR-PU (Wang et al., 2022b). For the distantly supervised setting, we only compare against SSR-PU, as it is shown to be better than ATLOP. Furthermore, DREAM and KD-DocRE both primarily use knowledge distillation to achieve their improvements over ATLOP, only using the distantly supervised data to create the teacher logits to supervise the student model. Therefore, we believe that our "fine-tuning on both the distantly and fully supervised training data" would not maintain the spirit of the baseline method. Finally, since the main focus of many of the previous baselines combines distantly supervised work with human annotations, we also evaluate all models on the combined human+distantly supervised datasets.
For ChemDisGene, we compared against baselines BRAN (Verga et al., 2018), PubmedBert (Gu et al., 2021), PubmedBert + BRAN (Zhang et al., 2022b), ATLOP, and Wang et al. (2022b) with (Positive-negative, positive-unlabelled, and the final SSR-PU variants).
In our experiments, we use precision, recall, and F1 scores as the evaluation metrics for the performance comparison. All standard deviations were calculated with 5 runs with different random seeds. More details about these evaluation metrics can be found in Appendix B.

heading: 5 Results
text: Main Results Table 3 shows the main results of our experiments. We see here that the Human Annotated Only dataset performs on the same level as the current SOTA results, DREEAM and SSR-PU. However, we see that on the 2 other settings, Distantly Supervised and the combined Human + Distantly Supervised, TTM-RE outperforms other methods by a significant margin, when considering the standard deviation (+9 F1 and +3 F1 respectively). This implies that our model is much more effective with larger-scale training data even with noise in it,
such as the distantly supervised training datasets. This intuitively makes sense since the memory tokens are initialized from scratch, and would benefit much more from larger-scale training data. Further research should seek to improve the initialization of the memory tokens, which could lead to faster training and further performance gains.
Finally, we observe that other baselines do not generally improve significantly even after training with distantly supervised and human-annotated data, which could be caused by architectural limits. Notably, this implies that TTM-RE’s memory module adds processing capability that is actually significantly useful (Section 5 shows us an example where adding more parameters does not help).
ChemDisGene Results Table 4 shows that TTM-RE does indeed translate to other domains beyond the general task, with a 5 F1 point improvement over the best baseline. We observe TTM-RE performs well on the human-annotated training data. This is presumably because ChemDisGene has a larger dataset for training, so the memory tokens can be learned more effectively and, therefore it does not negatively affect performance as compared to the ReDocRED fully supervised setting.
Classifying Frequent/Infrequent Labels Previous work has shown that adding a memory com-
ponent yields better performance on long-tailed or imbalanced class classification problems. We do generally see this phenomenon in Table 5, as the difference is around 4 F1 and 4.5 F1 on the top 10 labels and the rest of the data. This difference is slightly more pronounced in the top 5, as we see a difference of 3.5 F1 and 5.5 F1 respectively. This indicates that baseline models perform slightly worse on the infrequent classes, whereas TTM-RE’s memory component can help alleviate this performance drop.
Extremely Unlabeled Setting Wang et al. (2022b) introduced an "extremely unlabeled" scenario, that reduced the training labels to a mere 19% of the original labels. We also evaluate our model on the extremely unlabeled setting (19%) of the original training triples in ReDocRED (Wang et al., 2022b) (Table 6). We again see that TTM-RE does not work better than baselines on fully supervised, yet increases to 12 F1 points over the best baseline when allowed to train on distantly supervised data. We hypothesize that this is due to the better learning of the infrequent classes as shown in Table 5.
Memory Token Size In Figure 4, we generally see an improvement in model performance (F1, Precision, and Ign F1) when increasing the READ module as well as the memory token size of the Token Turing Machine. While we halted at 4 layers and 200 tokens due to computational constraints, this trend is promising as it suggests that there are potential performance improvements awaiting
exploration in future endeavors, albeit requiring increased computational resources.
Using DebertaV3 as the Base Model: Interestingly, all baselines generally rely on Robertalarge as the base model. We also explored using DebertaV3-large, which is presented as a more recent and powerful model due to its larger parameter count and higher performance on the GLUE benchmark (improvements include disentangling attention, an enhanced decoding layer (He et al., 2020), and Electra-style pretraining (He et al., 2021)). However, from Table 7, we see that for document RE, it surprisingly does not improve performance. Because of this observation, the TTM-RE also uses
Roberta-large. Additionally, this demonstrates a case where naively adding parameters does not help improve relation classification performance, whereas adding the memory mechanism does.

heading: 6 Discussion
text: Analysis of Memory Module The Token Turing Machine (TTM) memory module performs best with a large training set. This leads to an important question: what other applications could provide such extensive training data? While it is true that TTM benefits from a large dataset, potential applications include any field within NLP or computer vision. Moreover, these datasets do not necessarily need to be human-annotated. For instance, the DocRED dataset was distantly supervised using Wikipedia and spaCy for NER and relation linking via Wikidata. Similar methods can be employed to create datasets for specific tasks. TTM-RE has demonstrated improved performance with these distantly supervised datasets compared to other baselines.
Dataset Size We hypothesize that performance is correlated with the size of the dataset. Notably, TTM-RE outperforms baselines in ChemDisGene without finetuning on a related dataset. Given that the human-annotated ReDocRED dataset contains only 3,053 documents, compared to 101,873 for distantly supervised datasets and 76,942 for ChemDisGene, the memory mechanism may require a certain amount of finetuning before it is fully effective. This suggests further research is needed to find more efficient ways to optimize the memory mechanism. However, future work is required to fully investigate this phenomenon.

heading: 7 Conclusion
text: In this paper, we investigated TTM-RE, integrating TTMs relation classification and evaluated our model on ReDocRED and ChemDisGene RE datasets. To summarize our contributions, no previous work has explored memory in this distantly supervised setting. As such, TTM-RE demonstrates a completely new way of increasing performance for the difficult task of document-level relation extraction as opposed to previous work, which mainly improved the loss function (Zhou et al., 2021a; Wang et al., 2022b). For this work, we found compelling results by performing ablations that showed that adding (and increasing) the number of memory tokens/layers helped performance over
baselines, compared to simply using larger models like Deberta V3. Table 5 also demonstrates the improvement in the less-represented labels as opposed to the top 10 labels (which comprise 62% of the dataset). Additionally, we show that in Figure 4, performance continues to improve as we add more memory tokens.
We also observed that TTMs necessitate either fine-tuning on a large distantly-labeled training dataset or a significantly large human-annotated training dataset (ChemDisGene) to optimize memory vector initialization. We believe that this work lends itself to future work in this exciting area, and we hope that our findings will pave the way for future exploration of memory-augmented techniques in large language models for information extraction tasks.
Limitations
Although we investigated multiple different LLMs and parameters and the type relation distribution for relation prediction as well as addressing the false positives, the performance we attained is still limited compared to supervised methods on the same task. Relation prediction still requires a large amount of data, despite TTM-RE’s ability to use distantly supervised data. Future work should seek to tackle this approach that combines labeled data creation with SOTA document relation extraction models for maximum efficiency on human annotators.
Ethical Statement
Based on the methodology we have currently employed, we do not foresee any significant ethical concerns. All the documents and models utilized in our study were obtained from open-source domains, ensuring a transparent and accessible source of information. Additionally, TTM-RE is trained on purely open-source document relation extraction data, eliminating the risk of privacy leakage. Additionally, the task of relation extraction is a widely recognized and well-studied problem across various natural language processing applications.
However, it is crucial to acknowledge a minor factor, namely the presence of potential hidden biases within the pre-trained language models used in our analysis. These biases may stem from the data on which the models were trained, which could have inadvertently introduced implicit human biases. While our usage of these pre-trained language
models enables us to identify relationships between arbitrary entities, it is conceivable that biases may emerge if one were to explore sensitive relation classes and entities.
ChatGPT and Grammarly were used for parts of the writing. In total, training took more than 75 hrs on NVIDIA RTX A6000 for pretraining in total. The main roadblock was the distantly supervised finetuning portion for all of the models, due to the size of the dataset. Derivatives of data accessed for research purposes should not be used outside of research contexts. Code will be released at https://github.com/chufangao/TTM-RE.

heading: A Parameter Settings
text: All models were run on an NVIDIA A6000 with 48 gigabytes of VRAM. Still, around 10 days were required to fully run the experiments. For particularly expensive computations, like LogitsSR, only the fastest model–UnifiedQA-large–could be feasibly run.
All models were downloaded from Huggingface (Wolf et al., 2019). We used the default setup of the pre-trained models and did not do further finetuning. All the step mentioned in the methodology section works on the output of the pre-trained models.
Supervised results DREEEAM (Ma et al., 2023) and KD-DocRE (Tan et al., 2022a) were taken from the original source papers.

heading: B Evaluation Metrics
text: To keep in tradition with existing document relation extraction work, we report both F1 and Ign_F1 as computed by the official metrics from ReDocRED. F1 refers to micro-averaged F1 score that combines precision P and recall R
F1 = 2PR
P +R
P = length of correct (h,t,rel) preds
length of all (h,t,rel) preds
R = length of correct (h,t,rel) preds
length of correct (h,t,rel)
Where (h,t,rel) denotes a tuple of the predicted head, tail, and relation. Ign_F1 is computed similarly to above but ignores the samples in the DocRED’s distantly supervised training set. (Note that we do not use any distantly labeled data).

heading: C Ablation Tables
text: Tables 8 and 9 for Figure 4.

heading: D PCA Plots of Memory Tokens
text: Shown in Figure 5, we see that the token embeddings lie in a scattered space around the head entities. This makes intuitive sense as the prototypical tokens should capture a diverse set of different types of head tokens.

heading: E Literature Review Continued
text: Pre-trained language models, such as BERT-based architectures (Xu et al., 2021), have shown considerable efficacy in document-level relation extraction. BERT-based methodologies have integrated approaches like hierarchical inference networks (Tang et al., 2020), enhanced co-reference reasoning (Ye et al., 2020), and adaptive thresholding. Furthermore, graphical neural networks (GNNs) (Zeng et al., 2020) have been leveraged for document-level relation extraction, employing techniques such as feature learning on a coreference graph (Sahu et al., 2019), edge-oriented learning strategies (Christopoulou et al., 2019), attention mechanisms (Guo et al., 2019), and iterative refinement methods to aggregate multi-hop information (Nan et al., 2020).
E.1 Weakly Supervised Document Relation Extraction
Past studies in document-level relation extraction have heavily depended on human annotation to create training datasets, a process known for its time-consuming and labor-intensive nature. There has been minimal exploration into document relation extraction methods that do not necessitate human annotation.
Various weakly supervised methods have been extensively investigated for relation extraction (Jiang, 2009; Huang and Wang, 2017; Qu et al., 2018; Wang et al., 2018; Li et al., 2018). For instance, Huang and Wang (2017) employed residual connections and convolutional neural networks (CNNs) to identify pertinent candidates, thereby enhancing supervised relation classification. In a similar vein, Qu et al. (2018) extracted textual patterns from initial examples to offer supplementary supervision. Introducing a ranking-based approach for seed selection, Phi et al. (2018) improved bootstrapping and distantly supervised relation extraction. Additionally, Sainz et al. (2021) proposed representing each relation class using a label verbalizer and tackled the relation extraction task with a textual entailment model.
Moreover, Wang et al. (2022b) showed an "extremely unlabeled" scenario wherein each relation type comprised only one instance, consequently reducing the training set to a smaller number of labeled relation triplets. However, this scenario does not help in improving performance on the fully supervised test data overall.
Qu et al. (2018) derived textual patterns from initial samples and employed them as weak signals for relation extraction. Gao et al. (2023) investigated purely weakly-supervised prompting methods devoid of human labels and revealed significant limitations in relying solely on weak supervision, particularly in handling a high incidence of hallucinations when predicting no-relation entity-entity pairs.

heading: F SSR-PU Loss
text: F.1 Class-shift Adjusted Positive Unlabeled Loss Function (SSR-PU)
Previous supervised document-level RE methods only treated unlabeled relations as negative samples, which may lead to low recall in the presence of a large number of false negatives. To address this problem, we adopt PU learning with prior shift sim-
ilar to Wang et al. (2022b) for each class (Plessis et al., 2015; du Plessis et al., 2014).
PU learning assumes that unlabeled data can reflect the true overall distribution, that is, pUi(x) = pi(x). The expected classification risk formulation can be defined in a form that can be approximated using the data like so:
R̂PU(f) =
K∑
i=1
( πi nPi
nPi∑
j=1
ℓ(fi(x Pi j ),+1)
+ max(0, [ 1
nUi
nUi∑
j=1
ℓ(fi(x Ui j ),−1)
− πi nPi
nPi∑
j=1
ℓ(fi(x Pi j ),−1)]))
(1)
where πi = p(yi = +1) denotes probability of positive prior for relation class i. nPi are the number of positive and nUi are the unlabelled samples of class i, respectively. ℓ is a convex loss function, and fi(·) is a score function that predicts class i. xPij and xUij denotes that the j-th sample of class i is positive and unlabeled as class i respectively. Note that without the max function, the second term in Eq.1 can be negative and can be prone to overfitting (and therefore highly negative) when using a highly flexible model. Thus, a non-negative risk component (Kiryo et al., 2017) is used to solve the overfitting problem. Note that nUi is essentially a hyperparameter that one assumes before model training (nPi , πi can be directly calculated by counting, and everything else is learned via backprop).
While the original method additionally corrected for the heavy class imbalance problem via multiplying γi = 1−πiπi )
0.5 before positive risk estimations as the class weight, we found that this was unnecessary and that tuning the other hyper-parameters was sufficient in reproducing the original paper results.
Prior Shift: Ordinary PU learning requires an assumption that the overall distribution needs to be the same as the distribution of the unlabeled data. In contrast, with the document-level RE dataset constructed by a recommend-revise scheme, where Wang et al. (2022b) found that there existed a prior shift in the unlabeled data of the training data vs the test data. When these priors are different, ordinary PU learning will yield a biased result.
To address this problem, inspired by the method (Charoenphakdee and Sugiyama, 2019) for handling a prior shift between the test set and the training set, a correction term is added. For each class,
assume that the original prior πi = p(yi = +1). Let πlabeled,i = p(si = +1) and (1 − πlabeled,i) = (1 − p(si = +1)) = p(si = −1) where si = +1 or si = −1 mean that the i-th class is labeled or unlabeled, respectively.
The conditional probability of a positive sample under unlabeled data is:
πu,i = p(yi = 1 | si = −1)
= p(yi = 1, si = −1)
p(si = −1)
= p(yi = 1)− p(si = +1)
p(si = −1)
= πi − πlabeled,i 1− πlabeled,i
(2)
where Step 3 is true because positive samples are assumed to be a superset of the unlabelled data and negative data has no overlaps with the labeled data. I.e. p(yi = −1, si) = 0. Finally, the non-negative risk estimator (Kiryo et al., 2017) under class prior shift of training data is obtained as follows:
R̂S−PU(f) = K∑
i=1
( πi nPi
nPi∑
j=1
ℓ(fi(x Pi j ),+1)
+ max(0, [ 1
nUi 1− πi 1− πu,i
nUi∑
j=1
ℓ(fi(x Ui j ),−1)
− 1 nPi πu,i − πu,iπi 1− πu,i
nPi∑
j=1
ℓ(fi(x Pi j ),−1)]))
(3)
The proof is shown in Theorem 1 in Wang et al. (2022b).

heading: G Distribution of Labels
text: We visualize a distribution of labels to see which relations have the least amount of occurrences in Figure 6. We see that "unemployment rate", "sister city", and "separated from" are the smallest by a few orders of magnitude. This uneven label distribution makes document relation classification even harder.

heading: H Case Study On Rare Events
text: From Appendix G, we see that many labels are quite rare. What if we restrict our labels to only those with less than the 25th quartile? Then, we obtain a dataset that is 0.027 of the total relation labels in the test set. However, whenever TTM-RE predicts one of these relations, it has an 80% chance to get it right.
E.g. We were able to predict that "Republic of China is "territory claimed by" "People’s Republic of China" (a relation in the 6th percentile) from this paragraph: The " March of the Volunteers " is
the national anthem of the People ’s Republic of China , including its special administrative regions of Hong Kong and Macau . Unlike most previous Chinese state anthems , it is written entirely in the vernacular , rather than in Classical Chinese . Its lyrics were composed as a dramatic poem by the poet and playwright , the Japan - educated Tian Han in 1934 and set to music by Nie Er from Yunnan Province the next year for the film Children of Troubled Times . It was adopted as the PRC ’s provisional anthem in 1949 in place of the " Three Principles of the People " of the Republic of China and the Communist " Internationale " . When Tian Han was imprisoned during the Cultural Revolution in the 1960s , the march was briefly and unofficially replaced by " The East Is Red " , then played without words , then played with altered words . Restored to its original version , the " March of the Volunteers " was raised to official status in 1982 , adopted by Hong Kong and Macau upon their restorations to China in 1997 and

heading: 1999 , respectively , and included in the Chinese Constitution ’s Article 136 in

heading: 2004 .
text: However, we note that this fact is not explicitly mentioned in this. It is possible that the memory mechanism enabled this prediction. Although we hypothesize that adding the memory gives it a larger context to be able to compare incoming entities, further research is needed to fully investigate the rare relation performance.

heading: Future Work
text: We believe that this work lends itself to future work in this exciting area, and we hope that our findings will pave the way for future exploration of memory-augmented techniques in large language models for information extraction tasks. Limitations
Although we investigated multiple different LLMs and parameters and the type relation distribution for relation prediction as well as addressing the false positives, the performance we attained is still limited compared to supervised methods on the same task. Relation prediction still requires a large amount of data, despite TTM-RE’s ability to use distantly supervised data. Future work should seek to tackle this approach that combines labeled data creation with SOTA document relation extraction models for maximum efficiency on human annotators