Learning Chain Of Thoughts Prompts for Predicting Entities, Relations, and even Literals on Knowledge Graphs

1. Problem Statement

Traditional Knowledge Graph Embedding (KGE) models (e.g., TransE, DistMult) represent entities and relations as continuous vectors to perform link prediction. While effective for transductive settings, these models face three critical limitations:

1. Zero-Shot Generalization: They struggle with unseen entities or relations (inductive settings) because they rely on learned embeddings for specific IDs.
2. Literal Reasoning: They fail to effectively reason over literals (numerical values, text descriptions), which are pervasive in real-world graphs.
3. Static Nature: They require retraining to incorporate new data or structural changes.

The authors propose shifting the paradigm from continuous vector-parameterized scoring functions ($\phi_\Theta$) to string-parameterized scoring functions ($\phi_{\beta \le c}$). The goal is to leverage Large Language Models (LLMs) and automated prompt optimization to predict missing triples (entities, relations, or literals) without gradient-based fine-tuning, enabling reasoning over unseen elements and complex logical expressions.

2. Methodology: The RALP Framework

The authors introduce RALP (Reasoning-Augmented Link Prediction), a framework that treats link prediction as a prompt learning problem.

Core Concept

Instead of learning vector embeddings, RALP learns a Chain-of-Thought (CoT) prompt ($\beta$) consisting of up to $c$ tokens. This prompt conditions an LLM to score the plausibility of a triple $(h, r, t)$ or predict missing components.

Key Components

• Scoring Function: The model defines a scoring function $\phi_{\beta \le c}: E \times R \times E \to \mathbb{R}$, where the parameter $\beta$ is a discrete string (prompt) rather than a continuous vector.
• Optimization via MIPRO: Since the search space of strings is discrete and non-differentiable, the authors use MIPRO (a Bayesian Optimization algorithm) to find the optimal prompt.
  • Process: MIPRO iteratively selects candidate prompts, evaluates them on a validation set (using metrics like Cross-Entropy or MRR), and updates a probabilistic model to select the next promising candidate.
  • Efficiency: This allows the system to find high-performing prompts using fewer than 30 training examples without accessing model gradients.
• Two Variants:
  1. RALP: Uses a default, handcrafted prompt with in-context learning (few-shot examples).
  2. RALP†: Uses MIPRO to automatically learn and optimize the prompt specifically for the target Knowledge Graph and task.

Specific Tasks Addressed

1. Triple Completion (Entity/Relation Prediction): Given $(s, r)$, the LLM predicts the object $o$ and a confidence score.
2. Numerical Literal Prediction: For data properties (e.g., GDP, population), the model retrieves relevant context (other entities with the same relation) and performs CoT reasoning to predict a point estimate and a confidence interval (e.g., $[y_{min}, y_{max}]$).
3. OWL Instance Retrieval: The model predicts instances satisfying complex class expressions (e.g., $\exists hasChild.Female$) by interpreting logical syntax (Manchester or Description Logic) and reasoning over the graph.
4. Knowledge Graph Enrichment: An algorithm (Algorithm 1) uses RALP to query existing $(s, r)$ pairs, identify high-confidence missing triples, and add them to the training set to augment traditional KGE models.

3. Key Contributions

• Novel Formulation: First approach to treat numerical link prediction and complex logical reasoning as a prompt-based reasoning problem using LLMs directly, bypassing regression layers or binning.
• Gradient-Free Optimization: Demonstrates that effective scoring functions can be learned via Bayesian Optimization (MIPRO) with minimal data ($\le 30$ examples) and no gradient access.
• Generalization: Successfully handles unseen entities, relations, and literals, overcoming the zero-shot limitations of traditional KGE.
• Data Augmentation: Shows that RALP-predicted triples can be used to augment training data, significantly boosting the performance of standard KGE models (e.g., DistMult, KeCI).
• Open Source: The implementation, training, and evaluation pipeline are released as open source.

4. Experimental Results

The authors evaluated RALP on Countries (transductive), LitWD1K (numerical literals), and Father (OWL reasoning) datasets.

• Transductive Link Prediction:
  • RALP† achieved 1.000 MRR and Hits@1 on the Countries-S1 dataset, setting a new state-of-the-art (SOTA).
  • When used as a data augmentor, RALP improved SOTA KGE models by over 5% MRR across datasets. For example, DistMult improved from 0.783 to 0.990 MRR on S1.
• Numerical Literal Prediction:
  • Achieved high accuracy in predicting numerical values (e.g., GDP per capita) with low Mean Squared Error (MSE).
  • Interval Coverage Rate (ICR) exceeded 75% for all tested properties, indicating well-calibrated uncertainty estimates.
  • The model successfully performed CoT reasoning (e.g., calculating GDP per capita by dividing total GDP by population) to derive predictions.
• OWL Instance Retrieval:
  • Achieved >88% Jaccard similarity on complex class expressions.
  • Manchester Syntax outperformed Description Logic (DL) syntax, suggesting LLMs align better with human-readable logical structures.
  • Perfect performance (1.000 Jaccard) was achieved on atomic concepts and "at least" restrictions.

5. Significance and Implications

• Paradigm Shift: The paper challenges the dominance of vector-based embeddings, proposing that symbolic, string-based reasoning via LLMs is a viable and often superior alternative for dynamic and heterogeneous graphs.
• Scalability & Flexibility: By avoiding retraining embeddings, RALP offers a flexible solution for dynamic graphs where entities and relations change frequently.
• Interpretability: Unlike "black box" embeddings, the CoT prompts provide explainable reasoning (e.g., "Gabon's GDP is X, population is Y, therefore per capita is Z"), which is crucial for trust in AI systems.
• Future Directions: The work highlights the potential of prompt engineering as a learnable parameter, bridging the gap between symbolic AI (logic) and neural AI (LLMs).

Limitations Noted

• Outliers: Numerical predictions can be skewed by extreme outliers in the context (e.g., Venezuela's hyperinflation), though pre-processing can mitigate this.
• Token Constraints: The approach is limited by the LLM's context window, making it challenging to scale to massive graphs without sophisticated sampling strategies.
• Semantic Dependency: The method assumes entity/relation strings are semantically meaningful; it struggles if identifiers (like QIDs) are not mapped to natural language labels.

In conclusion, RALP demonstrates that learning to prompt is a powerful mechanism for knowledge graph completion, offering superior generalization, interpretability, and performance on complex reasoning tasks compared to traditional embedding-based methods.