SEPAL: Scalable Feature Learning on Huge Knowledge Graphs

Authors: Félix Lefebvre, Gael Varoquaux
Primary area: learning on graphs and other geometries & topologies
Keywords: Knowledge graph, Scalable, Feature learning

Abstract:
Knowledge graphs accumulate information about more and more entities of the world. Much research is conducted to improve embedding models that capture this information and give useful node features in many downstream applications. However, most current methods are hard to scale to large knowledge graphs, partly because GPU memory is too small to hold the embeddings of many entities --YAGO4 has 67M entities. To scale existing embedding models on modest hardware, we introduce SEPAL: Scalable Embedding Propagation Algorithm for Large knowledge graphs.
The key idea of SEPAL to reduce compute is to only optimize embeddings on a core subset of entities, those that come with much more information than others. Then SEPAL propagates these embeddings to the rest of the graph with message passing, but no explicit optimization.
To enable efficient message passing, we break down large graphs into well-connected subgraphs that fit in GPU memory using a new algorithm called BLOCS: Balanced Local Overlapping Connected Subgraphs.
We evaluate SEPAL on five different knowledge graphs for four downstream regression tasks. We show that SEPAL outperforms alternative on downstream tasks, while providing a $43\times$ speedup to its base embedding algorithm.
Moreover, outside the core subgraph, embeddings obtained by message passing are not degraded compared to traditional methods, demonstrating the validity of SEPAL's propagation.

Decision: reject
Reviewer scores: 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730269766877

thread_type:
official_review

<thread_object>
creator: reviewer_qyks

<content>
summary:
Knowledge graphs accumulate information about relations of entities of the world. Producing high quality knowledge graph embeddings are important to many downstream applications such as knowledge graph completion, knowledge graph based entity retrieval, etc. However, existing methods hardly scale to large knowledge graphs. Works like PyTorch-BigGraph and DGL-KE leverage graph partition and distributed training to scale knowledge graph embedding (KGE) training. This paper tackles the scalability issue of KGE training from a new way. Specifically, it trains knowledge graph embeddings on a core subset of entities, which greatly reduces the computation cost and memory requirements. After training the embeddings of core entities and relations, it propagates these embeddings to the rest of the graph with message passing which does not require further embedding training. The resulting system is called SEPAL (Scalable Embedding Propagation Algorithm for Large knowledge graph).

strengths:
The paper proposes a new way to train knowledge graph embedding. Instead of training embeddings on the entire KG, it only trains embeddings on a small subset of core entities of a KG, which can reduces the computation cost and memory requirements.

weaknesses:
The method of constructing the core subgraph cannot guarantee that all the relation types are included in the Core subgraph. Furthermore, there is no explanation of how to select \eta to minimize the training cost while maximize the model quality (embedding quality).

The diffuse step of BLOCS can easily cause the size of a subgraph to grow significantly if the subgraph has some super nodes. How to guarantee the size of subgraph will not exceed the upper bound is not discussed.

The paper claims that the proposed method can easily be adapted to most knowledge-graph embedding method. However, in the evaluation, it only shows the training efficiency and model performance of DistMult model. The performance (system and model) of other KGE models like TransE, RotatE should be presented.

The evaluation only uses YAGO datasets. The author should also include other large scale dataset like Freebase. Both PyTorch-BigGraph and DGL-KE can train KGE on Freebase.

One important task for KGE is knowledge completion. However, the paper does not have any evaluation on KG completion tasks.

The ablation study of how different hyper-parameters, I.e., \eta, m (L259), h (L308) of SEPAL impact KGE training speed and performance is missing.

The execution time breakdown of SEPAL is missing in the evaluation. How many time it takes to produce the core subgraph? How many time it takes to training the KGE on core subgraph? How many time it takes to run BLOCS and how many time it takes to do embedding propagation.

questions:
How do you define NORMALIZE in section 3.2?
How do you decide K for different datasets? Is it a hyper-parameter?
How does R2 score define?
How do you train the KGE models using PyTordh-BigGraph and DGL-KE, How do you split the training set? Have you included all the entities in Table 5 in the training set?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730344994365

thread_type:
official_review

<thread_object>
creator: reviewer_c7mu

<content>
summary:
This paper introduces SEPAL  to scale knowledge graph embedding models efficiently on limited hardware resources. SEPAL reduces computational demands by optimizing embeddings only for core subgraphs.   To facilitate efficient message passing, the authors developed BLOCS which partitions large graphs into smaller, well-connected subgraphs that fit within GPU memory. Evaluated on five knowledge graphs across four regression tasks, SEPAL demonstrates speedup over traditional embedding methods and achieves strong performance on downstream tasks.

strengths:
1.This paper proposes SEPAL, a method designed to reduce computational costs, achieving faster processing times and lower memory usage for knowledge graph embeddings.

2.Although SEPAL generates embeddings from subgraphs, which could be suboptimal, the approach demonstrates strong performance on downstream tasks in experimental evaluations.

weaknesses:
1.The work presents an interesting approach, though some aspects require further clarification. For instance, how specifically does the model reduce time and memory costs? The proposed method involves constructing core subgraphs and performing message passing on these subsets. In terms of memory, it appears that a substantial portion of memory usage stems from storing embeddings. However, it seems embeddings still need to be allocated for all entities and relations, raising the question of how memory savings are achieved. Regarding time reduction, is it primarily due to message passing on smaller subgraphs? Otherwise, this appears similar to conventional neighborhood-based message passing.

2.In the experimental section, the regression task setup lacks clarity. What specific values are being predicted? Additionally, why not evaluate the method on the knowledge graph completion task, which is a fundamental evaluation for knowledge graph embeddings? Including this evaluation would provide a more comprehensive assessment of the method’s performance.

questions:
Please refer to the weekness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730515219879

thread_type:
official_review

<thread_object>
creator: reviewer_zda3

<content>
summary:
This paper introduces Scalable Embedding Propagation Algorithm (SEPAL) for scalable and cost-effective training of knowledge graph (KG) embeddings over large KGs with a single GPU. Specifically, SEPAL only optimizes the embeddings for a subset of core entities, and it computes the embeddings for the rest entities with message passing. To address the scalability challenge of message passing over a large KG, it proposes an algorithm called Balanced Local Overlapping Connected Subgraphs (BLOCS) to break the KG into smaller connected subgraphs. Empirical studies demonstrate the effectiveness and efficiency of SEPAL on KGs of up to about 67M entities.

strengths:
**S1.** The problem formulation and methodology design are well motivated.

**S2.** Empirical studies demonstrate the effectiveness and efficiency of SEPAL for the experiments considered.

**S3.** The paper is easy to follow.

weaknesses:
**W1.** In order to defend the novelty and principled benefit of core set embedding + embedding propagation, the current discussion and comparison against previous bag-of-entities approaches, particularly NodePiece, is insufficient. Notably, NodePiece also embeds a subset of entities and the authors should discuss how this is different from SEPAL.

**W2.** Current ablation studies are insufficient. This is important in particular as both claimed contributions have counterparts explored in the existing literature, i.e., core entity embedding + embedding propagation vs bag-of-entities approaches like NodePiece and BLOCS vs graph sampling/partition like METIS. The authors should consider respectively fixing one contribution at a time and replacing the other contribution with existing approaches. 

**W3.** The experiment settings considered are kind of different from the established ones without sufficient justification. See the questions below.

questions:
**Q1.** Did you try evaluation for knowledge graph completion?

**Q2.** Did you try on common KGs like Freebase and Wikidata?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733226583680

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global summary of the discussions and new elements

comment:
We thank all reviewers for their careful reading of our work and their thoughtful feedback. All reviewers acknowledged that our work reduces the computation cost and memory requirements of training embedding models on huge knowledge graphs while demonstrating strong performance on downstream tasks.

Our paper provides a novel perspective for knowledge-graph embedding: optimizing the embeddings only on a small core of entities, and then using message passing to mimic the underlying embedding model and efficiently propagate to the rest of the graph. It also contributes a technical solution, BLOCS, that allows this approach to be scaled to huge graphs within the GPU RAM limits.

---------------

Reviewers posed insightful questions and constructive comments, which we have addressed in detail in our individual responses. These have helped us improve the quality of our work, and we sincerely thank the reviewers for that.

We provided a new version of the paper with the following main additions (in blue in the paper):
- New experiments on **Freebase**, a dataset with 85M entities and 338M triples, show that **on huge graphs SEPAL creates the best embeddings for downstream tasks and is fastest** (Figure 4, page 15).
- We showed empirically that SEPAL adapts well to **other embedding models (TransE and RotatE), making them faster and better performing** (Figure 4, page 15).
- We implemented a new core selection method that **guarantees by design that all relation types are included in the core**, addressing reviewer QYKs’ concern (Appendix A.5, pages 18-19).
- We added experimental results:
  - breakdown of execution time (Figure 5, Appendix A.2, page 16)
  - effect of hyperparameters (Figures 8 and 9, Appendix D, pages 22-23)
  - ablation study replacing BLOCS with METIS (Figure 6, Appendix A.3, page 17)
  - and comparison to NodePiece (Figure 7, Appendix A.4, pages 17-18).
- We expanded Appendix C on downstream tasks (pages 21-22), clarifying our experimental setup.
- And many minor additions or clarifications suggested by the reviewers.

The discussion with reviewer c7mu led us to analyze in detail the **origin of time and memory gains** of SEPAL, and we will include this analysis in the final version of the paper.

----------------

Finally, we wanted to discuss here the question of **knowledge graph completion**, as it was raised by the three reviewers.

Our paper focuses on feature learning, an active research field [Grover & Leskovec 2016, Cvetkov-Iliev 2023, Robinson 2024]. While knowledge graph (KG) completion is a popular use case for embeddings, it is fundamentally different from feature enrichment. Notably, embeddings optimized for link prediction may not perform well for feature enrichment, and vice versa. These are different tasks; for instance, FastRP—a strong baseline for feature learning—cannot be used for KG completion as it ignores relation types. There is a blind spot in the literature, as few works have studied KG embeddings for downstream tasks, despite their demonstrated utility.

Our paper shows that SEPAL consistently outperforms other methods for feature enrichment on 6 knowledge graphs. This stands independently of its performance on link prediction. We believe that having models performing well on feature enrichment is valuable for the community, and broadens the applicability of knowledge-graph embeddings.

Nevertheless, **we conducted additional experiments to test SEPAL on knowledge graph completion**. These **preliminary results show that SEPAL can also perform well for link prediction**. We detail them in the response to reviewer ZdA3.


[Grover & Leskovec 2016] Grover, A., & Leskovec, J. (2016). node2vec: Scalable feature learning for networks. KDD

[Cvetkov-Iliev 2023] Cvetkov-Iliev, A., Allauzen, A., & Varoquaux, G. (2023). Relational data embeddings for feature enrichment with background information. Machine Learning, 112(2)

[Robinson 2024] Robinson, J., Ranjan, R., Hu, W., Huang, K., Han, J., Dobles, A., ... & Leskovec, J. (2024). Relbench: A benchmark for deep learning on relational databases. NeurIPS.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a Scalable Embedding Propagation Algorithm (SEPAL) for scalable and cost-effective training of knowledge graph (KG) embeddings over large KGs with a single GPU. The key idea is to only optimize embeddings on a core subset of entities and compute the embeddings for the rest of the entities with message passing. To address the scalability challenge of message passing over a large KG, an algorithm called Balanced Local Overlapping Connected Subgraphs (BLOCS) is proposed to break the KG into smaller connected subgraphs.

The initial submission required extensive content additions to address the reviewers' concerns, and the revised version of the paper is substantially different from the initial version; thus, another round of reviews is needed to decide whether the revised paper is ready for publication. For example, Reviewer QYKs expressed remaining concerns regarding "The method of constructing the core subgraph cannot guarantee that all relation types are included in the core subgraph", and the authors added "a new core selection strategy", which should be one of the main components of the algorithm. Such a revision changes the core algorithm and, thus, should be described in the main section instead of the appendix. Also, the additional experiments added during the rebuttal and discussion phases had to be included in the initial submission. While the authors' efforts in improving the paper are appreciated, the initial submission is incomplete, and the revised algorithm and added experiments should be re-evaluated.