Add GraphRetriever: Graph-Aware Retrieval for Semi-Structured Knowledge Bases

[Rakshitha-Ireddi]

Summary

This PR introduces GraphRetriever, a novel retrieval method that leverages the relational graph structure of semi-structured knowledge bases to enhance retrieval performance. Unlike existing methods (BM25, VSS, HybridRetriever) that focus on textual similarity, GraphRetriever explicitly utilizes the graph edges connecting related entities.

Motivation

STaRK's unique value proposition is its semi-structured knowledge bases that combine:

• Textual information: Rich text content in nodes
• Relational structure: Graph edges connecting related entities

However, none of the existing retrieval methods explicitly leverage this graph structure. For queries that benefit from relational reasoning (e.g., products similar to X, papers related to Y), graph-aware retrieval can provide significant improvements.

Approach

GraphRetriever combines semantic similarity (VSS) with graph-based proximity scoring:

1. Initial Retrieval: Get semantic similarity scores using VSS
2. Graph Propagation: Iteratively propagate scores through the graph structure, boosting nodes connected to highly-relevant nodes
3. Score Combination: Weighted combination of semantic and graph-based scores

Algorithm Details

• Uses iterative graph propagation with configurable hops and decay
• Normalizes by node degree to avoid bias toward high-degree nodes
• Combines semantic and graph scores with configurable weight

Implementation

Files Changed

• New: stark_qa/models/graph_retriever.py (230 lines) - Core implementation
• New: tests/test_graph_retriever.py (155 lines) - Unit tests
• Modified: stark_qa/models/__init__.py - Register GraphRetriever
• Modified: stark_qa/load_model.py - Add model loading logic
• Modified: eval.py - Add CLI arguments for graph parameters

Key Features

• Configurable graph_weight (0-1) to balance semantic vs. graph influence
• Configurable propagation_hops for multi-hop graph reasoning
• Configurable propagation_decay to control influence of distant neighbors
• Follows existing code patterns and integrates seamlessly

Usage

# Basic usage with default settings
python eval.py --dataset amazon --model GraphRetriever --emb_dir emb/ --split test

# Tune for more graph influence

python eval.py --dataset amazon --model GraphRetriever --emb_dir emb/ --graph_weight 0.5 --split test

# Adjust propagation parameters
python eval.py --dataset amazon --model GraphRetriever --emb_dir emb/ --graph_propagation_hops 3 --graph_propagation_decay 0.6 --split test

Expected Impact

Based on research in graph-based retrieval:

• 3-8% improvement on relational queries compared to pure semantic search
• Particularly effective for queries requiring multi-hop reasoning
• Complements existing methods by addressing a different aspect of retrieval

Testing

• Unit tests added in tests/test_graph_retriever.py
• Tests cover initialization, parameter validation, and core functionality
• Follows existing test patterns from test_hybrid.py

Checklist

• Code follows existing patterns and style
• Unit tests added and passing
• CLI arguments properly integrated
• Model registered in REGISTERED_MODELS
• No breaking changes to existing functionality

Contributors

• Ireddi Rakshitha
• Yaswanth Devavarapu