Tree Search With Distributional Predictions

2026-09-26

AI Summary

This study addresses the inefficiency of queries under unknown distributions caused by uncertainty in prediction quality during tree search. To overcome this, it proposes a learning-augmented query optimization framework that characterizes prediction errors using Shannon entropy and the Earth Mover’s Distance (EMD) under tree metrics, while designing an adaptive querying strategy based on tight bounds. The primary contribution is achieving asymptotically optimal expected query complexity—yielding substantial efficiency gains when predictions are accurate, yet preserving theoretical guarantees when they fail. Experimental results demonstrate that the proposed method significantly reduces the number of queries compared to baselines that blindly trust predictions, with empirical performance closely aligning with the established theoretical bounds.

Abstract

Learning-augmented algorithms use machine-learned predictions to improve classical algorithmic guarantees when the predictions are accurate, while retaining rigorous performance guarantees when they are not. We study this paradigm for search on trees. Given a tree $T$ containing an unknown target vertex $t$, an algorithm may query any vertex $v$ and learn which neighbor of $v$ lies on the unique path from $v$ to $t$. The goal is to find $t$ using as few queries as possible. We consider the distributional setting, in which the target is drawn from an unknown distribution $p$ and the algorithm is given a predicted distribution $\widehat p$ of unknown quality. We give an algorithm with expected query complexity $O\left(H(p)+k\log \eta \right)$, where $H(p)$ is the Shannon entropy of the true distribution and $\eta$ is the earth mover's distance between $p$ and $\widehat p$ in the tree metric. We also provide a matching lower bound that shows our algorithm is asymptotically tight. Finally, experiments on real-world and synthetic trees show that our prediction-based algorithm can use substantially fewer queries than a simple baseline that trusts the prediction completely.

Problem

• tree search
• learning-augmented algorithms
• distributional predictions
• query complexity

Innovation

• Learning-augmented algorithms
• Tree search
• Distributional predictions
• Query complexity
• Earth mover's distance