SEGO : Sequential Subgoal Optimization for Mathematical Problem-Solving (2024)

Abstract: Large Language Models (LLMs) have driven substantial progress in artificial intelligence in recent years, exhibiting impressive capabilities across a wide range of tasks, including mathematical problem-solving. Inspired by the success of subgoal-based methods, we propose a novel framework called SE quential sub G oal O ptimization (SEGO) to enhance LLMs’ ability to solve mathematical problems. By establishing a connection between the subgoal breakdown process and the probability of solving problems, SEGO aims to identify better subgoals with theoretical guarantees. Addressing the challenge of identifying suitable subgoals in a large solution space, our framework generates problem-specific subgoals and adjusts them according to carefully designed criteria. Incorporating these optimized subgoals into the policy model training leads to significant improvements in problem-solving performance. We validate SEGO’s efficacy through experiments on two benchmarks, GSM8K and MATH, where our approach outperforms existing methods, highlighting the potential of SEGO in AI-driven mathematical problem-solving.

## introduction in recent years, the emergence of large language models (llms) has marked a significant milestone in the field of artificial intelligence. models such as chatgpt and llama have demonstrated remarkable capabilities across diverse tasks. within this context, addressing mathematical problems has attracted considerable interest from researchers, as it serves as a prominent showcase of the reasoning capabilities inherent in llms. reasoning involves a multitude of aspects, among which the ability to decompose the overall problem into smaller, more manageable subproblems (i.e., subgoals) is particularly essential for effective problem-solving. in this paper, we draw inspiration from the successful application of subgoal-based methods in both rl and llms @xcite and introduce a novel framework called sego (sequential subgoal optimization). intuitively, a good subgoal should serve as a bridge to solving a bigger problem, such that breaking down the problem into these subgoals makes the subproblems easier to solve, thereby increasing the likelihood of solving the entire problem. sego quantifies this intuition by establishing a theoretical connection between the subgoal breakdown process and the probability of solving the problem @xcite . concretely, we construct a lower bound on the probability of solving the complete problem using a proposal distribution considering a specific subgoal. we then employ a method inspired by annealed importance sampling @xcite to efficiently navigate through vast search spaces, seeking the subgoal corresponding to the theoretically optimal proposal distribution, while ensuring the process doesn't get trapped in suboptimal subgoals ( §3.2). by incorporating these sequentially optimized subgoals into the training of the policy model, we achieve significant improvements in solving mathematical problems. to empirically validate the efficacy of sego, we conducted experiments on two primary benchmarks: gsm8k @xcite and math @xcite . our approach demonstrated marked superiority against existing methods with comparable model sizes, highlighting the potential of sego in advancing the field of ai-driven mathematical problem-solving. we hope that our findings can open up new avenues for future research on the applicability of llms to complex tasks in diverse domains @xcite . ## method this work addresses challenges in subgoal-based rl, focusing on the suboptimality of generated subgoals and their selection process's lack of theoretical guarantees. we introduce the sego framework, which innovates beyond the traditional "generateselect" pipeline. sego employs a "generate-(sequentially) optimize-select" approach (figure 1 ), encompassing a policy network, subgoal generator, subgoal optimizer, reward network, and value network. notably, only the policy network is used in the testing phase. the "generate-(sequentially) optimize-select" pipeline starts with initial subgoal generation, followed by sequential optimizations. in each iteration, a new subgoal is proposed and evaluated for its increased likelihood of achieving the goal. improved subgoals are retained for further refinement. this results in a collection of refined subgoals, from which the most suitable are selected based on specific criteria. sego presents substantial advantages: (1) its sequential optimization aligns generated subgoals more closely with an optimal subgoal distribution. next, we detail subgoal-based fine-tuning in §3.2, the core of our framework. this section explains the "generate-(sequentially) optimize-select" process and how the resultant data updates various component parameters. the overall algorithm is outlined in §3.3. ## theoretical analysis we begin by constructing a lower bound on the probability of successfully solving the complete problem. this is done through the consideration of a proposal distribution focused on a specific subgoal. letting p π(•|•,g) (g | s) represent the probability of achieving a goal g from a state s under policy π(• | •, g), we have the following proposition: proposition 4.1. the objective defined below constitutes a lower bound on the probability of reaching the goal g from state s: we provide the proof in appendix a.1. next, we derive the analytical solution for the optimal subgoal distribution and obtain the following proposition. proposition 4.2. the optimal subgoal distribution satisfies the following condition: we provide the proof in appendix a.2. proposition 4.2 reveals that the optimal subgoal should not only be reachable from the starting point but also aid in ultimately reaching the final goal. we further investigate the ability of sego to provide an unbiased estimate of the z. inspired by annealed proposition 4.3. let ᾱ be defined as ᾱ = 1 n n @xmath0, wherein each α (i) adheres to the definition in eq.5. it follows that ᾱ constitutes an unbiased estimator of z. we provide the full proof of the unbiasedness in appendix a.3. proposition 4.3 reveals that the training objective for the value network can be approximated as a proportional estimate of the probability of attaining the goal g from state s following the current policy π. ## conclusion in conclusion, this work presents sego, an innovative framework aimed at improving llms' mathematical problem-solving abilities. drawing inspiration from subgoal-based rl, sego establishes a theoretical link between subgoal decomposition and the probability of solving problems. it enhances llms' performance by generating and refining problem-specific subgoals using theoretically defined criteria. empirical evaluations on benchmark datasets gsm8k and math demonstrate sego's ability to outperform existing approaches of comparable model sizes.