summaries: In this work, the authors suggest a new point of view on generalization through the lens of the distribution of the per-sample gradients. The authors consider the variance and mean of the per-sample gradients for each parameter of the model and define for each parameter the Gradient Signal to Noise ratio (GSNR). The GSNR of a parameter is the ratio between the mean squared of the gradient per parameter per sample (computed over the samples) and the variance of the gradient per parameter per sample (also computed over the samples). The GSNR is promising as a measure of generalization and the authors provide a nice leading order derivation of the GSNR as a proxy for the measure of the generalization gap in the model. After the derivation, experimental results on MNIST are presented and suggest that empirically there is a relation between the generalization gap of neural network trained by gradient descent and the GSNR quantity given in the paper. Next the author analyze the GNSR of DNNs as opposed with shallow models or other learning techniques and observe that the GSNR differs when using random labels (lower GSNR) as compared with true labels and exhibits different behavior along training for DNNs and gradient descent.

summaries: This paper introduces a quantity termed the "one-step generalization ratio" . They derive approximate relations between OSGR and GSNR then show experimental results demonstrating the validity of these approximations, thus linking GSNR and a quantity related to generalization. They investigate the empirical value of GSNR during training of a neural network on Cifar10 with real labels vs. random labels. The final section derives a relation which attempts to explain the correlation between the size of the expected gradient and the learning of features.

abstractText: As deep neural networks (DNNs) achieve tremendous success across many application domains, researchers tried to explore in many aspects on why they generalize well. In this paper, we provide a novel perspective on these issues using the gradient signal to noise ratio (GSNR) of parameters during training process of DNNs. The GSNR of a parameter is defined as the ratio between its gradient’s squared mean and variance, over the data distribution. Based on several approximations, we establish a quantitative relationship between model parameters’ GSNR and the generalization gap. This relationship indicates that larger GSNR during training process leads to better generalization performance. Moreover, we show that, different from that of shallow models (e.g. logistic regression, support vector machines), the gradient descent optimization dynamics of DNNs naturally produces large GSNR during training, which is probably the key to DNNs’ remarkable generalization ability.

authors:
Madhu S Advani
Andrew M Saxe

title: High-dimensional dynamics of generalization error in neural networks

venue: arXiv preprint arXiv:1710.03667,

authors:
Sanjeev Arora
Rong Ge
Behnam Neyshabur
Yi Zhang

title: Stronger generalization bounds for deep nets via a compression approach, 2018

authors:
Peter L Bartlett
Shahar Mendelson

title: Rademacher and gaussian complexities: Risk bounds and structural results

venue: Journal of Machine Learning Research,

authors:
Laurent Dinh
Razvan Pascanu
Samy Bengio
Yoshua Bengio

title: Sharp minima can generalize for deep nets

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Gintare Karolina Dziugaite
Daniel M Roy

title: Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data

venue: arXiv preprint arXiv:1703.11008,

title: Agenerating sequences with recurrent neural networks, 2013. arXiv:1308.0850v5

authors:
Moritz Hardt
Benjamin Recht
Yoram Singer

title: Train faster, generalize better: Stability of stochastic gradient descent

venue: arXiv preprint arXiv:1509.01240,

authors:
Elad Hoffer
Itay Hubara
Daniel Soudry

title: Train longer, generalize better: closing the generalization gap in large batch training of neural networks

authors:
Daniel Jakubovitz
Raja Giryes
Miguel RD Rodrigues

title: Generalization error in deep learning, 2019

authors:
Kenji Kawaguchi
Leslie Pack Kaelbling
Yoshua Bengio

title: Generalization in deep learning

venue: arXiv preprint arXiv:1710.05468,

authors:
Nitish Shirish Keskar
Dheevatsa Mudigere
Jorge Nocedal
Mikhail Smelyanskiy
Ping Tak Peter Tang

title: On large-batch training for deep learning: Generalization gap and sharp minima

venue: arXiv preprint arXiv:1609.04836,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Behnam Neyshabur
Srinadh Bhojanapalli
David McAllester
Nati Srebro

title: Exploring generalization in deep learning

venue: In Advances in Neural Information Processing Systems,

authors:
Roman Novak
Yasaman Bahri
Daniel A Abolafia
Jeffrey Pennington
Jascha Sohl-Dickstein

title: Sensitivity and generalization in neural networks: An empirical study

authors:
Tom Rainforth
Adam R Kosiorek
Tuan Anh Le
Chris J Maddison
Maximilian Igl
Frank Wood
Yee Whye Teh

title: Tighter variational bounds are not necessarily better

venue: arXiv preprint arXiv:1802.04537,

title: Stiffness: A new perspective on generalization in neural networks, 2019

authors:
Vladimir N Vapnik
A Ja Chervonenkis

title: The necessary and sufficient conditions for consistency of the method of empirical risk

venue: Pattern Recognition and Image Analysis,

authors:
Chiyuan Zhang
Samy Bengio
Moritz Hardt
Benjamin Recht
Oriol Vinyals

title: Adadelta: An adaptive learning rate method

venue: Understanding

heading: 1 INTRODUCTION

text:
Deep neural networks typically contain far more trainable parameters than training samples, which seems to easily cause a poor generalization performance. However, in fact they usually exhibit remarkably small generalization gaps. Traditional generalization theories such as VC dimension (Vapnik & Chervonenkis, 1991) or Rademacher complexity (Bartlett & Mendelson, 2002) cannot explain its mechanism. Extensive research focuses on the generalization ability of DNNs (Neyshabur et al., 2017; Arora et al., 2018; Keskar et al., 2016; Dinh et al., 2017; Hoffer et al., 2017; Novak et al., 2018; Dziugaite & Roy, 2017; Jakubovitz et al., 2019; Kawaguchi et al., 2017; Advani & Saxe, 2017).
Unlike that of shallow models such as logistic regression or support vector machines, the global minimum of high-dimensional and non-convex DNNs cannot be found analytically, but can only be approximated by gradient descent and its variants (Zeiler, 2012; Kingma & Ba, 2014; Graves, 2013). Previous work (Zhang et al., 2016; Hardt et al., 2015; Dziugaite & Roy, 2017) suggests that the generalization ability of DNNs is closely related to gradient descent optimization. For example, Hardt et al. (2015) claims that any model trained with stochastic gradient descent (SGD) for reasonable epochs would exhibit small generalization error. Their analysis is based on the smoothness of loss function. In this work, we attempt to understand the generalization behavior of DNNs through GSNR and reveal how GSNR affects the training dynamics of gradient descent. Stanislav Fort (2019) studied a new gradient alignment measure called stiffness in order to understand generalization better and stiffness is related to our work.
The GSNR of a parameter is defined as the ratio between its gradient’s squared mean and variance over the data distribution. Previous work tried to use GSNR to conduct theoretical analysis on deep learning. For example, Rainforth et al. (2018) used GSNR to analyze variational bounds in
∗corresponding author
unsupervised DNNs such as variational auto-encoder (VAE). Here we focus on analyzing the relation between GSNR and the generalization gap.
Intuitively, GSNR measures the similarity of a parameter’s gradients among different training samples. Large GSNR implies that most training samples agree on the optimization direction of this parameter, thus the parameter is more likely to be associated with a meaningful “pattern” and we assume its update could lead to a better generalization. In this work, we prove that the GSNR is strongly related to the generalization performance, and larger GSNR means a better generalization.
To reveal the mechanism of DNNs’ good generalization ability, we show that the gradient descent optimization dynamics of DNN naturally leads to large GSNR of model parameters and therefore good generalization. Furthermore, we give a complete analysis and a detailed interpretation to this phenomenon. We believe this is probably the key to DNNs remarkable generalization ability.
In the remainder of this paper we first analyze the relation between GSNR and generalization (Section 2). We then show how the training dynamics lead to large GSNR of model parameters experimentally and analytically in Section 3.

heading: 2 LARGER GSNR LEADS TO BETTER GENERALIZATION

text: In this section, we establish a quantitative relation between the GSNR of model parameters and generalization gap, showing that larger GSNR during training leads to better generalization.

heading: 2.1 GRADIENTS SIGNAL TO NOISE RATIO

text:
Consider a data distribution Z = X × Y , from which each sample (x, y) is drawn; a model ŷ = f(x, θ) parameterized by θ; and a loss function L.
The parameters’ gradient w.r.t. L and sample (xi, yi) is denoted by
g(xi, yi, θ) or gi(θ) := ∂L(yi, f(xi, θ))
∂θ (1)
whose j-th element is gi(θj). Note that throughout this paper we always use i to index data examples and j to index model parameters.
Given the data distribution Z , we have the (sample-wise) mean and variance of gi(θ). We denote them as g̃(θ) = E(x,y)∼Z(g(x, y, θ)) and ρ2(θ) = Var(x,y)∼Z(g(x, y, θ)), respectively.
The gradient signal to noise ratio (GSNR) of one model parameter θj is defined as:
r(θj) := g̃2(θj)
ρ2(θj) (2)
At a particular point of the parameter space, GSNR measures the consistency of a parameter’s gradients across different data samples. Figure 1 intuitively shows that if GSNR is large, the parameter gradient space tends to be distributed in the similar direction and if GSNR is small, the gradient vectors are then scatteredly distributed.

heading: 2.2 ONE-STEP GENERALIZATION RATIO

text:
In this section we introduce a new concept to help measure the generalization performance during gradient descent optimization, which we call one-step generalization ratio (OSGR). Consider training set D = {(x1, y1), ..., (xn, yn)} ∼ Zn with n samples drawn from Z , and a test set D′ = {(x′1, y′1), ..., (x′n′ , y′n′)} ∼ Zn ′ . In practice we use the loss on D′ to measure generalization. For simplicity, we assume the sizes of training and test datasets are equal, i.e. n = n′. We denote the empirical training and test loss as:
L[D] = 1
n n∑ i=1 L(yi, f(xi, θ)), L[D ′] = 1 n n∑ i=1 L(y′i, f(x ′ i, θ)), (3)
respectively. Then the empirical generalization gap is given by L[D′]− L[D].
In gradient descent optimization, both the training and test loss would decrease step by step. We use ∆L[D] and ∆L[D′] to denote the one-step training and test loss decrease during training, respectively. Let’s consider the ratio between the expectations of ∆L[D′] and ∆L[D] of one single training step, which we denote as R(Z, n).
R(Z, n) := ED,D ′∼Zn(∆L[D
′])
ED∼Zn(∆L[D]) (4)
Note that this ratio also depends on current model parameters θ and learning rate λ. We are not including them in the above notation as we will not explicitly model these dependencies, but rather try to quantitatively characterize R for very small λ and for θ at the early stage of training (satisfying Assumption 2.3.1).
Also note that the expectation of ∆L[D′] is over D and D′. This is because the optimization step is performed on D. We refer to R(Z, n) as OSGR of gradient descent optimization. Statistically the training loss decreases faster than the test loss and 0 < OSGR(t) < 1 (Middle panel of Figure 2), which usually results in a non-zero generalization gap at the end of training. If OSGR(t) is large (≈ 1) in the whole training process (Right panel of Figure 2), generalization gap would be small when training completes, implying good generalization ability of the model. If OSGR(t) is small (= 0), the test loss will not decrease while the training loss normally drops (Left panel of Figure 2), corresponding to a large generalization gap.

heading: 2.3 RELATION BETWEEN GSNR AND OSGR

text:
In this section, we derive a relation between the OSGR during training and the GSNR of model parameters. This relation indicates that, for the first time as far as we know, the sample-wise gradient distribution of parameters is related to the generalization performance of gradient descent optimization.
In gradient descent optimization, we take the average gradient over training set D, which we denote as gD(θ). Note that we have used gi(θ) to denote gradient evaluated on one data sample and g̃(θ) to denote its expectation over the entire data distribution. Similarly we define gD′(θ) to be the average gradient over test set D′.
gD(θ) = 1
n n∑ i=1 g(xi, yi, θ) = ∂L[D] ∂θ , gD′(θ) = 1 n n∑ i=1 g(x′i, y ′ i, θ) = ∂L[D′] ∂θ (5)
Both the training and test dataset are randomly generated from the same distribution Zn, so we can treat gD(θ) and gD′(θ) as random variables. At the beginning of the optimization process, θ is randomly initialized thus independent of D, so gD(θ) and gD′(θ) would obey the same distribution. After a period of training, the model parameters begin to fit the training dataset and become a function of D, i.e. θ = θ(D), therefore distributions of gD(θ(D)) and gD′(θ(D)) become different. However we choose not to model this dependency and make the following assumption for our analysis:
Assumption 2.3.1 (Non-overfitting limit approximation) The average gradient over the training dataset and test dataset gD(θ) and gD′(θ) obey the same distribution.
Obviously the mean of gD(θ) and gD′(θ) is just the mean gradient over the data distribution g̃(θ). ED∼Zn [gD(θ)] = ED,D′∼Zn [gD′(θ)] = g̃(θ) (6)
We denote their variance as σ2(θ), i.e.
VarD∼Zn [gD(θ)] = VarD,D′∼Zn [gD′(θ)] = σ 2(θ) (7)
It is straightforward to show that:
σ2(θ) = VarD∼Zn [ 1
n n∑ i=1 gi(θ)] = 1 n ρ2(θ) (8)
where σ2(θ) is the variance of the average gradient over the dataset of size n, and ρ2(θ) is the variance of the gradient of a single data sample.
In one gradient descent step, the model parameter is updated by ∆θ = θt+1 − θt = −λgD(θ) where λ is the learning rate. If λ is small enough, the one-step training and test loss decrease can be approximated by
∆L[D] ≈ −∆θ · ∂L[D] ∂θ +O(λ2) = λgD(θ) · gD(θ) +O(λ2) (9)
∆L[D′] ≈ −∆θ · ∂L[D ′]
∂θ +O(λ2) = λgD(θ) · gD′(θ) +O(λ2) (10)
Usually there are some differences between the directions of gD(θ) and gD′(θ), so statistically ∆L[D] tends to be larger than ∆L[D′] and the generalization gap would increase during training. When λ → 0, in one single training step the empirical generalization gap increases by ∆L[D] − ∆L[D′], for simplicity we denote this quantity as ▽:
▽ := ∆L[D]−∆L[D′] ≈ λgD(θ) · gD(θ)− λgD(θ) · gD′(θ) (11) = λ(g̃(θ) + ϵ)(g̃(θ) + ϵ− g̃(θ)− ϵ′) (12) = λ(g̃(θ) + ϵ)(ϵ− ϵ′) (13)
Here we replaced the random variables by gD(θ) = g̃(θ) + ϵ and gD′(θ) = g̃(θ) + ϵ′, where ϵ and ϵ′ are random variables with zero mean and variance σ2(θ). Since E(ϵ′) = E(ϵ) = 0, ϵ and ϵ′ are independent, the expectation of ▽ is
ED,D′∼Zn(▽) = E(λϵ · ϵ) +O(λ2) = λ ∑ j σ2(θj) +O(λ 2) (14)
where σ2(θj) is the variance the of average gradient of the parameter θj .
For simplicity, when it involves a single model parameter θj , we will use only a subscript j instead of the full notation. For example, we use σ2j , rj , and gD,j to denote σ
2(θj), r(θj), and gD(θj) respectively.
Consider the expectation of ∆L[D] and ∆L[D′] when λ → 0 ED∼Zn(∆L[D]) ≈ λED∼Zn(gD(θ) · gD(θ)) = λ ∑ j ED∼Zn(g 2 D,j) (15)
ED,D′∼Zn(∆L[D ′]) = ED,D′∼Zn(∆L[D]−▽) (16) ≈ λ ∑ j (ED∼Zn(g 2 D,j)− σ2j ) (17)
= λ ∑ j (ED∼Zn(g 2 D,j)− ρ2j/n) (18)
Substituting (18) and (15) into (4) we have: R(Z, n) = 1− ∑ j ρ 2 j
n ∑
j ED∼Zn(g 2 D,j)
(19)
Although we derived eq. (19) from simplified assumptions, we can empirically verify it by estimating two sides of the equation on real data. We will elaborate on this estimation method in section 2.4.
We can rewrite eq. (19) as:
R(Z, n) = 1− 1 n ∑ j ED∼Zn(g 2 D,j)∑ j′ ED∼Zn(g 2 D,j′) ρ2j ED∼Zn(g2D,j)
(20)
= 1− 1 n ∑ j ED∼Zn(g 2 D,j)∑ j′ ED∼Zn(g 2 D,j′) 1 rj + 1 n
(21)
where ED∼Zn(g2D,j) = V arD∼Zn(gD,j) + E 2 D∼Zn(gD,j) = 1 nρ 2 j + g̃ 2 j .
We define ∆Lj [D] to be the training loss decrease caused by updating θj . We can show that when λ is very small ∆Lj [D] = λg2D,j +O(λ 2). Therefore when λ → 0, we have
R(Z, n) = 1− 1 n ∑ j Wj 1 rj + 1 n , where Wj := ED∼Zn(∆Lj [D]) ED∼Zn(∆L[D]) with ∑ j Wj = 1 (22)
Eq. (22) shows that the GSNR rj plays a crucial role in the model’s generalization ability—the onestep generalization ratio in gradient descent equals one minus the weighted average of 1
rj+ 1 n over all model parameters divided by n. The weight is proportional to the expectation of the training loss decrease resulted from updating that parameter. This implies that larger GSNR of model parameters during training leads to smaller generalization gap growth thus better generalization performance of the trained model. Also note when n → ∞, we have R(Z, n) → 1, meaning that training on more data helps generalization.

heading: 2.4 EXPERIMENTAL VERIFICATION OF THE RELATION BETWEEN GSNR AND OSGR

text:
The relation between GSNR and OSGR, i.e. eq. (19) or (22) can be empirically verified using any dataset if: (1) The dataset includes enough samples to construct many training sets and a large enough test set so that we can reliably estimate ρ2j , ED∼Zn(g 2 D,j) and OSGR. (2) The learning rate is small enough. (3) In the early training stage of gradient descent.
To empirically verify eq. (19), we show how to estimate its left and right hand sides, i.e. OSGR by definition and OSGR as a function of GSNR. Suppose we have M training sets each with size n, and a test set of size n′. We initialize a model and train it separately on the M training sets and test it with the same test set. For the t-th training iteration, we denote the training loss and test loss of the model trained on the m-th training dataset as L(m)t and L ′(m) t , respectively. Then the left hand
side, i.e. OSGR by definition, of the t-th iteration can be estimated by Rt(Z, n) ≈ ∑M m=1 L ′(m) t+1 − L′ (m) t∑M
m=1 L (m) t+1 − L (m) t
(23)
For the model trained on the m-th training set, we can compute the t-th step average gradient and sample-wise gradient variance of θj on the corresponding training set, denoted as gm,j,t and ρ2m,j,t, respectively. Therefore the right hand side of eq. (19) can be estimated by
ED∼Zn(g 2 D,j,t) ≈
1
M M∑ m=1 g2m,j,t, ρ 2 j,t ≈ 1 M M∑ m=1 ρ2m,j,t (24)
We performed the above estimations on MNIST with a simple CNN structure consists of 2 Conv-Relu-MaxPooling blocks and 2 fully-connected layers. First, to estimate eq. (24) with M = 10, we randomly sample 10 training sets with size n and a test set with size 10,000. To cover different conditions, we (1) choose n ∈ {1000, 2000, 4000, 6000, 8000, 10000, 15000}, respectively; (2) inject noise by randomly changing the labels with probability prandom ∈ {0.0, 0.1, 0.2, 0.3, 0.5}; (3) change the model structure by varying number of channels in the layers, ch ∈ {6, 8, 10, 12, 14, 16, 18, 20}. See Appendix A for more details of the setup. We use the gradient descent training (not SGD), with a small learning rate of 0.001. The left and right hand sides of 19 at different epochs are shown in Figure 3, where each point represents one specific choice of the above settings.
At the beginning of training, the data points are closely distributed along the dashed line corresponding to LHS=RHS. This shows that eq. (19) fits quite well under a variety of different settings. As training proceeds, the points become more scattered as the non-overfitting limit approximation no longer holds, but correlation between the LHS and RHS remains high even when the training converges (at epoch 2,500). We also conducted the same experiment on CIFAR10 A.2 and a toy dataset A.3 observed the same behavior. See Appendix for these experiments.
The empirical evidence together with our previous derivation of eq. (19) clearly show the relation between GSNR and OSGR and its implication in the model’s generalization ability.

heading: 3 TRAINING DYNAMICS OF DNNS NATURALLY LEADS TO LARGE GSNR

text:
In this section, we analyze and explain one interesting phenomenon: the parameters’ GSNR of DNNs rises in the early stages of training, whereas the GSNR of shallow models such as logistic regression or support vector machines declines during the entire training process. This difference gives rise to GSNR’s large practical values during training, which in turn is associated with good
generalization. We analyze the dynamics behind this phenomenon both experimentally and theoretically.

heading: 3.1 GSNR BEHAVIOR OF DNNS TRAINING

text:
For shallow models, the GSNR of parameters decreases in the whole training process because gradients become small as learning converges. But for DNNs it is not the case. We trained DNNs on the CIFAR datasets and computed the GSNR averaged over all model parameters. Because ED∼Zn(g 2 D,j) = 1 nρ 2 j + g̃ 2 j and we assume n is large, ED∼Zn(g 2 D,j) ≈ g̃2j . In the case of only one large training datasets, we estimate GSNR of t-th iteration by rj,t ≈ g2D,j,t/ρ2D,j,t (25)
As shown in Figure 4, the GSNR starts out low with randomly initialized parameters. As learning progresses, the GSNR increases in the early training stage and stays at a high level in the whole learning process. For each model parameter, we also computed the proportion of the samples with the same gradient sign, denoted as psame_sign. In Figure 4c, we plot the mean of time series of this proportion for all the parameters. This value increases from about 50% (half positive half negetive due to random initialization) to about 56% finally, which indicates that for most parameters, the gradient signs on different samples become more consistent. This is because meaningful features begin to emerge in the learning process and the gradients of the weights on these features tend to have the same sign among different samples.
Previous research (Zhang et al., 2016) showed that DNNs achieved zero training loss by memorizing training samples even if the labels were randomized. We also plot the average GSNR for model trained using data with randomized labels in Figure 4 and find that the GSNR stays at a low level throughout the training process. Although the training loss of both the original and randomized labels go to zero (not shown), the GSNR curves clearly distinguish between these two cases and reveal the lack of meaningful patterns in the latter one. We believe this is the reason why DNNs trained on real and random data lead to completely different generalization behaviors.

heading: 3.2 TRAINING DYNAMICS BEHIND THE GSNR BEHAVIOR

text:
In this section we show that the feature learning ability of DNNs is the key reason why the GSNR curve behavior of DNNs is different from that of shallow models during the gradient descent training. To demonstrate this, a simple two-layer perceptron regression model is constructed. A synthetic dataset is generated as following. Each data point is constructed i.i.d. using y = x0x1 + ϵ, where x0 and x1 are drawn from uniform distribution [−1, 1] and ϵ is drawn from uniform distribution [−0.01, 0.01]. The training set and test set sizes are 200 and 10,000, respectively. We use a very simple two-layer MLP structure with 2 inputs, 20 hidden neurons and 1 output.
We randomly initialized the model parameters and trained the model on the synthetic training dataset. As a control setup we also tried to freeze model weights in the first layer to prevent it from learning features. Note that a two layer MLP with the first layer frozen is equivalent to a linear regression model. That is, regression weights are learned on the second layer using fixed features extracted by the first layer. We plot the average GSNR of the second layer parameters for both the frozen and non-frozen cases. Figure 5 shows that in the non-frozen case, the average GSNR over parameters of
the second layer shows a significant upward process, whereas in the frozen case the average GSNR decreases in the beginning and remains at a low level during the whole training process.
In the non-frozen case, GSNR curve of individual parameters of the second layer are shown in Figure 5. The GSNR for some parameters show a significant upward process. To measure the quality of these features, we computed the Pearson correlation between them and the target output y, both at the beginning of training and at the maximum point of their GSNR curves. We can see that the learning process learns “good” features (high correlation value, i.e. with stronger correlation with y) from random initialized ones, as shown in Table 1. This shows that the GSNR increasing process is related to feature learning.

heading: 3.3 ANALYSIS OF TRAINING DYNAMICS BEHIND DNNS’ GSNR BEHAVIOR

text:
In this section, we will investigate the training dynamics behind the GSNR curve behavior. In the case of fully connected network structure, we can analytically show that the numerator of GSNR, i.e. the squared gradient mean of model parameters, tends to increase in the early training stage through feature learning.
Consider a fully connected network, whose parameters are θ = {W(1),b(1), ...,W(lmax),b(lmax)}, where W(1),b(1) are the weight matrix and bias of the first layer, and so on. We denote the activations of the l-th layer as a(l) = {a(l)s (θ(l−))}, where s is the index for nodes/channels of this layer, and θ(l−) is the collection of model parameters in the layers before l, i.e. θ(l−) = {W(1),b(1), ...,W(l−1),b(l−1)}. In the forward pass on data sample i, {als(θ(l−))} is multiplied by the weight matrix W(l):
o (l) i,c = ∑ s W (l)c,sa (l) i,s(θ (l−)) (26)
where o(l) = {o(l)i,c} is the output of the matrix multiplication, for the i-th data sample, on the l-th layer, c = {1, 2, ..., C} is the index of nodes/channels in the (l + 1)-th layer. We use g(l)D to denote the average gradient of weights of the l-th layer W(l), i.e. g(l)D = 1 n ∑n i=1 ∂Li ∂W(l)
, where Li is the loss of the i-th sample.
Here we show that the feature learning ability of DNNs plays a crucial role in the GSNR increasing process. More precisely, we show that the learning of features a(l)(θ(l−)), i.e. the learning of parameters θ(l−) tends to increase the absolute value of g(l)D . Consider the one-step change of gradient mean ∆g(l)D = g (l) D,t+1 − g (l) D,t with the learning rate λ → 0. In one training step, θ is updated by ∆θ = θt+1 − θt = −λgD(θ). Using linear approximation with λ → 0, we have
∆g (l) D,s,c ≈ ∑ j ∂g (l) D,s,c ∂θj ∆θj = ∑ θj∈θ(l−) ∂g (l) D,s,c ∂θj ∆θj + ∑ θj∈θ(l+) ∂g (l) D,s,c ∂θj ∆θj (27)
where θ(l−) and θ(l+) denote model parameters before and after the l-the layer (including the l-th), respectively.
We focus on the first term of eq. (27), i.e. the one-step change of g(l)D caused by learning θ (l−). Substituting g(l)D = 1 n ∑n i=1 ∂Li ∂W(l) and ∆θj = (−λ 1n ∑n i=1 ∂Li ∂θj ) into eq. (27), we have
∆g (l) D,s,c = −
λ
n2 ∑ θj∈θ(l−) W(l)s,c( n∑ i=1 ∂Li ∂o (l) i,c ∂a (l) i,s ∂θj )2 + other terms (28)
The detailed derivation of eq. (28) can be found in Appendix B. We can see the first term (which is a summation over parameters in θ(l−)) in eq. (28) has opposite sign with W(l)s,c. This term will make ∆g
(l) D,s,c negatively correlated with W (l) s,c. We plot the correlation between ∆g (l) D,s,c with W (l) s,c for a model trained on MNIST for 200 epochs in Figure 6a. In the early training stage, they are indeed negatively correlated. For top-10% weights with larger absolute values, the negative correlation is even more significant.
Here we show that this negative correlation between ∆g(l)D,s,c and W (l) s,c tends to increase the absolute value of g(l)D through an interesting mechanism. Consider the weights W (l) s,c with {W(l)s,c > 0,g (l) D,s,c < 0}. Learning θl− would decrease g (l) D,s,c and thus increase its absolute value because the first term in eq. (28) is negative. On the other hand, learning W(l)s,c would increase W (l) s,c and its absolute value because ∆W(l)s,c = −λg(l)D,s,c is positive. This will form a positive feedback process, in which the numerator of GSNR, (g(l)D,s,c) 2, would increase and so is the GSNR. Similar analysis can be done for the case with {W(l)s,c < 0,g(l)D,s,c > 0}.
On the other hand, when {W(l)s,cg(l)D,s,c > 0}, we show that the weights tend to change into the earlier case, i.e. {W(l)s,cg(l)D,s,c < 0} during training. Consider the case of {W (l) s,c > 0,g (l) D,s,c > 0}, the first term in eq. (28) is negative, learning θ(l−) tends to decrease g(l)D,s,c or even change its sign. Another posibility is that learning W(l)s,c changes the sign of W (l) s,c because ∆W (l) s,c = −λg(l)D,s,c is negative. In both cases the weights change into the earlier case with {W(l)s,cg(l)D,s,c < 0}. Similar analysis can be done for the case of {W(l)s,c < 0,g(l)D,s,c < 0}.
Therefore {W(l)s,cg(l)D,s,c < 0} is a more stable state in the training process. For a simple model trained on MNIST, We plot the proportion of weights satisfying {W(l)s,cg(l)D,s,c < 0} in Figure 6b and find that there are indeed more weights with {W(l)s,cg(l)D,s,c < 0} than the opposite. Because weights with small absolute value easily change sign during training, we also plot this proportion for the top-10% weights with larger absolute values. We can see that for the weights with large absolute values, nearly 80% of them have opposite signs with their gradient mean, confirming our earlier analysis. For these weights, the numerator of GSNR, (g(l)D,s,c)
2, tends to increase through the positive feedback process as discussed above.

heading: 4 SUMMARY

text: In this paper, we performed a series of analysis on the role of model parameters’ GSNR in deep neural networks’ generalization ability. We showed that large GSNR is a key to small generalization gap, and gradient descent training naturally incurs and exploits large GSNR as the model discovers useful features in learning.

heading: A APPENDIX A

text:
A.1 MODEL STRUCTURE IN SECTION 2.4
As shown in Table 2, all models in the experiment consist of 2 Conv-Relu-MaxPooling blocks and 2 fully-connected layers, but they are different in the number of channels. We choose the number of channels p from {6, 8, 10, 12, 14, 16, 18, 20}.
A.2 EXPERIMENT ON CIFAR10
Different from the experiment on MNIST, we use a deeper network on CIFAR10. We also include the Batch Normalization (BN) layer, because we find that it’s difficult for the network to converge in the absence of it. The network consists of 4 Conv-BN-Relu-Conv-BN-Relu-MaxPooling blocks and 3 fully-connected layers. More details are shown in Table 3.
The experiment is conducted under a similar setting as that of MNIST in section 2.4. We choose n ∈ {2000, 4000, 6000, 8000, 10000}, prandom ∈ {0.0, 0.2, 0.4}, ch ∈ {6, 8, 10, 12, 14, 16, 18}. We use the gradient descent training (Not SGD), with a small learning rate of 0.001. The left and right hand sides of 19 at different epochs are shown in Figure 7, where each point represents one specific combination of the above settings. Note that at the evaluation step of every epoch, we use
the same mean and variance inside the BN layers as the training dataset. That’s to ensure that the network and loss function are consistent between training and test.
At the beginning of training, compared to that of MNIST, the data points no longer perfectly resides on the diagonal dashed line. We suppose that’s beacuse of the presence of BN layer, whose internal parameters, i.e. running mean and running variance, are not regular learnable parameters in the optimization process, but change their values in a different way. Their change affects the OSGR, yet we could not include them in the estimation of OSGR. However, the strong positive correlation between the left and right hand sides of eq. (19) can always be observed until the training begins to converge.
A.3 EXPERIMENT ON TOY DATASET
In this section we show a simple two-layer regression model consists of a FC-Relu structure with only 2 inputs, 1 hidden layer with N neurons and 1 output. A similar synthetic dataset with the training data used in the experiment of Section 3.2 is generated as follows. Each data point is constructed i.i.d. using y = x0x1 + ϵ, where x0 and x1 are drawn from uniform distribution of [−1, 1] and ϵ is drawn from uniform distribution of [−ηnoise, ηnoise]. To estimate eq. (24), we randomly generate 100 training sets with n samples each, i.e. M=100, and a test set with 20,000 samples. To cover different conditions, we (1) choose n ∈ {50, 100, 300, 600, 1000, 2000, 6000}; (2) inject noise with ηnoise ∈ {0.2, 2, 4, 6, 8}; (3) perturb model structures by choosing N ∈ {6, 8, 10, 12, 14, 16, 18, 20}. We use gradient descent with learning rate of 0.001.
Figure 8 shows a similar behavior as Fig. 3. During the early training stages, the LHS and RHS of eq. (19) are very close. Their highly correlated relation remains until training converges, whereas the RHS of eq. (19) decreases significantly.

heading: B APPENDIX B

text:
Derivation of eq. (28)
∆g (l) D,s,c = ∑ θj∈θ(l−) ∂g (l) D,s,c ∂θj ∆θj + other terms (29)
= ∑
θj∈θ(l−)
∂( 1n ∑n i=1 ∂Li
∂W (l) s,c
)
∂θj (−λ 1 n n∑ i=1 ∂Li ∂θj ) + other terms (30)
= ∑
θj∈θ(l−)
∂( 1n ∑n i=1 ∂Li ∂o
(l) i,c
∂o (l) i,c
∂W (l) s,c
)
∂θj (−λ n n∑ i=1 ∑ s′,c′ ∂Li ∂o (l) i,c′ ∂o (l) i,c′ ∂a (l) i,s′ ∂a (l) i,s′ ∂θj ) + other terms (31)
= − λ n2 ∑ θj∈θ(l−)
∂( ∑n
i=1 ∂Li ∂o
(l) i,c
a (l) i,s)
∂θj ( n∑ i=1 ∑ s′,c′ ∂Li ∂o (l) i,c′ W (l) s′,c′ ∂a (l) i,s′ ∂θj ) + other terms (32)
= − λ n2 ∑ θj∈θ(l−) n∑ i=1 ( ∂Li ∂o (l) i,c ∂a (l) i,s ∂θj + ∂2Li ∂o (l) i,c∂θj a (l) i,s)( ∑ s′,c′ W (l) s′,c′ n∑ i=1 ∂Li ∂o (l) i,c′ ∂a (l) i,s′ ∂θj )
+ other terms (33)
Above we used ∂o
(l) i,c′
∂a (l) i,s′ = W
(l) s′,c′ and
∂o (l) i,c
∂W (l) s,c
= a (l) i,s that can both be derived from eq. (26). Consider
the first term of eq. (33). When s′ = s, c′ = c, we have
∆g(l)s,c = − λ
n2 ∑ θj∈θ(l−) W(l)s,c( n∑ i=1 ∂Li ∂o (l) i,c ∂a (l) i,s ∂θj )2 + other terms (34)
Note that the term related to ∂ 2Li
∂o (l) i,c∂θj
a (l) i,s and the terms when s ′ ̸= s or c′ ̸= c in eq. (33) are merged into other terms of eq. (34).
C APPENDIX C Notations
Z A data distribution satisfies X × Y s or (x, y) A single data sample
D Training set consists of n samples drawn from Z D′ Test set consists of n′ samples drawn from Z θ Model parameters, whose components are denoted as θj gs(θ) or gi(θ) Parameters’ gradient w.r.t. a single data sample s or (xi, yi)
g̃(θ) Mean values of parameters’ gradient over a total data distribution, i.e., Es∼Z(gs(θ))
gD(θ) Average gradient over the training dataset, i.e., 1n ∑n i=1 gi(θ)
gD′(θ) Average gradient over the test dataset, i.e., 1n′ ∑n′ i=1 g ′ i(θ). Note
that, in eq. (5), we assume n′ = n
gD,j Same as gD(θj)
ρ2(θ) Variance of parameters’ gradient of a single sample, i.e., Vars∼Z(gs(θ))
ρ2j Same as ρ 2(θj)
σ2(θ) Variance of the average gradient over a training dataset of size n, i.e., VarD∼Zn [gD(θ)]
σ2j Same as σ 2(θj)
rj or r(θj) Gradient signal to noise ratio (GSNR) of model parameter θj L[D] Empirical training loss, i.e., 1n ∑n i=1 L(yi, f(xi, θ))
L[D′] Empirical test loss, i.e., 1n′ ∑n′ i=1 L(y ′ i, f(x ′ i, θ)))
∆L[D] One-step training loss decrease
∆Lj [D] One-step training loss decrease caused by updating one parameter θj
R(Z, n) One-step generalization ratio (OSGR) for the training and test sets of size n sampled from data distribution Z , i.e., ED,D′∼Zn (∆L[D ′])
ED∼Zn (∆L[D])
λ Learning rate
▽ One-step generalization gap increment, i.e., ∆L[D] - ∆L[D′] ϵ Random variables with zero mean and variance σ2(θ)
W(l) and b(l) Model parameters (weight matrix and bias) of the l-th layer
θ(l−) Collection of model parameters over all the layers before the l-th layer
g (l) D Average gradient of W (l) over the training dataset
θ(l+) Collection of model parameters over all the layers after the l-th layer, including the l-th layer
a(l) = {a(l)s (θ(l−))} Activations of the l-th layer, where s = {1, 2, ..., S} is the index of nodes/channels in the l-th layer.
o(l) = {o(l)c } Outputs of matrix multiplication of the l-th layer, where c = {1, 2, ..., C} is index of nodes/channels in the (l + 1)-th layer.
a (l) i,s and o (l) i,c a (l) s and o (l) c evaluated on data sample i

title: UNDERSTANDING WHY NEURAL NETWORKS GENER-


summaries: The paper introduces an approach to answering queries on knowledge graphs, called Query2Box. The idea is to work with the embeddings of the vertices of the knowledge graph as if they were kind of sets. In this way, from a set, called box, of entities embeddings it is possible to project them to find other boxes using the relations specified by the query (these boxes contain the embeddings of the entities linked with those of the previous box by the relation specified in the query), or to intersect them to find the common entities.

summaries: This paper proposes a method to answer complex logical queries in large incomplete knowledge bases (KB). Specifically it considers the class of existential first-order logical queries (EPFO) which includes the logical and, or and existential operator.  The key contribution of this paper is to represent sets of entities via regions, more specifically as boxes or hyper-rectangles. This is well motivated because such logical queries often involves working over sets of entities at once and involves applying set based operators. Previous work which represented queries as a point in vector space are not well suited for these queries.

abstractText: Answering complex logical queries on large-scale incomplete knowledge graphs (KGs) is a fundamental yet challenging task. Recently, a promising approach to this problem has been to embed KG entities as well as the query into a vector space such that entities that answer the query are embedded close to the query. However, prior work models queries as single points in the vector space, which is problematic because a complex query represents a potentially large set of its answer entities, but it is unclear how such a set can be represented as a single point. Furthermore, prior work can only handle queries that use conjunctions (∧) and existential quantifiers (∃). Handling queries with logical disjunctions (∨) remains an open problem. Here we propose QUERY2BOX, an embedding-based framework for reasoning over arbitrary queries with ∧, ∨, and ∃ operators in massive and incomplete KGs. Our main insight is that queries can be embedded as boxes (i.e., hyper-rectangles), where a set of points inside the box corresponds to a set of answer entities of the query. We show that conjunctions can be naturally represented as intersections of boxes and also prove a negative result that handling disjunctions would require embedding with dimension proportional to the number of KG entities. However, we show that by transforming queries into a Disjunctive Normal Form, QUERY2BOX is capable of handling arbitrary logical queries with ∧, ∨, ∃ in a scalable manner. We demonstrate the effectiveness of QUERY2BOX on three large KGs and show that QUERY2BOX achieves up to 25% relative improvement over the state of the art.

authors:
Carl Allen
Ivana Balazevic
Timothy M Hospedales

title: On understanding knowledge graph representation

venue: arXiv preprint arXiv:1909.11611,

authors:
Ben Athiwaratkun
Andrew Gordon Wilson

title: Hierarchical density order embeddings

venue: In International Conference on Learning Representations (ICLR),

authors:
Dzmitry Bahdanau
Kyunghyun Cho
Yoshua Bengio

title: Neural machine translation by jointly learning to align and translate

authors:
Jon Louis Bentley
Jerome H Friedman

title: Data structures for range searching

venue: ACM Computing Surveys (CSUR),

authors:
Kurt Bollacker
Colin Evans
Praveen Paritosh
Tim Sturge
Jamie Taylor

title: Freebase: a collaboratively created graph database for structuring human knowledge

venue: In ACM SIGMOD international conference on Management of data (SIGMOD),

authors:
Antoine Bordes
Nicolas Usunier
Alberto Garcia-Duran
Jason Weston
Oksana Yakhnenko

title: Translating embeddings for modeling multi-relational data

venue: In Advances in Neural Information Processing Systems (NeurIPS),

authors:
Nilesh Dalvi
Dan Suciu

title: Efficient query evaluation on probabilistic databases

title: The dichotomy of probabilistic inference for unions of conjunctive queries

venue: Journal of the ACM (JACM),

authors:
Rajarshi Das
Arvind Neelakantan
David Belanger
Andrew McCallum

title: Chains of reasoning over entities, relations, and text using recurrent neural networks

venue: In European Chapter of the Association for Computational Linguistics (EACL),

authors:
Brian A Davey
Hilary A Priestley

title: Introduction to lattices and order

venue: Cambridge university press,

title: Logical and relational learning

venue: Springer Science & Business Media,

title: Relational data mining

venue: In Data Mining and Knowledge Discovery Handbook,

title: Representing words as regions in vector space

venue: In Proceedings of the Thirteenth Conference on Computational Natural Language Learning,

authors:
Kelvin Guu
John Miller
Percy Liang

title: Traversing knowledge graphs in vector space

venue: In Empirical Methods in Natural Language Processing (EMNLP),

authors:
Will Hamilton
Payal Bajaj
Marinka Zitnik
Dan Jurafsky
Jure Leskovec

title: Embedding logical queries on knowledge graphs

authors:
Shizhu He
Kang Liu
Guoliang Ji
Jun Zhao

title: Learning to represent knowledge graphs with gaussian embedding

venue: In Proceedings of the 24th ACM International on Conference on Information and Knowledge Management,

authors:
Piotr Indyk
Rajeev Motwani

title: Approximate nearest neighbors: towards removing the curse of dimensionality

venue: In Proceedings of the thirtieth annual ACM symposium on Theory of computing,

authors:
Daphne Koller
Nir Friedman
Sašo Džeroski
Charles Sutton
Andrew McCallum
Avi Pfeffer
Pieter Abbeel
Ming-Fai Wong
David Heckerman
Chris Meek

title: Introduction to statistical relational learning

venue: MIT press,

authors:
Denis Krompaß
Maximilian Nickel
Volker Tresp

title: Querying factorized probabilistic triple databases

venue: In International Semantic Web Conference,

authors:
Alice Lai
Julia Hockenmaier

title: Learning to predict denotational probabilities for modeling entailment

venue: In Annual Meeting of the Association for Computational Linguistics (ACL),

authors:
Xiang Li
Luke Vilnis
Andrew McCallum

title: Improved representation learning for predicting commonsense ontologies

venue: arXiv preprint arXiv:1708.00549,

authors:
Xiang Li
Luke Vilnis
Dongxu Zhang
Michael Boratko
Andrew McCallum

title: Smoothing the geometry of probabilistic box embeddings

authors:
Tomas Mikolov
Kai Chen
Greg Corrado
Jeffrey Dean

title: Efficient estimation of word representations in vector space

authors:
Maximilian Nickel
Kevin Murphy
Volker Tresp
Evgeniy Gabrilovich

title: A review of relational machine learning for knowledge graphs

venue: Proceedings of the IEEE,

authors:
Zhiqing Sun
Zhi-Hong Deng
Jian-Yun Nie
Jian Tang

title: Rotate: Knowledge graph embedding by relational rotation in complex space

authors:
Kristina Toutanova
Danqi Chen

title: Observed versus latent features for knowledge base and text inference

venue: In Proceedings of the 3rd Workshop on Continuous Vector Space Models and their Compositionality,

title: The nature of statistical learning theory

venue: Springer science & business media,

authors:
Ivan Vendrov
Ryan Kiros
Sanja Fidler
Raquel Urtasun

title: Order-embeddings of images and language

authors:
Luke Vilnis
Andrew McCallum

title: Word representations via gaussian embedding

authors:
Luke Vilnis
Xiang Li
Shikhar Murty
Andrew McCallum

title: Probabilistic embedding of knowledge graphs with box lattice measures

authors:
Wenhan Xiong
Thien Hoang
William Yang Wang

title: Deeppath: A reinforcement learning method for knowledge graph reasoning

title: DETAILS OF NELL995 DATASET Here we detail our pre-processing of the NELL995

heading: 1 INTRODUCTION

text:
Knowledge graphs (KGs) capture different types of relationships between entities, e.g., Canada citizen−−−−→ Hinton. Answering arbitrary logical queries, such as “where did Canadian citizens with Turing Award graduate?”, over such KGs is a fundamental task in question answering, knowledge base reasoning, as well as AI more broadly.
First-order logical queries can be represented as Directed Acyclic Graphs (DAGs) (Fig. 1(A)) and be reasoned according to the DAGs to obtain a set of answers (Fig. 1(C)). While simple and intuitive, such approach has many drawbacks: (1) Computational complexity of subgraph matching is exponential in the query size, and thus cannot scale to modern KGs; (2) Subgraph matching is very sensitive as it cannot correctly answer queries with missing relations. To remedy (2) one could impute missing relations (Koller et al., 2007; Džeroski, 2009; De Raedt, 2008; Nickel et al., 2016) but that would only make the KG denser, which would further exacerbate issue (1) (Dalvi & Suciu, 2007; Krompaß et al., 2014).
Recently, a promising alternative approach has emerged, where logical queries as well as KG entities are embedded into a low-dimensional vector space such that entities that answer the query are embedded close to the query (Guu et al., 2015; Hamilton et al., 2018; Das et al., 2017). Such approach robustly handles missing relations (Hamilton et al., 2018) and is also orders of magnitude faster, as answering an arbitrary logical query is reduced to simply identifying entities nearest to the embedding of the query in the vector space.
∗Equal contributions. Project website with data and code: http://snap.stanford.edu/ query2box
However, prior work embeds a query into a single point in the vector space. This is problematic because answering a logical query requires modeling a set of active entities while traversing the KG (Fig. 1(C)), and how to effectively model a set with a single point is unclear. Furthermore, it is also unnatural to define logical operators (e.g., set intersection) of two points in the vector space. Another fundamental limitation of prior work is that it can only handle conjunctive queries, a subset of first-order logic that only involves conjunction (∧) and existential quantifier (∃), but not disjunction (∨). It remains an open question how to handle disjunction effectively in the vector space. Here we present QUERY2BOX, an embedding-based framework for reasoning over KGs that is capable of handling arbitrary Existential Positive First-order (EPFO) logical queries (i.e., queries that include any set of ∧, ∨, and ∃) in a scalable manner. First, to accurately model a set of entities, our key idea is to use a closed region rather than a single point in the vector space. Specifically, we use a box (axis-aligned hyper-rectangle) to represent a query (Fig. 1(D)). This provides three important benefits: (1) Boxes naturally model sets of entities they enclose; (2) Logical operators (e.g., set intersection) can naturally be defined over boxes similarly as in Venn diagrams (Venn, 1880); (3) Executing logical operators over boxes results in new boxes, which means that the operations are closed; thus, logical reasoning can be efficiently performed in QUERY2BOX by iteratively updating boxes according to the query computation graph (Fig. 1(B)(D)).
We show that QUERY2BOX can naturally handle conjunctive queries. We first prove a negative result that embedding EPFO queries to only single points or boxes is intractable as it would require embedding dimension proportional to the number of KG entities. However, we provide an elegant solution, where we transform a given EPFO logical query into a Disjunctive Normal Form (DNF) (Davey & Priestley, 2002), i.e., disjunction of conjunctive queries. Given any EPFO query, QUERY2BOX represents it as a set of individual boxes, where each box is obtained for each conjunctive query in the DNF. We then return nearest neighbor entities to any of the boxes as the answers to the query. This means that to answer any EPFO query we first answer individual conjunctive queries and then take the union of the answer entities.
We evaluate QUERY2BOX on three standard KG benchmarks and show: (1) QUERY2BOX provides strong generalization as it can answer complex queries; (2) QUERY2BOX can generalize to new
logical query structures that it has never seen during training; (3) QUERY2BOX is able to implicitly impute missing relations as it can answer any EPFO query with high accuracy even when relations involving answering the query are missing in the KG; (4) QUERY2BOX provides up to 25% relative improvement in accuracy of answering EPFO queries over state-of-the-art baselines.

heading: 2 FURTHER RELATED WORK

text:
Most related to our work are embedding approaches for multi-hop reasoning over KGs (Bordes et al., 2013; Das et al., 2017; Guu et al., 2015; Hamilton et al., 2018). Crucial difference is that we provide a way to tractably handle a larger subset of the first-order logic (EPFO queries vs. conjunctive queries) and that we embed queries as boxes, which provides better accuracy and generalization.
Second line of related work is on structured embeddings, which associate images, words, sentences, or knowledge base concepts with geometric objects such as regions (Erk, 2009; Vilnis et al., 2018; Li et al., 2019), densities (Vilnis & McCallum, 2014; He et al., 2015; Athiwaratkun & Wilson, 2018), and orderings (Vendrov et al., 2016; Lai & Hockenmaier, 2017; Li et al., 2017). While the above work uses geometric objects to model individual entities and their pairwise relations, we use the geometric objects to model sets of entities and reason over those sets. In this sense our work is also related to classical Venn Diagrams (Venn, 1880), where boxes are essentially the Venn Diagrams in vector space, but our boxes and entity embeddings are jointly learned, which allows us to reason over incomplete KGs.
Box embeddings have also been used to model hierarchical nature of concepts in an ontology with uncertainty (Vilnis et al., 2018; Li et al., 2019). While our work is also based on box embeddings we employ them for logical reasoning in massive heterogeneous knowledge graphs.

heading: 3 QUERY2BOX: LOGICAL REASONING OVER KGS IN VECTOR SPACE

text:
Here we present the QUERY2BOX, where we will define an objective function that allows us to learn embeddings of entities in the KG, and at the same time also learn parameterized geometric logical operators over boxes. Then given an arbitrary EPFO query q (Fig. 1(A)), we will identify its computation graph (Fig. 1(B)), and embed the query by executing a set of geometric operators over boxes (Fig. 1(D)). Entities that are enclosed in the final box embedding are returned as answers to the query (Fig. 1(D)).
In order to train our system, we generate a set of queries together with their answers at training time and then learn entity embeddings and geometric operators such that queries can be accurately answered. We show in the following sections that our approach is able to generalize to queries and logical structures never seen during training. Furthermore, as we show in experiments, our approach is able to implicitly impute missing relations and answer queries that would be impossible to answer with traditional graph traversal methods.
In the following we first only consider conjunctive queries (conjunction and existential operator) and then we extend our method to also include disjunction.

heading: 3.1 KNOWLEDGE GRAPHS AND CONJUNCTIVE QUERIES

text:
We denote a KG as G = (V,R), where v ∈ V represents an entity, and r ∈ R is a binary function r : V × V → {True,False}, indicating whether the relation r holds between a pair of entities or not. In the KG, such binary output indicates the existence of the directed edge between a pair of entities, i.e., v r−→ v′ iff r(v, v′) = True. Conjunctive queries are a subclass of the first-order logical queries that use existential (∃) and conjunction (∧) operations. They are formally defined as follows.
q[V?] = V? . ∃V1, . . . , Vk : e1 ∧ e2 ∧ ... ∧ en, (1) where ei = r(va, V ), V ∈ {V?, V1, . . . , Vk}, va ∈ V, r ∈ R,
or ei = r(V, V ′), V, V ′ ∈ {V?, V1, . . . , Vk}, V 6= V ′, r ∈ R,
where va represents non-variable anchor entity, V1, . . . , Vk are existentially quantified bound variables, V? is the target variable. The goal of answering the logical query q is to find a set of entities JqK ⊆ V such that v ∈ JqK iff q[v] = True. We call JqK the denotation set (i.e., answer set) of query q.
As shown in Fig. 1(A), the dependency graph is a graphical representation of conjunctive query q, where nodes correspond to variable or non-variable entities in q and edges correspond to relations in q. In order for the query to be valid, the corresponding dependency graph needs to be a Directed Acyclic Graph (DAG), with the anchor entities as the source nodes of the DAG and the query target V? as the unique sink node (Hamilton et al., 2018).
From the dependency graph of query q, one can also derive the computation graph, which consists of two types of directed edges that represent operators over sets of entities:
• Projection: Given a set of entities S ⊆ V , and relation r ∈ R, this operator obtains ∪v∈SAr(v), where Ar(v) ≡ {v′ ∈ V : r(v, v′) = True}. • Intersection: Given a set of entity sets {S1, S2, . . . , Sn}, this operator obtains ∩ni=1Si.
For a given query q, the computation graph specifies the procedure of reasoning to obtain a set of answer entities, i.e., starting from a set of anchor nodes, the above two operators are applied iteratively until the unique sink target node is reached. The entire procedure is analogous to traversing KGs following the computation graph (Guu et al., 2015).

heading: 3.2 REASONING OVER SETS OF ENTITIES USING BOX EMBEDDINGS

text:
So far we have defined conjunctive queries as computation graphs that can be executed directly over the nodes and edges in the KG. Now, we define logical reasoning in the vector space. Our intuition follows Fig. 1: Given a complex query, we shall decompose it into a sequence of logical operations, and then execute these operations in the vector space. This way we will obtain the embedding of the query, and answers to the query will be entities that are enclosed in the final query embedding box.
In the following, we detail our two methodological advances: (1) the use of box embeddings to efficiently model and reason over sets of entities in the vector space, and (2) how to tractably handle disjunction operator (∨), expanding the class of first-order logic that can be modeled in the vector space (Section 3.3).
Box embeddings. To efficiently model a set of entities in the vector space, we use boxes (i.e., axis-aligned hyper-rectangles). The benefit is that unlike a single point, the box has the interior; thus, if an entity is in a set, it is natural to model the entity embedding to be a point inside the box. Formally, we operate on Rd, and define a box in Rd by p = (Cen(p),Off(p)) ∈ R2d as:
Boxp ≡ {v ∈ Rd : Cen(p)− Off(p) v Cen(p) + Off(p)}, (2) where is element-wise inequality, Cen(p) ∈ Rd is the center of the box, and Off(p) ∈ Rd≥0 is the positive offset of the box, modeling the size of the box. Each entity v ∈ V in KG is assigned a single vector v ∈ Rd (i.e., a zero-size box), and the box embedding p models {v ∈ V : v ∈ Boxp}, i.e., a set of entities whose vectors are inside the box. For the rest of the paper, we use the bold face to denote the embedding, e.g., embedding of v is denoted by v.
Our framework reasons over KGs in the vector space following the computation graph of the query, as shown in Fig. 1(D): we start from the initial box embeddings of the source nodes (anchor entities) and sequentially update the embeddings according to the logical operators. Below, we describe how we set initial box embeddings for the source nodes, as well as how we model projection and intersection operators (defined in Sec. 3.1) as geometric operators that operate over boxes. After that, we describe our entity-to-box distance function and the overall objective that learns embeddings as well as the geometric operators.
Initial boxes for source nodes. Each source node represents an anchor entity v ∈ V , which we can regard as a set that only contains the single entity. Such a single-element set can be naturally modeled by a box of size/offset zero centered at v. Formally, we set the initial box embedding as (v,0), where v ∈ Rd is the anchor entity vector and 0 is a d-dimensional all-zero vector. Geometric projection operator. We associate each relation r ∈ R with relation embedding r = (Cen(r),Off(r)) ∈ R2d with Off(r) 0. Given an input box embedding p, we model the projection by p + r, where we sum the centers and sum the offsets. This gives us a new box with the translated center and larger offset because Off(r) 0, as illustrated in Fig. 2(A). The adaptive box size effectively models a different number of entities/vectors in the set.
Geometric intersection operator. We model the intersection of a set of box embeddings {p1, . . . ,pn} as pinter = (Cen(pinter),Off(pinter)), which is calculated by performing attention
over the box centers (Bahdanau et al., 2015) and shrinking the box offset using the sigmoid function: Cen(pinter) = ∑ i ai Cen(pi), ai = exp(MLP(pi))∑ j exp(MLP(pj)) ,
Off(pinter) = Min({Off(p1), . . . ,Off(pn)}) σ(DeepSets({p1, . . . ,pn})), where is the dimension-wise product, MLP(·) : R2d → Rd is the Multi-Layer Perceptron, σ(·) is the sigmoid function, DeepSets(·) is the permutation-invariant deep architecture (Zaheer et al., 2017), and both Min(·) and exp(·) are applied in a dimension-wise manner. Following Hamilton et al. (2018), we model all the deep sets by DeepSets({x1, . . . ,xN}) = MLP((1/N) · ∑N i=1 MLP(xi)), where all the hidden dimensionalities of the two MLPs are the same as the input dimensionality. The intuition behind our geometric intersection is to generate a smaller box that lies inside a set of boxes, as illustrated in Fig. 2(B).1 Different from the generic deep sets to model the intersection (Hamilton et al., 2018), our geometric intersection operator effectively constrains the center position and models the shrinking set size.
Entity-to-box distance. Given a query box q ∈ R2d and an entity vector v ∈ Rd, we define their distance as
distbox(v;q) = distoutside(v;q) + α · distinside(v;q), (3) where qmax = Cen(q) + Off(q) ∈ Rd, qmin = Cen(q) − Off(q) ∈ Rd and 0 < α < 1 is a fixed scalar, and
distoutside(v;q) = ‖Max(v − qmax,0) + Max(qmin − v,0)‖1, distinside(v;q) = ‖Cen(q)−Min(qmax,Max(qmin,v))‖1.
As illustrated in Fig. 2(C), distoutside corresponds to the distance between the entity and closest corner/side of the box. Analogously, distinside corresponds to the distance between the center of the box and its side/corner (or the entity itself if the entity is inside the box).
The key here is to downweight the distance inside the box by using 0 < α < 1. This means that as long as entity vectors are inside the box, we regard them as “close enough” to the query center (i.e., distoutside is 0, and distinside is scaled by α). When α = 1, distbox reduces to the ordinary L1 distance, i.e., ‖Cen(q)− v‖1, which is used by the conventional TransE (Bordes et al., 2013) as well as prior query embedding methods (Guu et al., 2015; Hamilton et al., 2018).
Training objective. Our next goal is to learn entity embeddings as well as geometric projection and intersection operators.
Given a training set of queries and their answers, we optimize a negative sampling loss (Mikolov et al., 2013) to effectively optimize our distance-based model (Sun et al., 2019):
L = − log σ (γ − distbox(v;q))− k∑
i=1
1 k log σ (distbox(v ′ i;q)− γ) , (4)
1One possible choice here would be to directly use raw box intersection, however, we find that our richer learnable parameterization is more expressive and robust
where γ represents a fixed scalar margin, v ∈ JqK is a positive entity (i.e., answer to the query q), and v′i /∈ JqK is the i-th negative entity (non-answer to the query q) and k is the number of negative entities.

heading: 3.3 TRACTABLE HANDLING OF DISJUNCTION USING DISJUNCTIVE NORMAL FORM

text:
So far we have focused on conjunctive queries, and our aim here is to tractably handle in the vector space a wider class of logical queries, called Existential Positive First-order (EPFO) queries (Dalvi & Suciu, 2012) that involve ∨ in addition to ∃ and ∧. We specifically focus on EPFO queries whose computation graphs are a DAG, same as that of conjunctive queries (Section 3.1), except that we now have an additional type of directed edge, called union defined as follows:
• Union: Given a set of entity sets {S1, S2, . . . , Sn}, this operator obtains ∪ni=1Si.
A straightforward approach here would be to define another geometric operator for union and embed the query as we did in the previous sections. An immediate challenge for our box embeddings is that boxes can be located anywhere in the vector space, so their union would no longer be a simple box. In other words, union operation over boxes is not closed.
Theoretically, we prove a general negative result that holds for any embedding-based method that embeds query q into q and uses some distance function to retrieve entities, i.e., dist(v;q) ≤ β iff v ∈ JqK. Here, dist(v;q) is the distance between entity and query embeddings, e.g., distbox(v;q) or ‖v − q‖1, and β is a fixed threshold. Theorem 1. Consider any M conjunctive queries q1, . . . , qM whose denotation sets Jq1K, . . . , JqM K are disjoint with each other, ∀ i 6= j, JqiK ∩ JqjK = ∅. Let D be the VC dimension of the function class {sign(β − dist(·;q)) : q ∈ Ξ}, where Ξ represents the query embedding space and sign(·) is the sign function. Then, we need D ≥M to model any EPFO query, i.e., dist(v;q) ≤ β ⇔ v ∈ JqK is satisfied for every EPFO query q.
The proof is provided in Appendix A, where the key is that with the introduction of the union operation any subset of denotation sets can be the answer, which forces us to model the powerset {∪
qi∈S JqiK : S ⊆ {q1, . . . , qM}} in a vector space.
For a real-world KG, there are M ≈ |V| conjunctive queries with non-overlapping answers. For example, in the commonly-used FB15k dataset (Bordes et al., 2013), derived from the Freebase (Bollacker et al., 2008), we find M = 13,365, while |V| is 14,951 (see Appendix B for the details). Theorem 1 shows that in order to accurately model any EPFO query with the existing framework, the complexity of the distance function measured by the VC dimension needs to be as large as the number of KG entities. This implies that if we use common distance functions based on hyper-plane, Euclidean sphere, or axis-aligned rectangle,2 their parameter dimensionality needs to be Θ(M), which is Θ(|V|) for real KGs we are interested in. In other words, the dimensionality of the logical query embeddings needs to be Θ(|V|), which is not low-dimensional; thus not scalable to large KGs and not generalizable in the presence of unobserved KG edges.
To rectify this issue, our key idea is to transform a given EPFO query into a Disjunctive Normal Form (DNF) (Davey & Priestley, 2002), i.e., disjunction of conjunctive queries, so that union operation only appears in the last step. Each of the conjunctive queries can then be reasoned in the low-dimensional space, after which we can aggregate the results by a simple and intuitive procedure. In the following, we describe the transformation to DNF and the aggregation procedure.
Transformation to DNF. Any first-order logic can be transformed into the equivalent DNF (Davey & Priestley, 2002). We perform such transformation directly in the space of computation graph, i.e., moving all the edges of type “union” to the last step of the computation graph. Let Gq = (Vq, Eq) be the computation graph for a given EPFO query q, and let Vunion ⊂ Vq be a set of nodes whose in-coming edges are of type “union”. For each v ∈ Vunion, define Pv ⊂ Vq as a set of its parent nodes. We first generate N = ∏ v∈Vunion |Pv| different computation graphs Gq(1) , . . . , Gq(N) as follows, each with different choices of vparent in the first step. 1. For every v ∈ Vunion, select one parent node vparent ∈ Pv .
2For the detailed VC dimensions of these function classes, see Vapnik (2013). Crucially, their VC dimensions are all linear with respect to the number of parameters d.
2. Remove all the edges of type ‘union.’ 3. Merge v and vparent, while retaining all other edge connections.
We then combine the obtained computation graphs Gq(1) , . . . , Gq(N) as follows to give the final equivalent computation graph.
1. Convert the target sink nodes of all the obtained computation graphs into the existentially quantified bound variables nodes. 2. Create a new target sink node V?, and draw directed edges of type “union” from all the above variable nodes to the new target node.
An example of the entire transformation procedure is illustrated in Fig. 3. By the definition of the union operation, our procedure gives the equivalent computation graph as the original one. Furthermore, as all the union operators are removed from Gq(1) , . . . , Gq(N) , all of these computation graphs represent conjunctive queries, which we denote as q(1), . . . , q(N). We can then apply existing framework to obtain a set of embeddings for these conjunctive queries as q(1), . . . ,q(N).
Aggregation. Next we define the distance function between the given EPFO query q and an entity v ∈ V . Since q is logically equivalent to q(1) ∨ · · · ∨ q(N), we can naturally define the aggregated distance function using the box distance distbox:
distagg(v; q) = Min({distbox(v;q(1)), . . . ,distbox(v;q(N))}), (5) where distagg is parameterized by the EPFO query q. When q is a conjunctive query, i.e., N = 1, distagg(v; q) = distbox(v;q). For N > 1, distagg takes the minimum distance to the closest box as the distance to an entity. This modeling aligns well with the union operation; an entity is inside the union of sets as long as the entity is in one of the sets. Note that our DNF-query rewriting scheme is general and is able to extend any method that works for conjunctive queries (e.g., (Hamilton et al., 2018)) to handle more general class of EPFO queries.
Computational complexity. The computational complexity of answering an EPFO query with our framework is equal to that of answering the N conjunctive queries. In practice, N might not be so large, and all the N computations can be parallelized. Furthermore, answering each conjunctive query is very fast as it requires us to execute a sequence of simple box operations (each of which takes constant time) and then perform a range search (Bentley & Friedman, 1979) in the embedding space, which can also be done in constant time using techniques based on Locality Sensitive Hashing (Indyk & Motwani, 1998).

heading: 4 EXPERIMENTS

text: Our goal in the experiment section is to evaluate the performance of QUERY2BOX on discovering answers to complex logical queries that cannot be obtained by traversing the incomplete KG. This means, we will focus on answering queries where one or more missing edges in the KG have to be successfully predicted in order to obtain the additional answers.

heading: 4.1 KNOWLEDGE GRAPHS AND QUERY GENERATION

text:
We perform experiments on three standard KG benchmarks, FB15k (Bordes et al., 2013), FB15k-237 (Toutanova & Chen, 2015), and NELL995 (Xiong et al., 2017) (see Appendix E for NELL995 pre-processing details). Dataset statistics are summarized in Table 5 in Appendix F.
u
u
u
u
3i2i3p2p1p ip pi 2u up
Training Conjunctive Queries Unseen Conjunctive Queries Union Queries
Figure 4: Query structures considered in the experiments, where anchor entities and relations are to be specified to instantiate logical queries. Naming for each query structure is provided under each subfigure, where ‘p’, ‘i’, and ‘u’ stand for ‘projection’, ‘intersection’, and ‘union’, respectively. Models are trained on the first 5 query structures, and evaluated on all 9 query structures. For example, “3p” is a path query of length three, and “2i” is an intersection of cardinality two.
Dataset 1p 2p 3p 2i 3i ip pi 2u up FB15k 10.8 255.6 250.0 90.3 64.1 593.8 190.1 27.8 227.0 FB15k-237 13.3 131.4 215.3 69.0 48.9 593.8 257.7 35.6 127.7 NELL995 8.5 56.6 65.3 30.3 15.9 310.0 144.9 14.4 62.5
Table 1: Average number of answer entities of test queries with missing edges grouped by different query structures (for a KG with 10% edges missing).
We follow the standard evaluation protocol in KG literture: Given the standard split of edges into training, test, and validation sets, we first augment the KG to also include inverse relations and effectively double the number of edges in the graph. We then create three graphs: Gtrain, which only contains training edges and we use this graph to train node embeddings as well as box operators. We then also generate two bigger graphs: Gvalid, which contains Gtrain plus the validation edges, and Gtest, which includes Gvalid as well as the test edges. We consider 9 kinds of diverse query structures shown and named in Fig. 4. We use 5 query structures for training and then evaluate on all the 9 query structures. We refer the reader to Appendix D for full details on query generation and Table 6 in Appendix F for statistics of the generated logical queries. Given a query q, let JqKtrain, JqKval, and JqKtest denote a set of answer entities obtained by running subgraph matching of q on Gtrain, Gvalid, and Gtest, respectively. At the training time, we use JqKtrain as positive examples for the query and other random entities as negative examples. However, at the test/validation time we proceed differently. Note that we focus on answering queries where generalization performance is crucial and at least one edge needs to be imputed in order to answer the queries. Thus, rather than evaluating a given query on the full validation (or test) set JqKval (JqKtest) of answers, we validate the method only on answers that include missing relations. Given how we constructed Gtrain ⊆ Gvalid ⊆ Gtest, we have JqKtrain ⊆ JqKval ⊆ JqKtest and thus we evaluate the method on JqKval\JqKtrain to tune hyper-parameters and then report results identifying answer entities in JqKtest\JqKval. This means we always evaluate on queries/entities that were not part of the training set and the method has not seen them before. Furthermore, for these queries, traditional graph traversal techniques would not be able to find the answers (due to missing relations).
Table 1 shows the average number of answer entities for different query structures. We observe that complex logical queries (especially 2p, 3p, ip, pi, up) indeed require modeling a much larger number of answer entities (often more than 10 times) than the simple 1p queries do. Therefore, we expect our box embeddings to work particularly well in handling complex queries with many answer entities.3

heading: 4.2 EVALUATION PROTOCOL

text:
Given a test query q, for each of its non-trivial answers v ∈ JqKtest\JqKval, we use distbox in Eq. 3 to rank v among V\JqKtest. Denoting the rank of v by Rank(v), we then calculate evaluation metrics for answering query q, such as Mean Reciprocal Rank (MRR) and Hits at K (H@K):
Metrics(q) = 1 |JqKtest\JqKval| ∑
v∈JqKtest\JqKval
fmetrics(Rank(v)), (6)
3On the simple link prediction (1p query) task, box embeddings provide minor empirical performance improvement over TransE, possibly because simple link prediction does not require modeling large sets of entities, as shown in Table 1. See Appendix C for full experimental results on link prediction.
where fmetrics(x) = 1x for MRR, and fmetrics(x) = 1[x ≤ K] for H@K.
We then average Eq. 6 over all the queries within the same query structure,4 and report the results separately for different query structures. The same evaluation protocol is applied to the validation stage except that we evaluate on JqKval\JqKtrain rather than JqKtest\JqKval.

heading: 4.3 BASELINE AND MODEL VARIANTS

text:
We compare our framework QUERY2BOX against the state-of-the-art GQE (Hamilton et al., 2018). GQE embeds a query to a single vector, and models projection and intersection operators as translation and deep sets (Zaheer et al., 2017), respectively. The L1 distance is used as the distance between query and entity vectors. For a fair comparison, we also compare with GQE-DOUBLE (GQE with doubled embedding dimensionality) so that QUERY2BOX and GQE-DOUBLE have the same amount of parameters. Refer to Appendix G for the model hyper-parameters used in our experiments. Although the original GQE cannot handle EPFO queries, we apply our DNF-query rewriting strategy and in our evaluation extend GQE to handle general EPFO queries as well. Furthermore, we perform extensive ablation study by considering several variants of QUERY2BOX (abbreviated as Q2B). We list our method as well as its variants below.
• Q2B (our method): The box embeddings are used to model queries, and the attention mechanism is used for the intersection operator. • Q2B-AVG: The attention mechanism for intersection is replaced with averaging. • Q2B-DEEPSETS: The attention mechanism for intersection is replaced with the deep sets. • Q2B-AVG-1P: The variant of Q2B-AVG that is trained with only 1p queries (see Fig. 4); thus,
logical operators are not explicitly trained. • Q2B-SHAREDOFFSET; The box offset is shared across all queries (every query is represented
by a box with the same trainable size).

heading: 4.4 MAIN RESULTS

text:
We start by comparing our Q2B with state-of-the-art query embedding method GQE (Hamilton et al., 2018) on FB15k, FB15k-237, and NELL995. As listed in Tables 2, our method significantly and consistently outperforms the state-of-the-art baseline across all the query structures, including those not seen during training as well as those with union operations. On average, we obtain 9.8% (25% relative), 3.8% (15% relative), and 5.9% (24% relative) higher H@3 than the best baselines on FB15k, FB15k-237, and NELL995, respectively. Notice that naïvely increasing embedding dimensionality in GQE yields limited performance improvement. Our Q2B is able to effectively model a large set of entities by using the box embedding, and achieves a significant performance gain compared with GQE-DOUBLE (with same number of parameters) that represents queries as point vectors. Also notice
4Note that our evaluation metric is slightly different from conventional metric (Nickel et al., 2016; Hamilton et al., 2018; Guu et al., 2015), where average is taken over query-answer pairs. The conventional metric is problematic as it can be significantly biased toward correctly answering generic queries with huge number of answers, while dismissing fine-grained queries with a few answers. Here, to treat queries equally regardless of the number of answers they have, we take average over queries.
that Q2B performs well on new queries with the same structure as the training queries as well as on new query structures never seen during training, which demonstrates that Q2B generalizes well within and beyond query structures.
We also conduct extensive ablation studies (Tables 3). We summarize the results as follows:
Importance of attention mechanism. First, we show that our modeling of intersection using the attention mechanism is important. Given a set of box embeddings {p1, . . . ,pn}, Q2B-AVG is the most naïve way to calculate the center of the resulting box embedding pinter while Q2B-DEEPSETS is too flexible and neglects the fact that the center should be a weighted average of Cen(p1), . . . ,Cen(pn). Compared with the two methods, Q2B achieves better performance in answering queries that involve intersection operation, e.g., 2i, 3i, pi, ip. Specifically, on FB15k-237, Q2B obtains more than 4% and 2% absolute gain in H@3 compared to Q2B-AVG and Q2B-DEEPSETS, respectively.
Necessity of training on complex queries. Second, we observe that explicitly training on complex logical queries beyond one-hop path queries (1p in Fig. 4) improves the reasoning performance. Although Q2B-AVG-1P is able to achieve strong performance on 1p and 2u, where answering 2u is essentially answering two 1p queries with an additional minimum operation (see Eq. 5 in Section 3.3), Q2B-AVG-1P fails miserably in answering other types of queries involving logical operators. On the other hand, other methods (Q2B, Q2B-AVG, and Q2B-DEEPSETS) that are explicitly trained on the logical queries achieve much higher accuracy, with up to 10% absolute average improvement of H@3 on FB15k.
Adaptive box size for different queries. Third, we investigate the importance of learning adaptive offsets (box size) for different queries. Q2B-SHAREDOFFSET is a variant of our Q2B where all the box embeddings share the same learnable offset. Q2B-SHAREDOFFSET does not work well on all types of queries. This is most likely because different queries have different numbers of answer entities, and the adaptive box size enables us to better model it. In fact, we find that box offset varies significantly across different relations, and one-to-many relations tend to have larger offset embeddings (see Appendix H for the details).

heading: 5 CONCLUSION

text: In this paper we proposed a reasoning framework called QUERY2BOX that can effectively model and reason over sets of entities as well as handle EPFO queries in a vector space. Given a logical query, we first transform it into DNF, embed each conjunctive query into a box, and output entities closest to their nearest boxes. Our approach is capable of handling all types of EPFO queries scalably and accurately. Experimental results on standard KGs demonstrate that QUERY2BOX significantly outperforms the existing work in answering diverse logical queries.

heading: ACKNOWLEDGMENTS

text:
We thank William Hamilton, Rex Ying, and Jiaxuan You for their helpful discussion. W.H is supported by Funai Overseas Scholarship and Masason Foundation Fellowship. J.L is a Chan Zuckerberg Biohub investigator. We gratefully acknowledge the support of DARPA under Nos. FA865018C7880 (ASED), N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR); Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Chan Zuckerberg Biohub, JD.com, Amazon, Boeing, Docomo, Huawei, Hitachi, Observe, Siemens, UST Global.
The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ARO, or the U.S. Government.

heading: A PROOF OF THEOREM 1

text:
Proof. To model any EPFO query, we need to at least model a subset of EPFO queriesQ = {∨ qi∈S qi : S ⊆ {q1, . . . , qM}}, where the corresponding denotation sets are {∪qi∈S JqiK : S ⊆ {q1, . . . , qM}}. For the sake of modeling Q, without loss of generality, we consider assigning a single entity embedding vqi to all v ∈ JqiK, so there are M kinds of entity vectors, vq1 , . . . ,vqM . To model all queries in Q, it is necessary to satisfy the following.
∃vq1 , . . . ,∃vqM ,∀S ⊆ {q1, . . . , qM},∃qS ∈ Ξ, such that dist(vqi ;qS) { ≤ β if qi ∈ S, > β if qi /∈ S.
(7)
where qS is the embedding of query ∨qi∈Sqi. Eq. 7 means that we can learn the M kinds of entity vectors such that for every query in Q, we can obtain its embedding to model the corresponding set using the distance function. Notice that this is agnostic to the specific algorithm to embed query ∨
q∈Sq into qS; thus, our result is generally applicable to any method that embeds the query into a single vector.
Crucially, satisfying Eq. 7 is equivalent to {sign(β − dist(·;q)) : q ∈ Ξ} being able to shutter {vq1 , . . . ,vqM}, i.e., any binary labeling of the points can be perfectly fit by some classifier in the function class. To sum up, in order to model any EPFO query, we need to at least model any query in Q, which requires the VC dimension of the distance function to be larger than or equal to M .

heading: B DETAILS ABOUT COMPUTING M IN THEOREM 1

text:
Given the full KG Gtest for the FB15k dataset, our goal is to find conjunctive queries q1, . . . , qM such that Jq1K, . . . , JqM K are disjoint with each other. For conjunctive queries, we use two types of queries: ‘1p’ and ‘2i’ whose query structures are shown in Figure 4. On the FB15k, we instantiate 308,006 queries of type ‘1p’, which we denote by S1p. Out of all the queries in S1p, 129,717 queries have more than one answer entities, and we denote such a set of the queries by S′1p. We then generate a set of queries of type ‘2i’ by first randomly sampling two queries from S′1p and then taking conjunction; we denote the resulting set of queries by S2i.
Now, we use S1p and S2i to generate a set of conjunctive queries whose denotation sets are disjoint with each other. First, we prepare two empty sets Vseen = ∅, andQ = ∅. Then, for every q ∈ S1p, if Vseen ∩ JqK = ∅ holds, we letQ ← Q∪ {q} and Vseen ← Vseen ∪ JqK. This procedure already gives us Q, where we have 10, 812 conjunctive queries whose denotation sets are disjoint with each other. We can further apply the analogous procedure for S2i, which gives us a further increased Q, where we have 13, 365 conjunctive queries whose denotation sets are disjoint with each other. Therefore, we get M = 13, 365.

heading: C EXPERIMENTS ON LINK PREDICTION

text: In Table 4, we report the link prediction performance (no multi-hop logical reasoning required) following the conventional metrics (taking average over the triples of head, relation, and tail). Here query2box is trained on all five query structures as shown in Figure 4, and query2box-1p is only trained on simple 1p queries. We found that our query2box is comparable or slightly better than TransE on simple link prediction. Note that in the case of simple link prediction, we do not expect a huge performance gain by using box embeddings as link prediction does not involve logical reasoning nor handling a large set of answer entities. Also, we see that even if we train query2box over diverse queries, its performance on link prediction is still comparable to TransE and query2box-1p, which are trained solely on the link prediction task.

heading: D DETAILS ON QUERY GENERATION

text:
Given Gtrain, Gvalid, and Gtest as defined in Section 4.1, we generate training, validation and test queries of different query structures. During training, we consider the first 5 kinds of query structures. For evaluation, we consider all the 9 query structures in Fig. 4, containing query structures that are both seen and unseen during training time. We instantiate queries in the following way.
Given a KG and a query structure (which is a DAG), we use pre-order traversal to assign an entity and a relation to each node and edge in the DAG of query structure to instantiate a query. Namely, we start from the root of the DAG (which is the target node), we sample an entity e uniformly from the KG to be the root, then for every node connected to the root in the DAG, we choose a relation r uniformly from the in-coming relations of e in the KG, and a new entity e′ from the set of entities that reaches e by r in the KG. Then we assign the relation r to the edge and e′ to the node, and move on the process based on the pre-order traversal. This iterative process stops after we assign an entity and relation to every node and edge in DAG. The leaf nodes in the DAG serve as the anchor nodes. Note that during the entity and relation assignment, we specifically filter out all the degenerated queries, as shown in Fig. D. Then we perform a post-order traversal of the DAG on the KG, starting from the anchor nodes, to obtain a set of answer entities to this query.
When generating validation/test queries, we explicitly filter out trivial queries that can be fully answered by subgraph matching on Gtrain/Gvalid.

heading: E DETAILS OF NELL995 DATASET

text: Here we detail our pre-processing of the NELL995 dataset, which is originally presented by Xiong et al. (2017). Following Allen et al. (2019), we first combine the validation and test sets with the training set to create the whole knowledge graph for NELL995. Then we create new validation and test set splits by randomly selecting 20,000 triples each from the whole knowledge graph. Note that we filter out all the entities that only appear in the validation and test sets but not in the training set.

heading: F DATASET STATISTICS

text: Table 5 summarizes the basic statistics of the three datasets used in our experiments. Table 6 summarizes the basic statistics of the generated logical queries.

heading: G HYPER-PARAMETERS

text: We use embedding dimensionality of d = 400 and set γ = 24, α = 0.2 for the loss in Eq. 4. We train all types of training queries jointly. In every iteration, we sample a minibatch size of 512 queries for each query structure (details in Appendix D), and we sample 1 answer entity and 128 negative entities for each query. We optimize the loss in Eq. 4 using Adam Optimizer (Kingma & Ba, 2015) with learning rate = 0.0001. We train all models for 250 epochs, monitor the performance on the validation set, and report the test performance.

heading: H ANALYSIS OF LEARNED BOX OFFSET SIZE

text: Here we study the correlation between the box size (measured by the L1 norm of the box offset) and the average number of entities that are contained in 1p queries using the corresponding relation. Table 7 shows the top 10 relations with smallest/largest box sizes. We observe a clear trend that the size of the box has a strong correlation with the number of entities the box encloses. Specifically, we see that one-to-many relations tend to have larger offset embeddings, which demonstrates that larger boxes are indeed used to model sets of more points (entities).

heading: I MRR RESULTS