Deliberate Reasoning in Language Models as Structure-Aware Planning with an Accurate World Model

TL;DR

This paper introduces SWAP, a novel structure-aware reasoning framework for language models that combines structured knowledge, symbolic verification, and diverse candidate exploration to improve multi-step reasoning accuracy.

Contribution

SWAP integrates entailment graphs, a policy-driven graph construction, and diversity-based sampling with contrastive ranking, advancing reasoning capabilities beyond prior natural language approaches.

Findings

SWAP outperforms baseline models on math, logic, and coding benchmarks.

Structured knowledge and symbolic verification enhance reasoning consistency.

Diversity sampling and contrastive ranking improve candidate plausibility and ranking accuracy.

Abstract

Enhancing the reasoning capabilities of language models (LMs) remains a key challenge, especially for tasks that require complex, multi-step decision-making where existing Chain-of-Thought (CoT) approaches struggle with consistency and verification. In this paper, we propose a novel reasoning framework, referred to as Structure-aware Planning with an Accurate World Model (SWAP), that integrates structured knowledge representation with learned planning. Unlike prior methods that rely purely on natural language reasoning, SWAP leverages entailment graphs to encode structured dependencies and enable symbolic verification of intermediate steps. To systematically construct and update the graph, SWAP employs a policy model to propose candidate expansions and a world model to predict structural updates. To improve accuracy, the world model generates multiple alternative updates, and a discriminator re-ranks them based on plausibility. To encourage diverse exploration, we introduce Diversity-based Modelling (DM), which samples candidates from the remaining probability mass after removing previously sampled candidates from the original policy distribution. Additionally, SWAP improves the discrimination accuracy through Contrastive Ranking (CR), which directly compares candidates within prompts and incorporates meta-knowledge to improve ranking quality. We evaluate SWAP across diverse reasoning-intensive benchmarks including math reasoning, logical reasoning, and coding tasks. Extensive experiments demonstrate that SWAP significantly improves upon the base models and consistently outperforms existing reasoning methods.