Proof Search Augmented Language Models

I have read all the reviews so far. Since the authors have not addressed any of the concerns or answered any questions, I have lowered my rating slightly to reflect that.

Thanks for your patience - we've now updated the PDF and posted a general response describing our revisions!

W1: See our general response on evaluation scope.

W2: While the ‘transformer’ was originally introduced as an encoder-decoder architecture, in contemporary practice it is frequently used interchangeably for encoder-only, decoder-only, or encoder-decoder variants – see e.g. the RoBERTa paper (Liu et al, 2019) p2: “BERT uses the now ubiquitous transformer architecture [...]”

Q1 and Q2: See our general response on evaluation scope and our new LLM comparison.

W1/Q2 (Comparison with LLMs): See the results of our new comparison with GPT-4o and o1-preview in the general response.

W2 (Unfair comparison): See the additional results we describe in our general response. We can now outperform o1 with pure end-to-end training.

Q1 (Other benchmarks): See our general response on evaluation scope.

Thanks for your careful consideration and in-depth questions!

W1: See our new results in the general response; we compare to GPT-4o and o1-preview and have adopted new NTP training techniques which allow PSALMs to learn SimpleLogic robustly from only end-to-end supervision.

W2: We comment on these results in the first paragraph of Section 6.

Q1 - Could you explain the reasons why loss $L_{rule}$ is minimized much smaller than the other two?: There is inherent uncertainty in $L_{demo}$, as multiple rules may be valid choices to discharge a particular subgoal; the model has no way of knowing which will occur in the gold demonstration, which leads to a nonzero loss floor for $L_{demo}$. The basic version of $L_{E2E}$ does not converge as it gets stuck in a bad local minimum at initialization. However, we have now fixed that with a modification to the method, as described in our general response. You are right that many entries in the target matrix are zeroes, however, the full $L_{rule}$ rebalances the labels in the objective so that the model does not resort to trivially predicting 0 for every score; see Appendix A.1.

Q2 - Should different weights improve the performance when combining losses?: Yes, these losses do have different scales, and reweighting them when combining them is potentially beneficial; however, we did not pursue a detailed hyperparameter search over loss weights in favor of attempting to mitigate the issues with L_E2E alone.

Q3 - Could you quantify the impact of independent rule encoding?: In our preliminary experiments, encoding rules jointly allowed the model to learn trivial rule systems consisting of only a single fact, whose unification score with the query encoded the predicted score directly resulting in no actual proof search occurring. This does not happen with $L_{rule}$, as adjacent rules are independent so no shortcuts are learned regardless of their availability; we kept rule encoding independent in the $L_{rule}$ configuration for simplicity.

Q4 - Could you explain how the larger range for $L_{rule}$ affect its generalization and separation between positive and negative ones?: The wider score separation for $L_{rule}$ is a reflection of the fact that the model is better able to minimize $L_{rule}$, resulting in more confident unification score predictions. We have updated Figure 3 to show the score separations learned by the new $L_{E2ER}$ configuration, which learns an even cleaner separation.

Q5 - Could you provide the key reasons underlying the big differences between the last two rows of Table 1 in terms of OOD accuracy and OOD soundness?: The OOD split has two major differences with the ID split: gold proofs are deeper, and the label imbalance is different - there are now more negatives than positives. Including the original version of $L_{E2E}$ in the loss mixture in the second-to-last row of Table 1 allows the model to learn reasonable score separations for the ID split, but this solution still has many incorrect unification scores (as evidenced by the poor ID soundness). Transferring to the OOD split increases the rate of incidence of these false-negative/false-positive steps per example (as examples are deeper).

Q6 - Could you provide a baseline or baselines using state of the art LLMs?: See our general response on our new LLM comparisons. For PSALM itself, larger LLMs are not needed to address the settings in this paper, but will generally be needed in the future as the architecture is scaled to higher complexity of text and more domains.

Q7 - Could you explain how the proposed method overcomes the exploding computational cost compared to other methods?: We have not undertaken a theoretical analysis of the impact of pruning or parallelization, but empirically they both confer a major increase in the maximum proof depths we are able to reach in a reasonable amount of time. Pruning in particular is essential, as completing one proof allows the elimination of the bulk of the fringe. In principle, the algorithm is still exponential, as it is possible to construct rule sets for which search continues to find higher-scoring partial proofs very deep in the tree, but in practice examples where this occurs are rare. While pruning is a well-known technique, prior work describing its application in NTPs either only uses it for fact unification (i.e. nearest-neighbor, Minervini et al. 2021), or uses a static threshold, which doesn’t result in as powerful of an effect.

Q8 - Could you comment on the complexity of the proposed system with respect to the depth?: See our response to Reviewer xqTT; we have included a breakdown of time by depth in Figure 6 in the appendix. The SimpleLogic depth analysis algorithm (used to generate and bucket examples) is also exponential in complexity by depth - quantifying example depths higher than 6 or 7 becomes impractically slow, so we have limited our analysis to depth 6.

W1/Q2: See our general response on evaluation scope.

W2: Our new results include additional analysis of median elapsed time for varying search depths, showing that dynamic pruning prevents search costs from exploding.