Paper: COINs: Model-based Accelerated Inference for Knowledge Graphs (ICLR-2024)

Reviewer comment (Experiments):
More datasets should be evaluated, especially bigger ones like WikiKG90M-LSC.

Author response:
- We agree on the relevance of experiments with larger datasets. We selected the three we use as the most classical ones for the field, and ones with the most available results from previous work to compare with as baselines. In fact, the WikiKG90M-LSC dataset does not have a publicly available test set with ground truth yet, as it was used for a recent competition. Nevertheless, we have started COINs-RotatE runs with larger knowledge graphs such as ogbl-biokg, but due to the greater number of edges, we require more time to achieve convergence. We would have no problem releasing these results when they are ready.

Reviewer comment (Experiments):
The trade-off between error, speedup and number of clusters should be investigated.

Author response:
- Thank you for the interesting analysis suggestion, we have managed to perform these experiments and obtain results. The full discussion is available in appendix sections C1 and C2, however here we summarize. We implemented the METIS min-cut algorithm (used for partitioning by DistDGL (Zheng et al. 2020)) as well as random uniform community assignment (used for partitioning by Pytorch-BigGraph (Lerer et al. 2019)), as alternatives to Leiden. We analyzed the cut sizes produced by all 3 (large range of resolution values for Leiden) in Figure 4. We then end-to-end trained and evaluated COINs-RotatE with all 3 community methods (a couple of resolution values for Leiden in ascending order) and plotted the relative error-scalability trade-off in Figure 6. The results can be summarized as follows: although the METIS and random communities might yield higher speed-ups, there are a lot of Leiden community assignments resulting in similar performance, and equally significant speed-ups but with significantly fewer parameters.

Reviewer comment (Experiments):
Static min-cut algorithms (outside the end-to-end training) could be compared.

Author response:
- Thank you for the interesting analysis suggestion, we have managed to perform these experiments and obtain results. The full discussion is available in appendix sections C1 and C2, however here we summarize. We implemented the METIS min-cut algorithm (used for partitioning by DistDGL (Zheng et al. 2020)) as well as random uniform community assignment (used for partitioning by Pytorch-BigGraph (Lerer et al. 2019)), as alternatives to Leiden. We analyzed the cut sizes produced by all 3 (large range of resolution values for Leiden) in Figure 4. We then end-to-end trained and evaluated COINs-RotatE with all 3 community methods (a couple of resolution values for Leiden in ascending order) and plotted the relative error-scalability trade-off in Figure 6. The results can be summarized as follows: although the METIS and random communities might yield higher speed-ups, there are a lot of Leiden community assignments resulting in similar performance, and equally significant speed-ups but with significantly fewer parameters.

Reviewer comment (Evaluation):
What's the relationship between cut size, speedup and performance?

Author response:
- Thank you for the interesting analysis suggestion, we have managed to perform these experiments and obtain results. The full discussion is available in appendix sections C1 and C2, however here we summarize. We implemented the METIS min-cut algorithm (used for partitioning by DistDGL (Zheng et al. 2020)) as well as random uniform community assignment (used for partitioning by Pytorch-BigGraph (Lerer et al. 2019)), as alternatives to Leiden. We analyzed the cut sizes produced by all 3 (large range of resolution values for Leiden) in Figure 4. We then end-to-end trained and evaluated COINs-RotatE with all 3 community methods (a couple of resolution values for Leiden in ascending order) and plotted the relative error-scalability trade-off in Figure 6. The results can be summarized as follows: although the METIS and random communities might yield higher speed-ups, there are a lot of Leiden community assignments resulting in similar performance, and equally significant speed-ups but with significantly fewer parameters.

Reviewer comment (Novelty):
Minimum cut literature studies the problem of reducing the number of inter-cluster relations.

Author response:
- Thank you for your helpful suggestions. We have added a comment on tf-GNN in our introduction and corrected the error w.r.t min-cut literature (mention of the min-cut problem and reference to the METIS algorithm added in Section 2.3.3 Edge locality, which as mentioned above, we now consider METIS as an alternative method for obtaining the community partition and for which we obtained some initial experiment results).

Reviewer comment (Novelty):
tf-GNN is another scalable GNN framework which uses sampling for large datasets.

Author response:
- Thank you for your helpful suggestions. We have added a comment on tf-GNN in our introduction and corrected the error w.r.t min-cut literature (mention of the min-cut problem and reference to the METIS algorithm added in Section 2.3.3 Edge locality, which as mentioned above, we now consider METIS as an alternative method for obtaining the community partition and for which we obtained some initial experiment results).

Reviewer comment (Experiments):
Did you consider producing different clusters for different relation types? Are the "optimal" clusters correlated?

Author response:
- No, as we believe this would cause too large of an increase in the number of parameters to be appropriate for our scalability goals in any case ($|R|$ new sets of 1 community embedder, $K$ intra-community embedders and 1 inter-community embedder), when the number of possible relations $|R|$ is large, but most relations could be underrepresented and we cannot assume any correlation between relation frequencies between community subgraphs. Instead, we consider handling the adaptivity to different community structures per relation by allowing each of the new COINs embedders to keep separate embeddings of each relation. This yields an addition of only $|R|(K + 2)$ new embeddings, and since $|R| \ll |V|$, the scalability effects for relation embeddings would be minuscule compared to those of node embeddings we analyzed thoroughly.

Reviewer comment (Evaluation):
What's the breakdown of speedup for training and inference?

Author response:
- The speed-up for inference for each dataset is given in Table 2 in the Acceleration column, based on the ratio between the baseline complexity $N |V|$ and our complexity detailed in Proposition 1. For training, one can observe that without COINs in a worst-case training batch, one would have to compute all $|V|$ node embeddings and their gradients, while for COINs this worst-case number is instead (previously in our response to Reviewer 2jBS's related question there was an erroneous factor of $B$): $K+\max{(\max_{k=1}^{K}{|C_k|},|V^*|+1)}$, where similarly to the inference cost it depends on quantities that we all aim to optimize with better community partitions. The specific reduction in node embeddings per mini-batch for every dataset would be: 1.0075x for FB15k-237, 3.1450x for WN18RR, and 3.9284x for NELL-995. For FB15k-237 the number of inter-community nodes is what kills the training gains (due to this graph's high edge density), while for the other two, the size of the largest community dominates.

Reviewer comment (Theory):
The proposition does not prove the author's remark. For proposition 1 (equation 2), proving the lower bound on runtime does not prove that the algorithm is better. A lower bound of zero would also say nothing. A tighter upper bound for the optimized cluster size is needed. The current upper bound is trivial, and it seems possible to make the lower bound an upper bound by substituting values. For proposition 2, equation 3 simply moves the left-hand side to the right-hand side. The scenario must be specified and the "expected time to a correct answer" quantified in the proposition. Clarification is needed on why the derivation is related to Prop 1, especially given a missing constant 2. The derivation should state something like "Our method is better when ratio A is better than ratio B."

Author response:
**"First, The proposition doesn't prove the author's remark. For proposition 1 (equation 2), proving the lower bound on runtime didn't prove that your algorithm is better. I may show a lower bound of zero, and it says nothing."**
        
- We do not wish to claim that the proven lower bound is sufficient evidence for superiority. We do not base our further analyses on it. It was presented as simply the limit of what can be possible with COINs. It is very important that we could achieve it and thus the lower-bound is tight. As a result, we do not need to worry about a vacuous lower-bound of 0. 
                
**"The current upper bound is simply trivial, and I see it's possible to make your lower bound an upper bound (just substitute the values)."**

- Yes, we agree that the proof of the bounds is less rigorous than desirable. Thus, we amended this error the following way: we now show first with Proposition 3 that an extreme value of the evaluation cost is only achieved if the community sizes and/or evaluation edge counts per community are uniformly distributed. Then, we show with Proposition 4 that the value of $K$ is what decides whether the extreme point is a maximum or minimum.

**"For proposition 2, equation 3 simply moves the left-hand side to the right-hand side. You need to specify the scenario and quantify the "expected time to a correct answer" in your proposition. Also, I need clarification on why the derivation is related to Prop 1 (with a missing constant 2). It should be simply "Our method is better when ratio A is better than ratio B" in the derivation.""**

- Thank you for the helpful suggestion, we have both simplified and clarified better the reasoning behind Proposition 2.

Reviewer comment (Experiments):
It is not possible to control the cluster size of the community detection algorithm. Tuning the resolution parameter changes the number of clusters, but cluster size depends on the graph structure and cannot be easily homogenized. There may be communities of only a few nodes, as well as a community consisting of one-quarter of the nodes. Therefore, the analysis represents "acceleration at the best case." Practically, showing good acceleration results is acceptable.

Author response:
**"Second, you cannot control the cluster size of the community detection algorithm. Tuning the resolution parameter changes the number of clusters, but the size of each cluster depends on the graph structure and cannot be easily homogenized. There might be communities of a few nodes, and there might be a community consisting of 1/4 of the nodes. So, the analysis is actually "acceleration at the best case." The result is okay from a practical perspective, and showing good acceleration results is good enough."**

- Yes we agree on the validity of your claim and counterexample. But to clarify, Proposition 1 holds for *any* disjoint partition $c$ and does not assume that Leiden community detection was used to obtain it. One can easily assign the groups arbitrarily focusing only on their size. E.g., with $c$=`np.random.choice(`$K$`, size=`$|V|$`)` just like what Lerer et al. in 2019 employ for PyTorch-BigGraph, one achieves pretty uniformly sized communities with low enough variance in the sizes. The idea is that due to this lack of correlation in such an assignment, the community prediction step for us would be extremely difficult.

Reviewer comment (Experiments):
The experiments do not support the claims in the introduction. All three traditional datasets used can be run within hours or minutes on a single desktop. While the experiments showed acceleration in terms of vector evaluations, the results do not demonstrate scalability comparable to the experimental scales in DistDGL or SMORE.

Author response:
**"Finally, the experiments don't support your claims in the introduction. All the 3 traditional datasets the authors tried can be run within hours or minutes on a single desktop. The experiments did show acceleration (in terms of vector evaluations), but the result doesn't support scalability compared to the scales in experiments from DistDGL or SMORE."**

- We chose those datasets as most, if not all, previous work has produced baselines for us to compare with. More importantly, we do not claim to replace DistDGL or SMORE in all application settings. The key gaps in the research that we address is the *requirement of multiple machines* for lossless scalable prediction, as well as the lack of scalability improvements for knowledge graph model *inference* and not simply faster training data sampling and model optimization. If one has at their disposal compute for parallelization and has to guarantee no deviation from the baseline prediction power, the other frameworks are viable. As a side note, we utilize SMORE's improved training query and answer sampling too, while also noting the potential in parallelizing the intra-community embedding of nodes for future work.

Reviewer comment (Reproducibility):
(Algorithm 2 line 11) The $L$ function is implicitly parameterized by the graph structure and the negative-sampling methods. There are now three graphs: the community graph, the intra-community graph, and the inter-community graph. It is unclear when each is used in the loss function and what the differences in sampling are. For example, node $\omega$ will not appear in the testing set but has many edges—how are these integrated?

Author response:
**"(Algorithm 2 [sic] line 11) The function is actually implicitly parameterized by the graph structure and the negative-sampling methods. However, there are now three graphs: the community graph the intra- and inter-community graph. So when will each be used in the loss function? And what's the difference in sampling? For example, the node won't appear in the testing set but has many edges. How are they integrated?"**

- We apologize for the lack of clarity. Thus, we give a bit more implementation details about Algorithm 1: to learn better community representations negative community examples have to be sampled according to the community graph as well. Using these and the positive community triplets from the batch first the loss terms $\mathcal{L}(e^{(c)}_h,e^{(c)}_r,e^{(c)}_t,y)$ are computed in a supervised manner according to the contrastive loss over the community graph. Then the triplet batch is split into sub-batches according to the community membership of the entities and the appropriate intra- or inter-community embedder for each sub-batch computes the temporary triplet embeddings (superscripted with $^*$). Negative node samples are sampled from the graph appropriate for the target community of each sub-batch (that community's subgraph). Temporary triplet embeddings are then simply concatenated across the batch dimension before going to step 8.

Reviewer comment (Theory):
(item 4 in page 5) Usually, the loss function is not convex with respect to the embeddings, and no information is given on the refinement used.

Author response:
**"Usually, the loss function is not convex to the embeddings. And there's no info on the refinement used."**

- We apologize for the imprecise statement. The idea is that if a given model's embedding scoring function $\mathcal{S}$ is convex w.r.t. the embeddings (like for TransE), then the weighted averaging operations (steps 8-10 of Algorithm 1, which are the final refinement steps) will preserve the convexity of $\mathcal{S}$ w.r.t. the embeddings, and the contrastive loss is always convex.

Reviewer comment (Theory):
(item 2 in page 5) The big-O notation provided is incorrect, as all embeddings always need to be swept through.

Author response:
**"The big-O notation is wrong. We always need to sweep through all embeddings."**

- Yes, but that would be with respect to epochs, what we have is with respect to mini-batches. If we assume mini-batch stochastic optimization, COINs will have at most $K + B \max(\max_{k=1}^{K}{|C_k|},|V^*|+1)$ node embeddings or their gradients in memory for a given mini-batch size of $B$.

Reviewer comment (Writing):
(Sec 3.1) The purpose of adding the $\omega$ node is unclear.

Author response:
**"What's the purpose of adding the $\omega$ node?"**

- It is a padding node (embedding) representing all of the nodes in $V \setminus V^*$ in the eyes of the inter-community embedder. The idea is that if the knowledge graph model requires the sampling of random walk contexts, k-hop neighbors, multi-hop query structure etc. around an inter-community edge then the entity sampling could in these cases produce results outside of $V^*$, which $f_*$ would not know how to deal with. Through the inclusion of $\omega$ this is prevented for any possible scenario.