Efficient Hyper-parameter Search for Knowledge Graph Embedding (2022)

Abstract: While hyper-parameters (HPs) are important for knowledge graph (KG) learning, existing methods fail to search them efficiently. To solve this problem, we first analyze the properties of different HPs and measure the transfer ability from small subgraph to the full graph. Based on the analysis, we propose an efficient two-stage search algorithm KGTuner, which efficiently explores HP configurations on small subgraph at the first stage and transfers the top-performed configurations for fine-tuning on the large full graph at the second stage. Experiments show that our method can consistently find better HPs than the baseline algorithms within the same time budget, which achieves 9.1% average relative improvement for four embedding models on the large-scale KGs in open graph benchmark. Our code is released in https://github.com/AutoML-Research/KGTuner .

## introduction knowledge graph (kg) is a special kind of graph structured data to represent knowledge through entities and relations between the entities @xcite . learning from kg aims to discover the latent properties from kgs to infer the existence of interactions among entities or the types of entities @xcite . kg embedding, which encodes entities and relations as low dimensional vectors, is an important technique to learn from kgs @xcite . the existing models range from translational distance models @xcite , tensor factorization models @xcite @xcite , neural network models @xcite , to graph neural networks @xcite . hyper-parameter (hp) search @xcite is very essential for kg learning. in this work, we take kg embedding methods @xcite , as a good example to study the impact of hps to kg learning. as studied, the hp configurations greatly influence the model performance @xcite ). an improper hp configuration will impede the model from stable convergence, while an appropriate one can make considerable promotion to the model performance. indeed, studying the hp configurations can help us make a more scientific understanding of the contributions made by existing works @xcite . in addition, it is also important to search for an optimal hp configuration when adopting kg embedding methods to the real-world applications @xcite @xcite . algorithms for hp search on general machine learning problems have been well-developed @xcite . as shown in figure 1(a) , the search algorithm selects a hp configuration from the search space in each iteration, then the evaluation feedback obtained by full model training is used to update the search algorithm. the optimal hp is the one achieving the best performance on validation data in the search process. representative hp search algorithms are within sample-based methods like grid search, random search @xcite , and sequential model-based bayesian optimization (smbo) methods like hyperopt @xcite , smac @xcite , spearmint @xcite as well as bore @xcite , etc. recently, some subgraph-based methods @xcite are proposed to learn a predictor with configurations efficiently evaluated on small subgraphs the predictor is then transferred to guide hp search on the full graph. however, these methods fail to efficiently search a good configuration of hps for kg embedding models since the training to address the limitations of existing hp search algorithms, we carry a comprehensive understanding study on the influence and correlation of hps as well as their transfer ability from small subgraph to full graph in kg learning. from the aspect of performance, we classify the hps into four different groups including reduced options, shrunken range, monotonously related and no obvious patterns based on their influence on the performance. by analyzing the validation curvature of these hps, we find that the space is rather complex such that only tree-based models can approximate it well. in addition, we observe that the consistency between evaluation on small subgraph and that on the full graph is high, while the evaluation cost is significantly smaller on the small subgraph. above understanding motivates us to reduce the size of search space and design a two-stage search algorithm named as kgtuner. as shown in figure 1(b), kgtuner explores hp configurations in the shrunken and decoupled space with the search algorithm rf+bore @xcite on a subgraph in the first stage, where the evaluation cost of hps are small. then in the second stage, the configurations achieving the top10 performance at the first stage are equipped with large batch size and dimension size for fine-tuning on the full graph. within the same time budget, kgtuner can consistently search better configurations than the baseline search algorithms for seven kg embedding models on wn18rr @xcite and fb15k-237 @xcite . by applying kgtuner to the large-scale benchmarks ogbl-biokg and ogbl-wikikg2 @xcite , the performances of embedding models are improved compared with the reported results on ogb link prediction leaderboard. besides, we justify the improvement of efficiency via analyzing the design components in kgtuner. ## background: hps in kg embedding we firstly revisit the important and common hps in kg embedding. following the general framework @xcite , the learning problem can be written as regularization. to balance the expressiveness and complexity, and to avoid unbounded embeddings, the regularization techniques can be considered, such as regularizers like frobenius norm (fro) @xcite , nuclear norm (nuc) @xcite as well as dura @xcite , and dropout on the embeddings @xcite . optimization. to optimize the embeddings, important optimization choices include the optimizer, such as sgd, adam (kingma and ba, 2014) and adagrad @xcite , learning rate, initializers, batch size, embedding dimension size, and add inverse relation @xcite or not. ## defining the search problem denote an instance @xmath0, x 2 . . . , x n ), which is called an hp configuration, in the search space x . let f (p , x) be an embedding model with model parameters p and hps x, we define m(f (p , x), d val ) as the performance measurement (the larger the better) on validation data d val and l(f (p , x), d tra ) as the loss function (the smaller the better) on training data d tra . we define the problem of hp search for kg embedding models in definition 1. the objective is to search an optimal configuration x * ∈ x such that the embedding model f can achieve the best performance on the validation data d val . definition 1 (hyper-parameter search for kg embedding). the problem of hp search for kg embedding model is formulated as definition 1 is a bilevel optimization problem @xcite , which can be solved by many conventional hp search algorithms. the most common and widely used approaches are samplebased methods like grid search and random search @xcite , where the hp configurations are independently sampled. to guide the sampling of hp configurations by historical experience, smbo-based methods @xcite learn a surrogate model to select configurations based on the results that have been evaluated. then, the model parameters p are optimized by minimizing the loss function l on d tra in eq. ( 3 ). the evaluation feedback m of x on the validation data d val is used to update the surrogate. there are three major aspects determining the efficiency of definition 1: (i) the size of search space x , (ii) the validation curvature of m(•, •) in eq. ( 2 ), and (iii) the evaluation cost in solving arg min p l in eq. ( 3 ). however, the existing methods @xcite ) directly search on a huge space with commonly used surrogate models and slow evaluation feedback from the full kg due to the lack of understanding on the search problem, leading to low efficiency. ## understanding the search problem to address the mentioned limitations, we measure the significance and correlation of each hp to determine the feasibility of the search space x in section 4.1. in section 4.2, we visualize the hps that determine the curvature of eq. ( 2 ). to reduce the evaluation cost in eq. ( 3 ), we analyze the approximation methods in section 4.3. following @xcite , the experiments run on the seven embedding models in section 2 and two widely used datasets wn18rr @xcite and fb15k-237 @xcite . the experiments are implemented with pytorch framework @xcite , on a machine with two intel xeon 6230r cpus and eight rtx 3090 gpus with 24 gb memories each. we provide the implementation details in the appendix d.1. ## related works in analyzing the performance of kg embedding models, @xcite pointed out that the earlier works in kg embedding only search hps in small grids. by searching hundreds of hps in a unified framework, the reproduced performance can be significantly improved. similarly, ali et al. (2020) proposed another unified framework to evaluate different models. rossi et al. (2021) evaluated 16 different models and analyzed their properties on different datasets. all of these works emphasize the importance of hp search, but none of them provide efficient algorithms to search hps for kg learning. autosf (zhang et al., 2020a) evaluates the bilinear scoring functions and set up a search problem to design bilinear scoring functions, which can be complementary to kgtuner. understanding the hps in a large search space is non-trivial since many hps only have moderate impact on the model performance @xcite and jointly evaluating them requires a large number of experiments @xcite . considering the huge amount of hp configurations (with 10 5 categorical choices and 5 continuous values), it is extremely expensive to exhaustively evaluate most of them. hence, we adopt control variate experiments to efficiently evaluate hps' properties instead of the quasi-random search in @xcite . technically, we are similar to autone @xcite and e-autogr @xcite by leveraging subgraphs to improve search efficiency on graph learning. since they do not target at kg embedding methods, directly adopt them is not a good choice. besides, based on the understanding in this paper, we demonstrate that transferring the surrogate model from subgraph evaluation to the full graph is inferior to directly transferring the top configurations for kg embedding models. ## conclusion in this paper, we analyze the hps' properties in kg embedding models with search space size, validation curvature and evaluation cost. based on the observations, we propose an efficient search algorithm kgtuner that efficiently explores configurations in a reduced space on small subgraph and then fine-tunes the top configurations with increased batch size and dimension size on the full graph. empirical evaluations show that kgtuner is robuster and more efficient than the existing hp search algorithms and achieves competing performance on large-scale kgs in open graph benchmarks. in the future work, we will understand the hps in graph neural network based models and apply kgtuner on them to solve the scaling limitations in hp search.