Tree Search With Distributional Predictions

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 $\hat{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 $\hat{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.