Proof Search Augmented Language Models

Kaj Bostrom, Swarat Chaudhuri, Greg Durrett

Keywords: reasoning, transformers, neural theorem proving, neural network architectures, differentiable algorithms

TL;DR: We add differentiable proof search to transformers to improve generalization across problem distributions

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.

Primary Area: neurosymbolic & hybrid AI systems (physics-informed, logic & formal reasoning, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 12133