GPU-Accelerated A* Search with Deep Neural Network Heuristics

Abstract

Heuristic search is a foundational technique in artificial intelligence. Recent work has demonstrated that powerful heuristic functions represented as deep neural networks (DNNs) can be learned with machine learning. This combination of machine learning and heuristic search has been successfully applied to puzzle solving, quantum circuit synthesis, and chemical reaction mechanism pathfinding. This success has been accompanied by batched heuristic search algorithms, such as batch A* search, which take advantage of parallelism provided by graphics processing units (GPUs) for DNNs. However, the rest of the heuristic search algorithm remains on the central processing unit (CPU). This results in the often costly transfer of data from the CPU to the GPU for heuristic computation and leaves other GPU-based acceleration options unexplored. To address this, we investigate the benefit of performing heuristic search entirely on the GPU. Using the Rubik's Cube as our case study, we systematically profile the speedup achieved by each component of batch A* search when moved to the GPU. Our results show that GPU execution yields about 12.03× speedup over the CPU baseline that only uses the GPU for the heuristic function. These findings suggest that fully GPU-native heuristic search with DNNs is a practical and effective strategy for scaling learned heuristics to demanding combinatorial search tasks.