DistRAG: Towards Distance-Based Spatial Reasoning in LLMs

TL;DR

DistRAG enhances large language models with a distance-based retrieval mechanism that encodes geodesic distances in a graph, improving their ability to perform spatial reasoning tasks involving distances.

Contribution

The paper introduces DistRAG, a novel method that enables LLMs to retrieve relevant spatial information from a graph to improve distance-based reasoning.

Findings

LLMs can answer distance-based questions more accurately with DistRAG.

DistRAG effectively encodes geodesic distances for spatial reasoning.

The approach provides a step towards a world model for LLMs.

Abstract

Many real world tasks where Large Language Models (LLMs) can be used require spatial reasoning, like Point of Interest (POI) recommendation and itinerary planning. However, on their own LLMs lack reliable spatial reasoning capabilities, especially about distances. To address this problem, we develop a novel approach, DistRAG, that enables an LLM to retrieve relevant spatial information not explicitly learned during training. Our method encodes the geodesic distances between cities and towns in a graph and retrieves a context subgraph relevant to the question. Using this technique, our method enables an LLM to answer distance-based reasoning questions that it otherwise cannot answer. Given the vast array of possible places an LLM could be asked about, DistRAG offers a flexible first step towards providing a rudimentary `world model' to complement the linguistic knowledge held in LLMs.