Editing Language Model-based Knowledge Graph Embeddings

Early KG embedding approaches primarily focused on deriving embeddings from structured information alone. Most existing KG embedding methods utilize the translation-based paradigm such as TransE (Bordes et al. 2013), TransR (Lin et al. 2015), TransH (Wang et al. 2014), and RotatE (Sun et al. 2019) or semantic matching paradigms, including DistMult (Yang et al. 2015), RESCAL (Nickel, Tresp, and Kriegel 2011), and HolE (Nickel, Rosasco, and Poggio 2016). Additionally, there has been significant interest in explicitly utilizing structural information through graph convolution networks (Kipf and Welling 2017; Velickovic et al. 2018; Vashishth et al. 2020; Liu et al. 2021b; Zhang et al. 2022b). Since pre-trained language models (PLMs) (Zhao et al. 2023) have made waves for widespread tasks, framing KG embedding via language models is an increasing technique that has led to empirical success (Pan et al. 2023). On the one hand, several works have emerged to leverage PLMs for KG embeddings. Some studies (Yao, Mao, and Luo 2019; Zhang et al. 2020c; Wang et al. 2021c, 2022a) leverage finetuning the PLMs, for example, KG-BERT (Yao, Mao, and Luo 2019), which takes the first step to utilize BERT (Devlin et al. 2019) for KG embeddings by regarding a triple as a sequence and turning link prediction into a sequence classification task. StAR (Wang et al. 2021a) proposes a structure-augmented text representation approach that employs both spatial measurement and deterministic classifier for KG embeddings, respectively. LMKE (Wang et al. 2022b) adopts language models which utilize description-based KG embedding learning with a contrastive learning framework. On the other hand, some studies (Lv et al. 2022) adopt the prompt tuning (Liu et al. 2022; Chen et al. 2022b; Zhang, Li, and et al. 2022) with language models. PKGC (Lv et al. 2022) converts triples into natural prompt sentences for KG embedding learning. There are also some other studies (Xie et al. 2022; Saxena, Kochsiek, and Gemulla 2022; Chen et al. 2022a) formulate KG embedding learning as sequence-to-sequence generation with language models. Nonetheless, KG embeddings in previous studies are usually deployed as static artifacts whose behavior is challenging to modify after deployment. To this end, we propose a new task of editing language model-based KG embeddings regarding the correction or changes for facts of existing KGs, which is the first work focusing on this to the best of our knowledge.

Editable training (Sinitsin et al. 2020; Yao et al. 2023) represents an early, model-agnostic attempt to facilitate rapid editing of trained models. Subsequently, numerous reliable and effective approaches have been proposed to enable model editing without the necessity for resource-intensive re-training (Cohen et al. 2023; Geva et al. 2023; Hase et al. 2023; Han et al. 2023). These methods can be broadly classified into two distinct categories: external model-based editor and additional parameter-based editor. The first category, external model-based editors, employs extrinsic model editors to manipulate model parameters. For instance, KnowledgeEditor (KE) (Cao, Aziz, and Titov 2021) adopts a hyper-network-based approach to edit knowledge within language models. MEND (Mitchell et al. 2022) utilizes a hypernetwork (MLP) to predict the rank-1 decomposition of fine-tuning gradients. The predicted gradients are used to update a subset of the parameters of PLMs. The second category, additional parameter-based editors, involves using supplementary parameters to adjust the final output of the model, thereby achieving model editing. Building upon prior research (Dai et al. 2022a), which suggests that Feed-Forward Networks (FFNs) in PLMs may store factual knowledge, CALINET (Dong et al. 2022) enhances a specific FFN within the PLM by incorporating additional parameters composed of multiple calibration memory slots. Existing approaches predominantly focus on modifying knowledge within pre-trained language models, constraining the manipulation and interpretation of facts. In contrast, by editing knowledge in KG embeddings, our proposed tasks enable adding and modifying knowledge, thereby extending the applicability of KG embeddings across various downstream tasks. Previous methods for updating knowledge graph embeddings (KGE), such as OUKE (Fei, Wu, and Khan 2021) and RotatH (Wei et al. 2021), mainly focused on modifying the entity and relation vectors in score function-based KGE models. These methods employed different dataset versions as snapshots to represent changes within the knowledge graph and validated these changes directly using hit@k. However, this approach falls short of effectively measuring the efficacy of edits. Our paper improves these methods by creating datasets and introducing metrics for a more accurate evaluation of KGE updates.

A KG can be represented as $\mathcal{G}=(\mathcal{E},\mathcal{R},\mathcal{T})$, where $\mathcal{E}$, $\mathcal{R}$ and $\mathcal{T}$ are sets of entities, relation types, and triples. Each triple in $\mathcal{T}$ takes the form $(h,r,t)$, where $h,t\in\mathcal{E}$ are the head and tail entities. For the EDIT task, knowledge requiring edits is defined as $(h,r,y,a)$ or $(y,r,t,a)$, where $y$ is an incorrect/outdated entity. We denote the original input entity as $x$ and the predicted target as $y$. $a$ points to the desired edited entity (change from $y$ to $a$). E.g., <Donald Trump, president_of, U.S.A.> is an outdated triple, which should be changed to <Joe Biden, president_of, U.S.A.> to update the KG embeddings. Apart from outdated information, numerous new facts may be absent in current KGs. Consequently, it becomes essential to flexibly incorporate these new triples into the KG embeddings. Thus, we introduce the ADD task to integrate new knowledge seamlessly. Note that the ADD task is similar to the inductive setting in KG completion but without re-training the model. Formally, we define the task of editing language model-based KG embeddings as follows:

$$p(a\mid\mathcal{\tilde{W}},\phi,x)\leftarrow p(y\mid\mathcal{W},\phi,x)$$

(1)

where $\mathcal{W}$ and $\mathcal{\tilde{W}}$ denote the original and edited parameters of KG embeddings, respectively. In this paper, $x$ refers to the original input (e.g., $(?,h,r)$), output entity ($y$), and desired edited entity ($a$). $\phi$ represents the external parameters of the editor network. At the same time, we need to strive to maintain the stability of the model for other correct knowledge, as described below:

$$p(y^{\prime}\mid\mathcal{\tilde{W}},\phi,x^{\prime})\leftarrow p(y^{\prime}\mid\mathcal{W},\phi,x^{\prime})$$

(2)

where $x^{\prime}$ and $y^{\prime}$ represent the input and label of the factual knowledge stored in the model, respectively, with $x^{\prime}\neq x$.

We construct four datasets for EDIT and ADD tasks, based on two benchmarks FB15k-237 (Toutanova et al. 2015), and WN18RR (Dettmers et al. 2018). After initially training KG embeddings with language models, we sample challenging triples as candidates, as detailed below. Recent research indicates that pre-trained language models can learn significant factual knowledge (Petroni et al. 2019; Cao et al. 2021), and predict the correct tail entity with a strong bias. This may hinder the proper evaluation of the editing task. Thus, we exclude relatively simple triples inferred from language models’ internal knowledge for precise assessment. Specifically, we select data with link prediction ranks >2,500. For EDIT, the Top-1 facts from the origin model replace the entities in the training set of FB15k-237 to build the pre-train dataset (original pre-train data), while golden facts serve as target edited data. For ADD, we leverage the original training set of FB15k-237 to build the pre-train dataset (original pre-train data) and the data from the inductive setting as they are not seen before. Unlike EDIT, for ADD, since new knowledge is not present during training, we directly use the training set to evaluate the ability to incorporate new facts. We adopt the same strategy to construct two datasets from WN18RR. To assess knowledge locality, we further create a reference dataset (L-Test) based on the link prediction performance with a rank value below $k$ (the same $k$ as in the RK metric). Five human experts scrutinize all datasets for ethical concerns. Ultimately, we establish four datasets: E-FB15k237, A-FB15k237, E-WN18RR, and A-WN18RR. Table 1 provides details of the datasets. The creation of the datasets for the EDIT task is shown in Figure 2.

This section introduces the technical background of language model-based KG embeddings. Specifically, we illustrate two kinds of methodologies to leverage language models, namely: finetuning (FT-KGE) and prompt tuning (PT-KGE). Note that our task is orthogonal to language-based KG embeddings, performance differences between FT-KGE and PT-KGE are discussed in the next section.

FT-KGE methods, such as KG-BERT (Yao, Mao, and Luo 2019), which regard triples in KGs as textual sequences. Specifically, they are trained with the description of triples that represent relations and entities connected by the [SEP] and [CLS] tokens and then take the description sequences as the input for finetuning. Normally, they use the presentation of [CLS] token to conduct a binary classification, note as:

$$p(\mathbf{y}_{\tau}|x)=p(\texttt{[CLS]}|x;\mathcal{W})$$

(5)

where $\mathbf{y}_{\tau}\in[0,1]$ is the label of whether the triple is positive.

PT-KGE methods like PKGC (Lv et al. 2022), utilize natural language prompts to elicit knowledge from pre-trained models for KG embeddings. Prompt tuning uses PLMs as direct predictors to complete a cloze task, linking pre-training and finetuning. In this paper, we implement PT-KGE with entity vocabulary expansion, treating each unique entity as common knowledge embeddings. Specifically, we consider entities $e\in\mathcal{E}$ as special tokens in the language model, turning link prediction into a masked entity prediction. Formally, we can obtain the correct entity by ranking the probability of each entity in the knowledge graph as follows:

$$p(\mathbf{y}|x)=p(\texttt{[MASK]}=y|x;\mathcal{W})$$

(6)

where $p(\mathbf{y}|x)$ is denoted as the probability distribution of the entity in the KG.

With the pre-trained KG embeddings via the sections mentioned above, we introduce several baselines to edit KG embeddings. Concretely, we divide baseline models into two categories. The first is the external model-based editor, which uses an extrinsic model editor to control the model parameters such as KE and MEND. Another is the additional parameter-based editor, which introduces additional parameters to adjust the final output of the model to achieve the model editing. There are also some non-editing methods in the baseline: Finetune and K-Adapter (Wang et al. 2021b).

Note that external model-based edit methods are flexible but have to optimize many editor parameters. In contrast, the additional parameter-based method only tunes a small number of additional parameters (0.9M) but encounters poor empirical performance due to the challenge of manipulating knowledge in PLMs. Intuitively, we capitalize on the advantages of both approaches and propose a simple yet robust baseline KGEditor, as shown in Figure 3, which employs additional parameters through hypernetwork. We construct an additional layer with the same architecture of FFN and leverage its parameters for knowledge editing. Moreover, we leverage an additional hypernetwork to generate the additional layer. With KGEditor, we can optimize fewer parameters while keeping the performance of editing KG embeddings. Specifically, when computing the output of FFN, we add the output of the additional FFNs to the original FFN output as an adjustment term for editing:

$$FFN^{\prime}(H)=FFN(H;\mathcal{W})+\bigtriangleup FFN(H;\mathcal{\tilde{W}}+\bigtriangleup\mathcal{\tilde{W}})$$

(7)

where $\mathcal{W}$ is denoted as the origin FFNs’ weight, and $\mathcal{\tilde{W}}$ is the weight of additional FFNs. $\bigtriangleup\mathcal{\tilde{W}}$ is the parameters’ shift generated by the hypernetwork. Inspired by KE (Cao, Aziz, and Titov 2021), we build the hypernetworks with a bidirectional LSTM. We encode $\langle x,y,a\rangle$ and then concatenate them with a special separator to feed the bidirectional LSTM. Then, we feed the last hidden states of bidirectional LSTM into the FNN to generate a single vector $h$ for knowledge editing. To predict the shift for the weight matrix $\mathcal{\tilde{W}}^{n\times m}$, we leverage FNN conditioned on $h$ which predict vectors $\alpha,\beta\in\mathbb{R}^{m},\gamma,\delta\in\mathbb{R}^{n}$ and a scalar $\eta\in\mathbb{R}$. Formally, we have the following:

	$\displaystyle\bigtriangleup\mathcal{\tilde{W}}=\sigma(\eta)\cdot\left(\hat{\alpha}\odot\nabla_{\mathcal{W}}\mathcal{L}(\mathcal{W};x,a)+\hat{\beta}\right)\;,$		(8)
	$\displaystyle\text{with}\quad\hat{\alpha}=\hat{\sigma}(\alpha)\;\!\gamma^{\top}\quad\text{and}\quad\hat{\beta}=\hat{\sigma}(\beta)\;\!\delta^{\top}\;,$

where $\hat{\sigma}$ refers to the Softmax function (i.e., $x\mapsto\exp(x)/\sum_{i}\exp(x_{i})$) and $\sigma$ refers to the Sigmoid function (i.e., $x\mapsto(1+\exp(-x))^{-1}$). $\nabla_{\mathcal{W}}\mathcal{L}(\mathcal{W};x,a)$ presents the gradient, which contains rich information regarding how to access the knowledge in $\mathcal{W}$. Note that the hypernetwork $\phi$ parameters can linearly scale with the size of $\mathcal{\tilde{W}}$, making the KGEditor efficient for editing KG embeddings in terms of computational resources and time.

We employ the constructed datasets, E-FB15k237, A-FB15k237, E-WN18RR, and A-WN18RR for evaluations. Initially, we utilize the pre-train datasets (the identical training set for the ADD task and the corrupted training set for the EDIT task) to initialize the KG embeddings with BERT. We then train the editor using the training dataset and evaluate using the testing dataset. Regarding the ADD task, the evaluation is conducted based on the performance of the training set, as they are not present in the pre-training dataset. We adopt PT-KGE as the default setting for the main experiments. Additionally, We supply the datasets, pre-trained KG embeddings, and a leaderboard for reference.

We compare KGEditor with several baselines, including MEND, KE, CALINET, as well as some variants: KGE_FT, which involves directly fine-tuning all parameters; KGE_ZSL, which entails directly inferring triples without adjusting any parameters; K-Adapter, which employs an adaptor to fine-tune a subset of parameters. Our experiments aim to 1) evaluate the efficacy of various approaches for editing KG embeddings (Table 2,3), 2) investigate the performance when varying the number ($n$) of target triples (Figure 5), and 3) examine the impact of distinct KGE initialization approaches on performance(Figure 4).