Efficient Relation-aware Scoring Function Search for Knowledge Graph Embedding

Generally, Neural Architecture Search (NAS) [21, 24, 22] is formed as a bi-level optimization problem [30] where we need to update the neural architectures on the upper-level and train the model parameters in the lower-level. Subsequently, three important aspects should be considered in NAS:

• •

  Search space: it defines what network architectures in principle should be searched, e.g., Convolutional Neural Networks (CNNs) [31] and Recurrent Neural Networks (RNNs) [32]. A well-defined search space should be expressive enough to enable powerful models to be searched. But it cannot be too large to search.

• •

  Search algorithm: it aims to efficiently search in the above space, e.g., bayesian optimization [33], reinforcement learning [24], evolution algorithm [34]. A search algorithm is required to perform an efficient search over the search space and be able to find architectures that achieve good performance.

• •

  Evaluation mechanism: it determines how to evaluate the searched architectures in the search strategy. Fast and accurate evaluation of candidate architectures can significantly boost the search efficiency.

Unfortunately, classic NAS [34, 24] methods are computationally consuming since candidate architectures are trained by the stand-alone way, i.e., many architectures are trained from scratch to convergence and are evaluated separately.

More recently, One-shot Architecture Search (OAS) methods [35, 36, 37], have been proposed to significantly reduce the search cost in NAS. OAS first represents the whole search space by a supernet [35], which is formed by a directed acyclic graph (DAG), where the nodes are the operations in neural networks (e.g., $3\times 3$ conv in CNNs). Every neural architecture in the space can be represented by a path in the DAG. Then, instead of training independent model weights of each candidate architecture like the stand-alone approach, OAS keeps weights for the supernet and forces different architectures to share the same weights if they share the same edges in the DAG (i.e., parameter-sharing [35, 38]). In this way, architectures can be searched by training the supernet once (i.e., the one-shot manner), which makes NAS much faster.

Generally, OAS methods unify the search space with the form of DAG but differ in their way to search the optimal subgraph of DAG. Sampling OAS (e.g., ENAS [35]) employs a controller to sample architectures and search an optimal subgraph of the whole DAG, which maximizes the expected reward of the subgraph on the validation set. Instead of involving controllers, differentiable OAS (e.g., DARTS [36] and NASP [37]) relaxes the search space to be continuous so that the architectures can be optimized by gradient descent. However, differentiable OAS may not able to derive an architecture that results in high evaluation performance when the evaluation metric is not differentiable. In comparison, sampling OAS is more suitable for the non-differentiable scenario since it utilizes the policy gradient [39] to optimize the controller.

Motivated by the expressiveness guarantee and computational efficiency of BLMs, AutoSF proposes to partition embeddings $\bm{h},\bm{r},\bm{t}\in\mathbb{R}^{d}$ into $M$ splits (e.g., $\bm{h}=[\bm{h}_{1};\cdots;\bm{h}_{M}]$ where $\bm{h}_{i}\in\mathbb{R}^{d/M}$), and represents scoring functions as:

where $\bm{o}\in\mathcal{O}$ with $\mathcal{O}\equiv\{\bm{0},\pm\bm{r}_{1},\cdots,\pm\bm{r}_{M}\}$. Note that $\langle\bm{h}_{i},\bm{o},\bm{t}_{j}\rangle$ computes the triple-dot product and is named as the multiplicative item. Then previous outstanding scoring functions [40, 15, 18, 16, 17] can be unified in $f$ with different choices of $\bm{o}$ [14]. This indicates that $f$ is general enough to represent good scoring functions which are designed manually. In this way, AutoSF generalizes from human wisdom and allows the discovery of better scoring functions, which are not visited in the literature. Subsequently, the search problem is defined as:

Given the embeddings $\bar{\bm{\omega}}$ learned on the training data $S_{\text{tra}}$, AutoSF aims to search for a better scoring function $f$ which leads to higher performance on the validation set $S_{\text{val}}$. Hence, AutoSF can find task-aware scoring functions that can achieve impressive performance on different KG tasks. However, it is non-trivial to efficiently search a proper scoring function from the AutoSF’s search space due to its size $O((2M+1)^{M^{2}})$.

Since a large number of possible scoring functions exist in the unified scoring function search space, AutoSF then proposes a progressive greedy search algorithm to find a proper scoring function according to an inductive rule:

where $b$ is the burget of non-zero multiplicative terms (i.e., $\bm{o}\neq\mathbf{0}$) in (1). The intuition behind (2) is to gradually add nonzero multiplicative terms $\langle\bm{h}_{i},\bm{o},\bm{t}_{j}\rangle$ to achieve the final desired $f$. Each greedy step in the search algorithm mainly contains two parts: sampling scoring functions and evaluate the sampled scoring functions. We summarize the search algorithm in Algorithm 1.

However, there are several issues in AutoSF. First, the evaluation mechanism in AutoSF is inefficient. In every greedy step, AutoSF trains all candidate scoring functions under budget $b$ to convergence for performance evaluation, i.e., step 5 of Algorithm 1. Then well-trained embeddings of all scoring functions will be discarded in the next greedy search. It wastes a lot of effort to train KG embeddings for performance evaluation. Moreover, within the budget $b$, the predictor in AutoSF search algorithm can only leverage the prior experience that is smaller than the budget $b$, i.e., step 4 of Algorithm 1, which may also bring unnecessary KG embeddings training due to inaccurate prediction. Second, AutoSF pursuits a universal scoring function over a given KG data. A universal scoring function that can learn certain relationships does not necessarily mean that this scoring function can perform well on them [28, 27]. This will be further discussed in Section III-A.

To handle the non-differentiable MRR, we first formulate the search problem of $\{f_{n}\}$ as a multi-step decision problem. Then we adopt reinforcement learning to solve (3).

Recall that each $f_{n}$ is a summation of multiplication terms $\langle\bm{h}_{i},\bm{o},\bm{t}_{j}\rangle$ in (1). We can sequentially determine which operation $\bm{o}\in\mathcal{O}$ is selected for the $(i,j)$-th multiplicative item in $f_{n}$. Let v denote the index of all multiplicative items in $\{f_{n}\}$ and $\alpha_{\emph{v}}$ denote the operation selected for the v-th multiplicative item. Then, as shown in Figure 1 (a), the search process of $\{f_{n}\}$ can be viewed as a multi-step decision problem: a list of tokens $\{\alpha_{\emph{v}}\}_{\emph{v}=1}^{V}$ ($V=NM^{2}$) that needs to be predicted. It is intuitive to adopt reinforcement learning to solve this problem. Let $\bm{A}\in\{0,1\}^{NM^{2}\times(2M+1)}$, where $A_{\emph{v}k}=1$ if $\alpha_{\emph{v}}$ chooses $k$-th operation $\bm{o}_{k}$ in $\mathcal{O}$ otherwise $A_{\emph{v}k}=0$. Then (3) can be reformulated as:

where $\pi(\bm{A};\theta)$ is a policy parameterized by $\theta$ for generating $\bm{A}$, and $Q\left(\bm{A},\bm{B},\bm{\omega};S_{\text{val}}\right)$ measures MRR performance as:

Note that $q(\cdot)$ measures MRR of a triplet and $f_{n}$ is the $n$-th scoring function based on $\bm{A}$. Then, the gradient of $\theta$ in (6) can be computed by REINFORCE gradient [39] as

where $Q_{u}$ denotes $Q\left(\bm{A}^{u},\bm{B},\bm{\omega};S_{\text{val}}\right)$, $\bm{A}^{u}$ is $u$-th sampled architecture from $\pi(\bm{A};\theta)$, $b$ is a moving average baseline to reduce variance, and $U$ denote the number of sampled scoring functions. We can see that solving (6) has been converted to optimize $\theta$ in (7), and whether $Q$ is differentiable does not affect the gradient computation w.r.t $\theta$.

Inspired by [24, 35], we also adopt a Long Short-term Memory (LSTM) [43] to parameterize $\theta$ for learning the policy $\pi(\bm{A};\theta)$. Specifically, the controller samples decisions in an autogressive way: the decided operation $\alpha_{\emph{v}}$ in the previous multiplicative item is carried out and then fed into the next to predict $\alpha_{\emph{v}+1}$ (see Figure 1 (a)). Finally, (7) is used to update the LSTM policy network $\theta$.

To fully evaluate the search space, we expect that every relation segmentation $\bm{r}_{i}\in\mathcal{O}\backslash\mathbf{0}$ must be selected at least once in the searched scoring functions, named as the exploitative constraint. Thus, we constrain $\bm{A}$ with this exploitative constraint. If the sampled $\bm{A}$ does not satisfy it, we will directly set the reward $Q$ to 0.

To enable fast evaluation, we propose a supernet view of the relation-aware scoring function search space. Specifically, the $f_{n}$ is represented as:

Recall that v is the index of multiplicative items in $\{f_{n}\}$ and $k$ is the index of operations in $\mathcal{O}$. Then, as shown in Figure 1 (b), we can take $\bm{A}$ as the adjacency matrix of a bipartite graph (i.e., supernet) from above (8), where multiplicative items and operations are nodes, and $A_{\emph{v}k}$ records the edge weight between multiplicative items and operations. Based on the supernet design, Figure 1 (c) illustrates that any relation-aware $\{f_{n}\}$ can be realized by taking subgraphs of the supernet. Then ERAS forces all subgraphs to share embeddings, thereby different scoring functions can be evaluated based on the same KG embedding. It enables us to evaluate candidate scoring functions faster by avoiding repetitive embedding training.

Given the fixed controller’s policy $\pi(\bm{A};\theta)$ and relation assignment $\bm{B}$, we propose to solve $\mathcal{M}_{\text{tra}}$ in objective (4) by minimizing the expected loss $\mathcal{L}$ on the training data, such as $\mathbb{E}_{\bm{A}\sim\pi(\bm{A};\theta)}[\mathcal{L}\left(\bm{A},\bm{B},\bm{\omega};S_{\text{tra}}\right)]$. Then stochastic gradient descent (SGD) can be performed to optimize $\bm{\omega}$. We approximate the gradient $\nabla_{\bm{\omega}}$ according to:

where $U$ is the number sampled scoring functions, and

Note that $\ell(\cdot)$ is the multiclass log-loss [19] and $f_{n}$ is the $n$-th scoring function based on sampled $\bm{A}^{u}$. Hence, (9) can be represented as:

Recall that $Q_{u}$ denotes the reward of the $u$-th sampled scoring function in (7). In the supernet, every sampled scoring function can be regarded as a subgraph of the supernet where some edges are activated. Hence, we can evaluate $Q_{u}$ based on the sampled scoring functions by activating the subgraph on the supernet.

To compare ERAS with the pioneering work AutoSF, we summarize them from the perspective of NAS principles in Table V. First, to capture the inherent properties of relation patterns in KGs, ERAS targets to solve the relation-aware scoring function search problem in this paper. This leads to the larger search space in ERAS than that in AutoSF. Second, ERAS proposes an alternative minimization way to solve the non-differentiable problem of $\mathcal{M}_{\text{val}}$. Finally, AutoSF evaluates every scoring function based on their stand-alone performance, which requires repeatedly training KG embeddings hundreds of times. On the contrary, ERAS shares KG embeddings across different scoring functions, which extremely reduces time expense for performance evaluation of candidate scoring functions.

Inspired by parameter sharing in OAS works, we propose to share KG embedding in the supernet, thereby ERAS avoids repeatedly training KG embedding hundreds of times. However, there exist some concerns about parameter-sharing in OAS [44, 38]. Specifically, while parameter sharing enables an alternative way to train all architectures in the search space in a cheaper way, it can lead to a biased evaluation problem: the evaluation of candidate architectures in the search phase (i.e., $\mathcal{M}_{\text{val}}$) is inconsistent with the stand-alone training phase. Especially, the correlation between one-shot and stand-alone performance will probably be unstable as the supernet goes deep and complex.

Therefore, to make parameter sharing work for our problem, we first construct a search space, which can be represented by a supernet in (8). Then, to avoid the supernet being deep and complex, we design a shallow and simple supernet with the form of a bipartite graph, which differs from the complex DAG supernet in classic OAS works. We demonstrate that the ERAS’s supernet design does not suffer from a biased evaluation problem in Section V-E1. In summary, we leverage the domain-knowledge of KG embedding to make embedding sharing works.

In our experiments, we mainly conduct experiments on five public benchmarks data sets: WN18 [7], WN18RR [12], FB15k [7], FB15k237 [41], and YAGO3-10 [12], that have been employed to compare KG embedding models in [7, 17, 18, 16, 15]. Note that WN18RR and FB15k237 remove duplicate and inverse duplicate relations from WN18 and FB15k, respectively. The statistics of five data sets are summarized in Table VII.

The hyperparameters in this work can be mainly categorized into searching and evaluation parameters. To fairly compare existing scoring functions, including human-designed and searched ones, we search a stand-alone parameter set on the SimplE with the help of HyperOpt, a hyperparameter optimization framework [33]. The tuned parameter set includes: learning rate, L2 penalty, decay rate, batch size, embedding dimensions. Then we compare the stand-alone performance of different scoring functions on the tuned parameter set. Besides, the searching parameters of ERAS are the number of segments $M$, relation groups $N$, sampled scoring functions $U$ in (7) and (9). Moreover, we optimize embeddings $\bm{\omega}$ with Adagrad [48] algorithm and the controller $\theta$ with Adam [49] algorithm.

We first test the performance of our proposed method on the link prediction task. This is the test bed to measure KG embedding models and works as an important task in KG completion. Given the triplet $(h,r,t)\in S_{\text{val}}\cup S_{\text{test}}$, the KG embedding model obtains the rank of $h$ through computing the score of $(h^{\prime},r,t)$ for all entities, and the same for $t$. We adopt the classic metrics [7, 8]:

• •

  Mean Reciprocal Ranking (MRR): $\nicefrac{{1}}{{|S|}}\sum_{i=1}^{|S|}\nicefrac{{1}}{{\text{rank}_{i}}}$, where $\text{rank}_{i}$ is the ranking result; and

• •

  Hit@1, i.e., $\nicefrac{{1}}{{|S|}}\sum_{i=1}^{|S|}\mathbb{I}(\text{rank}_{i}\!\leq\!1)$, and Hit@10, i.e., $\nicefrac{{1}}{{|S|}}\sum_{i=1}^{|S|}\mathbb{I}(\text{rank}_{i}\!\leq\!10)$, where $\mathbb{I}(\cdot)$ is the indicator function.

Note that the higher MRR, Hit@1 and Hit@10 values mean better embedding quality. We compare the proposed ERAS (Algorithm 2) with the popular KG embedding models mentioned in Section I:

• •

  Translational models (TDMs): TransE [7], TransH [8], and RotatE [45];

• •

  Neural network models (NNMs): NTN [11], ConvE [12], and HypER [25];

• •

  Tensor-based models (TBMs): TuckER [20], HolEX [29], QuatE [46], DistMult [15], ComplEx [16], Analogy [18] and SimplE [17];

• •

  The rule-based model: AnyBURL [50];

• •

  The scoring function search method: AutoSF [14].

The comparison of the global effectiveness between ERAS and other methods is in Table VI. Firstly, it is clear that traditional scoring functions, such as TDMs, NNMs, and TBMs, are not task-aware since no scoring functions can perform consistently over the benchmark data sets. This indicates that a single scoring function is hard to adapt to different KGs even though it is a universal scoring function like, as discussed in Table III. The task-aware method AutoSF can search the KG-dependent scoring functions on five data sets and perform consistently better than traditional scoring functions. Then, we compare AutoSF with a variant of ERAS, i.e., $\text{ERAS}^{N=1}$, that aims to search a universal scoring function since all relations are assigned into one group (i.e., only task-aware as AutoSF). For five benchmark data set, $\text{ERAS}^{N=1}$ shows almost the same performance with AutoSF. Moreover, as a task-aware and relation-aware method, ERAS performs better than AutoSF and other manually designed scoring functions.

As discussed in Section III-A, the existing scoring functions may achieve unsatisfactory performance on specific relation types of certain KG data. Corresponding to Table III, we investigate the performance of $\text{ERAS}^{N=1}$ and ERAS at the relation type level in Table VIII. Obviously, the relation-aware method ERAS can consistently achieve outstanding performance on various relation types of any KG data. Especially, ERAS improves the performance of symmetric relations in FB15k and FB15k237. However, since the fact of symmetric relation only accounts for 3% of the test data in FB15k and FB15k237 [27], the global improvement is not so notable.

To further demonstrate the effectiveness of ERAS, we also conduct the triplet classification experiments on FB15k, WN18RR, and FB15k237, where the positive and negative triplets are provided. We compare our methods with those that have reported results in public papers. This task aims to answer whether a given $(h,r,t)$ exists or not. In this task, we utilize the accuracy to evaluate the scoring functions. We set the same decision rule of classification in literature [14]: predicting a $(h,r,t)$ is positive if $f(h,r,t)>\theta_{r}$ otherwise negative, where $\theta_{r}$ is a relation-specific threshold and is inferred by maximizing the accuracy on $S_{\text{val}}$. As shown in Table X, the scoring function searched by relation-aware ERAS consistently outperforms other BLMs or searched scoring functions.

As discussed in Section IV-D2, the deep and complex supernet design will probably lead to the biased evaluation problem. In Figure 5(a), we first demonstrate the correlation between stand-alone validation MRR with one-shot validation MRR (i.e., $\mathcal{M}_{\text{val}}$) of various scoring functions in ERAS. It is obvious that one-shot validation MRR has near positive correlation with the stand-alone validation MRR. Therefore, the simplified design of supernet makes embedding sharing work, and there is no biased evaluation problem.

To further investigate the impact of $\mathcal{M}_{\text{val}}$ and optimization algorithms, we compare ERAS with following variants:

• •

  $\text{ERAS}^{\text{los}}$ utilizes the validation loss $\mathcal{L}$ to replace $\mathcal{M}_{\text{val}}$. Other steps are same with ERAS.

• •

  $\text{ERAS}^{\text{dif}}$ first replaces $\mathcal{M}_{\text{val}}$ with $\mathcal{L}$ as $\text{ERAS}^{\text{los}}$ does. Then the differentiable measurement $\mathcal{L}$ enables the differentiable optimization algorithm [36, 37] for search. The detailed implementing $\text{ERAS}^{\text{dif}}$ is presented in Appendix.

We show the correlation between stand-alone validation MRR with different $\mathcal{M}_{\text{val}}$ settings in Figure 5. Obviously, compared with using MRR as $\mathcal{M}_{\text{val}}$ in Figure 5 (a), Figure 5 (b) shows that using $\mathcal{L}$ as $\mathcal{M}_{\text{val}}$ has a low correlation with stand-alone validation MRR. This indicates that validation loss in search strategy cannot well evaluate the stand-alone performance of scoring functions. Subsequently, in Table XI, we can observe that the performance of ERAS using MRR as $\mathcal{M}_{\text{val}}$ is better than that of two variants, i.e., $\text{ERAS}^{\text{los}}$ and $\text{ERAS}^{\text{dif}}$. Moreover, $\text{ERAS}^{\text{los}}$ is worse than $\text{ERAS}^{\text{dif}}$ because optimizing loss with the RL approach could not make use of the differentiable nature of $\mathcal{L}$.

In this paper, Definition 2 formulates the problem with a bi-level optimization objective. As stated in Section III-B, bi-level optimization can benefit ERAS by updating the scoring functions and embeddings separately. To investigate the impact of optimization level, we add another variant of ERAS as $\text{ERAS}^{\text{sig}}$, which utilizes the training set to update (3) in Definition 2 (i.e., single-level problem). In Table XI, compared with $\text{ERAS}^{\text{sig}}$, ERAS demonstrates the bi-level optimization is needed because optimizing scoring functions on the validation set encourages ERAS to select scoring functions that generalize well rather than scoring functions that overfit on the training data.

To explore more about the influence of different grouping approaches, we set two variants of ERAS as follows:

• •

  $\text{ERAS}^{\text{pde}}$ does not update $\bm{B}$ during search. Instead, it fixes groupings based on the embedding trained from SimplE.

• •

  $\text{ERAS}^{\text{smt}}$ groups relations based on their semantic meanings (i.e., symmetric, anti-symmetric, asymmetric and inverse).

We compare these two variants with ERAS as shown in Table XI. Generally, the performance of $\text{ERAS}^{\text{smt}}$ is unsatisfactory due to the bias between human understanding of relation groups with the proper groups derived from data. In short, the performance comparison also indicates the importance of dynamically assigning relation groups in the search strategy, which encourages relations to be assigned to the appropriate scoring function.

To investigate the impact of grouping numbers in ERAS, we summarize the performance of searched scoring functions by AutoSF and different settings of ERAS in Figure 6, i.e., $\text{ERAS}^{N}$ ($N\in\{1,2,\cdots,5\}$) on WN18RR and FB15k237. Generally, the more groups there are, the longer the running time is. And ERAS achieves the best performance when $N=3$ or $4$. Comparing AutoSF and $\text{ERAS}^{N=1}$, they have a similar learning curve since their model complexities are the same.

AutoSF fixes the block number $M$ to 4 due to the efficiency issue. Once $M$ is changed (e.g., $M=3$ or $M=5$), all design details in AutoSF must be re-done. On the contrary, the efficiency of ERAS allows more flexible settings of $M$. Here we try $M\in\{3,4,5\}$ in order to learn more about how the block number $M$ influences the ERAS performance. As shown in Figure 7, we can observe that $M=4$ does have excellent performance among $\{3,4,5\}$.