You Don’t Have to Go the Distance

Neural models use large external memories accessed through learned representations. The cost of inference depends on the length of search for each query. Average recall gives equal weight to all recovered neighbors, while their recovery requires significantly different amounts of computation. We propose a learned search policy, TraCANN, that takes advantage of this asymmetry to reallocate the computation under a shared recall target. TraCANN's training trajectories provide the costs and neighbor recovery statistics used to fit continuation actions jointly, and it selects an action based on one policy evaluation from each query's early search progress. By learning the relative costs of recovery, TraCANN can stop promising searches before exhausting their candidate sets. Compared to eleven baselines, TraCANN's throughput is the highest at all evaluated recall thresholds on four standard ANN benchmarks. Our analysis demonstrates cases when such reallocations reduce computation and bounds the difference between discrete and fractional continuation assignments.