Thoughts-as-Planning: Latent World Models for Chain-of-Thoughts Optimization via Reinforcement Planning
About
The success of large language models (LLMs) across diverse NLP tasks has elevated the importance of reasoning chain optimization as a critical step in aligning model behavior with task objectives. Existing reasoning chain tuning methods often rely on black-box heuristics or gradient-free search, which lack interpretability, generalization, and sample efficiency. In this work, we introduce \textbf{Thoughts-as-Planning}, a novel framework that formalizes reasoning chain optimization as a sequential decision-making process over a latent semantic space. We model the LLM as a partially observable environment and learn a latent world model that simulates the effect of reasoning chain edits on downstream outputs. A proximity-preserving embedding space is constructed to encode reasoning chain-response dynamics, enabling planning via gradient descent or reinforcement learning. Our method supports multi-scale abstraction, allowing reasoning chain edits at token, segment, and instruction levels to be integrated into a unified planner. Through extensive experiments on language understanding and generation tasks, we demonstrate that Thoughts-as-Planning outperforms state-of-the-art reasoning chain tuning baselines in efficiency, robustness, and generalization, while offering interpretability through its structured planning trajectory. Our code is available at https://github.com/FastLM/Thoughts-as-Planning.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Commonsense Reasoning | StrategyQA (test) | Accuracy37.1 | 119 | |
| Physical Commonsense Reasoning | PIQA | Accuracy (PIQA)81.5 | 99 | |
| Multi-hop Reasoning | StrategyQA | Accuracy33.9 | 50 | |
| Logical reasoning | LogiQA | Accuracy21.6 | 8 | |
| Reasoning Chain Optimization | Reasoning tasks | Query Count47 | 3 | |
| Text Generation Evaluation | AlpacaEval | Performance Score76.1 | 3 | |
| Question Answering | PIQA | Accuracy79.8 | 1 | |
| Summarization | Reddit TL;DR | Performance Score21.8 | 1 |