Predicate-Conditional Conformalized Answer Sets for Knowledge Graph Embeddings (2025)

Abstract: Uncertainty quantification in Knowledge Graph Embedding (KGE) methods is crucial for ensuring the reliability of downstream applications. A recent work applies conformal prediction to KGE methods, providing uncertainty estimates by generating a set of answers that is guaranteed to include the true answer with a predefined confidence level. However, existing methods provide probabilistic guarantees averaged over a reference set of queries and answers (marginal coverage guarantee). In high-stakes applications such as medical diagnosis, a stronger guarantee is often required: the predicted sets must provide consistent coverage per query (conditional coverage guarantee). We propose CondKGCP, a novel method that approximates predicate-conditional coverage guarantees while maintaining compact prediction sets. CondKGCP merges predicates with similar vector representations and augments calibration with rank information. We prove the theoretical guarantees and demonstrate empirical effectiveness of CondKGCP by comprehensive evaluations.

## introduction knowledge graph embeddings (kge) encode entities and predicates as numerical vectors, enabling reasoning by exploiting similarities and analogies between entities and relations @xcite . while kge methods have demonstrated effectiveness in various downstream tasks such as link prediction @xcite and question answering @xcite , there remains uncertainty regarding the reliability of their predictions. specifically, kge models fail to identify when the answers to a query are uncertain @xcite . conformal prediction is a framework to quantify uncertainty by providing a prediction set-a set of possible solutions for a given task-that is guaranteed to cover the ground truth solution with a predefined confidence level @xcite . by assigning a score to each possible solution, the method defines a threshold to choose the minimum number of elements for the prediction set to provide the coverage guarantee. thus, the size of the predicted set reflects the uncertainty of the predictions, with larger sets indicating higher uncertainty. recently, @xcite introduced a method, conformalized knowledge graph embedding (kgcp), which applies conformal prediction to quantify uncertainty in the predictions from kge models. they show that kgcp provides marginal coverage guarantees, ensuring that the prediction sets meet the desired confidence level on average across all queries. however, predictive uncertainty may vary substantially across predicates, necessitating tailored coverage guarantees conditioned on predicates (predicate-conditional coverage guarantees). such conditional guarantees are especially crucial for real-world applications where specific subgroups demand reliable uncertainty estimates. for instance, in a medical diagnosis system leveraging kge, predicates like "contraindicated_for" (indicating that a treatment is not recommended for certain patients) and "has_symptom" (indicating that a specific disease or condition is associated with certain symptoms) may require different thresholds to achieve prediction sets with the desired confidence level. a shared threshold might fail to cover the true answer for the predicate "con-traindicated_for", as it is often associated with fewer triples and demands a higher threshold. conditional coverage guarantee can be achieved by performing conformal prediction at the subgroup level @xcite . however, the highly imbalanced distribution of triples across predicates in kgs @xcite poses challenges, often resulting in prediction sets that are either overly large or fail to cover the true answer @xcite . to address this limitation, we propose condkgcp, a method de-signed to approximate predicate-conditional guanratee while maintaining compact prediction sets. the key components of condkgcp are as follows: (1) it merges predicates with similar vector representation to increase the number of calibration triples available for reliable subgroup-level conformal prediction, and (2) it introduces a dual calibration schema that combines score calibration with rank calibration to exclude noisy answer entities, thereby reducing the size of prediction sets. we provide theoretical guarantees that cond-kgcp achieves conditional coverage probabilities tightly centered around the desired confidence level and that the dual calibration schema reduces expected prediction set sizes under certain conditions. empirically, we demonstrate that condkgcp outperforms five baseline methods, achieving a superior trade-off between conditional coverage probability and prediction set size across commonly used benchmark datasets. ## related work the majority of kge methods aim to improve model performance by capturing relational patterns through more expressive embedding spaces, such as complex @xcite , hyperbolic @xcite , or probabilistic spaces @xcite . by enabling richer representations, these methods have shown strong performance across downstream tasks including query answering @xcite @xcite , recommendation @xcite , and image classification @xcite . despite these successes, uncertainty quantification in kge remains largely underexplored. most uncertainty quantification methods for kge calibrate the plausibility scores generated by the models @xcite . however, these methods lack formal guarantees for the resulting probabilities. in contrast, @xcite introduce an approach that provides formal statistical guarantees. it is well-established that no prediction interval can achieve a conditional coverage guarantee in a finite sample without additional assumptions about the data distribution @xcite . consequently, many of the existing works provide coverage guarantees conditioned on specific subgroups, such as class-conditional coverage guarantees @xcite . two main strategies have been proposed to im-prove conditional coverage probabilities. the first involves modifying the nonconformity measure. for instance, @xcite enhance conditional coverage by defining cumulative probability of ground truth as nonconformity score, though their approach often results in larger prediction sets. to reduce the size of prediction sets, angelopoulos et al. (2021) introduce a regularization term in the nonconformity score. the second strategy leverages subgroup-level conformal prediction. @xcite propose mondrian conformal prediction, which performs conformal prediction within specific subgroups. building on this, @xcite cluster calibration points based on the distribution of nonconformity scores, balancing the trade-off between conditional coverage probability and the size of the prediction sets. additionally, @xcite further optimize prediction set sizes by incorporating rank information during the calibration step. however, these methods focus on classification, while our approach targets kge-based link prediction, which is more challenging due to the large number of potential answers and the highly imbalanced triple distribution across predicates. ## conditional conformal prediction for knowledge graph embedding (condkgcp) the goal of this paper is to approximate predicateconditional coverage guarantee. given a pair (q, e) ∈ t test , and an arbitrary predicate @xmath0, equation (4) defines the predicate-conditional coverage guarantee, which ensures that the true answer e is included in the prediction set of query q with a probability of at least 1ϵ. equation (4) does not necessarily hold for kgcp because the predictive uncertainty and the nonconformity score distribution can vary dramatically across predicates, which violates the i.i.d assumption in theorem 1. to have the guarantee in equation (4), a method called mondrian conformal prediction (mcp) @xcite performs conformal prediction separately for each predicate. given a subset a ⊆ r, let t cal [a] be the query-answer subset of t cal such that @xmath1. the prediction set for the method mcp is defined as: however, it is well known that most predicates in kgs are associated with very few triples @xcite , resulting in small sets t cal @xcite . this leads to unstable thresholds ŝϵ (t cal [{r}]), which in turn causes prediction sets to become overly large or fail to cover the ground truth. to address these issues, we propose: (1) merging predicates to increase the number of triples in the calibration set, and (2) augmenting the calibration process with rank information to reduce the size of prediction sets. ## coverage & size guarantees in this section, we will show the conditional coverage guarantee and size reduction guarantee of condkgcp. all proofs are in appendix a. ## experiment we evaluate the condkgcp empirically and demonstrate its effectiveness in balancing the tradeoff between predicate-conditional coverage probability and the size of the prediction sets. ## discussion and conclusion in this paper, we introduces condkgcp, a novel method that addresses the limitations of existing conformalized kge method by approximating predicate-conditional coverage guarantees while maintaining compact prediction sets. we theoretically prove that the deviation from the desired confidence level is bounded and empirically demonstrate the effectiveness of condkgcp across six kge methods and two benchmark datasets. our method offers a useful uncertainty quantification tool for high-stakes applications, such as medical diagnosis, and can be easily adapted to quantify uncertainty under other types of conditions, such as entity-type. additionally, it can be seamlessly extended to other tasks, including embedding-based query answering @xcite and probabilistic reasoning over kg @xcite .