SparseTransX: Efficient Training of Translation-Based Knowledge Graph Embeddings Using Sparse Matrix Operations

SparseTransX: Efficient Training of Translation-Based Knowledge Graph Embeddings Using Sparse Matrix Operations

Speakers: Md Saidul Hoque Anik, PhD Student, Texas A&M University; Ariful Azad

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into SparseTransX, a novel approach presented at MLSys 2025 by Md Saidul Hoque Anik and Ariful Azad, addressing the pervasive inefficiency in training Knowledge Graph Embedding (KGE) models. Knowledge graphs (KGs), structured collections of facts represented as triplets (head, relation, tail), have become increasingly vital, particularly with the rise of large language models (LLMs) and techniques like Retrieval-Augmented Generation (RAG) where they serve as crucial external knowledge sources. Despite their growing importance, the training of KGE models has traditionally lagged behind the computational efficiencies seen in other deep learning paradigms.

The core problem tackled by SparseTransX is the reliance of KGE training on shallow lookup tables and inefficient gather/scatter operations, which prevent these models from fully leveraging high-performance linear algebra libraries. Anik and Azad demonstrate that by reformulating the KGE training process, specifically for translation-based models, into a sparse-dense matrix multiplication (SPMM) problem, significant performance gains can be achieved. This paradigm shift not only accelerates training but also drastically reduces memory footprint without compromising model accuracy, offering a path to more scalable and resource-efficient KGE development.

The work introduces a generalized framework that converts the discrete index lookups of KGE triplets into a highly regular sparse matrix representation. This allows KGE training to tap into the mature and highly optimized ecosystems of sparse matrix computation, including advanced hardware acceleration and distributed processing capabilities. SparseTransX represents a critical step forward in bringing KGE training into the mainstream of high-performance computing, making it more accessible and efficient for a broader range of applications and researchers working at the intersection of HPC and graph neural networks.

Background

▶ Watch: Introduction, talk outline, and what is a knowledge graph? (0:00)

Knowledge graphs are specialized graph structures that capture relationships between entities in the real world. A fundamental unit in a knowledge graph is a "fact" or "triplet," comprising a head entity, a relation, and a tail entity (e.g., "person hunts bear"). The proliferation of large language models has significantly amplified the relevance of knowledge graphs, often employed as external knowledge sources for fine-tuning or in RAG systems to ground LLM responses with factual accuracy.

In the broader machine learning landscape, the efficiency of training models often hinges on the effective use of high-performance linear algebra libraries. Deep neural networks, for instance, extensively utilize dense matrix multiplication (GEMM) for their feed-forward and backpropagation computations. Similarly, Graph Neural Networks (GNNs) leverage sparse-dense matrix multiplication (SPMM) to handle the sparse adjacency matrices inherent in graph structures. However, KGE training has historically diverged from this paradigm. Instead of matrix operations, KGE models typically rely on shallow lookup tables to retrieve embedding vectors for entities and relations.

The traditional KGE training process involves several computationally intensive steps. For each mini-batch of triplets, the model must:

1. Gather the embedding vectors corresponding to the head, relation, and tail entities from a large embedding matrix.
2. Perform computations (e.g., vector addition, subtraction) to calculate a "score" for the triplet.
3. In the backward pass, scatter the computed gradients back to the respective positions in the embedding matrix to update the embeddings.

As highlighted by the speakers, profiling existing KGE models reveals that these gather and scatter operations constitute a significant bottleneck. For a typical mini-batch, computing both positive and negative scores (necessary for contrastive learning in KGEs) can necessitate up to six gather and six scatter operations per mini-batch. These operations are inherently memory-bound and do not efficiently utilize the computational parallelism offered by modern GPUs or specialized linear algebra accelerators, leading to substantial slowdowns and high memory consumption. The motivation behind SparseTransX is to bridge this gap, bringing the performance benefits of matrix multiplication to the domain of knowledge graph embedding training.

Key Findings

▶ Watch: Bottleneck: Many scatter and gather operations in KGE (4:00)

The central discovery of this research is the successful reformulation of translation-based Knowledge Graph Embedding (KGE) training into an efficient sparse-dense matrix multiplication (SPMM) problem. This paradigm shift fundamentally alters how KGE models compute scores and gradients, yielding substantial performance improvements and memory efficiencies.

Key findings and contributions include:

• Novel Formulation of KGE Training: The core insight is to convert the traditional gather/scatter operations for triplets into a single SPMM. Instead of looking up and manipulating individual embeddings, triplets (head, relation, tail indexes) are transformed into a three-hot encoded sparse matrix. Each row in this sparse matrix corresponds to a triplet in a mini-batch, with exactly three non-zero entries: +1 at the head entity's column, -1 at the tail entity's column, and +1 at the relation's column (offset to distinguish from entity columns).
• Leveraging High-Performance Libraries: By formulating KGE training as SPMM, the approach immediately benefits from highly optimized linear algebra libraries (e.g., NVIDIA cuSPARSE, Intel MKL) that support SIMD vectorization and distributed SPMMs. This allows for efficient parallel computation across cores and multiple nodes.
• Unified Forward and Backward Propagation: A significant advantage is that the backward propagation (gradient computation) for SPMM is also another SPMM. This means that a single, high-quality SPMM kernel can accelerate both the forward pass (score computation) and the backward pass (gradient updates), maximizing hardware utilization.
• Regular Sparsity Pattern: Unlike general graph adjacency matrices, the constructed sparse matrix exhibits a highly regular sparsity pattern, with precisely three non-zero values per row. This regularity ensures predictable performance and efficiency, as the sparsity is not affected by the underlying density of the knowledge graph itself.
• Scalable Complexity: The computational complexity of the proposed SPMM formulation is dependent only on the mini-batch size and the embedding dimension, not on the total number of entities or relations in the knowledge graph. This makes the approach highly scalable to very large knowledge graphs.
• Generalizability: The technique is highly generalizable, successfully applied to 11 different KGE models. This includes translation-based models (e.g., TransE, TransR) directly, and even non-translation models that use element-wise operations (e.g., DistMult, ComplEx) by leveraging the semiring functionality available in advanced SPMM libraries.
• Significant Performance Gains: Experiments demonstrated substantial speedups: 2x to 5x on CPU and up to 4x on GPU (A100) compared to existing popular KGE frameworks like DGL-KE and PyTorch Geometric.
• Exceptional Memory Efficiency: The SparseTransX framework achieved up to 11x reduction in GPU peak memory usage. This is primarily due to the in-place accumulation feature of SPMM, eliminating the need to store intermediate head, tail, and relation embedding vectors individually.
• Preserved or Improved Accuracy: Crucially, these performance gains were achieved without sacrificing model accuracy. In Hits@k benchmarks, SparseTransX maintained or even slightly improved accuracy compared to non-sparse baselines. Any minor deviations in loss curves due to gradient smoothing were mitigated using an adaptive learning rate scheduler.
• Open-Source Framework: The entire approach is encapsulated in an open-source, PyTorch-based framework available on GitHub, designed to be SPMM-agnostic, allowing users to plug in different underlying SPMM libraries. It also supports native PyTorch distributed training features like DDP and FSDP.

Technical Deep Dive

▶ Watch: Proposed solution: Merging scatter/gather into sparse-dense matrix multiplica... (5:00)

The technical innovation of SparseTransX lies in its re-conceptualization of Knowledge Graph Embedding (KGE) training from a sequence of discrete lookups and updates into a unified sparse-dense matrix multiplication operation.

Traditional KGE Training Bottlenecks:

Consider a translation-based KGE model like TransE, which aims to enforce the relationship head_embedding + relation_embedding ≈ tail_embedding for valid triplets (h, r, t).

1. Initialization: Entities and relations are initialized with random embedding vectors, stored in a large embedding matrix E.
2. Triplet Processing: For a mini-batch of triplets, the model performs:

• Gather: Given a triplet's indices (h_idx, r_idx, t_idx), the corresponding embedding vectors e_h, e_r, e_t are gathered from E.
• Score Computation: A scoring function f(e_h, e_r, e_t) is applied (e.g., ||e_h + e_r - e_t||_L1/L2).
• Loss Calculation & Backpropagation: The loss is computed, and gradients are propagated back.
• Scatter: The computed gradients for e_h, e_r, e_t are scattered back to their original positions in the embedding matrix E to update the weights.

1. Mini-Batch Overhead: For contrastive learning, both positive and negative triplets are processed. This means for each mini-batch, there are typically six gather operations (positive head, relation, tail; negative head, relation, tail) and six corresponding scatter operations. These operations involve irregular memory access patterns, leading to frequent cache misses and poor utilization of parallel hardware, making them a significant bottleneck.

SparseTransX Formulation: The SPMM Approach

SparseTransX transforms these gather/scatter operations into a single sparse-dense matrix multiplication (SPMM): S @ E.

1. Sparse Matrix Construction (S):

• For each mini-batch of B triplets (h_idx, r_idx, t_idx), a sparse matrix S of dimensions B x (N_entities + N_relations) is constructed.
• This matrix S is a three-hot encoding of the triplets. For each triplet (row i in the mini-batch):
• A +1 coefficient is placed at S[i, h_idx].
• A -1 coefficient is placed at S[i, t_idx].
• A +1 coefficient is placed at S[i, r_idx + offset]. The offset is N_entities, ensuring relations are mapped to distinct columns beyond the entity embeddings.
• Crucially, each row of S will have exactly three non-zero entries. This creates a highly regular sparsity pattern, which is beneficial for SPMM performance.

1. Dense Embedding Matrix (E):

• The dense matrix E has dimensions (N_entities + N_relations) x D_embedding, where D_embedding is the embedding dimension. It concatenates all entity embeddings with all relation embeddings.

1. Score Computation via SPMM:

• The operation Scores = S @ E directly computes the combined head + relation - tail vector for each triplet in the mini-batch.
• The non-zero entries in S act as selectors. When S is multiplied by E, the +1 and -1 coefficients effectively select the corresponding head, relation, and tail embeddings and perform the required addition/subtraction.
• The result Scores is a B x D_embedding dense matrix, where each row represents the (e_h + e_r - e_t) vector for a triplet.
• This formulation allows for in-place accumulation, meaning intermediate head, relation, and tail embedding vectors do not need to be stored separately, significantly reducing memory overhead.

Benefits of SPMM Formulation:

• Performance: SPMM kernels are highly optimized for modern hardware (CPUs with SIMD, GPUs with CUDA/cuSPARSE). They exploit parallelism, memory locality, and specialized instructions far more effectively than scattered gather/scatter operations.
• Unified Forward/Backward Pass: The gradient computation for S @ E with respect to E is S^T @ G_scores, where G_scores are the gradients of the loss with respect to Scores. This is also an SPMM, meaning the same highly optimized SPMM kernel can be used for both forward and backward passes.
• Memory Efficiency: By directly computing the combined vector and accumulating results, SparseTransX avoids storing multiple copies of embedding vectors, leading to substantial memory savings (up to 11x observed).
• Scalability: The complexity O(B * D_embedding) means performance scales with mini-batch size and embedding dimension, not the total graph size, which is critical for large KGs. The non-square nature of S (mini-batch rows vs. total entities+relations columns) also facilitates distributed training by splitting the matrix across nodes.

Generalization to Diverse KGE Models:

SparseTransX is designed for broad applicability:

• Translation Models: For models like TransE, TransR, or RotatE, which rely on vector addition/subtraction (h + r - t), the direct SPMM formulation described above is used.
• Models with Scaling Factors: Some models might apply factors (e.g., α h + β r - γ * t). For these, the sparse matrix S can be constructed with α, β, γ as coefficients instead of just +1/-1.
• Models with h - t: If relations are handled separately or are very small, the SPMM can compute h - t, and r can be added later. This is particularly useful if the relation embeddings are small and can be cached more effectively.
• Non-Translation Models (Element-wise Operations): For models like DistMult or ComplEx that involve element-wise multiplication (e.g., h r t or complex number operations), the framework leverages semiring operations within SPMM libraries. A semiring defines the "addition" and "multiplication" operations used in matrix multiplication. By modifying the semiring, an SPMM library can perform element-wise multiplication instead of the default sum-product, extending the technique's applicability to a wider range of KGE architectures.

Framework Implementation:

The speakers developed a PyTorch-based framework that handles the conversion of input datasets into the sparse matrix format and integrates the SPMM operations into KGE models. The framework is SPMM-agnostic, allowing developers to use any high-performance SPMM library (e.g., PyTorch's native sparse operations, external CUDA/cuSPARSE bindings) underneath. This flexibility ensures compatibility and future-proofing as new SPMM optimizations emerge. It also seamlessly integrates with PyTorch's distributed training capabilities like DDP (DistributedDataParallel) and FSDP (FullyShardedDataParallel).

Experimental Setup & Results

▶ Watch: Benefits of the SPMM formulation for KGE training (6:00)

To validate the efficiency and effectiveness of SparseTransX, a comprehensive experimental evaluation was conducted against leading KGE frameworks, focusing on training time, memory consumption, and model accuracy.

Comparative Frameworks:

The SparseTransX framework was benchmarked against two widely used and optimized KGE libraries:

• DGL-KE: A KGE library built on Deep Graph Library (DGL).
• PyTorch Geometric (PyG): A popular library for geometric deep learning, including graph neural networks.

KGE Models Evaluated:

The experiments focused on four common KGE models, which are representative of translation-based and other common architectures. While not explicitly named in the transcript, the accompanying slides indicate models such as TransE, TransR, DistMult, and ComplEx. These models were chosen because they are available in most KGE frameworks and cover different types of scoring functions.

Datasets:

Seven diverse knowledge graph datasets were used for evaluation, ranging in size and complexity. One particular dataset was noted as being "so large that it has its own scale," indicating the testing of the framework's scalability with massive graphs. Common KGE benchmarks typically include datasets like FB15k-237, WN18, WN18RR, YAGO3-10, and others.

Hardware Configuration:

Experiments were performed on both CPU and GPU environments to assess performance across different hardware:

• CPU: Single CPU.
• GPU: Single NVIDIA A100 GPU.

Training Protocol:

For each model and dataset, the training loop was run for 200 epochs. Key metrics were collected during this process.

Key Performance Metrics:

• Training Time: Measured the total time required to complete the 200 training epochs.
• GPU Peak Memory: Monitored the maximum memory allocated on the GPU during training.
• Accuracy: Evaluated using standard KGE metrics, specifically Hits@k, which measures the proportion of correctly predicted entities among the top k ranked candidates.

Headline Results:

1. CPU Speedup: SparseTransX demonstrated significant speed improvements on CPU, achieving 2x to 5x faster training times compared to DGL-KE and PyTorch Geometric across the evaluated datasets. The bar charts presented clearly illustrate the slowdown factors of other frameworks relative to SparseTransX.
2. GPU Speedup: On the A100 GPU, SparseTransX achieved even more impressive gains, with up to 4x speedup in training time. This highlights the framework's ability to effectively leverage the parallel processing capabilities of modern GPUs through its SPMM formulation.
3. Memory Efficiency: One of the most striking results was the dramatic reduction in memory footprint. SparseTransX achieved up to 11x lower GPU peak memory usage compared to the baseline frameworks. This efficiency stems directly from the SPMM approach, which avoids storing numerous intermediate variables and performs in-place accumulation of scores and gradients. This enables training larger models or using larger mini-batches on hardware with limited memory.
4. Accuracy Preservation/Improvement: Crucially, the performance and memory benefits did not come at the cost of accuracy. The Hits@k scores for SparseTransX were consistently equal to or, in some cases, better than those achieved by the non-sparse approaches.

• The speakers noted that the use of matrix multiplication naturally smoothes some gradients, leading to a slightly different loss curve compared to fine-grained gather/scatter operations. This difference could potentially impact convergence or final accuracy.
• To counteract this, a learning rate scheduler was introduced. By dynamically adjusting the learning rate during training, SparseTransX was able to maintain or even surpass the accuracy levels of traditional methods, demonstrating that the reformulated training process is robust and effective. The stability of accuracy was further confirmed by running experiments nine times for one dataset and comparing the variation.

In summary, the experimental results unequivocally demonstrate that SparseTransX provides a highly efficient and memory-optimized method for training KGE models, delivering substantial speedups and memory reductions while preserving or even enhancing model accuracy.

Practical Implications

▶ Watch: Complexity analysis and regular sparsity pattern (6:40)

The SparseTransX framework and its underlying SPMM-centric approach carry significant practical implications for various stakeholders involved in AI/ML development and deployment.

For Practitioners (ML Engineers, Researchers):

• Faster Iteration Cycles: The 2x-5x CPU and up to 4x GPU speedups directly translate to significantly reduced training times. This allows practitioners to iterate faster on model design, hyperparameter tuning, and experimentation with larger datasets or more complex KGE models.
• Scalability to Larger KGs: The improved efficiency and memory footprint (up to 11x reduction) mean that knowledge graphs previously considered too large to train on available hardware can now be tackled. This democratizes access to KGEs for applications with extensive knowledge bases.
• Resource Optimization: Reduced training time and memory usage lead to lower computational costs, whether running on local workstations or cloud-based GPU instances. This is particularly relevant as KGEs become integral to LLM pipelines, where resource consumption is a major concern.
• Broader Model Applicability: The generalizability to 11 different KGE models, including non-translation and element-wise operation models via semirings, means practitioners can apply this optimization broadly across their KGE projects without being restricted to specific model architectures.

For Infrastructure Teams (MLOps, HPC Engineers):

• Leveraging Existing Optimizations: SparseTransX's reliance on SPMM allows infrastructure teams to leverage highly optimized, battle-tested linear algebra libraries (e.g., cuSPARSE, Intel MKL) that are already integrated into their HPC environments. This avoids the need for custom, potentially less efficient, implementations of gather/scatter operations.
• Simplified Distributed Training: The formulation makes KGE training more amenable to distributed setups. SPMMs are well-understood in distributed contexts, and the PyTorch-based framework already supports native PyTorch distributed features (DDP, FSDP), simplifying the scaling of KGE training across multiple GPUs and nodes.
• Efficient Hardware Utilization: By transforming memory-bound gather/scatter operations into compute-bound SPMMs, the approach ensures better utilization of GPU compute units, leading to more efficient use of expensive hardware resources.

For Model Builders and Deployers:

• Enabling Complex Models: The newfound efficiency might encourage the development of more sophisticated KGE models that were previously impractical to train due to computational constraints.
• Streamlined Deployment: Models trained with SparseTransX can potentially be more compact and efficient, leading to smaller deployment footprints or faster inference if the SPMM formulation can also be applied to inference. (The talk focuses on training, but the SPMM benefits could extend).
• Reduced Trade-offs: The ability to achieve significant performance gains and memory efficiency without sacrificing accuracy means model builders face fewer trade-offs between model complexity, dataset size, and training efficiency.

Tradeoffs and Limitations:

• Data Preprocessing Overhead: Implementing SparseTransX requires an initial preprocessing step to convert the raw triplet data into the sparse matrix format. While this is a one-time cost, it adds complexity to the data pipeline compared to direct index lookups.
• SPMM Library Dependence: The performance benefits are highly dependent on the quality and optimization of the underlying SPMM library. While the framework is SPMM-agnostic, users might need to ensure their chosen library has robust semiring support for non-translation models.
• Applicability Scope: While highly generalizable, the approach is most directly beneficial for KGE models whose scoring functions can be effectively mapped to linear algebraic operations, particularly addition, subtraction, or element-wise products. Highly complex, non-linear scoring functions might require more intricate mapping or may not benefit as much.
• Learning Curve: Existing KGE practitioners might face a learning curve in adapting their current pipelines to the sparse matrix formulation, although the open-source PyTorch framework aims to mitigate this.

Overall, SparseTransX offers a compelling solution to a long-standing efficiency problem in KGE training, promising to accelerate research and development in knowledge graph-powered AI systems.

Key Takeaways

• KGE Training Bottleneck Addressed: Traditional Knowledge Graph Embedding (KGE) training suffers from severe bottlenecks due to inefficient gather/scatter operations, which prevent leveraging high-performance linear algebra libraries.
• SPMM Reformulation: SparseTransX reformulates translation-based KGE training as a sparse-dense matrix multiplication (SPMM) problem, converting triplets into a three-hot encoded sparse matrix that multiplies with the dense embedding matrix.
• Significant Performance & Memory Gains: This reformulation yields substantial speedups (2x-5x on CPU, up to 4x on A100 GPU) and dramatic memory efficiency (up to 11x reduction in GPU peak memory) without compromising accuracy.
• Broad Generalizability: The approach is highly generalizable, applicable to 11 different KGE models, including non-translation models by utilizing semiring operations within SPMM libraries.
• Leveraging HPC Optimizations: The SPMM approach inherently benefits from highly optimized linear algebra libraries, SIMD vectorization, and distributed computing paradigms, making KGE training more scalable and resource-efficient.
• Open-Source and PyTorch-Native: The entire method is encapsulated in an open-source, PyTorch-based, SPMM-agnostic framework that supports native PyTorch distributed training features (DDP, FSDP).

About the Speaker(s)

Md Saidul Hoque Anik is a third-year PhD student at Texas A&M University. His research focuses on the intersection of High-Performance Computing (HPC) and Graph Neural Networks (GNNs), exploring methods to enhance the efficiency and scalability of graph-based machine learning models. This work on SparseTransX is a direct outcome of his research in this domain.

Ariful Azad is also listed as a speaker and co-author on this work. While the transcript does not provide a detailed bio for Dr. Azad, his involvement as a co-author indicates his contribution to the research.