Knowledge Graph Completion with Counterfactual Augmentation

TL;DR

This paper introduces a counterfactual augmentation approach for Knowledge Graph Completion using GNNs, which improves performance and interpretability by generating and incorporating hypothetical relations based on causal modeling.

Contribution

It proposes a novel counterfactual relation generation method grounded in causal inference to augment GNN training for knowledge graphs, achieving state-of-the-art results.

Findings

Outperforms existing KGC methods on benchmark datasets

Enhances interpretability of GNN-based models

Demonstrates the effectiveness of counterfactual augmentation

Abstract

Graph Neural Networks (GNNs) have demonstrated great success in Knowledge Graph Completion (KGC) by modeling how entities and relations interact in recent years. However, most of them are designed to learn from the observed graph structure, which appears to have imbalanced relation distribution during the training stage. Motivated by the causal relationship among the entities on a knowledge graph, we explore this defect through a counterfactual question: "would the relation still exist if the neighborhood of entities became different from observation?". With a carefully designed instantiation of a causal model on the knowledge graph, we generate the counterfactual relations to answer the question by regarding the representations of entity pair given relation as context, structural information of relation-aware neighborhood as treatment, and validity of the composed triplet as the outcome. Furthermore, we incorporate the created counterfactual relations with the GNN-based framework on KGs to augment their learning of entity pair representations from both the observed and counterfactual relations. Experiments on benchmarks show that our proposed method outperforms existing methods on the task of KGC, achieving new state-of-the-art results. Moreover, we demonstrate that the proposed counterfactual relations-based augmentation also enhances the interpretability of the GNN-based framework through the path interpretations of predictions.

Tables

Table 1. Table 1. KGC results. Results of NeuraLP and DRUM are taken from (Sadeghian et al . , 2019 ) . Results of RotatE, HAKE, and LowFER are taken from their original papers  (Sun et al . , 2019 ; Zhang et al . , 2020a ; Amin et al . , 2020 ) . Results of all GNNs-based models are reproduced with DGL  (Wang et al . , 2019 ) . Results of the other embedding methods are taken from (Sun et al . , 2019 ; Zhu et al . , 2021b ) .

Class	Method	FB15k-237					WN18RR
		MRR	MR	H@1	H@3	H@10	MRR	MR	H@1	H@3	H@10
Path-based	Path Ranking (Lao and Cohen, 2010)	0.174	3521	0.119	0.186	0.285	0.324	22438	0.276	0.360	0.406
	NeuralLP (Yang et al., 2017)	0.240	-	-	-	0.362	0.435	-	0.371	0.434	0.566
	DRUM (Sadeghian et al., 2019)	0.343	-	0.255	0.378	0.516	0.486	-	0.425	0.513	0.586
Embeddings	TransE (Bordes et al., 2013)	0.294	357	-	-	0.465	0.226	3384	-	-	0.501
	DistMult (Yang et al., 2015)	0.241	254	0.155	0.263	0.419	0.429	5110	0.391	0.442	0.490
	ComplEx (Trouillon et al., 2016)	0.247	339	0.158	0.275	0.428	0.438	5261	0.412	0.464	0.512
	RotatE (Sun et al., 2019)	0.338	177	0.241	0.375	0.553	0.476	3340	0.428	0.492	0.571
	HAKE (Zhang et al., 2020a)	0.346	-	0.250	0.381	0.542	0.497	-	0.452	0.516	0.582
	LowFER (Amin et al., 2020)	0.359	-	0.266	0.396	0.544	0.465	-	0.434	0.479	0.526
	TuckER (Balažević et al., 2019)	0.358	-	0.266	0.394	0.544	0.470	-	0.433	0.482	0.526
	ConvE (Dettmers et al., 2018)	0.319	276	0.232	0.351	0.492	0.462	4888	0.431	0.476	0.525
GNNs	RGCN (Schlichtkrull et al., 2018) + TransE	0.324	220	0.238	0.357	0.497	0.198	3272	0.066	0.263	0.478
	RGCN (Schlichtkrull et al., 2018) + DistMult	0.332	242	0.242	0.366	0.508	0.393	6980	0.350	0.413	0.481
	RGCN (Schlichtkrull et al., 2018) + ConvE	0.336	197	0.249	0.368	0.511	0.428	2875	0.370	0.460	0.527
	CompGCN (Vashishth et al., 2019) + TransE	0.335	205	0.247	0.369	0.511	0.206	3182	0.064	0.281	0.502
	CompGCN (Vashishth et al., 2019) + DistMult	0.342	200	0.252	0.372	0.520	0.430	4559	0.395	0.439	0.513
	CompGCN (Vashishth et al., 2019) + ConvE	0.351	245	0.254	0.386	0.535	0.469	3065	0.433	0.480	0.543
	LTE (Zhang et al., 2022) + TransE	0.334	182	0.241	0.370	0.519	0.211	3290	0.022	0.362	0.521
	LTE (Zhang et al., 2022) + DistMult	0.335	238	0.246	0.367	0.517	0.437	4485	0.403	0.447	0.517
	LTE (Zhang et al., 2022) + ConvE	0.352	249	0.262	0.385	0.533	0.472	3434	0.436	0.485	0.544
	NBFNet (Zhu et al., 2021b)	0.415	114	0.321	0.454	0.599	0.551	636	0.497	0.573	0.666
	KGCF (Ours)	0.421	124	0.326	0.461	0.603	0.553	628	0.503	0.577	0.670

Table 2. Table 2. Examples of path interpretations of predictions on FB15k-237. For a valid triplet as a query, we visualize the S ubstitution o f its C ounterfactual R elation, which is denoted as SoCR, as well as the top-1 path interpretations and their weights. Inverse relations are denoted with a superscript -1 .

Query	$⟨$$h_i,r_j,t_k$$⟩$: $⟨$Barzil, olympics, 1972 summer olympics$⟩$, $({𝒯}_{h_i,r_j,t_k}^F=1$, ${𝒯}_{h_i,r_j,t_k}^{CF}=0)$
SoCR	$⟨$$h_a,r_j,t_b$$⟩$: $⟨$Argentinia, olympics, Bids for the 1976 summer olympics$⟩$ $({𝒯}_{h_a,r_j,t_b}^F=0$, ${𝒯}_{h_a,r_j,t_b}^{CF}=0)$
0.350	$⟨$Barzil, medal, Gold Medal$⟩$ $({𝒯}_{h_i,r_j,t_k}^F=1$, ${𝒯}_{h_i,r_j,t_k}^{CF}=0)$
	$\wedge$ $⟨$Gold Medal, medal-1, 1972 summer olympics$⟩$ $({𝒯}_{h_i,r_j,t_k}^F=1$, ${𝒯}_{h_i,r_j,t_k}^{CF}=0)$
Query	$⟨$$h_i,r_j,t_k$$⟩$: $⟨$Bridesmaids, film release region, Republica Portuguesa$⟩$, $({𝒯}_{h_i,r_j,t_k}^F=0$, ${𝒯}_{h_i,r_j,t_k}^{CF}=1)$
SoCR	$⟨$$h_a,r_j,t_b$$⟩$: $⟨$The Hangover, film release region, Norwegen$⟩$ $({𝒯}_{h_a,r_j,t_b}^F=1$, ${𝒯}_{h_a,r_j,t_b}^{CF}=0)$
0.445	$⟨$Bridesmaids, film release region, Espagna$⟩$ $({𝒯}_{h_i,r_j,t_k}^F=0$, ${𝒯}_{h_i,r_j,t_k}^{CF}=1)$
	$\wedge$ $⟨$Espagna, Adjoins-1, Republica Portuguesa$⟩$ $({𝒯}_{h_i,r_j,t_k}^F=1$, ${𝒯}_{h_i,r_j,t_k}^{CF}=0)$

Table 3. Table 3. Significance test of KGCF over NBFNet on WN18RR.

Metric	MRR	H@1	H@3	H@10
p-value	P<0.005	P<0.001	P<0.001	P<0.005

Table 4. Table 4. Inference time comparison on FB15k-237 test set with a single GPU. We report the average time over 5 5 5 splits for each method here. s 𝑠 s denotes seconds and m 𝑚 m denotes minutes.

Methods	RGCN	CompGCN	LTE	NeuralLP	NBFNet	KGCF
Wall Time	13$s$	5$s$	11$s$	2.1$m$	4.0$m$	4.2$m$

Table 5. Table 5. Dataset statistics for KGC.

Dataset	#Entity	#Relation	#Triplet
			#Train	#Validation	#Test
FB15k-237 (Toutanova and Chen, 2015)	14,541	237	272,115	17,535	20,466
WN18RR (Dettmers et al., 2018)	40,943	11	86,835	3,034	3,134

Equations

$$\begin{gathered} \mathcal{G}=\{(e_i,r_j,e_k)\} \\ \subset (\mathcal{E}\times \mathcal{R}\times \mathcal{E}), \end{gathered}$$

$$\text{ITE}(\mathbf{z})=g(\mathbf{z},1)-g(\mathbf{z},0),\text{ATE}={\mathbb{E}}_{\mathbf{z}\sim \mathbf{Z}}\text{ITE}(\mathbf{z}).$$

$${\text{ITE}}_{(h,t)}=g({\mathbf{z}}_r(h,t),1)-g({\mathbf{z}}_r(h,t),0).$$

$$\displaystyle(h_{a},r_{j},t_{b})=\operatorname*{arg\,min}_{(h_{a},r_{j},t_{b})\in{\mathcal{E}}\times{\mathcal{E}}}\big{\{}$$

$$\displaystyle\text{ }{\mathcal{T}}^{F}_{(h_{a},r_{j},t_{b})}=1-{\mathcal{T}}^{F}_{(h_{i},r_{j},t_{k})}\big{\}},$$

$$\mathbf{M}\in {\mathbb{R}}^{\mathcal{E}\times d}=\sum _{r_j}^{\mathcal{R}}{\psi }_j{\mathbf{M}}_{r_j}.$$

$$\displaystyle(h_{a},r_{j},t_{b})=\operatorname*{arg\,min}_{(h_{a},r_{j},t_{b})\in{\mathcal{E}}\times{\mathcal{E}}}\big{\{}$$

$$\displaystyle\text{ }{\mathcal{T}}^{F}_{(h_{a},r_{j},t_{b})}=1-{\mathcal{T}}^{F}_{(h_{i},r_{j},t_{k})}\big{\}},$$

$${\mathcal{T}}^{CF}_{(h_{i},r_{j},t_{k})},{\mathcal{A}}^{CF}_{i,j,k}=\begin{cases}{\mathcal{T}}^{F}_{(h_{a},r_{j},t_{b})},{\mathcal{A}}^{F}_{a,j,b}&\text{, if }\exists$ $(h_{a},r_{j},t_{b})\in{\mathcal{E}}\times{\mathcal{E}}\\ &\quad\text{ satisfies \ref{eq:nearest2};}\\ {\mathcal{T}}^{F}_{(h_{i},r_{j},t_{k})},{\mathcal{A}}^{F}_{i,j,k}&\text{, otherwise}.\end{cases}$$

$$\text{vspace}-6mm{\mathbf{z}}_r^{(0)}(h,t)$$

$${\mathbf{z}}_r^{(l+1)}(h,t)$$

$$\displaystyle\qquad\qquad\qquad(h,r,e)\in\mathcal{N}_{in}(e)\Big{\}}\cup\Big{\{}{\bm{z}}^{(0)}_{r}(h,t)\Big{\}}\Big{)}\vspace{-6mm}$$

$${\hat{p}}^F(t|h,r)=g(\mathcal{Z},{\mathcal{T}}^F)=\text{MLP}([{\mathbf{z}}_r(h,t),{\mathcal{T}}_{(h,r,t)}^F]),$$

$${\hat{p}}^{CF}(t|h,r)=g(\mathcal{Z},{\mathcal{T}}^{CF})=\text{MLP}([{\mathbf{z}}_r(h,t),{\mathcal{T}}_{(h,r,t)}^{CF}]),$$

$${\mathcal{L}}_F=$$

$${\mathcal{L}}_{CF}=$$

$${\mathcal{L}}_{disc}=\text{disc}({\hat{P}}_f^F,{\hat{P}}_f^{CF}),\text{ where }\text{disc}(P,Q)=||P-Q|{|}_F,$$

$$\mathcal{L}={\mathcal{L}}_F+\alpha \cdot {\mathcal{L}}_{CF}+\beta \cdot {\mathcal{L}}_{disc},$$

$$\widehat{\text{ATE}}=$$

$$+(1_{\mathcal{E}\times \mathcal{E}}-{\mathcal{T}}^F)\odot ({\mathcal{A}}^{CF}-{\mathcal{A}}^F){\}}_{i,j}.$$

$${\mathcal{P}}_1,{\mathcal{P}}_2,...,{\mathcal{P}}_k=\underset{\mathcal{P}\in {\mathcal{P}}_{h,t}}{top-k}\frac{\partial {\hat{p}}^F(h,r,t)}{\partial {\hat{P}}^F},$$

$$\displaystyle\mathrm{MRR}=\frac{1}{2*|\mathcal{T}|}\sum_{(h,r,t)\in\mathcal{T}}\Big{(}\frac{1}{\mathrm{rank}_{h}(h,r,t)}+\frac{1}{\mathrm{rank}_{t}(h,r,t)}\Big{)},$$

$$\displaystyle\mathrm{MRR}=\frac{1}{2*|\mathcal{T}|}\sum_{(h,r,t)\in\mathcal{T}}\Big{(}\mathrm{rank}_{h}(h,r,t)+\mathrm{rank}_{t}(h,r,t)\Big{)},$$

Full text

causal inference, knowledge graph completion, graph augmentation, graph neural networks

††submissionid: 5698††copyright: acmcopyright††journalyear: 2023††copyright: acmlicensed††conference: Proceedings of the ACM Web Conference 2023; April 30-May 4, 2023; Austin, TX, USA††booktitle: Proceedings of the ACM Web Conference 2023 (WWW ’23), April 30-May 4, 2023, Austin, TX, USA††price: 15.00††doi: 10.1145/3543507.3583401††isbn: 978-1-4503-9416-1/23/04††ccs: Computing methodologies Reasoning about belief and knowledge

1. Introduction

Knowledge can be encoded in knowledge graphs (KGs), which store structured information of real-world entities as nodes and relations as edges. Real-world KGs are usually noisy and incomplete. Because of the nature of incompleteness of KGs, knowledge graph completion (KGC) is designed to predict missing relations (Bordes et al., 2013; Galárraga et al., 2015; Ho et al., 2018).

GNN-based KGC models have shown effects on the KGC task and attracted tremendous attention in recent years (Ji et al., 2021). Many popular GNNs follow an iterative message passing then aggregate scheme to adaptively learn the representations of nodes in graph (Zhang et al., 2020b; Chang et al., 2020, 2022; Guan et al., 2021). When learning on a homogeneous graph, the representation of the center node is updated by aggregating the message gathered from the neighbors of the node while preserving the local structure around the node during each iteration, While a homogeneous graph could be viewed as a special version of KGs, GNNs are not able to be directly employed on KGs and are usually modified to model the interactions between entities and multiple relations. However, as shown in Figure 2(b), the distribution of relation types is usually imbalanced on KGs. Most of the current GNN-based approaches fail to capture this information and could not generalize well. This failure calls for a better learning regime on KGC.

In this work, we propose to alleviate this failure with data augmentation on KGs, especially from the perspective of causal inference by answering the following counterfactual question:

Would the relation still exist if the neighborhood of entities became different from observation?

If a GNN-based KGC model learns this causal relationship via answering the counterfactual question above, such knowledge will help to discover causal effects on relation types and improve the accuracy of prediction as well as the generalization ability (Zhao et al., 2022). Note that the causality here falls into the scope of causal inference models, where we focus on taking advantage of existing basic yet effective causal models to enhance the task of KGC. This aligns with a line of research works (Johansson et al., 2016; Alaa and Van Der Schaar, 2019; Zhao et al., 2022), and is not related to the explicit semantic relations in KGs.

A counterfactual question is often associated with three factors: context (as a data point), treatment (as a type of manipulation), and outcome (Van der Laan and Petersen, 2007; Johansson et al., 2016; Zhao et al., 2022). Previous effort (Zhao et al., 2022) proposes to use whether two nodes live in the same neighborhood to define such treatment on a homogeneous graph. However, we argue that this definition has flaws in KGs since it ignores the specific relation connected entities. Taking Figure 1 as an example. Let the blue and orange ovals be two detected communities by treating the KG as a homogeneous graph. Hedwig and Ron Weasley obviously share the same neighborhood within the ovals if we ignore the relation type on edges. Then the treatment definition from (Zhao et al., 2022) would be confused and make no distinction between Ron Weasley, who is a person, and Hedwig, which is an owl. As a result, it is even harder to help predict the different relations "pet of" and "friend" that are associated with them and Harry Potter. Therefore, it is non-trivial to directly migrate the treatment definition from a homogeneous graph to KGs due to the multi-relation nature.

Based on the idea of considering relation types with causal learning on KGs, we propose an instantiation of causal models on KGs with a new definition of treatment for context. By answering the counterfactual questions, we discover the counterfactual relations for all observed contexts and calculate their counterfactual treatments accordingly. Then, we integrate the representation learning of a pair of entities given a query relation with both their factual and counterfactual treatments and propose a counterfactual relation augmented GNN-based framework (KGCF) for KGC. Specifically, since this counterfactual treatment is not available from the observed data (valid triplets on KGs), we match it with the nearest neighbor of an entity pair in hidden embedding space and use the neighbor’s factual treatment as a substitution. To further address the imbalanced distribution of relation types, we first embed the entities to an embedding space that is associated with each relation type, then find the nearest neighbors within this embedding space. After obtaining both the factual and counterfactual views of training data, we consociate them into an encoder-decoder style framework to learn better representations and predict missing relations.

The contributions of this paper are summarized as follows:

• •

Motivated by the need for data augmentation on KGs, we propose the first instantiation of the causal model for KGs via answering counterfactual questions and considering the relation types.

• •

We present KGCF that utilizes counterfactual relations to augment the representation learning on KGs with special consideration on the imbalanced relation distribution.

• •

Extensive experiments show the superiority of KGCF on KGC task. We further demonstrate that counterfactual relations also enhance the interpretation ability through path-based explanation of predictions.

2. Related Work

2.1. Knowledge Graph Completion (KGC)

Popular research KGC models mainly fall into three categories (Ji et al., 2021): embedding-based methods, relation path inference, and rule-based reasoning. While the last two research directions focus on exploring the multi-step relationships, we focus on the recently preliminary direction of KGC in this paper: embedding-based KGC methods with GNNs and cover some of the recent progress in this line of research due to a large amount of literature. Jung et al. (2020) propose a new GNN encoder, which can effectively capture the query-related information from the temporal KG and is interpretable in the reasoning process. Zhang et al. (2022) verify by experiments that the graph structure modeling in GCNs has no significant impact on the performance of the KGC model. Instead, they propose a simple and effective LTE-KGE framework by removing the GCN aggregation. Wang et al. (2022) propose a new model called CFAG, which uses a coarse-grained aggregator (CG-AGG) and a fine-grained generative adversarial net (FG-GAN) to resolve the emergence of new entities. Huang et al. (2022) propose a KG inference framework SS-AGA based on self-supervised graph alignment. Despite the progress of the above-mentioned models, most of them fail to consider the imbalanced distribution of relation types.

2.2. Causal Inference

Finding causal relationships between different variables from observational data is a long-standing scientific problem in the field of statistics and artificial intelligence (Yao et al., 2020). Causal inference for GNNs is an emerging field of research, which can be divided into three categories: counterfactual learning, causal graph-guided representation learning, and causal discovery. Lin et al. (2021) use causal inference to improve the explainability of GNNs. They transform the explainability problem of reasoning decisions in the GNNs into a causal learning task, and then a causal explanation model is trained based on the objective function of Granger causality. Ma et al. (2022) propose a new framework called causal inference under spillover effects, which solves the problem of individual treatment effect estimation with high-order interference on hypergraphs. Zhao et al. (2021) first propose to use causal inference to improve link prediction on homogeneous graphs, which is the most related one to this paper. They propose the concept of counterfactual link prediction (CFLP) and representations are learned from observed and counterfactual links which are used as augmented training data. In contrast, we focus on KGs in this paper, to which CFLP could not be trivially migrated as discussed in Sec. 1.

2.3. Data Augmentation on Knowledge Graph

Graph Data Augmentation (GraphDA) is to find a mapping function that enriches or changes the information in a given graph (Gu et al., 2020; Jin et al., 2021; Chang et al., 2021a, b), which is a relatively underexplored area for KGs. Chen et al. (2021) propose a rumor data augmentation method called graph embedding-based rumor data augmentation (GERDA) to solve the data imbalance of rumor detection, which uses KGs to simulate the rumor generation process from the perspective of knowledge. Tang et al. (2022b) propose positive-unlabeled learning with adversarial data augmentation (PUDA). PUDA alleviates the problem of false negatives through positive-unlabeled learning and positive sample sparsity through adversarial data augmentation. Different from them, we propose to use the causal model to generate counterfactual augmentations to facilitate representation learning on KGs.

3. Preliminaries

3.1. Knowledge Graphs (KGs)

A knowledge graph (KG) is denoted as a set of triplets:

[TABLE]

where ${\mathcal{E}}$ and ${\mathcal{R}}$ represent a set of entities (nodes) and relations (edges), respectively. $e_{i}$ and $r_{j}$ are the $i$-th entity and $j$-th relation, and the types of relations could be in great numbers within a KG. We also usually distinguish the entity pair $(e_{i},e_{k})$ with $(h_{i},t_{k})$, considering the direction of the two entities. Note that a homogeneous graph ${\mathcal{G}}=({\mathcal{V}},{\mathcal{E}})$ can be viewed as only one relation type for all edges versions of KGs. Throughout this paper, we use bold terms, ${\bm{W}}$ or ${\bm{e}}$, to denote matrix/vector representations for weights and entities, respectively. And we select italic terms, $w_{h}$ or $\alpha$, to denote scalars. We can also use a third-order binary tensor ${\mathcal{A}}\in\{0,1\}^{|{\mathcal{E}}|\times|{\mathcal{R}}|\times|{\mathcal{E}}|}$ to uniquely define a KG, which is also known as the adjacency tensor of $\mathcal{G}$. The $(i,j,k)$ entry ${\mathcal{A}}_{ijk}=1$ if $(h_{i},r_{j},t_{k})$ is valid or otherwise ${\mathcal{A}}_{ijk}=0$. The $j$-th frontal slice of ${\mathcal{A}}$ is the adjacency matrix of the $j$-th relation $r_{j}$, which we denote as ${\mathcal{A}}_{r_{j}}$ for specification. To ease the notation, we also use ${\mathcal{A}}_{(h,r,t)}$ to denote the entry of a given triplet $(h,r,t)$ in the adjacency tensor.

3.2. GNN-Based framework for KGC

GNN-Based framework dealing with KGC usually adopts an encoder-decoder style framework (Schlichtkrull et al., 2018), where GNNs perform as the encoder and embedding-based score functions (e.g., TransE, DistMult, RotatE and ConvE) perform as the decoder. A KG’s entities and relations are first represented as embeddings by the GNN encoder. After having the embeddings from GNN (encoder) for entities and relations, KGC models usually simulate how entities and relations interact with the help of a score function (decoder) $s:{\mathcal{E}}\times{\mathcal{R}}\times{\mathcal{E}}\rightarrow\mathbb{R}$. The decoder then predicts the entries in the adjacency tensors ${\mathcal{A}}$ utilizing the derived representations. The prediction could be viewed as a recovery of the original graph structures since adjacency tensors and graph structures are in bijection. In this way, the score function-based decoder can detect missing relations in the original KG while completing the KG by recovering the graph structure. It is worth mentioning that in some GNN-based frameworks for KGC, such as NBFNet, the aforementioned score functions are implicitly included in the procedure of message passing, and a simple feed-forward neural network is employed as decoder on the learned entity pair representations to predict the missing links given the head/tail entity.

4. Proposed Method

4.1. Incorporating GNN-based framework for KGC with Causal Model

Previous works on GNN-based framework for KGC (Ji et al., 2021) have shown that message passing and aggregating scheme is able to generate more structure-enriched representations. However, recalling the example in Figure 2(b), the distribution of relations appears to be quite imbalanced in KGs. This imbalanced distribution might lead the learning of GNNs to be trapped into the dominant relation types and fail to generalize to other relations during inference. Nevertheless, with the help of recent progress of data augmentation on graph learning (Zhu et al., 2021a; Li et al., 2022; Ding et al., 2022), we propose to enhance the generalization ability of current GNN-based learning frameworks on KGC from the perspective of a causal model.

Counterfactual causal inference seeks to determine the causal relationship between treatment and outcome by posing counterfactual questions. An illustration of the counterfactual question is “would the outcome be different if the treatment was different?” (Morgan and Winship, 2015). In order to address the counterfactual question, researchers would next develop a causal model and apply causal inference accordingly. For example, the left subfigure from Figure 3 is a typical triangle causal-directed acyclic graph (DAG) model. In this example, we denote the context (confounder) as ${\bm{Z}}$, the treatment as ${\bm{T}}$, and the outcome as ${\bm{Y}}$. Once three of them are decided, the counterfactual inference methods could determine the effect of treatment ${\bm{T}}$ on the outcome ${\bm{Y}}$ given context ${\bm{Z}}$. Two statistics, individual treatment effect (ITE) and its expectation averaged treatment effect (ATE) (Van der Laan and Petersen, 2007; Weiss et al., 2015; Zhao et al., 2022), are usually employed to measure this effect. Given a binary treatment variable $T=\{0,1\}$, we denote $g(\mathbf{z},T)$ as the outcome of $\mathbf{z}$, then

[TABLE]

Bigger ATE/LTE indicates a stronger causal relationship between treatment and outcome. Recent effort (Zhao et al., 2022) proposes to develop a causal model to enhance the representation learning on graphs. Through proper definitions of context, treatments, and their corresponding outcomes on graphs, promising performance on plain link prediction problems for homogeneous graphs is obtained.

On the contrary, in this work, we focus on a more comprehensive version of graphs, the KG, and focus on the KGC task. Given a pair of head entities and relation $(h,r,?)$, our model needs to predict the existence of a tail relation $t$. Hence, we aim to learn an effective representation ${\bm{z}}_{r}(h,t)$ of head-tail pair $(h,t)$ given a query relation $r$. As shown in the right subfigure from Figure 3, we define the context as the entity pair representation ${\bm{z}}_{r}(h,t)$, and the outcome in the adjacency tensor ${\mathcal{A}}_{(h,r,t)}$ is the validity of triplet $(h,r,t)$. Therefore, our causal DAG contains the following nodes:

• •

${\bm{z}}$ (Confounder; unobservable): latent entity pair representations;

• •

${\mathcal{T}}$ (Treatment; observable): KG structural information (e.g., whether two entities belong to the same neighborhood);

• •

${\mathcal{A}}$ (Outcome; observable): triplet existence.

And the edges in our DAG are: ${\bm{z}}$ -> ${\mathcal{T}}$, ${\bm{z}}$ -> ${\mathcal{A}}$ and ${\mathcal{T}}$ -> ${\mathcal{A}}$.

In this way, we instantiate the typical causal model on the KG by leveraging the estimated $\text{ITE}({\mathcal{A}}_{(h,r,t)}|{\mathcal{T}}_{(h,r,t)})$ to improve the learning of ${\bm{z}}_{r}(h,t)$. Here, the objective is different from classic causal inference, as well as the definition of the plain link prediction in (Zhao et al., 2022), considering the multiple types of relations. More specifically, for each pair of entities $(h,t)$ given a query relation $r$, the corresponding ITE can be estimated by

[TABLE]

In the next, we will incorporate the information of ITE from the counterfactual relations to augment the KG and improve the learning of ${\bm{z}}_{r}(h,t)$, accordingly. Both the entity pair representation ${\bm{z}}_{r}(h,t)$ and the outcome ${\mathcal{A}}_{(h,r,t)}$ need to consider the presence of multiple types of relations. Furthermore, the imbalanced distribution also makes the design of incorporation to the GNN-based embedding framework on KGC difficult.

As a preparation for the next step of incorporation, we denote the observed adjacency tensor by ${\mathcal{A}}^{F}$ as factual outcomes. In other words, ${\mathcal{A}}^{F}$ is constructed with all triplets from the training set. We also denote the unobserved tensor of the counterfactual relations by ${\mathcal{A}}^{CF}$ when the treatment is different from the counterfactual outcomes. For the definition of treatment, we denote ${\mathcal{T}}^{F}\in\{0,1\}^{|{\mathcal{E}}|\times|{\mathcal{R}}|\times|{\mathcal{E}}|}$ as the binary factual treatment tensor, where ${\mathcal{T}}^{F}_{(h,r,t)}$ indicates the treatment of the entity pair $(h,t)$ given a query relation $r$. Then we denote ${\mathcal{T}}^{CF}$ as the counterfactual treatment matrix where ${\mathcal{T}}^{CF}_{(h,r,t)}=1-{\mathcal{T}}^{F}_{(h,r,t)}$. We are interested in two aspects: 1) Estimating the counterfactual outcomes ${\mathcal{A}}^{CF}$ 2) learning from both factual (observed data) outcomes ${\mathcal{A}}^{F}$ and counterfactual (augmented data) outcomes ${\mathcal{A}}^{CF}$ to enhance the ability of GNN-based framework on KGC.

4.2. Treatment Variable on KG

In this section, we describe how we design to acquire the factual treatment variable ${\mathcal{T}}^{F}$. As in many GNN-based embedding frameworks for KGC, we assume the entities that share similar neighborhoods would tend to share similar relations as well. Therefore, finding the individual treatment for a pair of entities from their neighborhood under a special consideration to relation type is intuitively effective for downstream tasks. Meanwhile, as pointed out by recent works on homogeneous graphs (Cai et al., 2021; Tang et al., 2022a; Zhao et al., 2022), models may fail to identify more global but also important factors purely from the association between local structural information and link presence. Therefore, in this work, we use both local and global structural roles of each entity pair as its treatment.

In KGs, since the relations are divided into multiple types, we need to define treatment based on the adjacency tensor ${\mathcal{A}}_{r}$ w.r.t the query relation $r$, rather than simply dealing with a collapsed adjacency matrix as the regular practice in GNN-based approaches (Schlichtkrull et al., 2018; Shang et al., 2019; Vashishth et al., 2019). Note that though in this work, we define the treatment from the structural information w.r.t the view of each relation $r_{j}$ on KGs for illustration, the causal models shown in 3 do not limit the treatment to be structural roles on the graph. The treatment variable ${\mathcal{T}}^{F}_{(h,r,t)}$ can be defined from any property of triplet $(h,r,t)$.

Then for each adjacency tensor in ${\mathcal{A}}_{r_{j}}$, we define the binary treatment variable according to whether the two entities in a pair belong to the same community/neighborhood from a global perspective and calculate them separately. In this way, the relation types are explicitly considered in our treatment definition. We denote $c:\mathcal{V}\rightarrow\mathbb{N}$ as any graph clustering or community detection approach that assigns each entity with an index of the community/cluster/neighborhood that the entity belongs to. The treatment tensor ${\mathcal{T}}$ is defined as ${\mathcal{T}}^{F}_{(h,r,t)}=1$ if $c(h)=c(t)$ on adjacency tensor ${\mathcal{A}}_{r}$, and ${\mathcal{T}}^{F}_{(h,r,t)}=0$ otherwise. For the choice of $c$, without the loss of generality, an unsupervised approach K-core (Malliaros et al., 2020) that is widely used for clustering graph structure is utilized in our work as an example. K-core is strongly related to the concept of graph degeneracy, which has a long history in graph theory. Given a graph ${\mathcal{G}}$ and an integer $k$, a K-core of a graph ${\mathcal{G}}$ is a maximal connected subgraph of ${\mathcal{G}}$ in which all vertices have degrees at least $k$. Equivalently, it is the connected components that are left after all vertices of degree less than $k$ have been removed.

The reason that we choose the K-core algorithm falls into three aspects: 1) The core decomposition-based approach can be used to quantify node importance in many different domains efficiently and effectively. Precisely, the K-core algorithm for a graph ${\mathcal{G}}$ can be computed in linear time w.r.t the number of edges ${\mathcal{E}}$ of ${\mathcal{G}}$, which is of great importance considering the massive entities in a KG. 2) By assigning to each graph node an integer number $k$ (the core number), K-core could capture how well the given node is connected w.r.t its neighbors, which could also reflect the first-order proximity on graph (Tang et al., 2015), which is a desired structural property for neighborhood. 3) As one of the generally used clustering methods that consider the global graph structural information, the role of each entity pair from a global perspective could also be reserved to complement the local structure information.

4.3. Counterfactual Relations

Once we have defined the treatment from the observed space, we can discover the counterfactual treatment ${\mathcal{T}}^{CF}$ and outcome ${\mathcal{A}}^{CF}$ for each pair of entities given a query relation. However, this counterfactual relationship is not available from the observed data. As a substitution, we would like to find ${\mathcal{T}}^{CF}$ and ${\mathcal{A}}^{CF}$ from the nearest observed context as substitution  (Zhao et al., 2022). It is a common practice to estimate treatment effects from observational data using this form of matching on covariates (Johansson et al., 2016; Alaa and Van Der Schaar, 2019; Zhao et al., 2022). In this way, for a pair of entities $(h,t)$ given a relation $r$ as a candidate, we can match their counterfactual treatment ${\mathcal{T}}^{CF}_{(h,r,t)}$ to their nearest neighbor with the opposite treatment. Accordingly, the outcome of their nearest neighbor is treated as counterfactual relation.

Formally, $\forall(h_{i},r_{j},t_{k})\in{\mathcal{S}}$, where ${\mathcal{S}}$ is a selected set of entity pair candidates, its counterfactual relation $(h_{a},r_{j},t_{b})$ with the same type of relationship is

[TABLE]

where $d(\cdot,\cdot)$ is a distance measuring metric between a pair of entity pairs (a pair of contexts) in the factual adjacency tensor ${\mathcal{A}}^{F}_{r_{j}}$ given a query relation $r_{j}$. Here we omit the index in relation $r$ since the counterfactual relation should have the same relation type as the factual one. Nevertheless, there are still two challenges remaining for the above practice:

• •

It needs $O(N^{4})$ comparisons for finding the nearest neighbors if we choose the commonly used all $O(N^{2})$ entity pairs as ${\mathcal{S}}={\mathcal{E}}\times{\mathcal{E}}$, which is inefficient and infeasible in the application and hard to scale to large graphs. Considering the validity does not hold for all entity pairs, we design to set ${\mathcal{S}}$ as all entity pairs that have appeared in the dataset, but we eliminate the types of relations as well as the validity of the formed triplets, which prevents the risk of data leakage. Thus, the number of comparisons is greatly reduced by this selection of ${\mathcal{S}}$.

• •

Directly calculating the closest distance of entity pairs from ${\mathcal{A}}^{F}$ is still trapped in the observed space and could not bring extra counterfactual information. Therefore, we propose to find the nearest observed context of entity pairs from hidden low-dimensional embedding space. We take the unsupervised embedding method node2vec (Grover and Leskovec, 2016) to learn the embedding space for entities since it could effectively preserve the proximity in the graph. This aligns with our second reason for choosing K-core for defining factual treatment in the observed space.

After the aforementioned two challenges are solved, we propose to use an entity embedding matrix, which is denoted as ${\bm{M}}$, to find the counterfactual substitutes of triplets from candidates. Especially, taking the imbalanced distribution into account, we design to learn ${\bm{M}}$ w.r.t the proportion $\psi_{j}$ of a given relation $r_{j}$ in the distribution of all relation types. That is for every ${\mathcal{A}}^{F}_{r_{j}}$, we employ node2vec to get the embeddings accordingly as ${\bm{M}}_{r_{j}}$, and the overall embedding matrix is defined as:

[TABLE]

where $d$ is the dimension of embedding. Thus, $\forall(h_{i},r_{j},t_{k})\in{\mathcal{S}}$, we define its substitution for counterfactual relation $(h_{a},r_{j},t_{b})$ as

[TABLE]

where we specify $d(\cdot,\cdot)$ as the Euclidean distance on the hidden embedding space of ${\bm{M}}$, and ${\bm{m}}_{i}$ indicates the embedding of entity $h_{i}$ in the weighted node2vec space and so on. It is worth mentioning that in comparison with the approximated solution in  (Zhao et al., 2022), our introduced pair candidates could not only enhance the validity of potential triplets but also cast off the need for another hyperparameter for constraining the maximum distance.

To guarantee that all neighbors are sufficiently comparable (as substitutes) in the hidden embedding space, we do not assign counterfactual treatments for the given entity pair if there does not exist any entity pair satisfying Eq.(4.3), Thus, $\forall(h_{i},r_{j},t_{k})\in{\mathcal{S}}$ the counterfactual treatment matrix ${\mathcal{T}}^{CF}$ and the counterfactual adjacency tensor ${\mathcal{A}}^{CF}_{r_{j}}$ given a query relation $r_{j}$ are substituted as

[TABLE]

Note that the embeddings ${\bm{M}}$ from hidden space is only used once for finding the substitutes for counterfactual treatment and outcome. Both ${\bm{M}}$ and the nearest neighbors are also computed only once and remain unchanged during the learning process.

After the instantiation of the causal model on KGC, we can design our model on how to learn the embeddings for the KG from the counterfactual distributions, which we denote as KGCF. We denote $P^{F}$ as the factual distribution of the observed contexts and treatments, and $P^{CF}$ be the counterfactual distribution that is composed of the observed contexts and opposite treatments. For all possible triplets $(h_{i},r_{j},t_{k})$ and their treatments, the empirical factual distribution $\hat{P}^{F}\sim P^{F}$ is defined as $\hat{P}^{F}=\{((h_{i},r_{j},t_{k}),{\mathcal{T}}^{F}_{(h_{i},r_{j},t_{k})})\}^{\mathcal{E}}_{i,k=1}$, and define the empirical counterfactual distribution $\hat{P}^{CF}\sim P^{CF}$ as $\hat{P}^{CF}=\{((h_{i},r_{j},t_{k}),{\mathcal{T}}^{CF}_{(h_{i},r_{j},t_{k})})\}^{\mathcal{E}}_{i,k=1}$. While existing GNN-based KGC models only use the information from $\hat{P}^{F}$ as input, we benifit from the counterfactual distribution $\hat{P}^{CF}$ by augment the training data with the counterfactual relations.

4.4. Counterfactual Relations Augmented Learning Framework

After all the data preparation, as with most of the GNN-based embedding model, our model KGCF consists of two trainable components: a KG encoder $f$ and a relation decoder $g$.

Encoder. While we aim to learn a pairwise representation ${\bm{z}}_{r}(h,t)$, the encoder $f$ should capture the important neighborhood information between $h$ and $t$ w.r.t the query relation $r$, such as the proximity that is preserved by counting different types of random walks similar to the one we use during finding counterfactual relations. This kind of encoded local structure is often captured by path formulation, which is widely used in KG representation learning (Lao and Cohen, 2010; Gardner and Mitchell, 2015; Li et al., 2019; Zhu et al., 2021b) Without the loss of generality, we adopt the most advanced variant NBFNet (Zhu et al., 2021b) as the encoder $f$ in KGCF.

NBFNet is a path formulation-based graph representation learning framework. In NBFNet, the pair representations of nodes are interpreted as a generalized sum and each path representation is defined as the generalized product of the edge representations in the path. Then NBFNet solves the path formulation efficiently with learned operators with the help of the generalized Bellman-Ford algorithm, then reads out the pair representation as a source-specific message-passing process. In NBFNet, each layer of GNN is defined as

[TABLE]

where $l$ is the layer index. ${\bm{z}}^{(l)}_{r}(h,t)$ is the pair representation from $l$ layer for entity pair $(h,t)$ given a query relation $r$. The set $\mathcal{N}_{in}(e)$ contains all triplets $(h,r,e)$ such that $(h,r,e)$ is valid for a preset entity $e$. ${\bm{w}}(h,r,e)$ is the edge representation for the triplet $(h,r,e)$. The IND is the indicator function outputs $\text{\scriptsize{1}⃝}_{r}$ if $h=t$ and $\text{\scriptsize{0}⃝}_{r}$ otherwise, which aims to provide a non-trivial representation for the heat entity $h$ as the boundary condition. The AGG function is instantiated as natural sum, max or min in traditional methods followed by a linear transformation and a non-linear activation, or principal neighborhood aggregation (PNA) in recent efforts (Corso et al., 2020). For the MSG function, the NBFNet incorporates the traditional score functions (e.g., TransE, DistMult, or RotatE) as transformations correspond to the relational operators, in order to model the message passing procedure among entities and relations.

Decoder. For the relation decoder that predicts whether a triplet is valid between a pair of entities with a query relation $r$, we remain consistent with  (Zhao et al., 2022) and adopt a simple decoder based on multi-layer perceptron (MLP) to predict the conditional likelihood of the tail entity $t$. With the pair representation ${\bm{z}}_{r}(h,t)$ from the output of the last layer in encoder $f$ and the treatments of the triplet, the decoder $g$ is defined as

[TABLE]

where $[\cdot,\cdot]$ stands for the concatenation of vectors. The outputs are the conditional likelihood of the tail entity $t$: $\hat{p}^{F}(t|h,r)$ and $\hat{p}^{CF}(t|h,r)$ for both factual and counterfactual empirical distribution. The conditional likelihood of the head entity $h$ can be predicted by $\hat{p}^{F}(h|t,r^{-1})$ and $\hat{p}^{CF}(h|t,r^{-1})$ with the same model. Note that only the factual output $\hat{p}^{F}(t|h,r)$ will be used in inference.

Loss design. During the training process, data samples from the empirical factual/counterfactual distribution $\hat{P}^{F}$/$\hat{P}^{CF}$ are fed into decoder $g$ and optimized towards $p^{F}(t|h,r)$ and $p^{CF}(t|h,r)$, respectively. We follow the previous works (Bordes et al., 2013; Sun et al., 2019) to minimize the negative log-likelihood of positive and negative triplets from factual and counterfactual views. The negative samples are generated according to Partial Completeness Assumption (PCA) (Galárraga et al., 2013) as in NBFNet, which creates a negative sample by corrupting one of the entities in a positive triplet. The corresponding losses are defined as:

[TABLE]

where $n$ is the number of negative samples for a single positive sample and $(h_{i}^{\prime},r,t_{i}^{\prime})$ is the $i$-th negative sample. The negative samples remain the same for factual and counterfactual treatments.

Since the test data contains only observed (factual) samples, we force the distributions of representations of factual distributions and counterfactual distributions to be similar, which could protect the model from exposing to the risk of covariant shift (Johansson et al., 2016; Assaad et al., 2021). Similar to  (Zhao et al., 2022), we minimize the discrepancy distance (Mansour et al., 2009) between $\hat{P}^{F}$ and $\hat{P}^{CF}$ to regularize the representation learning:

[TABLE]

where $||\cdot||_{F}$ denotes the Frobenius Norm. $\hat{P}^{F}_{f}$ and $\hat{P}^{CF}_{f}$ are the representations of entity pairs given a query relation learned by graph encoder $f$ from factual distribution and counterfactual distribution, respectively. Therefore, the overall training loss of our proposed KGCF is

[TABLE]

where $\alpha$ and $\beta$ are hyperparameters to control the weights of counterfactual outcome estimation loss and discrepancy loss. Recall from Sec. 4.1 that our estimated $\widehat{\text{ATE}}$ is formulated as:

[TABLE]

We can find from loss Eq.(11) that $\widehat{\text{ATE}}$ is implicitly optimized during learning from the counterfactual relations. As a result, the learned causal relations help to improve the learning of pair representation.

Computational complexity. We first analyze the computational complexity of finding counterfactual relations. The complexity of the node2vec for hidden embedding space is bounded by $\mathcal{O}({\mathcal{E}}{\mathcal{R}})$. The step of finding counterfactual links with nearest neighbors in Eq.(4.3) has complexity in proportion to the size of the candidate set $|S|$ and entity pairs. Since the computation in 4.3 can be parallelized, the time complexity is $O(|S|{\mathcal{E}}^{2}/C)$ where $C$ is the number of processes. Then for the complexity of the training counterfactual learning model, the NBFNet encoder $f$ has a time complexity of $\mathcal{O}({\mathcal{E}}D_{f}+{\mathcal{H}}D_{f}^{2})$, where ${\mathcal{H}}$ is the number of training triplets and $D_{f}$ is the dimension of entity representations. While we sample the same number of non-existing triplets as that of observed triplets during training, the complexity of a two-layer MLP decoder $g$ is $O(((D_{f}+1)\cdot D_{g}+1\cdot D_{g}){\mathcal{H}})=O((D_{f}+2)D_{g}{\mathcal{H}})$, where $D_{g}$ is the number of neurons in the MLP hidden layer.

5. Experiments

5.1. Experimental Setup

We evaluate KGCF on the task KGC on the most popular datasets: FB15k-237 (Toutanova and Chen, 2015) and WN18RR (Dettmers et al., 2018). Both datasets are publicly accessible from DGL111https://github.com/awslabs/dgl-ke/tree/master/examples (Wang et al., 2019). We use the standard splits of both datasets. Statistics of datasets can be found in the Appendix.

Baselines. We compare KGCF against 21 baselines including three types: path-based methods, embedding methods, and GNNs with different scoring functions. Note that LTE (Zhang et al., 2022) is not exactly a GNN-based model but we include it here since it shares the same design framework with other GNN-based models.

Implementation Details.. Our implementation is similar with (Zhu et al., 2021b). For the encoder, we closely follow the official implementation for NBFNet and keep all architecture details as reported to maintain a fair comparison. In NBFNet, we set GNN layers as $6$ and hidden states as $32$. The negative ratio (#negative/#positive) is set to $32$. For other details please kindly refer to (Zhu et al., 2021b). For the decoder, we set MLP as two layers with $64$ hidden units and use ReLU as the activation function. The integer $k$ in K-core is set to $2$. The $\text{disc}(\cdot)$ function is selected as KL divergence. We use a public implementation for node2vec and the output dimension for node2vec is set to $32$. We decide the hyperparameters in loss through grid search within: $\alpha\in\{0.001,0.01,0.1,1\}$ and $\beta\in\{0.001,0.01,0.1,1\}$. The batch size is set to $32$. We use the Adam optimizer with a learning rate $5e-3$ for optimization. Our model is trained on a single Tesla V100 GPU for 20 epochs and reports the results from the model selected based on the best performance on the validation set.

5.2. KGC Results

Table 1 shows the performance of KGCF on the task of KGC in comparison with state-of-the-art baselines. We report the results regarding the filtered ranking protocol (Bordes et al., 2013). We have the following observations: 1) Our KGCF achieves the best performance across all measurements on both datasets, which aligns with our motivation that balancing the distribution of relation types with augmentation could contribute to better pair representation learning. 2) The path-based GNN approaches generally outperform other GNN baselines, which implies that the local subgraph structure that maintains the proximity between an entity pair is more essential to the massage-passing mechanism in the GNN-based KGC framework. This is consistent with the observation from LTE (Zhang et al., 2022) that the contribution of current convolutional GNN-based designs to KGC is somehow unnecessary and needs further exploration. 3) We further use a two-sided t-test to evaluate the significance of KGCF over NBFNet. Due to limited space here, we only show the result of the p-value on WN18RR here:

We can find that KGCF achieves strongly significant improvements (P < 0.005) over NBFNet across all metrics, which statistically confirms that KGCF’s scores in Table 1 are significantly better than NBFNet’s.

5.3. Ablation Studies

Visualization of prediction. With the help of the interpretation tool from NBFNet, we could interpret the KGC prediction $\hat{p}^{F}(h,r,t)$ of KGCF through paths from $h$ to $t$ that contribute most to the results. The importance of a path interpretation for $\hat{p}^{F}(h,r,t)$ is defined as its weight from a linear model, where the weight is approximated by the sum of the importance of edges in that path. In other words, the $\operatorname*{top-k}$ path interpretations are equivalent to the $\operatorname*{top-k}$ longest paths on the edge importance graph. The is computed by the partial derivative of the prediction w.r.t the path:

[TABLE]

where ${\mathcal{P}}_{1},{\mathcal{P}}_{2},...,{\mathcal{P}}_{k}$ is the $\operatorname*{top-k}$ paths acquired from ranking weights.

Table 2 shows the interpretation of two triplets examples with the path as well as their substitutions that are used for calculating counterfactual treatment SoCR. Note that SoCR does not have to be a valid triplet. Either way, its validity could augment the learning with a counterfactual view of information. This augmentation helps avoid inaccurate predictions through directly analogical reasoning, which is the common practice of previous efforts, from only observed factual neighborhoods.

Our observations from Table 2 are: 1) The reasoning over two entities is conducted over the paths among them with both information from factual and counterfactual treatments. Our method provides two additional views of the learned factual and counterfactual relations, which helps enhance the confidence level of reasoning by answering the counterfactual question. 2) The nearest observed context that is used as a substitution for counterfactual estimation is effective on the KG. We can find that similar entities (e.g., the films Bridesmaids and The Hangover) are close to each other in the proximity-preserving embedding space and with analogous context. 3) The interpretation answers our question from Section 1 that the relationship could not exist when the neighborhood information varies. In the first example, Argentina did not bid for the 1976 summer Olympics thus the relation is not valid in counterfactual treatment. This augmented counterfactual view of information helps the learning avoid false prediction because analogical reasoning from only observed factual neighborhoods might infer it as a valid one otherwise.

Running time comparison. In complementary of the complexity analysis from Section 4.4, we report the wall time on the FB15k-237 dataset for inference efficiency comparison over a single Tesla V100 GPU. It is obvious that the KGCF is very comparable to the NBFNet encoder, which indicates that our augmented framework brings no burden to the existing GNN-based method. Meanwhile, we can find that the convolutional GNN-based methods are the most efficient since they do not consider a path-based strategy, which results in worse KGC performance conversely.

Ablation study on $\mathcal{L}_{CF}$ and $\mathcal{L}_{disc}$. To investigate how $\mathcal{L}_{CF}$ and $\mathcal{L}_{disc}$ affect the overall KGC performance, we conduct sensitivity analysis on the performance of KGCF w.r.t different combinations of $\alpha$ and $\beta$. While we tune $\alpha$ and $\beta$ through grid search within the same set $\{0.001,0.01,0.1,1\}$, we search through all combinations on WN18RR and report the corresponding H@10 performance in Figure 4. We observe that the performance is relatively stable and KGCF is generally robust to the hyperparameters $\alpha$ and $\beta$ for finding the optimal values.

6. Conclusion

In this paper, we propose a novel counterfactual relation-based data augmentation framework KGCF for the task of KGC. Given a pair of entities and a query relation, we first explore the factual treatments for the observed context from the view of the proximity of entities. Considering the imbalance of the distribution of relation types, we define the counterfactual relations with the help of their nearest neighbors in the embedding space from relation type weighted node2vec. The counterfactual relations augment the training data with causal relationships and facilitate pair representation learning for KGC. Comprehensive experiments demonstrate that KGCF achieves the SOTA performance on benchmarks as well as endows the interpretability of the path-based GNN models on KGC from a new causal perspective.

Acknowledgments

This work was partially supported by NSFC Grant No. 62206067 and HKUST-GZU Joint Research Collaboration Fund No.GZU22EG05.

Appendix A Overall Algorithm of KGCF.

Alg. 1 summarizes the overall algorithm of KGCF during the training and inference stages. Note that only the output $p^{F}(t|h,r)$ from the factual treatments will be used during the inference stage for the task of KGC.

Appendix B Statistics of datasets

Dataset statistics of FB15k-237 and WN18RR for KGC are summarized in Table 5. We use the standard splits (Toutanova and Chen, 2015; Dettmers et al., 2018) for a fair comparison.

Appendix C Evaluation metrics

We follow the standard filtered ranking protocol (Bordes et al., 2013) and evaluate the performance of models by five metrics Mean Reciprocal Rank (MRR), Mean Rank (MR), and Hits@K (H@K). The detailed definitions of five metrics is introduced below:

MRR. For a test triplet $(h,r,t)$, we rank it against all negative triplets $(h^{\prime},r,t)$ or $(h,r,t^{\prime})$ obtained by corrupting the head (or the tail) of the relation. Let $\mathrm{rank}_{h}(h,r,t)$ be the ranking of $(h,r,t)$ among all head-corrupted relations, and let $\mathrm{rank}_{t}(h,r,t)$ denote a similar ranking with tail corruptions. MRR is the mean of the reciprocal rank:

[TABLE]

where $\mathcal{T}$ is the test set.

MR. MR is the mean of both ranks:

[TABLE]

Hits@K. Hits@K measures the proportion of triples in $\mathcal{T}$ that rank among top-K after corrupting both heads and tails.