# PIE: a Parameter and Inference Efficient Solution for Large Scale Knowledge Graph Embedding Reasoning

arXiv: https://arxiv.org/abs/2204.13957 · HTML: https://mlsanity.com/paper/2204.13957/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 50 references that are in the map.
- First posted on arXiv 29 Apr 2022.

## Builds on (most cited first, 30 of 31)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Geometric deep learning: going beyond Euclidean data](https://mlsanity.com/paper/1611.08097/) (2016)
- [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://mlsanity.com/paper/1412.6575/) (2014)
- [Complex Embeddings for Simple Link Prediction](https://mlsanity.com/paper/1606.06357/) (2016)
- [RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space](https://mlsanity.com/paper/1902.10197/) (2019)
- [Representation Learning on Graphs with Jumping Knowledge Networks](https://mlsanity.com/paper/1806.03536/) (2018)
- [Convolutional 2D Knowledge Graph Embeddings](https://mlsanity.com/paper/1707.01476/) (2017)
- [Circle Loss: A Unified Perspective of Pair Similarity Optimization](https://mlsanity.com/paper/2002.10857/) (2020)
- [OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs](https://mlsanity.com/paper/2103.09430/) (2021)
- [Inductive Relation Prediction by Subgraph Reasoning](https://mlsanity.com/paper/1911.06962/) (2019)
- [SimplE Embedding for Link Prediction in Knowledge Graphs](https://mlsanity.com/paper/1802.04868/) (2018)
- [Embedding Logical Queries on Knowledge Graphs](https://mlsanity.com/paper/1806.01445/) (2018)
- [CoDEx: A Comprehensive Knowledge Graph Completion Benchmark](https://mlsanity.com/paper/2009.07810/) (2020)
- [Question Answering Over Temporal Knowledge Graphs](https://mlsanity.com/paper/2106.01515/) (2021)
- [NodePiece: Compositional and Parameter-Efficient Representations of Large Knowledge Graphs](https://mlsanity.com/paper/2106.12144/) (2021)
- [DGL-KE: Training Knowledge Graph Embeddings at Scale](https://mlsanity.com/paper/2004.08532/) (2020)
- [PairRE: Knowledge Graph Embeddings via Paired Relation Vectors](https://mlsanity.com/paper/2011.03798/) (2020)
- [Relational Message Passing for Knowledge Graph Completion](https://mlsanity.com/paper/2002.06757/) (2020)
- [Orthogonal Relation Transforms with Graph Context Modeling for Knowledge Graph Embedding](https://mlsanity.com/paper/1911.04910/) (2019)
- [TorusE: Knowledge Graph Embedding on a Lie Group](https://mlsanity.com/paper/1711.05435/) (2017)
- [AutoSF: Searching Scoring Functions for Knowledge Graph Embedding](https://mlsanity.com/paper/1904.11682/) (2019)
- [Open Question Answering with Weakly Supervised Embedding Models](https://mlsanity.com/paper/1404.4326/) (2014)
- [Neural Fine-Grained Entity Type Classification with Hierarchy-Aware Loss](https://mlsanity.com/paper/1803.03378/) (2018)
- [LowFER: Low-rank Bilinear Pooling for Link Prediction](https://mlsanity.com/paper/2008.10858/) (2020)
- [Zero-Shot Open Entity Typing as Type-Compatible Grounding](https://mlsanity.com/paper/1907.03228/) (2019)

## Built on by (most cited first, 4 of 4)

- [BESS: Balanced Entity Sampling and Sharing for Large-Scale Knowledge Graph Completion](https://mlsanity.com/paper/2211.12281/) (2022)
- [From Wide to Deep: Dimension Lifting Network for Parameter-efficient Knowledge Graph Embedding](https://mlsanity.com/paper/2303.12816/) (2023)
- [Are We Wasting Time? A Fast, Accurate Performance Evaluation Framework for Knowledge Graph Link Predictors](https://mlsanity.com/paper/2402.00053/) (2024)
- [TIEG-Youpu Solution for NeurIPS 2022 WikiKG90Mv2-LSC](https://mlsanity.com/paper/2603.28512/) (2026)

## Abstract (arXiv)

Knowledge graph (KG) embedding methods which map entities and relations to unique embeddings in the KG have shown promising results on many reasoning tasks. However, the same embedding dimension for both dense entities and sparse entities will cause either over parameterization (sparse entities) or under fitting (dense entities). Normally, a large dimension is set to get better performance. Meanwhile, the inference time grows log-linearly with the number of entities for all entities are traversed and compared. Both the parameter and inference become challenges when working with huge amounts of entities. Thus, we propose PIE, a parameter and inference efficient solution. Inspired from tensor decomposition methods, we find that decompose entity embedding matrix into low rank matrices can reduce more than half of the parameters while maintaining comparable performance. To accelerate model inference, we propose a self-supervised auxiliary task, which can be seen as fine-grained entity typing. By randomly masking and recovering entities' connected relations, the task learns the co-occurrence of entity and relations. Utilizing the fine grained typing, we can filter unrelated entities during inference and get targets with possibly sub-linear time requirement. Experiments on link prediction benchmarks demonstrate the proposed key capabilities. Moreover, we prove effectiveness of the proposed solution on the Open Graph Benchmark large scale challenge dataset WikiKG90Mv2 and achieve the state of the art performance.