Generic Structure Extraction with Bi-Level Optimization for Graph Structure Learning

Currently, most Graph Structure Learning (GSL) methods, as a means of learning graph structure, improve the robustness of GNN merely from a local view by considering the local information related to each edge and indiscriminately applying the mechanism across edges, which may suffer from the local structure heterogeneity of the graph (i.e., the uneven distribution of inter-class connections over nodes). To overcome the drawbacks, we extract the graph structure as a learnable parameter and jointly learn the structure and common parameters of GNN from the global view. Excitingly, the common parameters contain the global information for nodes features mapping, which is also crucial for structure optimization (i.e., optimizing the structure relies on global mapping information). Mathematically, we apply a generic structure extractor to abstract the graph structure and transform GNNs in the form of learning structure and common parameters. Then, we model the learning process as a novel bi-level optimization, i.e., Generic Structure Extraction with Bi-level Optimization for Graph Structure Learning (GSEBO), which optimizes GNN parameters in the upper level to obtain the global mapping information and graph structure is optimized in the lower level with the global information learned from the upper level. We instantiate the proposed GSEBO on classical GNNs and compare it with the state-of-the-art GSL methods. Extensive experiments validate the effectiveness of the proposed GSEBO on four real-world datasets.

Introduction
Based on the homophily assumption of "like to associate with like" [1,2], the Graph Neural Network (GNN) [3] has become the promising solution for node classification. However, a large portion of edges are inter-class connections [4], and representation propagation over such connections can largely hinder the GNN from obtaining class-separated node representations, hurting the performance.
Existing GSL methods are roughly categorized into the attentive mechanism, noise detection, and probabilistic mechanism. Attentive mechanism calculates weights for edges to adjust the contribution of different neighbors during representation propagation [5][6][7][8][9]. These methods can hardly work well in practice for two reasons: (1) the mechanism may not generalize well to all nodes with different local structures (cf. Figure 1a); and (2) the attention cannot be easily trained well due to the limited labeled data [10]. Noise detection incorporates an edge classifier to estimate the probability of inter-class connection for each edge [11][12][13][14][15][16]. Although it can be better trained owing to the supervision of edge labels, the edge classifier also suffers from local structure heterogeneity and lacks consideration of global information. Probabilistic mechanism models connection from a global view, which assumes a prior distribution of edge and estimates GNN parameters with Bayesian optimization to overcome the impact of inter-class edges [17][18][19][20]. Although the edge specific parameterization largely enhances the model representation ability, it is hard to accurately access the prior distribution.  Figure 1. Examples for (a) the impact of local structure heterogeneity (node 1 and 9, the weights between nodes do not provide enough information for classes differentiation) and (b) the global view of messaging (the edge between node 3 and node 5, the deleted edge prevents the intra-class information transfer).
Despite the achievements of the existing methods, there still exists some common drawbacks: (1) Edge modeling method, the existing methods model edges with the parameter sharing mechanism, which may suffer from the local structure heterogeneity problem; (2) Local optimization, the local optimization problem focuses on optimizing the parameters with the information of neighbor nodes, which ignores the impact from the global view. Therefore, we come up with the key considerations for GSL: (1) modeling graph connection in an edge-specific manner instead of a shared mechanism; and (2) optimizing the corresponding parameters with a global objective of accurately classifying all nodes. The edge-specific modeling can overcome the local structure heterogeneity, i.e., handling nodes with different properties (e.g., node 1 and node 9 in Figure 1a) via different strategies. Besides, blindly removing the inter-class edges will increase the risk of misclassifying the target nodes (in dash circle) due to cutting off their connections to the labeled neighbors in the same class (e.g., edge between node 5 and node 9 in Figure 1b). Thus, it is necessary to optimize the graph structure from the global view instead of the local ones.
However, it is non-trivial to achieve the targets mentioned above due to the following challenges: (1) the graph structure is embedded into the GNN model, which affects the procedure of model parameter optimization once updated, requiring a careful design of the model training algorithm; (2) the calculation of the ideal global objective is intractable due to the limited labelled nodes, especially the hard semi-supervised setting.
In this work, we propose a new Generic Structure Extraction with Bi-level Optimization for Graph Structure Learning (GSEBO), which optimized the graph structure and learsn the node embeddings from the global view. In particular, we first devise a new generic structure extractor, which accounts for the graph structure with both the connectedness between nodes and the strength of connections. In addition to the adjacency matrix, the extractor adopts a learnable matrix to represent the graph structure and adjusts the representation propagation. Moreover, we design a bi-level optimization algorithm where the outer and inner optimizations update the structure and the parameters of the base graph convolution (vs feature mapping parameters). In this way, we decompose the hard optimization issue of GSEBO into two easy ones. In addition, we separate the training set into two parts (e.g., train_1, train_2), and set the objective of outer optimization as the train_2 loss to better approximate the ideal global objective. The proposed generic structure extractor can be extended to most existing graph convolution operators. We instantiate it on four representative GNN models (i.e., GCN [21], GAT [5], GraphSAGE [22], and JK-Net [23]) and compare GSEBO with state-of-the-art GSL methods, which are evaluated on four real-world node classification datasets. Extensive experiments justify the rationality, effectiveness and robustness of the proposed method. In summary, our main contributions are as follows: • We propose a novel GSEBO with bi-level optimization for edge-specific graph structure learning, which learns the graph structure from a global view by optimizing a global objective of node classification. • We devise a generic structure extractor, which parameterizes the strength of each edge during representation propagation. Besides, we summarize how the classical GNN methods are transferred in the form of a learnable graph structure with a generic structure extractor. • We evaluate the proposed GSEBO with classical GNNs as backbones and compare it with the state-of-the-art GSL methods. Extensive experiments on four real-world datasets show the superior learning ability of the proposed method compared to the existing methods.

Attentive Mechanism
The attentive mechanism methods adaptively learn the weights of edges and adjust the contributions of neighbor nodes. MAGNA [6] incorporates multi-hop context information into every layer of attention computation. IDGL [8] uses the multi-head self-attention mechanism to reconstruct the graph, which has the ability to add new nodes without retraining. HGSL [9] extends the graph structure learning to heterogeneous graphs, which constructs different feature propagation graphs and fuses these graphs together in an attentive manner. However, those methods suffer from different local structures and are difficult to train.

Noise Detection
The noise detection methods leverage the off-shelf pre-trained model to induce node embeddings or labels and incorporate an edge classifier to estimate the probability of each edge. NeuralSparse [11] considers the graph sparsification task by removing irrelevant edges. GAUG [12] utilizes a GNN to parameterize the categorical distribution instead of MLP in NerualSparse. PTDNet [13] prunes task-irrelevant edges by penalizing the number of edges in the sparsified graph with parameterized networks. Even though the noise detection methods can be well trained with the supervision of edge labels, the edge classifier also suffers from local structure heterogeneity and lacks consideration of global information.

Probabilistic Mechanism
This type of method assumes the prior distribution of graph or noise and estimates the parameters through observed values, then resamples the edges or noise to obtain a new graph. BGCN [17] estimates the parameter distribution of edges and communities by sampling edges from graph, and resamples new graphs with the estimated parameters for prediction. VGCN [18] trains a graph distribution parameter similar to the original structure through ELBO, and resamples graphs for prediction. However, both the BGCN and VGCN models are sampled from a noisy graph, and the estimated parameters also contain noise. DenNE [19] assumes the observed graph is composed of real values and noise and the prior distribution of features and that the noise is known. With a generative model, the likelihood is used to estimate the representation of nodes. However, this method highly relies on the priors of feature and noise, which is difficult to obtain accurately.

Bi-Level Optimization on GNN
LDS [15] jointly learns the graph structure and GNN parameters by solving a bi-level optimization issue that learns a discrete probability distribution for each edges. According to the learned distributions, LDS generates a new graph structure by sampling. Towards this end, LDS sets the objective of outer optimization as generating the observed edges, which clearly has a gap to the overall classification objective. Moreover, LDS needs to estimate N 2 distribution parameters at least, which is hard due to insufficient labels (|E | N 2 |. On the contrary, our method only activates a small portion of entries in Z, where A ij = 1, i.e., the number of estimated parameters is same as the number of edges in graph G (i.e., |E |).

Methodology
We first introduce the essential preliminaries for GNN, and then elaborate the graph convolution operator and bi-level optimization algorithm of the proposed GSEBO.

Preliminary
Let G = (V, E , X) represents a graph with N nodes and M edges, where V = {v 1 , v 2 , · · · , v N } and E = {e 1 , e 2 , · · · , e M } denote the set of nodes and edges, respectively.
This task aims to learn a classifier f (A, X; θ) from a set of labeled nodes to forecast the remaining nodes labels, where θ denotes model parameters. Assuming there are N labels, we index them from 1 to N without loss of generality. Formally, Y = [y 1 , y 2 , · · · , y N ] ∈ R N are the labels of the nodes, where y i ∈ R is the label of node i. The target is achieved by optimizing the model parameter θ with respect to the labeled nodes, which is formulated as: where l(·) is a classification loss and λ is a hyperparameter to adjust the strength of parameter regularization.

Generic Structure Extraction
To optimize the graph structure, the key consideration lies in (1) decoupling the graph structure from the GNNs to account for the edge-specific modeling and (2) learning the graph structure from the global information in θ.
Towards the first purpose, the core idea is to decompose the graph structure information into connectedness (the edges in the adjacency matrix) and the strength of connection (the latent variable). In general, there are two ways to model the connection strength regarding whether relying on the inductive bias of translation invariant or not. On the one hand, attentive mechanisms or noise detection models are translation invariant, which decode the connection strength of each edge from its local features. However, with the consideration of the local structure heterogeneity issue in most real-world graphs [23], it is risky to rely on the translation invariant bias. On the other hand, probabilistic mechanisms separately model the connection strength for each edge, where each edge corresponds to a specific distribution. However, it is non-trivial to set a proper prior in practice. According to these advantages and disadvantages, we summarize two considerations for extending the graph convolution: (1) edge-specific modeling; and (2) optimization without prior.
Towards this end, we first propose the generic structure extractor (GSE) to decouple the graph structure into the connectedness between nodes and the strength of the connections. To model the edges in an edge-specific way, we apply the bi-level optimization method with the inner optimization to update the common parameters and the outer optimization to optimize the weight of each edge. By introducing the bi-level optimization, we can learn the structure information from the global parameters, thus avoiding the prior assumptions. Specifically, we model the connection strength as a parameter matrix Z with the same size as A. Formally, the generic structure extractor (GSE) is abstracted as: Ã , where σ(·) is a non-negative activation function, since the value of strength is always positive (in this work, we use the min[max[0,x],1] function to restrict the value within [0, 1]). COM and AGG are the combination and aggregation functions, respectively. Noteworthy, different from GNNs, GSE decouples the graph structure from GNNs and treats it as a learnable objective. Besides, GSE is a generic extractor, which can be instantiated over most existing graph convolutions. In this way, as long as learning the connection strength is set properly, GSEBO can downweight the neighbors with inter-class connection during representation propagation, reducing the impact of inter-class with bi-level optimization.

Update Parameters with Bi-Level Optimization
To achieve the second purpose of learning graph structure from global information, it is essential to carefully design a proper training algorithm to optimize the connection strength matrix Z. Assume that we construct GSEBO with K layers, which is denoted as We have three main considerations for designing the training algorithm: • Connection strength is a relative value, which changes across different views. As shown in Figure 1b, the connection between node 3 and node 5 is weak from the local view, i.e., the edge is inter-class and should be assigned a low weight. However, this edge is essential for the classification of node 9, 10, 11, which deserves a high weight. Therefore, the optimization objective of Z should be the overall performance of node classification. • Z and W = {W (k) |k ∈ [1, K]} play different roles, but are closely related. The role of Z is close to A, which restricts the model space for the mapping from node feature to label, and the role of W includes the global mapping information for classification, which would relieve the cons of local optimization. That is, an update on Z will adjust W and its optimization procedure. Therefore, the optimization of Z and W are at two different but dependent levels. • Directly minimizing the objective function of Equation (1) to obtain the parameters Z and W is not able to achieve the desired purpose of learning the structure for the reason of over-fitting, which is shown in Figure 2.
Inner objective  Towards this end, we propose a bi-level optimization [24][25][26] with inner and outer optimizations to learn W and Z, respectively. Figure 3 shows the overall procedure of the algorithm, where the inner and outer optimization steps are iteratively executed until the stop condition.  Figure 3. The framework of GSEBO. The generic structure extractor (GSE) decouples the graph structure from GNNs to account for the edge-specific modeling; the inner and outer optimization steps are adopted for common parameters and graph structure optimization.

Inner Optimization for Common Parameters
This step is similar to the normal training of GNN for W optimization. In particular, we update W with a gradient descent based optimizer (e.g., Adam [27]) over the train_1 nodes by minimizing Equation (3) (i.e., set l inner = l( f (A, X), y; θ) + λ θ ). When calculating the gradient of W, we treat the connection strength matrix Z as constant.

Outer Optimization for Graph Structure
Similarly, we treat the sequence of W as constant to optimize the graph structure Z. Ideally, the objective should be the overall classification loss of all nodes in graph. Formally: Apparently, the calculation of the ideal objective is intractable. Similar to the global parameters optimization, we can approximate the ideal objective as the empirical risk over train_1 nodes. However, it can easily suffer from the over-fitting issue, which is shown in Figure 2. We believe the empirical risk over train_2 nodes is a better approximation of the ideal objective since it reflects to what extent the parameter generalizes well. In this light, we set l outer as the classification loss and optimize Z over the train_2 set V train2 by minimizing Equation (6). Assuming: We optimize the outer parameters Z by fixing the inner parameters W 1 , · · · , W τ . Formally, the derivative of the outer objective to the hyperparameter Z (hypergradient) is formulated as: For t = τ down to 1, we define the matrices: and the update of Z can be formulated as:

Training Process
The bi-level optimization cannot guarantee the convergence of the model; thus, we set the early stop condition as the additional requirement for training, which is shown in the second line of Algorithm 1. We set the early stop condition as: within 20 epochs, if the accuracy of the verification set is not improved, the model will stop training.
The process of updating Z is shown on line 8-11 in Algorithm 1. With the parameter Z updated by the graph structure optimization, we reoptimize the W with Z by global parameters optimization, and repeat this process until the early stopping is met. To summarize, Algorithm 1 shows the training procedure of GSEBO.

Complexity Analysis
The overall framework of GSEBO is illustrated in Algorithm 1. The computing complexity of our GSEBO mainly depends on two steps. Given a graph G, ||A|| 0 is the number of nonzeros in the adjacency matrix, d is the feature dimension, L is the layer number of GCN, and |V| is the number of nodes. In the inner optimization step, the graph convolutional network takes O(τL||A|| 0 d) computational time. In the outer optimization step, the computational time is O(τ||A|| 2 0 ). As a result, the total computational is O(τ||A|| 2 0 ) if Ld is smaller than ||A|| 0 .

Unifying Various GNN and Beyond
To extend the graph convolution operator to the most existing graph convolutions, we reformulate the classical GNN model into a unified form.

Gcn
The forward propagation of GCN is formulated as follows: where H (0) = X,Ã = A + I,D nn = ∑ iÃni , and W ∈ R D×C are the trainable parameters, I denotes the identity matrix, and σ is a nonlinear function. For better understanding of the propagation of GCN, we rewrite the embedding update in the following form: where W (l) reflects the trainable parameters on the l-th layer, σ 1 is the activate function, such as ReLU, and σ 2 = min[max[0,x],1] is the a non-negative activation function.

Gat
GAT assigns different weights to each neighbor node, and updates the node embeddings with weighted average of neighbors: whereÃ ij = 1, and j ∈ N (i), where W (l−1) k is the parameter of kth multi-head attention on (l − 1)th layer, So f tmax is the Softmax operation, and atten is the self-attention operation.

GraphSAGE
GraphSAGE is one representative of the spatial approaches, which learns how to aggregate feature information from a node's local neighborhood, and the reformulation is as follows:

JK-Net
JK-Net learns the node representation from deep layers by using a jumping network. There exist different ways to aggregate features from different layers; we take the concatenation as an example in our work (other aggregation methods are also applicable). The formulation of JK-Net is as follows: where FC is a fully connected layer and Concat is the concatenation operation.

Experiments
In this section, we conduct experiments on four datasets to answer the following research questions: We select four widely used real-world node classification benchmark datasets with graphs of citation networks (Cora and Citeseer [21]), social networks (Terrorist) [28], and air traffic (Air-USA) [29]. The statistics of the datasets are shown in Table 1. • Citation networks. Cora and Citeseer are citation networks, where the nodes are papers published in computer science, the features of each publication are described by a 0/1-valued word vector indicating the absence/presence of the corresponding word from the dictionary, the adjacency matrix is binary and undirected, which denotes the citation relation between papers, and the labels are the category of each paper. • Terrorist Attacks. Terrorist attacks describe the information related to terrorism attack entities: the attributes of the entities and the links that connect various entities together to form a graph structure. This dataset consists of 1293 terrorist attacks each assigned one of six labels indicating the type of the attack. Each attack is described as a 106-dimensional vector, with a 0 and 1 value indicating the absence and presence, respectively, of an entity. The edges are binary and undirected. • Air traffic network. Air-USA is the airport traffic network in the USA, where the nodes represent airports and the binary undirected edges indicate the existence of commercial flights between the airports. The features of nodes are one-hot degree vectors and the labels are generated based on the activity measured by people and flights passed the airports.
All the baselines and our proposed method can be applied to all types of networks. We adopt the same data split of Cora and Citeseer as [21], and a split of training, validation, and testing with a ratio of 10:20:70 on other datasets [12].

Implementation Details
In the experiments, the separate ratio of training data is set to 0.8 for the train_1 and train_2 datasets, and we optimize the outer and inner parameters on the train_1 and train_2 datasets, respectively. The latent dimension of all the methods is set to 16. The parameters for all baseline methods are initialized as the corresponding papers, and are carefully tuned to achieve optimal performances. The learning rate of inner and outer optimization is set to 0.01. The hyperparameter λ is set to 5 × 10 −4 , and we search for the inner learning depth τ over the range [5,10,15,20,25]. As bi-level optimization cannot guarantee the convergence, we set the patience of early stopping to 20 to terminate the optimization of GSEBO. To prevent overfitting, the dropout ratio is set to 0.5 for all the methods. The network architectures of all the methods are configured to be the same as described in the original papers. Our experiments are conducted with Tensorflow running on GPU machines (NVIDIA 2080Ti). For all the compared methods, we report the average accuracy on the test set over 10 runs. Table 2 shows the performance of GSEBO instantiate with classical GNNs, and Table 3 presents the results comparison of GSEBO and state-of-the-art GSL methods. From Tables 2 and 3, we have the following observations:

Improvement over Baselines
GSEBO outperforms the baselines in most cases. Considering the average performance over four datasets, the improvement of GSEBO over the baselines is in the range of 3.0-12.5%, which validates the effectiveness of the proposed method. In particular: • Probabilistic mechanisms. The performance gain of BGCN and VGCN over the vanilla GCN are limited, which might because of the unsatisfied assumption of the prior distribution. This result shows the rationality of relaxing the prior assumption and modeling the connection strength with parameters directly. • Connection modeling. DropEdge, AdaEdge, GAUG, and PTDNet achieve better performance than vanilla GNN, which modifies the structure of the graph from different perspectives, such as the smoothness [30] and robustness [31]. These results reflect the benefit of connection modeling. However, there is still a clear gap between these methods and GSEBO, which is attributed to the global objective for learning the edge strength. • Attention mechanism. The performance of MAGNA is inferior, which is consistent with the analysis in [10].

Effects across GNN Architectures
On the four GNN architectures, GSEBO achieves better performance than the vanilla version in all cases. In particular, GSEBO achieves an average improvement across datasets of 4.2%, 4.4%, 4.0%, and 5.74% over the vanilla GCN, GAT, GraphSAGE, and JK-Net, respectively. These results justify the effectiveness of structure learning of the proposed GSEBO. Across the four architectures, GSEBO achieves the largest improvement over JK-Net. We postulate the reason is that the jump connection in JK-Net makes it aggregate more hops of neighbors than the other GNNs. As the hops increase, the homophily ratio of the neighbors will decrease, i.e., more neighbors are in classes different from the target node. Therefore, optimizing the connection strength (i.e., Z) is more beneficial on JK-Net.

Effects across Datasets
From the perspective of the dataset, GSEBO consistently performs better than the vanilla version on the four datasets. Specifically, the average improvement over the four across classic GNN architectures achieved by GSEBO are 1.9%, 3.1%, 5.3%, and 8.0% on the four datasets, respectively. Considering that the four datasets come from different scenarios, these results are evidence for the potential of GSEBO to be widely applied in practice. Note that the trend of performance improvement is similar to the density of graph, where Air-USA is the most dense graph with the largest performance improvement. As the number of neighbors increases, the percentage of neighbors essential for the classification of the target node will decrease. This result can reflect the rationality of optimizing the connection strength according to the overall classification objective.
Moreover, the training loss and test accuracy of GNN methods and the proposed GSEBO are shown in Figure 4. We observe that the loss of GSEBO tend to be stable after 30 epochs and is smaller than GCN, GAT, and GraphSAGE, showing the empirical convergence of GSEBO. This is because simultaneously optimizing the parameters of W and Z would make the loss function more easily approach to the labels. However, the loss of JK-Net is smaller than GSEBO, which we attribute to that JK-Net is easier to trap in over-fitting.
GNN has achieved promising achievements with the assumption of homophily, while the noise of graphs breaks the assumption and hinders the node representation of GNN.

Robustness Analysis
We investigate the robustness of GSEBO under the different inter-class levels. In particular, we follow [13] and construct graphs based on Cora by randomly adding 1000, 3000, 5000, 10,000, and 20,000 inter-class edges, respectively. On the synthetic datasets, we compare GSEBO with Vanilla, GAUG, DropEdge, and AdaEdge. Figure 5 shows the performance on the five GNN architectures. From the figures, we have the following observations: • The margins between GSEBO and vanilla GNN on the synthetic datasets are larger than the original Cora. For example, when adding 20,000 inter-class edges, GSEBO improves the accuracy by 17.6%, 3.1%, 40.6%, and 4.4% compared to GCN, GAT, GraphSAGE, and JKNet. This result indicates the robustness of GSEBO's structure learning ability. • In most cases, GSEBO outperforms the baselines at different noisy levels, which further justifies its robustness. • GAUG and AdaEdge utilizes different strategies to update the structure of the graph, which also consistently perform better than vanilla GNN. However, their gaps to GSEBO on the synthetic data are larger than the original Cora. We postulate that the reason for this is that their objectives are affected by the intentionally added noise. • DropEdge shows worse performances than the vanilla GNN on the synthetic datasets. The comparison shows that randomly dropping edges fails to enhance GNN robustness when the noisy level is high.

Visualization of Denoising (RQ3)
To investigate the effect of the connection strength matrix Z, we further visually compare the initialization of Z (i.e., initial weight) and its value after optimization (i.e., denoised weight). As both of them are large and sparse matrix, we visualize part of the denoised weight with a region of 100 × 100, which is shown on Figure 6. We show the results on Core and Air-USA, where the improvement of GSEBO is the smallest and largest, respectively. In the figures, we use red pixels and blue pixels to represent intraclass and inter-class edges, respectively, where a deeper color represents larger values. From Figure 6a,b, we find that the denoised weight matrix mainly decreases the value corresponding to inter-class connections, which, thus, can downweight the corresponds neighbor during representation propagation. In some cases, the denoised weight increases the weight of intra-class connections, which also facilitates denoising. Apparently, most edges in Air-USA are intra-class, and GSEBO still downweights a portion of inter-class edges and increases the intra-class weights, as shown in Figure 6c,d. This is because the denoised weight is learned in order to reach the overall classification objective.

Hyperparametric Analysis (RQ4)
We then investigate how the hyperparameters of GSEBO affect its effectiveness. As introduced in Section 3, the most critical hyperparameter of GSEBO is the training depth of inner optimization step (i.e., τ), which also is the number of gradient update in each inner optimization step. For the consideration of computing cost, we cap the value of τ at 25 and changes its value with step of 5, i.e., setting τ ∈ {5, 10, 15, 20, 25}. Figure 7 shows the performance of GSEBO applied to the four classical GNN architectures on the four datasets as changing the value of τ. From the figures, we have the following observations: (1) in most cases, the overall performance of the proposed GSEBO on each dataset is relatively stable as increasing the value of τ. This result indicates the insensitivity of GSEBO to the hyperparameter. (2) While increasing the value of τ can slightly improve the performance of GSEBO, we still suggest setting it as a relatively small value. This is because enlarging τ will increase the memory and computation cost of GSEBO.

Conclusions
In this work, we propose a novel GSEBO, which utilizes the generic structure extractor to extract the graph structure as learnable parameters and learns the structure from the global view. To better optimize the structure and common parameters, we decompose it into two mutually constrained optimization objectives, i.e., a bi-level optimization, with inner and outer optimization for common parameters and graph structure optimization, respectively. Extensive experiments demonstrate the effectiveness and robustness of GSEBO on both benchmark and synthetic datasets.
Even though GSEBO achieved impressive results, there are still some limitations: GSEBO cannot be effectively applied to large graphs, which requires implementation of mini-batches. Besides, we evaluate GSEBO in the transductive setting; when new nodes are added to the graph after training, GSEBO has to retrain the entire model. For future research, we would like to explore solutions to the above limitations.


-48-Dishes Recognition System Based on Deep Learning

: In recent years, food image recognition has emerged as a research hotspot in the field of computer vision. This paper proposes a dish recognition system based on image recognition and deep learning image processing and recognition, which can replace manual dish recognition. Because the food image is complex and dynamic, the system employs a Convolutional Neural Network (CNN). The dish detection algorithm has been completed using CNN. Multiple sets of comparative experiments were designed for the data before and after the improvement, and the ionfluence of the improvement points on the test data results was analyzed.

Introduction
With the rapid expansion of the mobile Internet's scale, sensing technology and artificial intelligence have gradually matured and made significant progress. In the last two years, the payment method in the cafeterias represented by college cafeterias has remained inefficient. This type of canteen has a high volume of people in a short period of time, is noisy and chaotic, does not print receipts, and has management difficulties. It is prone to long checkout lines and erroneous checkout links [1].
At the same time, a dinner plate is frequently required at the checkout. The scene of checking out with a card or mobile phone in one hand, the sound of the dinner plate falling on the floor one after the other during the peak period of the meal. In China, several "unmanned supermarkets" have emerged, improving people's shopping efficiency and eliminating long lines at checkout. With rising labor costs and a trend toward faster-paced living, it is clear that the service models of "self-service" and "machine instead of labor" will replace traditional service methods on a large scale.

Principles of Convolutional Neural Networks
Convolutional Neural Network (CNN) is a feed-forward neural network. Artificial neurons can respond to surrounding units and can perform large-scale image processing [2]. Convolutional neural networks include convolutional layers and pooling layers. Convolutional neural networks are MLPs (multilayer perceptrons) inspired by biological thinking. They have different category levels, and the working methods and functions of each layer are also different.

CNN Network Hierarchy
The CNN network has a total of 5 hierarchies as follows: •Input layer

Convolutional Layer
Local perception: In the process of the human brain recognizing a picture, it does not recognize the entire picture at once, but first perceives each feature in the picture locally, and then performs comprehensive operations on the parts at a higher level to obtain global information [3].
Through the local perception characteristics, the calculation parameters of the model are greatly reduced. But there are still many parameters just like this. This has the weight sharing mechanism: for each layer, the weights corresponding to all neurons should be equal, for example, the parameter vector of a neuron is [w1,w2,...,w100], Then other neurons in the same layer are also [w1,w2,...,w100], which is weight sharing.
Why do we need weight sharing? The connection parameters of neurons in the same layer are only related to the feature extraction and have nothing to do with the specific location, so it can be guaranteed that the connections of all locations in the same layer are weight-shared. For example, the first layer of the hidden layer is generally used for edge detection, the second layer is to combine the edge curves learned in the first layer to obtain some features, such as angle, line shape, etc.; the third layer will learn more complex features, such as: "Eyes, eyebrows," etc. For the same layer, they extract features in the same way. The neurons in the third layer are all used to extract the features of the "eye", so the parameters that need to be calculated are the same.

Incentive Layer
The so-called excitation is a nonlinear mapping of the output result of the convolutional layer.
If the activation function is not used (in fact, the activation function is equivalent to f(x)=x), in this case, the output of each layer is the linear function of the input of the previous layer. It is easy to conclude that no matter how many neural network layers there are, the output is a linear combination of inputs. The effect is the same as without hidden layers. This is the most primitive perception [4].
Commonly used excitation functions are: Suggestion from the incentive layer: ReLU first, because the iteration speed is fast, but the effect may not increase. If ReLU fails, consider using Leaky ReLU or Maxout. In this case, the general situation can be solved. The Tanh function has better results in text and audio processing.
The structure of DNN is not fixed. A general neural network includes an input layer, a hidden layer, and an output layer. A DNN structure has only one input layer, one output layer, and there are hidden layers between the input layer and the output layer. Each layer of the neural network has several neurons. The neurons between the layers are connected, the neurons in the layer are not connected, and the neurons of the next layer are connected to all the neurons of the previous layer.
A neural network with more hidden layers is called a deep neural network (the number of network layers in DNN does not include the input layer). The expressive power of a deep neural network is stronger than that of a shallow network. A neural network with only one hidden layer can fit Any function, but it requires many, many neurons.
 Advantages: Since DNN can fit almost any function, the nonlinear fitting ability of DNN is very strong. Often deep and narrow networks need to be more resource-efficient.
 Disadvantages: DNN is not easy to train, requires a lot of data, and a lot of skills to train a deep network.

Sensor
DNN can also be called multi-layer perceptron (MLP). The network structure of DNN is too complex and the number of neurons is too large. To facilitate the explanation, we design the simplest DNN network structure-perceptron.  Output: y=f(w * x+b)y=f(w * x+b)

Data Collection&Data Augmentation
Comparing to some classical data set including ILSVRC or CIFAR-10/100, we only need to consider a relatively small set of 6 types of cuisines. Basically as the data set for a neural network enlarged, the test performance will be more accurate. Here, we create a small data set that contains 54000 pictures to verify that our algorithm works as expected, which is a good start point to test the codes and make an efficient classifier for dishes. By using web spider and paragraphing on our own, we derive 300 pictures for every 6 cuisines and assign a value as their labels( ). Then, we rotate each picture to a certain angle to gain 9000 pictures per dish. A Gaussian filter is also employed to make the pictures smooth. To avoid aliasing, we preserve the RGB channels while resizing the pictures to 32*32 pixels. Last, we transfer the data to numpy.array with a shape (pic_numbers,32,32,3) and divide the data into train set and test set (ratio = 4:1). The label according to each single picture is transfered to the one-hot-encoding format to compare with the result of classifying later on..

The Architecture
After data collection, we define the model to extract features from those numerical data. Our model contains 6 layers, with 4 convolutional and 2 fully-connected. The first convolutional layer filters the 32*32*3 pictures with 32 kernels of 3*3*3. The other 3 convolutional layers all have the same input and Briefly describe our model as above, it is plain to see that our model has a relatively single size as a dishes classifier only need to distinguish 6 cuisines. In fact, we have tried pictures ranging from 16*16 to 256*256 and found that there will be a prompt of test accuracy from 16 to 32, with negligible gains thereafter. We also notice that the training time is only 2-3 minutes on the 54000*32*32*3 data set, but the time is about 10 times longer when it comes to the same pictures of128*128. So we employ the current model, which is inexpensive and competitive when put into commercial use, since the kinds of cuisines might change frequently and the data will be enlarged. Along with the small imput size, the convolutional layers also have small-scale stride as 1 pixel and kernels. In this way, depth is increased for better performance.
We also employ some techniques to solve the problem of overfitting which will be introduced below.

Activation Function
A standard model to activate an output of an neuron is the function f(x) = tanh(x)or f(x) = (1 + e −x ) −1 . However, the gradient of this function lies between (0 , 0.25), which indicates the gradient of the loss function has the multiply of these minor gradients by 4 times and becomes negligible to back propagation. The chain rule also implies that the different gradient on each layer will lead to a huge discrepancy on different depth. To solve this problem called vanishing gradient, we tried the ReLU function: f(x) = max(0, x) and finally adopted the LeakyReLU: f(x) = { ax, x < 0 x, x ≥ 0 . These functions are one-side saturating, which is better for feature extraction. To simplify, when kernels distinguish features from inputs, a high value represents high level of matching and a low value describes obscure, while a negative value is meaningless. They should be undiscriminating to those mismatching. In spite of this, we still introduce a coordinate to the negative values to prevent dying. The neurons of negative values are part of back propagation and if we set their gradients to 0, the neuron chains containing them will no more able to update.
Using LeakyReLU is fast and accurate. It is 6-7 times faster to use LeakyReLU than tanh(x) unints in our model. Also,when we changed the activation function ReLU, the accuracy will drop by an average 0.5%.

Dropout
When training with a small dataset, overfitting occurs. To prevent overfitting，introducing different classifiers and using the average is a good approach, since they will counterbalance the errors to some degree. Dropout is a method to train different models while less time-consuming than training traditionally. Instead, it hides each neuron in fully-connected layers with equal probability p=0.4 . When a neuron is dropped, it no longer acts in propagation. Dropout only appears at training, and the weights will be divided by (1-p) at test. This method also reduces co-adaptations of neurons. With dropout, two neurons may not occur simultaneously, and each neuron have to learn more robust feature, which exists at other model with different dropout pattern. All models share weights and update one after another.

Training and Testing
We train our model with a batch size of 64, and the data set will be shuffled after certain times of iteration. To initialize, the weights in each layer obey Gaussian distribution with mean=0 and standard deviation of 0.1. Also, the biases are all of constant 1 in every 6 layers. We add regularization to our loss function to prevent overfitting. To specify, the loss function is: Where loss is the mean of softmax cross entropy between logits and labels, γ is the regularization coordinate, W fci and b fci are weights and biases in the ith fully-connected layer.
Once the loss function is derived, we adapt adamoptimizer to update the weights, which is robust and outperforms other stochastic optimizers. The updating rule is:  Where α is the stepsize, β 1 and β 2 are exponential decay rates, ϵ is 1e − 8, and gt = ∇ θft(θt−1). To compare several optimizers, we record the accuracy and iteration steps in fig1. However, compare to the 100% accurate performance on training data, the accuracy of test is considered as weak--less than 85%. By reducing the iteration step to 1400, the model has approximately 94% accuracy of test data but with great fluctuation on both train and test data, since the steps to converge is different at each time. So we use the average accuracy of last 10 batches and set the threshold to 98%.


Entity-aware Transformers for Entity Search

Pre-trained language models such as BERT have been a key ingredient to achieve state-of-the-art results on a variety of tasks in natural language processing and, more recently, also in information retrieval.Recent research even claims that BERT is able to capture factual knowledge about entity relations and properties, the information that is commonly obtained from knowledge graphs. This paper investigates the following question: Do BERT-based entity retrieval models benefit from additional entity information stored in knowledge graphs? To address this research question, we map entity embeddings into the same input space as a pre-trained BERT model and inject these entity embeddings into the BERT model. This entity-enriched language model is then employed on the entity retrieval task. We show that the entity-enriched BERT model improves effectiveness on entity-oriented queries over a regular BERT model, establishing a new state-of-the-art result for the entity retrieval task, with substantial improvements for complex natural language queries and queries requesting a list of entities with a certain property. Additionally, we show that the entity information provided by our entity-enriched model particularly helps queries related to less popular entities. Last, we observe empirically that the entity-enriched BERT models enable fine-tuning on limited training data, which otherwise would not be feasible due to the known instabilities of BERT in few-sample fine-tuning, thereby contributing to data-efficient training of BERT for entity search.

INTRODUCTION
Pre-trained language models (LMs) such as BERT [11] and its successors [27,38,52] learn rich contextual information about words Permission to make digital or hard copies of part or all of this work for personal or classroom use is granted without fee provided that copies are not made or distributed for profit or commercial advantage and that copies bear this notice and the full citation on the first page. Copyrights for third-party components of this work must be honored. For all other uses, contact the owner/author(s). SIGIR '22, July 11-15, 2022 from large-scale unstructured corpora and have achieved intriguing results on a variety of downstream tasks in natural language processing (NLP) [11,25,26] and information retrieval (IR) [32,46]. It is even shown that these language models have the capability of capturing a tremendous amount of world knowledge, including information about real-world entities otherwise found in knowledge graphs (KGs) [35,42]. For example, language models can predict masked objects in cloze sentences such as "The native language of Mammootty is _____" and "The Sharon Cuneta Show was created in _____", where each of them demonstrates a ⟨subject, relation, object⟩ triple in a knowledge graph.
Language models, however, fail to perform complex reasoning about entities, as they cannot capture sparse world facts [23,35,39]. In information retrieval systems, users ask complex queries about entities, such as "What is the second highest mountain in the world?" or "Most famous civic-military airports. " Answering such queries requires leveraging rich human curated information about entities in the knowledge graphs. To bridge the gap between LMs and KGs, recent works enhance language models with rich structured entity information from knowledge graphs, showing superior performance for knowledge-driven NLP tasks such as relation classification, entity typing, and cloze-style question answering [23,34,37,43,45,50,54]. Despite this success, no previous study has examined the effect of entity-enriched language models for answering entity-oriented queries in IR.
This work explores the benefit of enriching BERT-based retrieval models with auxiliary entity information from KGs for the entity retrieval task, where users' queries are better answered with entities rather than passages or documents [2]. It is shown that entity information improves performance of document and entity retrieval tasks [9,15,18,47,48]. Yet, these studies were performed using traditional retrieval models, without utilizing LMs such as BERT.
The development of BERT-based models for entity retrieval faces a major challenge: there is limited training data for the entity retrieval task, and fine-tuning of BERT in data-constraint regimes leads to instabilities in model performance [13,36,53]. More specifically, the official entity retrieval dataset, DBpedia-Entity v2 [21], contains 467 queries; considered a small dataset for fine-tuning of a large neural language model. When fine-tuning BERT (especially its large variant) on a target task, the training data should be large enough that the model parameters get close to the ideal setting of the target task, or otherwise it causes forgetting of what the model has already learned.
Against this background, our first and primary research question is RQ1: Can an entity-enriched BERT-based retrieval model improve the performance of entity retrieval? To address this question, we propose an entity-enriched BERT-based retrieval model, EM-BERT, where factual entity information is injected into the monoBERT model [32] in the form of Wikipedia2Vec [49] entity embeddings that are transformed into the BERT's word piece vector space [37]; see Figure 1. The EM-BERT model is first trained with the MS MARCO passage dataset and further fine-tuned on the DBpedia-Entity v2 collection (in a cross validation setup). Our experiments indicate that our EM-BERT model improves state-of-the-art entity retrieval results by 11% with respect to NDCG@10. We also make the intriguing finding that while fine-tuning of the plain BERTbased retrieval model (monoBERT) on DBpedia-Entity v2 is brittle and prone to degenerate performance, the same fine-tuning process using the entity-enriched BERT model (EM-BERT) results in a stable model and brings consistent and meaningful improvements. We posit that this is attributed to direct injection of entity information into the BERT model, which brings the input distribution of the model close to the output label space, thereby contributing to dataefficient training of BERT-based models for entity-oriented tasks.
After observing the effectiveness of our EM-BERT model, we focus on understanding the whens and whys of entity-enriched BERT models for information retrieval tasks. Our next research question is: RQ2: When and which queries are helped by the EM-BERT model? We observe that EM-BERT mainly helps queries that are annotated with at least one entity; plain and entity-enriched BERT models perform on par with each other for queries without linked entities. We further examine queries with linked entities and find that EM-BERT is most helpful when entity mentions are broken into multiple word pieces by the BERT tokenizer (e.g. mention "Tuvalu" in Figure 1). This indicates that directly injecting entity information into the BERT-based models is particularly important for less popular entities, which BERT does not recognize as a single word (or phrase) because they are less observed during pre-training.
In our third question, we investigate RQ3: Why does the EM-BERT model work and what does it learn during the fine-tuning stage? We approach this by comparing the embeddings produced at the final layer of the EM-BERT network for both entities and their mention tokens. We find that the learned entity representations, unlike their corresponding word piece representations, form clear clusters in the embedding space, capturing the relation between entities. We also study some examples, and show that attention weights of entity tokens are larger than other tokens for queries requesting lists of entities or queries related to less known entities.
Finally, we question RQ4: How does our entity-enriched BERTbased model perform on other ranking tasks? To address this research question, we apply the EM-BERT model to the passage ranking task, where entity information is known to be less important and a large amount of training data can be used. We observe that the entity information that is embedded in the plain BERT-based retrieval model is enough for addressing this task and auxiliary entity information from knowledge graphs does not bring significant improvements.
In summary, this work makes the following contributions: • We study the effect of enriching BERT with entity information for information retrieval tasks and propose an entityenhanced BERT-based re-ranker EM-BERT.
• We establish new state-of-the-art results for the entity retrieval task on DBpedia-Entity v2 collection.
• We show that our entity-enriched model, unlike its equivalent model based on plain BERT, is robust against instabilities of BERT in data-constraint regimes and introduce EM-BERT as a data-efficient BERT-based model for entity retrieval.
• We perform thorough analysis of our EM-BERT model and add to our understanding of the model. We unfold when and for which queries this model work, what it learns, and how it works. The resources developed within the course of this paper are available at https://github.com/informagi/EMBERT.

RELATED WORK
Graph Embeddings. After the introduction of word embeddings methods like Word2Vec [29], neural embedding methods became increasingly popular. Word2Vec uses a single-layer neural network to either predict a word depending on its context (Continuous Bag of Word) or the context depending on a word (Skip-Gram). The resulting hidden layer of this network captures a representation of words and can be used for other downstream tasks. After the success of these, embedding methods were expanded to capture other types of knowledge as well, for example knowledge graphs. These so-called graph embeddings aim to embed the nodes of a knowledge graph in a way that entities and their relations are mapped accordingly into the vector space. One of the earliest introduced methods was TransE [4], which optimizes embeddings for head-relation-tail triples. The graph embedding method utilized in this paper is Wikipe-dia2Vec [49]. This method is similar to Word2Vec, but is extended to be used on text-heavy knowledge bases like Wikipedia. Instead of solely relying on words or edges between entities, it uses a combination of words, entities, and anchor texts. Using either Skip-Gram or Continuous Bag of Words method like in Word2Vec, it employs a single layer neural network to predict neighboring words to capture word similarities, entities relations (based on neighboring entities), and word-entity relations (based on surrounding anchor-text). Through this combination, words and entities get embedded into the same vector space.
Entity Retrieval. Graph embeddings have been used to assist in entity-oriented information retrieval. Gerritse et al. [16] uses Wiki-pedia2Vec embeddings for entity retrieval on the DBpedia-Entity v2 collection [21]. This is done by first tagging entities in the queries by using the TAGME entity linker [14]. Retrieved entities are then re-ranked by computing the cosine similarity score between the retrieved entities and the tagged entities, improving upon the state of the art for entity retrieval. KEWER [31] is another method for entity retrieval based on graph embeddings. It learns joint embeddings for both words and entities, combining this with entities annotated queries using SMAPH [7]. KEWER also improved previous state of the art for entity retrieval. The most recent graph-based entity ranking method, ENT Rank [12], utilizes entity-neighbor-text relations in a learning-to-rank model. ENT Rank incorporates a variety of entity features based on text, co-occurring entities, and entity neighbor relations, and performs competitively on entity ranking tasks. The paper, however, uses a different version of the DBpedia-Entity v2 collection and calculates the evaluation metrics only over judged entities, making it incomparable to our and the aforementioned entity retrieval methods.
The BEIR [40] benchmark provides a number of zero-shot retrieval models for several IR databases, including the DBpedia-Entity v2 entity retrieval collection. The best performing entity  Figure 1: Illustration of the EM-BERT model. Entity annotated query and documents are tokenized and mapped to their corresponding vector representations using and − functions.
retrieval model in BIER, BM25 + CE, re-ranks the top 100 entities retrieved by BM25, and then uses the MiniLM cross-encoder model [44], trained with the knowledge distillation setup provided by Hofstätter et al. [22].
Entity Linking. Utilizing entity information for retrieval models requires identifying entity information from documents and queries [3,17,19]. This is performed by annotating text with entity linking toolkits [6,7,14,20,41]. In this paper, we use the REL entity linker [41], which is the state-of-the-art open source entity linking toolkit. REL detects mentions using Flair [1] embeddings. REL performs candidate selection based on Wikipedia2Vec embeddings, and entity disambiguation based on latent relations between entity mentions in the text.
Transformer-based Rankers. Bidirectional Encoder Representations from Transformers (BERT) [11] is a language representation model, pre-trained to capture bi-directional relations between words. The pre-trained BERT model can be fine-tuned for various NLP tasks. BERT can be used for binary relevance classification by feeding it a sentence A and sentence B, and using the BERT classifier token (further denoted as [CLS]) to predict relevance. This setup has been shown to be effective for question answering.
BERT-based models are effective for information retrieval as well, and due to their relatively low number of input tokens became especially popular for passage retrieval. Several BERT-based methods have made enormous improvements on the previous state of the art [32,46]. MonoBERT [32] is a point-wise re-ranker, in which the BERT model is used as a binary relevance classifier. MonoBERT, at the time of release, obtained state-of-the-art results on the MS MARCO and TREC CAR datasets. After pre-training, the ranking is done by first making an initial ranking with a baseline method like BM25, and then re-ranking a query and passage by feeding a pair of sentences A (query) and B (passage) to BERT.
Entities and Transformers. During pre-training, BERT does not explicitly get structured data about entities. Some researchers have claimed that adding this may not be necessary since BERT captures this information implicitly [35,42]. Others set out to enrich BERT with entity information. One of the first works combining entities with transformers, ERNIE [54], is an enhanced language representation model, where the combination of entities and text is used to fine-tune the model. In KnowBert [34], contextual word embeddings are also enhanced by integrating information about entities. This is done by introducing a Knowledge Attention and Recontextualization component, which takes information in one layer of a BERT network, computes the enhanced entity information of that layer, and feeds it to the next layer. This method improves on BERT for entity-related tasks like entity linking. E-BERT [37] is another method to enhance BERT with entity information. The main point which distinguishes E-BERT from the previously mentioned work is that it does not require any additional training of the network; it only requires the computation of one large transformation matrix. This allows the method to be applied to any fine-tuned transformer model. Since Wikipedia2Vec embeds both words and entities in the same vector space, the embeddings of words in Wikipedia2Vec can be used to align a mapping. E-BERT has been shown to improve performance on unsupervised QA tasks like LAMA, as well as downstream tasks like relation classification.
Other works combining language models and knowledge graphs include [5], which enhances BERT with entities by using BERT itself for entity linking and seeing how much information of knowledge graphs is already contained in BERT. Next, BERT classifies texts to predict whether certain tokens belong to an entity. The authors show that specifically fine-tuning BERT on entity information greatly improves their scores on entity linking, thus showing that pre-trained BERT does not embed all information about entities yet, claiming that additional entity information does not help if either entities are too scarce in the data or if the task does not require entity knowledge.
In this section, we first provide a brief overview of knowledge graph embeddings and Wikipedia2Vec (Section 3.1), and then describe the EM-BERT model, which is an entity-enhanced version of a BERTbased retrieval model (Section 3.2), followed by a description of EM-BERT query-document input representation.

Background
Knowledge graph embeddings provide vector representations of entities in a knowledge graph, projecting entity properties and relations into a continuous vector space. These embeddings have proven effective not only for entity-related tasks such as entity linking [51], but also for general tasks such as question answering [31]. Knowledge graph embeddings can be constructed purely based on graph topology (i.e., entities and their relations) [4], or combining graph topology with additional textual information such as the entity descriptions [51]. In this paper, we use Wikipedia2Vec embeddings [49], which extends the skip-gram model [29] with entity relations and descriptions from Wikipedia and maps words and entities into a shared vector space. Following [37], we transform Wikipedia2Vec entity embeddings into the BERT word piece embedding space and use them as input to the BERT model, similar to a BERT word piece vector.
Formally, Wikipedia2Vec embeds a list of words L and a list of entities L into a vector space R Wikipedia , where Wikipedia is dimensions of the Wikipedia embeddings. The embeddings are trained by fitting a one layer neural network with the loss: L = L + L + L , which consist of three cross-entropy loss functions: L : predicting context words of a given word. L : predicting neighboring entities of a given entity. L : predicting anchor text of an entity to combine words and entities in a single embedding space.
The weight layer of this neural network results in the function Wikipedia : L ∪ L → R Wikipedia , which embeds both words and entities in the Wikipedia vector space. BERT, on the other hand, tokenizes text into the word piece dictionary L . We use the lookup function : L → R to transform word pieces into the -dimensional native BERT word piece vector space, the resulting vector being the input for the BERT model.

EM-BERT
Our proposed model, referred to as EM-BERT, incorporates entity embeddings into a point-wise document ranking approach. In its essence, EM-BERT combines E-BERT [37] with monoBERT [32] and predicts a retrieval score for a query-document pair, each represented by a sequence of words and entities. The model takes as input the concatenation of the query tokens 1 , ..., and the document tokens 1 , ..., , where each token is either a BERT native word piece token or an entity token.
Entity-enriched BERT. The off-the-shelf BERT model only accepts its native word pieces. To incorporate entity embeddings into BERT, we need to align entity vectors with word piece vectors.
Following [37], this is done by a linear transformation of entity vectors to BERT-like vectors. Since L does not contain any entities, the Wikipedia2Vec word dictionary L is used to obtain the linear transformation W ∈ R × Wikipedia , learned from the Wikipedia2Vec word vector space Wikipedia [L ] and BERT word piece space [L ]: Here, the intersection between L and L is taken to ensure that the selected words have embeddings in both Wikipedia and BERT vector spaces. The equation computes matrix W, in which the distance between W · Wikipedia ( ) and ( ) is minimal for all . Using W, we can then construct the function − , which maps both entities and word tokens to the BERT input vector space: The aligned entity vectors are then fed into BERT when an entity is mentioned in the input text. The entities are obtained by annotating the text with an entity; e.g., given the text "Who produced films starring Natalie Portman", the annotated sentence becomes "Who produced films starring Natalie Portman ENTITY/NataliePortman". This text is then tokenized as "who produced films starring natalie port ##man /ENTITY/Natalie_Portman," where the entity ID (in bold face) is a single token, embedded by − (Eq. 2). We note that the transformation matrix W can be fit into any pre-trained or fine-tuned BERT model, thus making E-BERT embeddings usable for every available BERT model.

Retrieval.
The retrieval phase is based on monoBERT [33], which is a multi-stage ranking method with BERT and has been shown to achieve competitive results on MS MARCO passage retrieval [30] and TREC-CAsT [8]. The simplicity of this model, coupled with its high effectiveness enables us to discern the effect of entity information on BERT-based retrieval models and obtain an understanding of the whys and whens of entity-enriched BERT models for information retrieval tasks.
In the multi-stage ranking method, for a query , documents are ranked with an initial ranking method; e.g., BM25. The topdocuments, denoted as = { 1 , 2 , . . . , }, are then passed to a second ranker for re-ranking; here a BERT-based retrieval model. Every query-document pair ⟨ , ∈ ⟩ is passed to BERT as two sentences and , with the separation token [SEP] in between. The BERT [CLS] classification vector is used as input to a single layer neural network to obtain the probability of the document being relevant to the query . To train this model, the cross-entropy loss is used: where and are indexes of relevant and non-relevant documents for all queries.
Putting all the pieces together, the training process of the EM-BERT model is summarized as follows. First, the transformation matrix W is trained using Wikipedia2Vec embeddings and BERT word pieces. All queries and documents are annotated with an entity linker, and the tagged queries and documents are tokenized and mapped to the corresponding vector representations using function − . This input is fed into the EM-BERT model and the model is first fine-tuned on MS MARCO passages and then on the entity retrieval collection.
Query-Document Representation. In the EM-BERT model (similar to monoBERT), queries are truncated to have a maximum token length of 64, and the combination of the query, passage, and separation tokens are truncated to have a maximum length of 512 tokens. Following [37], we concatenate mentions and entities, separated by the token '/'; i.e., "mention + '/' + EntityID.

EXPERIMENTAL SETUP
To address our research question, we compare the EM-BERT model with several baselines on the entity retrieval task and analyze the results. Due to limited training data for entity retrieval, we perform two-stage fine-tuning: we first fine-tune our models on the annotated MS MARCO passage dataset, and then continue fine-tuning on DBpedia-Entity v2 collection. In the following, we describe our experimental setup.

Entity Linking and Embeddings
We use the REL entity linking toolkit [41] to annotate texts with entities. REL is known to have high precision compared to other entity linking toolkits such as GENRE [6] and TAGME [14], ensuring that limited noise (incorrectly annotated entities) is added to the model. REL incorporates Wikipedia2Vec embeddings trained on the Wikipedia link graph. To avoid a mismatch between annotated entities and Wikipedia2Vec entity embeddings [12,16], we follow the REL Wikipedia2Vec training procedure, using the Wikipedia dump of 2019-07-01, setting min-entity-count parameter to zero, with a dimensionality of 500.

First Stage Fine-tuning
Collection. For the first stage fine-tuning of EM-BERT, we use the MS MARCO passage ranking collection, consisting of 8,841,823 passages, extracted from web documents retrieved by Bing. MS MARCO provides a training set with approximately 800k queries, a development set with 6980 queries with public relevance judgments, and an evaluation set with 6837 queries and private relevance judgments.
Although MS MARCO passage ranking is not an entity-related task, REL annotates 29% of queries in the development sets with at least one entity; e.g., "what were the brothers grimm names." The statistics of this collection (with REL annotations) are reported in Table 1. It shows that on average every query and document is annotated with 0.34 and 2.55 entities respectively, indicating that a fair amount of entities is present in this dataset.
Training. We fine-tune our EM-BERT model on the officially provided MS MARCO training set. This set provides triples of a query, a positive example, and a negative example, making the ratio of positive-negative documents equal during training. We refer to these triples as "sample." Starting from the monoBERT model provided by HuggingFace, 1 we fine-tune EM-BERT with randomly selected 300k samples (cf. §3.2); fine-tuning with more samples, up until 900k, did not give any further improvements on the development set. We use an accumulated batch size of 64, a learning rate of 10 −6 , 40k warm-up steps, and AdamW as optimizer. Fine-tuning for 300k samples took around 34 hours on a single GeForce RTX 3090 with 24 GB memory.

Second Stage Fine-tuning
Collection. For the second stage fine-tuning, we use the standard entity retrieval collection, DBpedia-Entity v2 [21]. The collection is based on DBpedia 2015-10 and consists of four different categories of queries: (i) SemSearch: named entity queries looking for particular entities; e.g., "mario bros", (ii) INEX-LD: keywordstyle entity-oriented queries; e.g., "bicycle benefits health", (iii) ListSearch: queries asking for a list of entities; e.g., "Hybrid cars sold in Europe", and (iv) QALD-2: natural language questions about entities; e.g., "Who produces Orangina?" Relevance judgements in this dataset are graded: Highly relevant (2), Relevant (1), and irrelevant (0). Documents in the DBpedia-Entity v2 collection are entities, each represented by an abstract and other metadata (stored in RDF format). These entity abstracts provide concise descriptions about entities and are comparable in length to MS MARCO passages. Both queries and entity abstracts are highly entity-oriented, reflected by the average number of entity annotations for queries and abstracts in Table 1.
Training. We consider entity abstracts as entity descriptions and use them as documents in monoBERT and EM-BERT. Since the DBpedia-Entity v2 collection provides a limited amount of training data (only 467 queries), we take the fine-tuned model on the MS MARCO passage dataset (Section 4.2) and continue fine-tuning for both monoBERT and EM-BERT models using the provided folds for 5-fold cross validation, each fold containing 48k query-entity pairs Table 2: Results on DBpedia-Entity v2 collection. Superscripts 1/2/3/4 denote statistically significant differences (better or worse) compared to base method/monoBERT (1st)/monoBERT/EM-BERT (1st), respectively. The highest value per column is marked with underline and in bold. The highest value per column and per block is marked with underline.

NDCG
SemSearch INEX-LD ListSearch QALD-2 Total @10 @100 @10 @100 @10 @100 @10 @100 @10 @100 ESIM [16] 0 In this process, both models are fine-tuned for one epoch using all training data of each fold, resulting in 5 different models each. The fine-tuned models are then used for ranking the test queries in the corresponding folds. Note that in this fine-tuning process, no hyper-parameter optimization is performed. In this process, both models are fine-tuned for one epoch using the training data of each fold, resulting in 5 different models each. The fine-tuned models are used for ranking the test queries in the corresponding folds. Note that in this fine-tuning process, no hyper-parameter optimization is performed. For both models an aggregated batch size of 64, learning-rate of 10 −6 , and AdamW optimizer is used. Warmup step is changed to 4000, following the 10% warmup step rule.

Baselines
We compare EM-BERT with the state-of-the-art entity retrieval methods using term-matching, neural, and BERT-based approaches. The baselines are: BM25: The BM25 run, based on the short abstracts of entities, with parameters = 0.9 and = 0.4. Similar to monoBERT and EM-BERT, this run uses only abstract information of entities (unlike other methods that use reach representation of entities). BM25F-CA [21]: The fielded variant of BM25, obtaining scores from 5-field representation of entities. Field weights are computed using Coordinate Ascent on each train fold, and then scored on the corresponding test fold. This is the best non-neural run and also the best baseline reported in [21]. ESIM [16]: ESIM ranks entities based on the aggregated cosine similarity scores between the Wikipedia2Vec embedding of each linked entity in the query and the target entity. Entities in this method are linked using TAGME [14]. The method re-ranks the BM25F-CA run. GEEER [16]: GEEER linearly combines the BM25F-CA score and entity-query similarities based on ESIM . The weights of this linear combination are computed on each train fold using Coordinate Ascent, and scored accordingly on each test fold. Similar to ESIM , GEEER re-ranks the BM25F-CA run. KEWER [31]: KEWER uses an embedding method with joint word and entity embeddings. A query-entity score is computed on all query terms by the weighted sum of the similarity between the target entity embedding and the term, re-ranking the BM25F-CA run. BM25F+KEWER [31]: The model linearly combines the KEWER score as listed above with BM25F-CA and re-ranks the BM25F-CA run. BLP-TransE [10]: BLP-TransE is a BERT-based method, combining a BERT cross-encoder architecture with TransE graph embeddings. It encodes queries and entities, uses the similarity between these encodings as query score, and then re-ranks and linearly combines the BM25F-CA scores. monoBERT: Our Pygaggle implementation of the monoBERT model. This implementation achieves slightly better performance compared to results reported in [33]: MRR@10 of 0.379 vs 0.372 on the MS MARCO development set. Following the multi-stage reranking method [33], we re-rank three entity retrieval runs using monoBERT: BM25, BM25F-CA, and GEEER (the best baseline run). The re-ranking is performed on the top-1000 entities of the BM25 and BM25F-CA, and top-100 entities of the GEEER run (as GEEER re-ranks the BM25F-CA run).

Evaluation Metrics
For evaluation of entity retrieval, we use the commonly reported metric in [21]; the Normalized Discounted Cumulative Gain (NDCG) at ranks 10 and 100. Statistical significant differences of NDCG@10 and NDCG@100 values are determined using the two-tailed paired t-test with p-value < 0.05.

RESULTS AND ANALYSIS
In the following, we evaluate our entity-enriched BERT-based retrieval model and answer our four research questions (Sections 5.1-5.4 ) listed in Section 1.

Entity Retrieval Results
In this section, we answer our first research question: RQ1: Can an entity-enriched BERT-based retrieval model improve the performance of entity retrieval?
We compare retrieval performance of EM-BERT with a variety of models on the DBpedia-Entity v2 collection in Table 2. The baseline (cf. Section 4.4) results are presented in the top part of the table. The next three compartments in the Table summarise results of re-ranking BM25, BM25F-CA, and GEEER runs with the following models: • monoBERT (1st): The monoBERT model after first stage finetuning; cf. Section 4.4.
• EM-BERT (1st): The EM-BERT model fine-tuned on the MS MARCO passage collection. Re-ranking setup is similar to the monoBERT runs; cf. Section 4.4.
• EM-BERT: The EM-BERT model after second stage fine-tuning. Setup is similar to the previous run.
As shown in Table 2, re-ranking of the GEEER run with EM-BERT markedly outperforms all baselines and establishes a new state-of-the-art result on the DBpedia-Entity v2 collection, with 11% improvements over the best run. When re-ranking the BM25 and BM25F-CA runs, we observe the same trend that EM-BERT outperforms all the corresponding BERT-based baselines. Comparing different query categories, EM-BERT improves over all runs by a large margin for ListSearch and most importantly for QALD-2 queries; improvements over GEEER are 20% and 24% with respect to NDCG@10 for ListSearch and QALD-2 queries, respectively. QALD queries are complex natural language queries that are hard to answer for most neural and non-neural models, reflected by the lowest NDCG scores on all baseline models. Table 2 also shows that monoBERT performance is decreased when fine-tuned with a limited amount of training data. The drop in the performance of monoBERT is expected due to the known instabilities of BERT with few-sample training [53], which can cause forgetting of what the model has already learned. The striking observation, however, is that EM-BERT results are improved when trained on the same data and using the same procedure. This indicates that entity-enriched BERT models enable fine-tuning on limited training data, thereby contributing to data-efficient training of BERT for entity-oriented tasks.
Based on these results, we can positively answer our first research question: our entity-enriched BERT-based model significantly improves state-of-the-art results on entity retrieval, and more interestingly, it is robust against instabilities of BERT when fine-tuned with limited training data.

INEX_LD-2010004
Indian food INEX_XER-127 german female politicians QALD2_te- 39 Give me all companies in Munich QALD2_tr- 26 Which bridges are of the same type as the Manhattan Bridge?

Query Analysis
In the following we answer our research question RQ2: When and which queries are helped by the EM-BERT model?
Entity annotations of queries. Knowing that not all queries in DBpedia-Entity v2 have a linked entity, we expect EM-BERT to perform better on queries with linked entities than on queries without them. To study this hypothesis, we compare NDCG@10 for two subsets of DBpedia-Entity v2 queries: with and without at least one linked entity. Table 3 shows the results. We see substantial improvements of EM-BERT over monoBERT and monoBERT (1st) for queries with linked entities, confirming our hypothesis.
Tokenization of entity mentions. To further our understanding of helped queries, we investigate the difference in EM-BERT performance for various types of linked queries by looking into their mentions tokenization. The BERT tokenizer splits words into word pieces, which can be categorized into beginning word pieces, indicating either the start of a word or a full word, and middle/end word pieces starting with a '##'. BERT employs around 30k tokens, consisting of around 24k beginning tokens and 6k middle/end tokens. Given that Wikipedia has around 12 million distinct entities, we expect only the most common entities to be present as a single word piece. We hypothesize that there is a performance difference between these very common entities, which are often seen in their natural form during the BERT pre-training process, and the rarer entities, where their mentions are broken into several word pieces. We, therefore, divide entity mentions into four categories: • One token: Mentions that are considered as a single token by the BERT tokenizer; e.g., "France" is tokenized to 'France'.
• Multiple tokens, no ##: Mentions that are tokenized to more than one token, where each token represents a single unbroken word; e.g., "Yoko Ono" is tokenized to 'Yoko', 'Ono'.
• One ##: Mentions that are split into multiple word pieces, of which exactly one is a middle/end word piece, e.g., "Weser" is be tokenized to 'wes', '##er' • Multiple ##: Mentions which gets split into multiple word pieces, of which more than one is a middle/end word piece, for example "Frisian" is tokenized to 'fr', '##isi', '##an'.
We categorize each query with a linked entity in one of the previously mentioned categories, prioritizing the most complicated entity; i.e., if a query has two mentions belonging to "One token" and "Multiple ##" categories, the query is categorized as "Multiple ##". Figure 2 compares the performance of monoBERT and EM-BERT models for these four mention tokenization categories. It shows that gains of EM-BERT are more pronounced for queries in "Multiple ##" and "One ##" categories than other categories. This supports our claim that queries with uncommon entity mentions are most helped by including their entity embeddings in the EM-BERT model. Additionally, we illustrate the relation between different tokenization categories and the query types mention (cf. Section 4.3) in Figure 3. The plot shows that 51% of INEX-LD queries have no linked entities, explaining marginal improvements of EM-BERT for these queries. This plot also describes why EM-BERT performs best on ListSearch and QALD-2 queries: most of the less known entities fall in "Multiple ##" and "One ##" categories.
These results provide an answer to RQ2: EM-BERT helps queries that are linked to at least one entity. Additionally, entity information provided by the EM-BERT model is most useful for less popular entities, which BERT cannot treat their mentions as a single token.

Model Analysis
To investigate RQ3: Why does the EM-BERT model work and what does it learn during the fine-tuning stage?, we examine the proximity of the embedding tokens compared to the mention tokens in the final layer of the BERT network, and also discuss several examples of helped queries.
Entities vs. mention embeddings. To get an understanding of what BERT learns from entity information, we compare entity and mention embeddings in the final layer of EM-BERT. We randomly choose eight queries with 100-200 relevant entities. For each of those documents, we take the first four linked entities and extract their entity and mention embeddings from the final layer of EM-BERT. The mention embeddings are taken from the first starting token before that entity. We use UMAP [28] with the default settings to plot these embeddings in 2-dimensional space; see Figure 4. As is evident from the plots, the entity embeddings form more pronounced clusters than the entity mention embeddings. For example, in the entity embedding plot, American bridges and rivers build a cluster together with other entities based in the US or Canada, while being close to the corresponding queries. This also indicates that the entity information injected through entity embeddings prevails into the final layers of the BERT model, keeping a better representation of entity information than provided with the default BERT word piece tokens.
Insightful Examples. We discuss in detail a few (hand-picked) example queries that are especially helped or hurt by re-ranking with the EM-BERT model. Table 4 visualizes the attention weights of the [CLS] token for the first attention head on the first layer, for sample query-entity pairs. In the query "Give me all movies directed by Francis Ford Coppola", monoBERT incorrectly ranks other filmmaker members of the Coppola family, matching the last name only. EM-BERT, in contrast, finds the correct entity, as movies directed by Francis Ford Coppola are situated close to him in the embedding space (as opposed to the movies by his daughter Sofia Coppola). Table 4 shows that EM-BERT gives high attention to the entity tokens relevant to the entity and query. It also shows high attention to the words "movies", which is essential here.
We can also see that through fine-tuning on DBpedia-entity, EM-BERT learns to retrieve specific entities for queries requesting a list. For example for the query "Give me the capitals of all countries in Africa", while monoBERT's highest ranked entities are general overview pages like 'List of African dependencies', all of EM-BERT's highest-ranked entities are capitals of African countries like 'Dakar' and 'Porto-Novo'.

Performance on Passage Retrieval
Finally, we turn to our last research question, RQ4: How does our entity-enriched BERT-based model perform on other ranking tasks? To further our understanding on the effect of entity-enriched BERT for IR tasks, we compare the monoBERT and EM-BERT models after the first fine tuning stage on the MS MARCO collection. We observe no significant differences between the two models for both queries with and without linked entities. A similar observation is made for the TREC Conversation Assistant Track (TREC-CAsT) [24]. We posit that the added value of enhancing BERT-based retrieval models with explicit entity representations is mainly pronounced for entity-oriented tasks.

CONCLUSION AND FUTURE WORK
In this research, we investigated the value of adding entity embeddings to BERT for entity search. We proposed an entity-enhanced BERT-based retrieval model, EM-BERT, and compared it with a regular BERT model. We found that EM-BERT improves substantially over the state-of-the-art results, showing that entity-oriented tasks benefit from entity embeddings in BERT, as is hypothesized in [5]. We found three categories where entities help the most: complex natural language queries, queries requesting a list of entities with a certain property, and queries with less known entities, in which their mentions get split in multiple tokens by the BERT tokenizer. We also showed that EM-BERT, unlike BERT-based retrieval models, can be fine-tuned with limited training data, thus being effective for lesser-resourced entity-related tasks.
For further research, the effect of the entity linking method can be investigated; here we used REL, which is an entity linker with high precision. It would be also interesting to evaluate the performance of other dense and sparse retrieval methods on entityoriented queries and tasks. This, however, may require additional training data for fine-tuning on entity-oriented queries, which would create a supplementary MS Marco subset to train and evaluate ranking systems.