Proof Search Augmented Language Models

Authors: Kaj Bostrom, Swarat Chaudhuri, Greg Durrett
Primary area: neurosymbolic & hybrid AI systems (physics-informed, logic & formal reasoning, etc.)
Keywords: reasoning, transformers, neural theorem proving, neural network architectures, differentiable algorithms

Abstract:
Transformer language models (TLMs) exhibit an impressively general range of capabilities. A growing body of work aims to harness these models for complex reasoning problems expressed in natural language. However, recent theoretical and empirical results have revealed limits to the algorithmic generalization of TLM reasoning. Transformers trained to solve deduction problems from one distribution fail to solve instances of the same problem type drawn from other distributions. We propose to improve the systematic reasoning capabilities of TLMs via a differentiable proof search module, yielding proof-search augmented language models (PSALMs).
In a PSALM, a Transformer is responsible for predicting rule and fact representations for a neural theorem prover (NTP). The NTP performs a backward-chaining search over proofs, scoring them based on a soft unification operation. Our results show that PSALMs successfully generalize in deduction tasks where vanilla transformers do not learn systematic behavior, can be adapted to more natural text with only label supervision, and robustly handle large examples where proprietary LLMs make mistakes.

Decision: reject
Reviewer scores: 5, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1729326525335

thread_type:
official_review

<thread_object>
creator: reviewer_td5p

<content>
summary:
Summary:

This paper applies an encoder-only model to represent the rule and statement in the proof. Then, the Neural Theorem Prover (NTP) utilizes the representation to perform a backward-chaning search of proofs and sort them.

Contributions:

This paper proposes a new method to improve the model reasoning abilities by semantics proof search. The experimental results on the SimpleLogic task are nice.

strengths:
- This paper proposes a new method to improve the model reasoning abilities by semantics proof search. The experimental results on the SimpleLogic task are nice.

weaknesses:
1. The experiment section is insufficient. First, only one dataset named SimpleLogic is used in the experiment. Second, the author only uses one model in the experiment. Last but not least, there are some work aimed to improve the proof reasoning abilities. However, I have not seen the gap between those work and the method proposed by the authors.

2. The terminology in this article is not to the standard. For instance, we typically refer "transformer" to encoder-decoder architecture models instead of encoder-only models.

questions:
1. Why are language models such as Llama not used in the experiments?
2. If language models are used, do we still need the complex procedure to make it work?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730660924423

thread_type:
official_review

<thread_object>
creator: reviewer_1xbq

<content>
summary:
The paper proposes a PSALM model that combines a transformer-based language model with a proof search system. Compared to the NTP introduced by Rocktaschel & Riedel in 2017, it uses pruning and parallel execution to improve scaling and throughputs. Using a new loss term, the system provides significant improvements over a vanilla model for (limited) out-of-distrbution generalization.

strengths:
The paper is well written and the key components are motivated and explained well. The experimental results seem to be systematic and convincing.

weaknesses:
The novelty introduced in the paper seems limited. Pruning is well known; while it improves scaling, in the worst case, the complexity remains the same. Also the improvements over the vanilla model do not reflect the improvements over the state of the art transformer models.

The experimental results are not fully explained. For example, the row of PSALM L_{rule} in Table 1is not explained and no comments are provided in the paper even though it is a key result.

questions:
- Could you explain the reasons why loss L_{rule} is minimized much smaller than the other two? It seems to me it shows that many entries in matrix T are zeroes.

- As a follow-up question, since the loss terms can be minimized to a different scale, should different weights improve the performance when combining them such as the results in Tables 1 and 2?

- Section 3.1 states that "Encoding rules independently prevents the TLM from “shortcutting” the NTP," could you provide corresponding empirical evidence and quantify the impact?

- Figure 3 shows the range of proof score is different from the case on the left and the case on the right. Could you explain how the larger range for L_{rule} affect its generalization and separation between positive and negative ones?

- 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?

- I assume the vanilla TLM model is the DeBerta model with 435M parameters. However, there are manyother larger LLMs developed. Could you provide a baseline or baselines using state of the art LLMs?

- Could you explain how the proposed method overcomes the exploding computational cost compared to other methods? As far as I could understand, the pruning and parallelization do not change the nature of exponential growth.

- Could you comment on the complexity of the proposed system with respect to the depth? In addition to depths 5-6, could you provide results for depths 7, 10, and 20?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702402694

thread_type:
official_review

<thread_object>
creator: reviewer_xqtt

<content>
summary:
The authors introduce Proof Search Augmented Language Models (PSALMs), a differentiable proof search module combined with a transformer. The authors propose an efficient hardware-aware method for proof search (at further depths than prior works) and pruning and performing ablations to identify the strengths of granular rule supervision.

strengths:
- The paper presents multiple training objectives and studies which impact performance.
- An efficient and hardware-aware algorithm has been proposed for proof search.
- The authors show evaluations of the SimpleLogic dataset and generalization capabilities.

weaknesses:
- There are limited empirical evaluations other than SimpleLogic.
- More investigation could be conducted into the scaling of the proof search.

questions:
- How does the approach work on tasks other than SimpleLogic?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714843690

thread_type:
official_review

<thread_object>
creator: reviewer_hfn3

<content>
summary:
The paper focuses on multi-step reasoning tasks that require a model to predict whether a statement is true given unification rules and facts in natural language. The authors first define several rule templates that contain different numbers of unification terms, then train a Transformer and cross-attention module to fill the term slots with entities and their features from a sentence. Lastly, the extracted rule is input to a neural theorem prover to obtain the proof and the truth prediction. The author uses three types of supervision: label supervision, proof supervision, and rule supervision. Experiments show training with rule supervision can obtain 96.7% accuracy on the OOD test split.

strengths:
The proposed method is effective. The experiments show the proposed method trained with rule labels can obtain nearly 100% accuracy on the OOD splits, which largely improves the previous methods.

weaknesses:
- The evaluation dataset and the compared method are too simple. Since the paper claims current TLMs have limits in reasoning, it should compare with recent SOTA LLMs and show their incapability in a complicated benchmark.

- The comparison is unfair. The proposed method with proof and label supervision achieves inferior performance to baseline TLM and only obtains nearly 100% accuracy with rule-level supervision. A more appropriate baseline should also have these labels. For example, an LLM trained to generate text in "fun :- happy kind" format and predict the label with CoT.

questions:
- Is there any harder QA benchmark that involves logic or multi-step reasoning suitable to evaluate the proposed method?

- How do current LLMs w/wo CoT perform in the SimpleLogic dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732767399603

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to reviewers

comment:
Thanks to all the reviewers for their feedback!

**New results:**

We have new results with a further modified NTP that is able to achieve 99% OOD accuracy using only end-to-end supervision. We achieve this by replacing the hard min/max in the NTP score aggregation with a soft relaxation, and injecting small amounts of noise in the unification score function to help escape poor local minima at initialization. This removes the previous need for additional supervision in the form of gold symbolic rule systems. We have revised our contributions (line 71), Table 1, discussion in Section 6, and conclusion to reflect these new results. We describe the new objective in Section 4.1 and Appendix A.2.

We have also included GPT-4o (80%) and o1-preview (96%) in our evaluation to provide a strong pre-trained LLM reference point; see Table 1, discussion at the start of Section 6, and Figure 9 in Appendix A.4.

Content changes in the updated PDF are highlighted in cyan.

**Scope of evaluation:**

Our main goal in this paper is to demonstrate a way to overcome the poor inductive bias of transformer models when trained to decide logical reasoning problems. We adopt SimpleLogic as a testbed because it was previously introduced specifically to highlight this shortcoming in vanilla transformers. Other, more difficult logical reasoning benchmarks exist, e.g. P-FOLIO, PrOntoQA, LogicBench, and MUSR. However, the logical structures backing these datasets are richer: P-FOLIO involves nested disjunction and negation, PrOntoQA involves quantification and degree properties, LogicBench involves quantification, negation, disjunction, and defeasibility, and MUSR involves both complex predicate-argument and temporal/causal structure, including implicit defeasible world knowledge. For PSALMs to succeed beyond propositional Horn clause systems (SimpleLogic), we would need to enrich the syntax of the NTP rule structures significantly. This is not an inherent or insurmountable limitation of the PSALM architecture, but rather a pragmatic limitation of the rule projection layer we use in the current version of the system. Changing this would increase the complexity beyond what we feel is appropriate in scope for a single ICLR paper with a focused contribution.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes Proof Search Augmented Language Models (PSALMs), which combine a transformer-based language model with a proof search system. The authors use three types of supervision: label, proof, and rule supervision. They conduct experiments on the SimpleLogic dataset to demonstrate the effectiveness of their method in overcoming the poor inductive bias of transformer models in logical reasoning problems. However, 1) the evaluation dataset and compared methods are too simple. the high accuracy of CoT-LLM shows the problem may not be challenging enough, and a harder benchmark is needed. 2) Limited novelty as pruning is well-known, and improvements over the vanilla model do not reflect improvements over SOTA transformer models.  Overall, the authors should verify their proposed method with more reasonable benchmarks with more experimental results as well as comprehensive analysis.