Conformalized Answer Set Prediction for Knowledge Graph Embedding

2024-08-15

arXiv.org

AI Summary

Knowledge graph embedding (KGE) link prediction lacks well-calibrated uncertainty estimates, limiting its deployment in high-stakes, reliability-critical applications. To address this, we introduce conformal prediction—a distribution-free, model-agnostic framework—for KGE for the first time, enabling statistically guaranteed coverage (e.g., 90%) of ground-truth answers via confidence sets. Our method generates answer sets per query without assuming data distribution or modifying base embeddings, and supports adaptive set-size control based on query difficulty. It is compatible with six major KGE models, including TransE and RotatE. Experiments across four benchmark datasets demonstrate strict adherence to user-specified coverage levels (100% compliance), reasonable average set sizes, and dynamic adjustment—shrinking for easy queries and expanding for hard ones. This yields substantially improved trustworthiness and practical utility of KGE-based inference.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Data Mining & Knowledge Management: Representing, Reasoning, and Using Provenance, Trust
• Reasoning under Uncertainty: Other Foundations of Reasoning under Uncertainty

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Knowledge Graph Embedding
• Probability Estimation
• Reliability Assessment

Innovation

• Consistent Prediction
• Knowledge Graph Embedding
• Probability Verification