Conformalized Answer Set Prediction for Knowledge Graph Embedding (2025)

Abstract: Knowledge graph embeddings (KGE) apply machine learning methods on knowledge graphs (KGs) to provide non-classical reasoning capabilities based on similarities and analogies. The learned KG embeddings are typically used to answer queries by ranking all potential answers, but rankings often lack a meaningful probabilistic interpretation - lower-ranked answers do not necessarily have a lower probability of being true. This limitation makes it difficult to quantify uncertainty of model’s predictions, posing challenges for the application of KGE methods in high-stakes domains like medicine. We address this issue by applying the theory of conformal prediction that allows generating answer sets, which contain the correct answer with probabilistic guarantees. We explain how conformal prediction can be used to generate such answer sets for link prediction tasks. Our empirical evaluation on four benchmark datasets using six representative KGE methods validates that the generated answer sets satisfy the probabilistic guarantees given by the theory of conformal prediction. We also demonstrate that the generated answer sets often have a sensible size and that the size adapts well with respect to the difficulty of the query.

## introduction knowledge graph embeddings (kge) map entities and predicates into numerical vectors, providing non-classical reasoning capabilities by exploiting similarities and analogies between entities and relations @xcite . kge models are typically evaluated through link prediction @xcite @xcite ). to answer queries in the form of ⟨head entity, predicate, ?⟩ or ⟨?, predicate, tail entity⟩, all possible entities are ranked according to their plausibility scores returned by the kge models. the higher the cor-rect answer is ranked, the more precise the model's prediction is for that particular query. however, high-ranked entities does not necessarily correlate with high likelihood of correctness @xcite . for instance, in a medical diagnosis setting, ranking "common cold" above "cancer x" may imply various scenarios: (1) the common cold is highly likely while cancer is improbable, (2) both conditions are likely but the common cold is more probable, or (3) neither is likely, but the common cold is relatively more probable. these interpretations can lead to vastly different treatments, making rankings alone insufficient for medical decisions. it is especially important for the patient to know when cancer can be confidently excluded from consideration. although uncertainty quantification is crucial for high-stakes applications like in medicine domain, where reliable predictions and risk assessment are essential, methods for quantifying uncertainty in kge model predictions remain underexplored. existing approaches often reply on off-the-shelf calibration techniques, such as platt scaling @xcite and isotonic regression @xcite , to map uncalibrated plausibility scores to the expected correctness of predictions @xcite . however, perfect calibration is impossible in practice @xcite , both platt scaling and isotonic regression are empirical calibration techniques that lack formal probabilistic guarantees and are highly sensitive to the calibration set. in contrast, our work diverges from these approaches by applying the theory of conformal prediction @xcite to quantify uncertainty with formal statistical guarantees. conformal prediction assigns a score to each candidate answer entity and defines a threshold to construct answer sets that include the truth answer with a desired confidence level. the size of the answer set reflects the uncertainty of the model's output for a given query. to the best of our knowledge, this is the first method that does not merely convert plausibility scores into probabilities but instead ensures statistical validity in the uncertainty quantification of the predictions within the context of kge. in our paper, we carefully design conformal predictors tailored to the link prediction task such that the answer sets are (1) probabilistically guaranteed to include the true answer entity at a specified confidence level, (2) tight, and (3) adaptive, providing smaller sets for easier queries than harder ones. we perform extensive experiments on commonly used benchmark datasets and a variety of kge methods. our empirical results show that: (1) conformal predictors satisfy the statistical guarantees in proposition 1 and produce tighter answer sets compared to other baselines (experiments 1); ## related work as highlighted in recent work @xcite , predictions from kge models exhibit substantial uncertainty, even among models with similar overall performance. this phenomenon, termed predictive multiplicity, arises from the highly non-convex loss functions employed by kge methods. it leads to multiple models that capture distinct patterns from the training kg and generalize differently. consequently, quantifying the uncertainty of kge methods is critical; however, this remains a relatively underexplored area of research. existing approaches incorporate uncertainty into kge by modeling entities and relations using probability distributions @xcite . however, these methods primarily focus on enhancing the performance of kge models through more expressive representations, without systematically analyzing or rigorously evaluating the quality of uncertainty in embeddings or predictions. furthermore, research by tabacof and costabello (2020) and @xcite applies off-theshelf calibration techniques, such as platt scaling and isotonic regression, to kge methods. these techniques aim to convert uncalibrated plausibility scores into probabilities by minimizing the negative log-likelihood on a validation set. however, these approaches are quite sensitive to the validation set and do not provide formal guarantees about the generated probabilities. this paper applies conformal prediction, which has its roots in online learning literature, is a method that produces predictive sets ensuring coverage guarantees @xcite . this approach has been successfully applied across various domains, including image classification @xcite , natural language processing @xcite and node classification/regression on graphs @xcite @xcite . however, to the best of our knowledge, it has not yet been applied to kge. ## kge-based answer set prediction in this section, we formally define the kge-based answer set prediction task and outline three key desiderata guiding the development of effective set predictors. we then introduce and discuss several basic set predictors. ## conformal prediction for kge-based answer set prediction to improve the basic set predictors, we apply conformal prediction, a general framework that requires adaptation to be effective in the context of kge. the two essential components in this design are the nonconformity measure and the method for constructing answer sets. in this section, we propose several kge-specific nonconformity measures and outline an efficient approach to constructing answer sets. ## experiments in this section, we present five experiments that evaluate the quality of the answer sets from (baseline) predictors in section 4 (naive, platt, topk) and conformal predictors (negscore, softmax, minmax) in section 5 based on coverage, size and adaptiveness desiderata. ## discussion our method predicts answer sets for queries with a guaranteed coverage of the true answer at a prespecified probability, such as 90%, while maintaining a small average size. unlike ranking-based outputs, our approach is particularly well-suited for decision-making in high-stakes domains, including medical diagnosis, drug discovery, and fraud detection. for instance, a doctor could use our method to automatically eliminate a large number of irrelevant diseases, thereby referring the patient to the most appropriate specialists. additionally, our method is easy to implement and is compatible not only with any kge models but also with embedding methods capable of answering more complex queries @xcite and approximate statistical reasoning in ontology @xcite . the adaptability of our answer sets to the uncertainty of queries also enables our method to quantify the predictive uncertainty of kge models. this feature broadens the applicability of our approach by systematically identifying hard or uncer-tain queries during testing. detecting such queries can help identify potential failure cases or outliers, alerting users when the model's predictions may be unreliable. ## limitations a limitation of our method is the requirement to divide the training set into two parts: one for training the model and another for calculating the nonconformity scores, due to the adoption of split conformal prediction. this division reduces the number of triples available for model training. however, this issue is mitigated by the fact that the validation set, typically reserved for hyperparameter tuning, can also serve as the calibration set. moreover, as demonstrated in experiment 4, even a small subset of the validation set is sufficient to produce nearly optimal answer sets. another limitation is that the probabilistic guarantee provided by theorem 1 and proposition 1 relies on the i.i.d. assumption, which may not hold under distribution shifts. we are currently extending our conformal predictors to covariant shift case, where only the input distribution p (x) changes while the conditional distribution p (y |x) remains the same. we begin with the simpler scenario where the likelihood ratio between the training and test distributions is known. following (tibshirani et al., 2019), we weight each nonconformity score proportionally to the likelihood ratio to ensure the probabilistic guarantee in proposition 1 holds beyond the i.i.d. assumption. pedro tabacof and luca costabello. 2020. probability calibration for knowledge graph embedding models. in iclr. openreview.net. ryan j tibshirani, rina foygel barber, emmanuel candes, and aaditya ramdas. 2019. conformal prediction under covariate shift. advances in neural information processing systems, 32. kristina toutanova and danqi chen. 2015. observed versus latent features for knowledge base and text inference. in proceedings of the 3rd workshop on continuous vector space models and their compositionality, pages 57-66. théo trouillon, johannes welbl, sebastian riedel, éric gaussier, and guillaume bouchard. 2016. complex embeddings for simple link prediction. in international conference on machine learning, pages 2071-2080. pmlr. vladimir vovk, alexander gammerman, and glenn shafer. 2005. algorithmic learning in a random world, volume 29. springer. quan wang, zhendong mao, bin wang, and li guo. 2017. knowledge graph embedding: a survey of approaches and applications. ieee transactions on knowledge and data engineering, 29(12):2724-2743. han xiao, minlie huang, and xiaoyan zhu. 2016. transg : a generative model for knowledge graph embedding. in proceedings of the 54th annual meeting of the association for computational linguistics (volume 1: long papers), pages 2316-2325, berlin, germany. association for computational linguistics. bishan yang, wen-tau yih, xiaodong he, jianfeng gao, and li deng. 2015. embedding entities and relations for learning and inference in knowledge bases. in iclr (poster). soroush h zargarbashi, simone antonelli, and aleksandar bojchevski. 2023. conformal prediction sets for graph neural networks. in international conference on machine learning, pages 12292-12318. pmlr. soroush h zargarbashi and aleksandar bojchevski. 2023. conformal inductive graph neural networks. in the twelfth international conference on learning representations. yuqicheng zhu, nico potyka, mojtaba nayyeri, bo xiong, yunjie he, evgeny kharlamov, and steffen staab. 2024a. predictive multiplicity of knowledge graph embeddings in link prediction. in findings of the association for computational linguistics: emnlp 2024, pages 334-354. yuqicheng zhu, nico potyka, bo xiong, trung-kien tran, mojtaba nayyeri, evgeny kharlamov, and steffen staab. 2024b. approximating probabilistic inference in statistical el with knowledge graph embeddings. corr, abs/2407.11821. yuqicheng zhu, nico potyka, bo xiong, trung-kien tran, mojtaba nayyeri, steffen staab, and evgeny kharlamov. 2023. towards statistical reasoning with ontology embeddings. in iswc (posters/demos/industry), volume 3632 of ceur workshop proceedings. ceur-ws.org.