KeeA*: Epistemic Exploratory A* Search via Knowledge Calibration

We sincerely appreciate the reviewers’ insightful comments and constructive suggestions. We will carefully revise the manuscript and refine the relevant details accordingly. Our detailed responses to each of the concerns are presented as follows.

Q1: Additional overhead is involved in assigning scores to clusters.

A1: We agree that calculating cluster reward increases the computational overhead per node expansion, because an additional presampling step is required to evaluate each cluster, along with computations for cluster rewards and path‑aware sampling. However, in practical scenarios, the computational overhead is usually dominated by calls to the policy and value functions as well as the realization of state transitions, and thus the additional cost of node selection is negligible. In particular, by improving search efficiency, our search algorithm is able to identify solutions with fewer node expansions, which in turn reduces the overall computation time. Taking retrosynthesis as an example, the primary bottleneck lies in the policy network’s single-step reaction prediction. On the USPTO benchmark, the average decision time of KeeA* is $38.70$s, which is only slightly higher than SeeA* (Cluster)'s $37.98$s.

Q2: Concerns on the experimental benchmarking, and comparisons of this work to SeeA* with UCT.

A2: In the logic synthesis task, the average ADP reduction rate of SeeA* (Cluster) is $22.5\\%$, while KeeA* achieves $25.2\\%$. In the retrosynthesis task, SeeA* (UCT) achieves an average success rate of $63.31\\%$, below KeeA*’s $63.76\\%$. SeeA* (UCT) is not included as a baseline in this paper due to its inferior performance compared to SeeA (Cluster). Although SeeA* (UCT) outperforms SeeA* (Cluster) on the USPTO test set, this benchmark comprises only $190$ carefully selected molecules, and the guiding heuristics are also trained on the USPTO dataset, raising concerns of potential overfitting. Across $4,909$ diverse test cases from seven benchmarks, SeeA* (Cluster) performs better than SeeA* (UCT). Moreover, the actual success rate of KeeA* is $98.9473\\%$, consistent with that of SeeA* (UCT), and the discrepancy arises from a rounding error. We appreciate your attention to this detail. Relevant results will be included in the revised paper.

Q3: Variability of KeeA* due to randomness.

A3: To quantify KeeA*’s variability induced by randomness, we executed three runs for each of ten random seeds on the USPTO test set. A t-test with the $0.05$ significance level is used to compute the error intervals. The average success rate over all $30$ runs is $98.53\\%$, which is comparable to the $98.84\\%$ reported in the paper. KeeA* exhibits robust performance across various random seeds, achieving more stable results compared to SeeA* (Cluster). SeeA* not only exhibits a larger variance across multiple runs but also shows greater performance fluctuations under different random seeds, highlighting the superiority of KeeA* over SeeA* (Cluster) with respect to stability.

Q4: Choice of score metric for the clusters, and the interplay with path-aware sampling

A4: We agree that the clustering score metric can be improved, and it is necessary to design more effective metrics. However, the mean-to-variance ratio employed in this paper aligns with a commonly used risk-adjusted evaluation metric, designed to provide solutions with higher returns relative to their variability, similar to the Sharpe Ratio. Clusters with reduced variance exhibit more concentrated and stable performance, leading to higher reliability in obtaining rewards and improved robustness under uncertainty. Clusters with high mean and low variance are preferred, while those with low mean and high variance are least desirable. In clusters with high mean and high variance, while the chance of sampling superior nodes increases, the risk of sampling inferior nodes also grows, thereby reducing performance stability. For clusters with low mean and low variance, the upper extremes of node quality are constrained, yet their stability increases the probability of achieving relatively consistent and satisfactory performance. The latter two scenarios involve a trade-off between potential gains and stability. Therefore, the cluster scoring method employed in this paper is reasonable. In three independent runs on the USPTO test set, using the mean alone yields an average success rate of $98.07\pm 1.64\\%$, whereas the result of MEAN/STD scoring is $98.53\pm 1.49\\%$, achieving higher average performance with lower variance and reflecting its greater robustness and overall superiority.

In the situation that deeper nodes may receive a higher reward $f$ than shallower ones, selecting nodes with lower $f$-values at shallower depths reflects the algorithm’s exploratory behavior, rather than greedy expansion based on maximal estimates $f$. Since $f$ is a biased estimate of the true return $f^\*$, the node with the highest $f$ does not necessarily correspond to the optimal node with the highest $f^\*$. Depth-penalty encourages selecting nodes beyond those with the highest $f$ values, introducing exploration into the search algorithm to escape local optima. Depth-based penalties have also been adopted in LevinTS [1] to encourage exploration. It is a classical exploitation-exploration trade-off in search algorithms, and is not necessarily detrimental to the search algorithm.

[1] Orseau, Laurent, et al. "Single-agent policy tree search with guarantees." Advances in Neural Information Processing Systems 31 (2018).

Q5: Performance improvement compared to SeeA*

A5: KeeA* demonstrates greater robustness than SeeA*, while consistently yielding solutions with significantly higher quality. Both SeeA* and KeeA* are sampling-based search algorithms. As shown in Table A3, KeeA* exhibits greater stability across different random seeds, achieving an average success rate of $98.53\pm 0.22\\%$ , outperforming SeeA* (Cluster), which achieves $97.30\pm 0.35\\%$. Furthermore, solution length is a critical indicator of solution quality. With $30$ test runs, KeeA* achieves a significantly shorter average path length of $6.15\pm0.06$ versus $6.78\pm0.10$ for SeeA* (Cluster) on USPTO benchmark . Similar advantages in solution quality are also observed on other test sets, as shown in Table 2 of the paper. The substantial difference in solution quality indicates that KeeA*’s advantage is not due to hyperparameter settings obtained by chance. In addition, we conducted a grid search over different hyperparameter settings, with three trials conducted for each configuration. The results are consistent with those in Appendix G: KeeA* demonstrates robust performance across a wide range of hyperparameter settings and outperforms $97.30\pm 2.36\\%$ of SeeA* (Cluster), for which the number of runs is also limited to three. It need to be noticed that the limited $3$ runs lead to higher variability than the $30$-run setting in t-test.

Q6: Theoretical contributions of this work.

A6: In this work, we theoretically demonstrate that cluster sampling outperforms uniform sampling in expectation for the first time. It is not readily intuitive since the advantage of cluster sampling does not universally hold. In uniform sampling, drawing $m$ samples from $N$ nodes yields a probability of $m/N$ for including the optimal node. Under cluster sampling, $m/K$ nodes are sampled from each of the $K$ clusters uniformly. If the optimal node resides in a large cluster with more than $N/K$ nodes, the probability of selecting the optimal node under cluster-based sampling will be lower than $m/N$ in uniform sampling, leading to reduced search efficiency. This paper offers a complementary perspective by first showing that cluster sampling increases the variance of collected candidates. Leveraging the extreme value theorem, a higher expected maximum is achieved by cluster sampling, providing a theoretical justification for SeeA* (Cluster). What's more, Theorem 5.1 establishes a theoretical connection between the predicted value $f$ and the ground-truth value $f^\*$, which is rarely investigated in neural-guided search algorithms.

We sincerely appreciate the reviewer’s insightful comments and constructive suggestions. We provide detailed responses to the concerns raised below, and the relevant revisions and clarifications will also be incorporated into the updated version of the manuscript.

Q1: Please clarify the advantages and disadvantages of cluster sampling.

A1: Following SeeA*, competitive learning as a distance-based clustering algorithm is employed to reduce computational overhead, where nodes are grouped by measuring their distance in the embedding space. A key advantage of this clustering algorithm is that it improves computational efficiency by avoiding re-clustering from scratch with each node generation. Newly generated nodes are incrementally assigned to the nearest cluster, with the cluster centers updated in turn. Meanwhile, besides the initialization of cluster centers, cluster sampling depends on the quality of feature vectors extracted by neural networks and the sequence of node expansions. As this study primarily focuses on tree search algorithms, the development of more advanced clustering methods is an essential avenue for future research.

Q2: Why keep A* in the retrosynthetic planning problems and the exclusion of state-of-the-art machine learning algorithms.

A2: This paper focuses on neural network‑guided heuristic search algorithms, aiming to enhance search efficiency when the estimated $f$-values deviate from the true optimal cost $f^\*$. A* search is the foundation of both SeeA* and KeeA*, and is retained as a baseline search algorithm. KeeA* is designed as a general‑purpose problem‑solving algorithm, not limited to retrosynthesis planning. SOTA general search algorithms, such as LevinTS and PHS, are also evaluated under the same heuristic function for a fair comparison.

As summarized in Appendix D, retrosynthetic planning aims to identify a synthetic route from available building block molecules to a target compound. One-step retrosynthesis model is employed to predict single-step chemical reactions that can generate the target molecule. If the reactants of a reaction are non-building-block molecules, one-step model is called recursively to find chemical reactions that can decompose these molecules until all molecules are reduced to building blocks. Machine learning algorithms are commonly employed to learn the one-step retrosynthesis model to guide the search process, and a synthesis route composed of multiple chemical reactions typically requires search algorithms for a solution. Directly generating the entire pathway using machine learning models remains highly challenging, and current SOTA retrosynthesis algorithms still rely heavily on search algorithms.

The following table is a comparative analysis on the USPTO benchmark against SOTA retrosynthesis algorithms. Notably, KeeA* attains comparable performance despite using less accurate heuristics, benefiting from search algorithm refinements without the need for expensive neural network training. We also test EG-MCTS [4] on the entire seven benchmarks, attaining an average success rate of $62.78\\%$, below KeeA*’s $63.76\\%$. Furthermore, both GRASP [1] and RetroGraph [2] are not open-source. PDVN [3] did not release their trained guiding heuristic models. Therefore, the results for [1,2,3] on all seven benchmarks could not be obtained. The latest DESP algorithm [5] requires processing a large set of starting molecules, leading to significant computational overhead and failing to yield feasible solutions within the 10-minute time constraint.

[1] Yu, Yemin, et al. "Grasp: Navigating retrosynthetic planning with goal-driven policy." Advances in Neural Information Processing Systems 35 (2022): 10257-10268.

[2] Xie, Shufang, et al. "Retrograph: Retrosynthetic planning with graph search." Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2022.

[3] Liu, Guoqing, et al. "Retrosynthetic planning with dual value networks." International conference on machine learning. PMLR, 2023.

[4] Hong, Siqi, et al. "Retrosynthetic planning with experience-guided Monte Carlo tree search." Communications Chemistry 6.1 (2023): 120.

[5] Yu, Kevin, et al. "Double-ended synthesis planning with goal-constrained bidirectional search." Advances in Neural Information Processing Systems 37 (2024): 112919-112949.

Q3: Investigations on the hyperparameter $K$.

A3: As indicated in Line 204 and Algorithm 1, $K$ represents the number of clusters, which is also a critical parameter of SeeA* (Cluster) and was comprehensively analyzed in the original study. KeeA* adopts the same $K$ value as SeeA* in the paper. To assess the influence of $K$ on KeeA*, three runs are performed for various $K$, and a t-test with the $0.05$ significance level is used to compute the error intervals. Because only three runs are conducted, the error interval is relatively large. KeeA* consistently outperforms SeeA* (Cluster) on the USPTO test set. An excessively small $K$ induces under-fitting, leading to a model of limited flexibility, while an overly large $K$ risks overfitting to collected states. Both extremes impair the efficacy of decision-making during search, highlighting the importance of appropriate cluster number selection to optimize the efficiency of SeeA* and KeeA*. We will include the relevant discussions and results in the revised manuscript to address this point.

We sincerely thank the reviewer for the valuable comments. We will carefully revise the manuscript to correct typos and clarify variable notations to improve readability. Our responses to the concerns are provided below.

Q1: Theorem 5.1 still rests on homoscedastic Gaussian error assumptions

A1: Due to the presence of prediction errors, for two nodes $n_1$ and $n_2$, we have that $f_1=f^\*_1+\epsilon_1$ and $f_2=f^\*_2+\epsilon_2$. Define $\Delta=f_2-f_1>0$ and $\eta=\epsilon_2-\epsilon_1$.

Since $P(\Delta>\eta|\Delta)$ is increasing with $\Delta$, $P(f^\*_2>f^\*_1|f_2-f_1=\Delta)$ also increases with the gap between $f_1$ and $f_2$. The monotonicity property established in Theorem 5.1 remains valid even when the noise distribution $\epsilon$ is non-Gaussian and node-dependent; however, Equation (11) can no longer be derived.

Q2: Grid search on both SeeA* and KeeA* and confidence intervals.

A2: KeeA* extends SeeA* by incorporating two additional hyperparameters, $\alpha$ and $\beta$. Other parameters remain consistent with the optimal settings provided in SeeA*. To demonstrate the robustness of KeeA*, an exhaustive grid search is conducted with three independent runs per configuration, and a t-test at the $0.05$ significance level is used to compute the error intervals. Because only three runs are conducted, the deviation is relatively large. Both the mean and the error interval are reported in the following table. KeeA* demonstrates robust performance, consistently outperforming SeeA*'s $97.30\pm 2.36\\%$ across a wide range of hyperparameter settings, with the highest success rate reaching $98.77\pm 0.62\\%$. This observation aligns with the findings shown in Appendix G, which provides the performance of KeeA* under various $\alpha$ and $\beta$.

To strengthen the statistical validity of Figures 4–6, confidence intervals (t-test, $\alpha$=0.05) based on empirical standard deviations can be computed. Revisions will be incorporated into the final manuscript. Due to time constraints, we have not completed the grid search and confidence interval calculations for the McNemar test across all seven datasets. We will include the comprehensive results in the final version of the paper.

We sincerely thank the reviewer for the insightful comment. The following content will be added to the revised manuscript to provide further clarification.

Q1: Discuss application scenarios where KeeA* might underperform compared to SeeA*.

A1: KeeA* and SeeA* are both search algorithms designed to address the discrepancy between the heuristic estimate $f$ and the true return $f^\*$. Compared to SeeA*, KeeA* introduces two key improvements:

(1) KeeA* allocates the number of sampled nodes from each cluster in proportion to the mean of the estimated $f$-values of their nodes, while SeeA* samples the same number of nodes from each cluster,

(2) KeeA* improves the probability of sampling nodes at shallower depths for exploration, instead of uniformly selecting candidates from the chosen cluster as in SeeA*

The improvement (1) exploits the observation that, despite potential bias, $f$ estimated by a well-trained heuristic remains informative. The improvement (2) introduces a depth-based penalty to encourage exploration. Given the exponential growth of the search tree, the number of nodes increases significantly at deeper levels; favoring shallower nodes helps escape from over-explored branches. The effectiveness of both improvements depends on the degree of bias between the heuristic estimate $f$ and the true evaluation $f^\*$. Therefore, there are two scenarios where KeeA* might underperform SeeA*:

Case (1): If the prediction quality of $f$ is very low, then the clusters with higher $f$ are more likely to have lower true returns $f^\*$. In such cases, KeeA* may over-sample from suboptimal clusters, reducing the likelihood of selecting the optimal node and thereby impairing search efficiency.

Case (2): If the prediction quality of $f$ is sufficiently high, greedy expansion based on $f$-values becomes an effective strategy to identify the node with the highest true return $f^\*$. The depth-based penalty used to encourage exploration introduces unnecessary node expansions, reducing the search efficiency of KeeA*.

Excluding the two aforementioned extreme cases, KeeA* leverages both exploitation of the heuristic function $f$ and exploration encouraged by the depth-based penalty, enhancing search efficiency under biased estimates by effectively balancing exploitation and exploration.