Kwai-STaR: Transform LLMs into State-Transition Reasoners

TL;DR

Kwai-STaR transforms large language models into state-transition reasoners to improve mathematical problem-solving by modeling reasoning as state transitions, leading to enhanced accuracy and efficiency.

Contribution

The paper introduces a novel framework that leverages state transition modeling to enhance LLMs' mathematical reasoning capabilities through a curriculum training strategy.

Findings

Significant performance improvements on GSM8K and GSM-Hard datasets.

Enhanced training and inference efficiency of LLMs using the state transition approach.

Effective generalization demonstrated across different LLM architectures.

Abstract

Mathematical reasoning presents a significant challenge to the cognitive capabilities of LLMs. Various methods have been proposed to enhance the mathematical ability of LLMs. However, few recognize the value of state transition for LLM reasoning. In this work, we define mathematical problem-solving as a process of transiting from an initial unsolved state to the final resolved state, and propose Kwai-STaR framework, which transforms LLMs into State-Transition Reasoners to improve their intuitive reasoning capabilities. Our approach comprises three main steps: (1) Define the state space tailored to the mathematical reasoning. (2) Generate state-transition data based on the state space. (3) Convert original LLMs into State-Transition Reasoners via a curricular training strategy. Our experiments validate the effectiveness of Kwai-STaR in enhancing mathematical reasoning: After training on the small-scale Kwai-STaR dataset, general LLMs, including Mistral-7B and LLaMA-3, achieve considerable performance gain on the GSM8K and GSM-Hard dataset. Additionally, the state transition-based design endows Kwai-STaR with remarkable training and inference efficiency. Further experiments are underway to establish the generality of Kwai-STaR.