Learning to Reason via Program Generation, Emulation, and Search

TL;DR

This paper introduces CoGEX, a method that extends language models' program synthesis capabilities to complex reasoning tasks by generating pseudo-programs, emulating their execution, and searching for optimal solutions, significantly improving performance on diverse tasks.

Contribution

The paper presents CoGEX, a novel approach combining program generation, emulation, and search to enable language models to handle complex reasoning tasks beyond traditional code-based problems.

Findings

CoGEX outperforms standard in-context learning on various algorithmic and reasoning tasks.

Emulating program execution helps fill knowledge gaps in language models.

Program search yields more accurate solutions across diverse datasets.

Abstract

Program synthesis with language models (LMs) has unlocked a large set of reasoning abilities; code-tuned LMs have proven adept at generating programs that solve a wide variety of algorithmic symbolic manipulation tasks (e.g. word concatenation). However, not all reasoning tasks are easily expressible as code, e.g. tasks involving commonsense reasoning, moral decision-making, and sarcasm understanding. Our goal is to extend an LM's program synthesis skills to such tasks and evaluate the results via pseudo-programs, namely Python programs where some leaf function calls are left undefined. To that end, we propose, Code Generation and Emulated EXecution (CoGEX). CoGEX works by (1) training LMs to generate pseudo-programs, (2) teaching them to emulate their generated program's execution, including those leaf functions, allowing the LM's knowledge to fill in the execution gaps; and (3) using them to search over many programs to find an optimal one. To adapt the CoGEX model to a new task, we introduce a method for performing program search to find a single program whose pseudo-execution yields optimal performance when applied to all the instances of a given dataset. We show that our approach yields large improvements compared to standard in-context learning approaches on a battery of tasks, both algorithmic and soft reasoning. This result thus demonstrates that code synthesis can be applied to a much broader class of problems than previously considered. Our released dataset, fine-tuned models, and implementation can be found at \url{https://github.com/nweir127/CoGEX}.