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 SEquential subGoal Optimization (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 problemsolving.

heading: 1 Introduction
text: 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.
∗ This work was done during an internship at Tencent AI Lab.
In this paper, we draw inspiration from the successful application of subgoal-based methods in both RL and LLMs (Zhang et al., 2020; Zhao et al., 2023) 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 (Eq. 6). 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 (Neal, 2001) 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 (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021). 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 (Yao et al., 2022; Liu et al., 2023).
7544

heading: 2 Preliminaries

heading: 2.1 Problem Formulation
text: This study focuses on a goal-conditioned reinforcement learning (RL) framework, consisting of goal space (G), state space (S), action space (A), transition probability (P), and reward function (R). The transition probability P(s′|s,a) indicates the probability of transitioning from a current state s to a new state s′ after an action a. The reward function R(s, g) gives a reward of 1 if the goal g is reached at state s, and 0 otherwise. The policy π(a|s, g) maps state-goal pairs to actions in A.
Building on the goal-conditioned RL framework, in mathematical problem-solving, an action denotes a step in the solution process, while the state comprises cumulative actions. The (sub-)goal is the specific problem targeted for resolution. The transition probability, P(s′|s, a), uniquely assigns a probability of 1 to the state s′ = [s; a] where [; ] represents sequence concatenation, and 0 to all others. The reward function R(s, g) evaluates whether state s correctly solves the goal g. In this work, we employ a program-aided approach (Gao et al., 2023; Chen et al., 2022; Drori et al., 2022) to form the solution. For illustration, consider the goal g as “Calculate sin(30◦)”. Here, a state s could be “import math; def solve(): angle = math.radians(30);”, and an action a, “return math.sin(angle)”. This work aims to create a policy network that predicts trajectories for new goals. It uses a demonstration dataset D = {τ : (g; s0,a0, . . . , sℓ,aℓ)}, where τ represents a trajectory of length ℓ with states st and actions at at every timestep. The special state, ŝ, consists solely of essential imports and function definitions. The task is to predict a trajectory, starting from ŝ and aligned with a given goal g. This is formulated as:
p(τ | ŝ, g) = ℓ∏
t=0
π(at | st, g) · P(st+1 | at, st),
with s0 = ŝ.
(1)

heading: 2.2 Subgoal-based Reinforcement Learning
text: The main idea behind subgoal-based RL involves decomposing a challenging task into two more manageable sub-tasks, each of which can be addressed by the existing policy (Li et al., 2022). In subgoal-based RL, a typical approach consists of three phases: subgoal collection, trajectory sampling, and training, which together form a cyclical process.
The subgoal collection phase is central to this framework and follows a “generate-select” pipeline. Specifically, for a challenging goal g, the process generates a variety of potential subgoals, each paired with its respective state. A suitable subgoal, gw, and its state, sw
1, are then selected based on criteria that vary among different algorithms (Li et al., 2022; Zhang et al., 2021; Chane-Sane et al., 2021). These criteria typically ensure that the chosen subgoal is attainable from the initial state and facilitates achieving the final goal. For example, the subgoal might be “Calculate the radian value of 30◦” with the state “import math; def solve(): angle = math.radians(30);”. In the trajectory sampling phase, trajectories τ1 and τ2 are drawn from the distributions p(τ |ŝ, gw) and p(τ |sw, g) respectively. The final training phase utilizes these trajectories to optimize the policy network, thereby enhancing the ability to achieve both subgoals and the ultimate goal.

heading: 3 Method
text: 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. (2) It accurately calculates subgoal weights based on an unbiased estimate of the probability of reaching a goal from a given state.
1In this work, the subscript “w” denotes “waypoint”, which is used interchangeably with “subgoal”.
Road Map. We start by discussing the initialization fine-tuning in §3.1, which includes setting up key components and preparing initial training data. 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.

heading: 3.1 Initialization Fine-tuning
text: The SEGO framework employs the following key components: policy network, subgoal generator, subgoal optimizer, reward network, and value network, all of which are implemented using large language models (LLMs) (Touvron et al., 2023a,b; Rozière et al., 2023). We defer the training details of these components after an overview of their implementation. More details about these components are provided in Appendix B.
Policy Network. The policy network π(a | s, g) processes the current state and goal, represented as token sequences, to predict actions. This network employs standard decoding methods like greedy search or top-k sampling (Holtzman et al., 2019).

heading: Subgoal Generator and Subgoal Optimizer.
text: The subgoal generator f , represented as sw, gw = f(s, g), breaks complex tasks into simpler subtasks, transforming the current state and goal into a subgoal and its associated state. This method
ensures manageable progression towards the ultimate goal. The subgoal optimizer h, denoted as s′w, g ′ w = h(sw, gw, s, g), refines these subgoals and states to improve goal decomposition efficiency.
Reward Network and Value Network. The reward network r(s, g) evaluates if the current state achieves the goal, acting as a surrogate for the actual reward function R. The value network vπ(s, g), a regression model, assesses the success probability from a given state under policy π.
Training of the SEGO components begins with a goal collection, Dg = {g}, and a trajectory dataset, D = {τ : (g; s0,a0, s1,a1, . . .)}, created using GPT-3.5-turbo.2 This dataset comprises various mathematical problem-solving trajectories. The policy network is trained with triplets (g, si,ai) from D to predict action ai for a given state si and goal g. For each trajectory, a state st is randomly chosen, and GPT-3.5-turbo predicts the intermediate subgoal gw and its state sw. This step trains the subgoal generator to predict subgoals and their states from (st, g). GPT-3.5-turbo also introduces slight modifications to gw and sw, producing g̃w and s̃w. The subgoal optimizer is trained to restore (gw, sw) from these corrupted versions, considering the current state st and goal g.
After initializing the policy network, it gener-
2Further details regarding the trajectory dataset are elaborated in Appendix C.
Algorithm 1 SEGO: Sequential Subgoal Optimization Requires: π, f , h, r, vπ: policy network, subgoal generator, subgoal optimizer, reward network, value network, respectively.
Kmax: maximum iterations for subgoal-based fine-tuning. Dg: a collection of goals (or mathematical problems).
1: Construct the trajectory dataset D using GPT-3.5-turbo and Dg . 2: Initialize and fine-tune π, f , h, r, and vπ with Dg and D. 3: k ← 0. 4: while k < Kmax do 5: Dp,Dv ← ∅. ▷ Prepare datasets Dp for policy and Dv for value network training. 6: for τ ∈ D do ▷ Each τ is a trajectory in the form (g; s0,a0, . . . , st,at, . . .). 7: Generate a diverse set of subgoals via Eq.2. 8: Optimize each subgoal using Eq.3 and Eq.4. 9: Select a subgoal via Eq.5. 10: Sample new trajectories τ1 and τ2 utilizing the selected subgoal. 11: Calculate ᾱ and form a triplet (g, st, ᾱ). ▷ ᾱ estimates the probability of achieving g from st under π. 12: Update Dp ← Dp ∪ {τ1, τ2}, Dv ← Dv ∪ {(g, st, ᾱ)}. 13: Train policy network π with Dp and value network vπ with Dv . 14: k ← k + 1.
ates trajectories for each goal in Dg, with goals linked to human-annotated answers. Trajectories leading to correct answers are positive examples, {τ : (g : s0,a0, . . . , sℓ,aℓ)}, and those missing the correct answers are negative examples, {τ : (g : s̃0, ã0, . . . , s̃ℓ, ãℓ)}. The reward network is trained to classify the final state-goal pair (sℓ, g) as positive and (s̃ℓ, g) as negative. Simultaneously, the value network trains to approximate to 1 for (st, g) and 0 for (s̃t, g), where st and s̃t are randomly selected from their respective sets.

heading: 3.2 Subgoal-based Fine-tuning
text: The policy network, when only fine-tuned at the initialization phase, struggles with complex problems (Luo et al., 2023). Inspired by recent advancements in subgoal-based RL (Li et al., 2022; Zhang et al., 2021; Chane-Sane et al., 2021) and annealed importance sampling (Neal, 2001), we introduce a fine-tuning stage that emphasizes decomposing tasks into subgoals. Additionally, it evolves from the traditional “generate-select” pipeline to a more advanced “generate-(sequentially)optimize-select” approach.
Subgoal Collection. For each trajectory τ : (g; s0,a0, s1,a1, . . .) ∈ D, this phase aims to generate a subgoal pair (sw, gw) that decompose g into more manageable subtasks. The procedure starts by generating N independent pairs of initial subgoals, denoted as {(s(i,1)w , g(i,1)w )}Ni=1. Subsequently, each pair (s(i,1)w , g (i,1) w ) proceeds through a sequential optimization process, which yields a sequence of subgoal pairs: {(s(i,1)w , g(i,1)w ), . . . , (s(i,η)w , g(i,η)w )}, where η represents the maximum number of iterations within the sequential optimization.
Within each trajectory τ , a state st is randomly selected from the set {s0, s1, . . .}. Subsequently, the subgoal generator is tasked with producing a series of subgoals, defined as follows:
s(i,1)w , g (i,1) w = f(st, g), for i = 1, . . . , N (2)
To ensure the generation of diverse subgoals, a top-k sampling strategy (Holtzman et al., 2019) is implemented.
The pipeline then progresses to a sequential optimization process. At the j-th iteration, the subgoal optimizer proposes a potentially improved subgoal pair, which is defined as:
s(i,j)w , g (i,j) w = h(s (i,j−1) w , g (i,j−1) w , st, g),
for i = 1, . . . , N (3)
To ensure the improvement of the new subgoal pair (s(i,j)w , g (i,j) w ) over its predecessor (s (i,j−1) w , g (i,j−1) w ), it is necessary to establish a rigorous criteria for evaluation. To do that, a criteria is derived from a theoretical perspective to guarantee an unbiased estimation, as detailed in Proposition 4.3. Formally, the criteria is defined as:
∆ =βj−1 log p(s
(i,j) w , g (i,j) w | st, g; f)
p(s (i,j−1) w , g (i,j−1) w | st, g; f)
+ (1− βj−1) log ( vπ(s (i,j) w , g)
vπ(s (i,j−1) w , g)
× v π(st, g (i,j) w )
vπ(st, g (i,j−1) w )
× exp
( r(s (i,j) w , g (i,j) w ) )
exp ( r(s (i,j−1) w , g (i,j−1) w ) )
 
+ log p(s
(i,j−1) w , g (i,j−1) w | s(i,j)w , g(i,j)w , st, g;h)
p(s (i,j) w , g (i,j) w | s(i,j−1)w , g(i,j−1)w , st, g;h)
(4)
where the sequence of weights βj satisfies 1 = β0 > β1 > . . . > βη = 0. If ∆ ≤ 0, the subgoal pair at the j-th step is redefined as (s(i,j−1)w , g (i,j−1) w ); otherwise, (s (i,j) w , g (i,j) w ) is maintained. Intuitively, as the coefficient βj approaches 0, the criteria increasingly emphasizes the comparison between the values of two subgoals within the optimal distribution (Proposition 4.2), represented in logarithmic form. Specifically, vπ(sw, g) and vπ(s, gw) serve as proxies for pπ(g | sw) and pπ(gw | s), respectively. This comparison favors the subgoal that better aligns with the optimal distribution, thus incrementally steering the subgoal optimization towards more theoretically effective choices. The final term in ∆ acts as a regularization factor.
The weight α(i) associated with each subgoal pair (s(i,η)w , g (i,η) w ) is defined as follows:
logα(i) =
η∑
j=1
[ (βj − βj−1) log p(s(i,j)w , g(i,j)w | st, g; f)
+ (βj−1 − βj) ( log vπ(s(i,j)w , g)
+ log vπ(st, g (i,j) w ) + r(s (i,j) w , g (i,j) w )
)]
(5)
Subsequently, the subgoal pair (sw, gw) is selected based on a softmax distribution over these weights, i.e., (sw, gw) ∼ Softmax(logα(i)). Trajectory Sampling and Component Training. Upon obtaining a trajectory τ and the predicted subgoal pair (sw, gw), the subsequent procedure involves generating two new trajectories through the policy network. These trajectories, denoted as τ1 and τ2, are sampled from p(τ | st, gw) and p(τ | sw, g) (defined in Eq.1), respectively. Within these trajectories, for each triplet (g, si,ai), the policy network is trained to predict the action ai given the state si and the goal g.
As a byproduct of this sequential optimization process, the average coefficient ᾱ = 1N ∑N i=1 α (i) acts as an unbiased estimator that correlates with the probability of successfully achieving the goal g from the state st when guided by the policy network π (see Proposition 4.3). Leveraging this byproduct, the value network is further trained to regress towards ᾱ, using the state st and the goal g as inputs.

heading: 3.3 SEGO: Sequential Subgoal Optimization
text: After completing the initialization phase, our approach involves repeated cycles of subgoal collec-
tion, trajectory sampling and component training. This procedure leads to the development of our final framework, SEGO, detailed in Algorithm 1.
Remarks. In this work, we concentrate on mathematical problem-solving, yet our proposed methodology serves as a universal framework for tackling a wide range of complex tasks that can be modeled as goal-conditioned reinforcement learning problems (see §2.1), including code generation (Chen et al., 2021) and commonsense reasoning (Clark et al., 2018). To do that, one only needs to customize the goal, action, and state space definitions to suit the task specifics and adjust prompts for trajectory generation and subgoal prediction using GPT-3.5-turbo, aligning them with the specific requirements of the task.

heading: 4 Theoretical Analysis
text: 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:
log pπ(·|·,g)(g | s) ≥ Eq(g w ,s w |g,s) [ log pπ(·|·,g)(g | sw)+ log pπ(·|·,g w)(gw | s) + r(sw, gw)− log q(sw, gw | s, g) ] .
(6)
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:
q⋆(sw, gw | s, g)
= pπ(·|·,g)(g | sw)pπ(·|·,g w)(gw | s)exp(r(sw, gw))
Z ,
where Z = ∫∫ pπ(·|·,g)(g | s′w)pπ(·|·,g ′ w )(g′w | s)
× exp(r(s′w, g′w))ds′wdg′w.
(7)
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
importance sampling (Neal, 2001), we arrive at the following proposition:
Proposition 4.3. Let ᾱ be defined as ᾱ = 1 N ∑N i=1 α
(i), 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 π.

heading: 5 Experiments

heading: 5.1 Dataset and Evaluation
text: Evaluation and Training Data. Our model is evaluated using two datasets: GSM8K (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021). GSM8K contains 8, 792 math word problems for elementary students, with 1, 319 reserved for testing. MATH, with 12, 500 problems (including 5, 000 for testing), focuses on advanced mathematics, featuring questions from competitions like the AMC and AIME. Data preprocessing follows the methodologies in the original papers to ensure consistent evaluation. We provide the details of the training data in Appendix D.
Evaluation Metric. We evaluate by comparing the results of the solution generated by the policy network in SEGO to the provided correct answers within the datasets. For evaluation, we report the
accuracy, specifically focusing on the rate at which the policy network correctly solves the problems on the first attempt.

heading: 5.2 Baselines
text: Due to space constraints, details on the baselines are available in Appendix D.

heading: 5.3 Main Results
text: As indicated in Table 1, our key findings include: (1) SEGO’s performance on the GSM8K and MATH datasets is notable. SEGO (7B) achieves 68.7% accuracy on GSM8K and 36.8% on MATH, while SEGO (13B) reaches 72.5% and 40.0%, respectively. These results surpass those of comparable models, underscoring SEGO’s effectiveness in mathematical problem-solving; and (2) The integration of finetuning and the Program of Thought (PoT) approach substantially enhances model performance, particularly in complex tasks. This is evident in SEGO and MetaMath, where finetuning aligns models with task specifics, and in comparisons involving CodeLLaMA and LLaMA2 on the MATH dataset, showcasing PoT’s efficiency. Additionally, incorporating Sequential Subgoal Optimization into SEGO underlines the significance of strategic planning in complex mathematical problem-solving, resulting in notably improved accuracy.

heading: 6 Analysis

heading: 6.1 Ablation Study
text: In our study, we conducted ablation experiments on 7B CodeLLaMA using SEGO and three variants to assess each component’s impact: (1) -Sequential: the sequential subgoal optimization is omitted. (2) -Sequential & Subgoal: the subgoal-based finetuning is omitted. (3) -Sequential & Subgoal & FT: both subgoal-based finetuning and initialization fine-tuning are omitted. Results in Table 2 show the crucial role of sequential subgoal optimization in SEGO, with its absence in the -Sequential variant leading to reduced accuracy. The significant performance drop in the -Sequential & Subgoal & FT variant, comparable to the base 7B CodeLLaMA, highlights the collective value of all components in enhancing SEGO’s mathematical problem-solving capabilities.

heading: 6.2 Analysis of Hyperparameters
text: In this section, we conduct a detailed examination of the hyperparameters N and η, where N represents the number of sequences and η denotes the length of each sequence, as defined in Proposition 4.3. All the experiments in this section are anchored on the 7B CodeLLaMA to ensure consistency in the results.
The balance between N and η. We begin by exploring various combinations of N and η, illustrated in Figure 2, to comprehend the synergistic effects of these parameters on the model’s performance. The results on GSM8K and MATH reveal that incrementing both N and η typically enhances the model’s accuracy, achieving 68.7% on GSM8K and 36.8% on MATH at N = 2 and η = 3. However, the enhancements appear to stabilize beyond certain thresholds, indicating optimal points for these parameters.

heading: In-depth analysis of Hyperparameters N and
text: η. We further conduct an in-depth analysis of the hyperparameters N and η, examining each one’s individual impact by holding one constant and varying the other. The results are illustrated in Figure 3. From the results, it is clear that when N = 2, the model achieves peak accuracy at η = 3 for both GSM8K and MATH, with no significant gains beyond this point. Similarly, with η = 3, optimal accuracy is reached at N = 2, remaining stable thereafter.

heading: 6.3 Analysis of Subgoal Evolution

heading: Validity and Progression of Subgoals. To
text: deepen our understanding of subgoals during the Reinforcement Learning phase, we analyze the evolution of subgoal validity and its correlation with the performance on the test set. A subgoal (i.e., gw and sw) is deemed valid if both τ1 and τ2, sampled with policies π(·|sw, g) and π(·|s, gw), yield correct solutions for goals g and gw respectively. Our findings, illustrated in Figure 4 (Left), reveal a positive correlation between the progression of training steps and the percentage of valid subgoals. This increase in valid subgoals is paralleled by improvements in accuracy on both GSM8K and MATH
datasets, suggesting that the validity of subgoals is a crucial factor in enhancing the model’s problemsolving capabilities.

heading: Hardness of Problems Yielding Valid Subgoals.
text: To further our understanding of subgoals, we delve into the relationship between the hardness of problems and the emergence of valid subgoals. This analysis aims to reveal any trends in the difficulty of problems that tend to yield valid subgoals, providing insights into the learning progression. The hardness of each problem is labeled by ChatGPT, with more details available in Appendix E. The results, shown in Figure 4 (Right), reveal a correlation between training progression and the model’s ability to formulate valid subgoals for increasingly intricate problems, underscoring its evolving sophistication and adaptability in problem-solving.

heading: 7 Related Works
text: Mathematical Reasoning with LLMs. Large Language Models’ (LLMs) advancement in mathematical reasoning is largely driven by datasets like GSM8K (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021), with additional resources like MAWPS (Koncel-Kedziorski et al., 2016) and MWPToolkit (Lan et al., 2022) enhancing the field. Research focuses on two main ar-
eas: prompting strategies, involving techniques like Chain-of-Thought (Wei et al., 2023), ProgressiveHint Prompting (Zheng et al., 2023), bi-modal behavioral alignment (Zhao et al., 2024) and learning with verifications, using methods like outcomebased verifiers (Cobbe et al., 2021). Our approach, orthogonal to these methods, emphasizes adaptive curricula with subgoals to improve LLMs’ mathematical reasoning. Concurrently, MAmmoTH (Yue et al., 2023) explores instruction finetuning in LLMs for math problem-solving, a concept related to our strategy. This can be considered as an implementation of the instruction finetuning stage within our framework.
Subgoal-based RL. In reinforcement learning, Subgoal Search is crucial for navigating complex tasks, offering insights into subgoal benefits (Zhai et al., 2022), hierarchical structures (Wen et al., 2020), option selection (Jinnai et al., 2019a), and temporal abstraction (Fruit et al., 2017). Research focuses on exploring efficient strategies (Jinnai et al., 2019b; Hartikainen et al., 2019; Pitis et al., 2020; OpenAI et al., 2021) and enhancing planning through various algorithms (Eysenbach et al., 2019; Parascandolo et al., 2020; Li et al., 2022; Moro et al., 2022; Chane-Sane et al., 2021). It also develops curricula for complex subgoals (Zhang
et al., 2020, 2021). Our work addresses subgoal learning in mathematical problem-solving, exploring optimal subgoal identification within expansive state spaces. Owing to space constraints, a detailed discussion of related works is provided in Appendix F.

heading: 8 Conclusion
text: 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.

heading: Ethical Considerations
text: In accordance with the established Code of Ethics, this research exclusively utilizes data and information that is publicly accessible, thereby ensuring that no private or confidential resources are engaged.

heading: Limitations
text: While SEGO represents a significant advancement in the realm of mathematical problem-solving, several limitations need further investigation to fully harness its potential. These limitations include aspects such as the efficiency of SEGO, the scope of problem difficulty it addresses, and potential framework extensions:
(1) While SEGO demonstrates enhanced efficacy in identifying subgoals compared to non-sequential methods, there is room for improvement in efficiency. This can be addressed through dynamic resource allocation, such as adjusting the annealing schedule in response to performance metrics or the complexity of the problem at hand, alongside the deployment of more sophisticated proposal distribution mechanisms that more accurately mirror the target distribution.
(2) Our evaluation benchmarks predominantly include elementary to middle school-level problems. Exploring more complex problems, such as those at the undergraduate level, is a promising future direction.
(3) In the current SEGO framework, only the policy network is retained during inference. An intriguing future direction involves integrating the subgoal generator/optimizer and the value network to recursively decompose complex problems into simpler subgoals.

heading: Acknowledgements
text: We would like to thank the HKU NLP group and the anonymous reviewers for their helpful suggestions, which greatly improved this work. We especially appreciated the valuable discussions with Shansan Gong. This work is partially supported by the joint research scheme of the National Natural Science Foundation of China (NSFC) and the Research Grants Council (RGC) under grant number N_HKU714/21.

heading: A Proofs

heading: A.1 Proof of proposition 4.1
text: In this subsection, we establish the proof of Proposition 4.1
Proof. We start by considering the joint distribution p(g, sw, gw | s), which can be factorized as pπ(·|·,g)(g | sw)pπ(·|·,gw)(gw | s)p(sw | gw).
The log-likelihood of reaching the goal g from s can be expressed as:
log pπ(·|·,g)(g | s) = logEq(gw,sw|g,s) [ p(g, sw, gw | s) q(gw, sw | g, s) ]
Expanding the expectation, we get:
log pπ(·|·,g)(g | s) = log ∫∫
q(gw, sw | g, s) p(g, sw, gw | s) q(gw, sw | g, s) dgwdsw
Utilizing Jensen’s inequality, we establish a lower bound for the log-likelihood as follows:
log pπ(·|·,g)(g | s) ≥ Eq(gw,sw|g,s) [ log pπ(·|·,g)(g | sw) + log pπ(·|·,gw)(gw | s)
+ log p(sw | gw)− log q(gw, sw | g, s) ]
Given that log p(sw | gw) = r(sw, gw) − log (∑ s′w exp(r(s′w, gw)) ) and that
log (∑
s′w exp(r(s′w, gw))
) can be absorbed into the lower bound as a constant term, which
does not affect the optimization process, the lower bound L can be written as:
L = Eq(gw,sw|g,s) [ log pπ(·|·,g)(g | sw) + log pπ(·|·,gw)(gw | s) + r(gw, sw)− log q(gw, sw | g, s) ] (8) This completes the proof of proposition 4.1. The underlying premise of this approach is predicated on the assumption that the ratio of exponentiated rewards, exp(r(s,g))exp(r(s′,g′)) , is equivalent to the ratio of the probabilities p(s|g)p(s′|g′) . In essence, this implies that the reward function r(s, g) is directly proportional to the conditional probability p(s | g).

heading: A.2 Proof of proposition 4.2
text: In this subsection, we establish the proof of Proposition 4.2
Proof. The optimization objective for finding q(gw, sw | g, s) is:
Eq(gw,sw|g,s) [ log pπ(·|·,g)(g | sw) + log pπ(·|·,gw)(gw | s) + r(gw, sw)− log q(gw, sw | g, s) ]
Introducing a Lagrange multiplier λ, the Lagrangian J is constructed as:
J = Eq(gw,sw|g,s) [ log pπ(·|·,g)(g | sw) + log pπ(·|·,gw)(gw | s) + r(gw, sw)
− log q(gw, sw | g, s) ] + λ (∫ q(gw, sw | g, s)dgwdsw − 1 )
Differentiating J with respect to q(gw, sw | g, s) and setting it to zero yields:
log pπ(·|·,g)(g | sw) + log pπ(·|·,gw)(gw | s) + r(gw, sw)− log q(gw, sw | g, s)− 1 + λ = 0
Simplifying, we get:
q(gw, sw | g, s) = exp(λ− 1)pπ(·|·,g)(g | sw)pπ(·|·,gw)(gw | s) exp(r(gw, sw))
To ensure q(gw, sw | g, s) is a valid probability distribution, it is normalized as:
q⋆(gw, sw | g, s) = pπ(·|·,g)(g | sw)pπ(·|·,gw)(gw | s) exp(r(gw, sw))∫∫
pπ(·|·,g)(g | s′w)pπ(·|·,g ′ w)(g′w | s)exp(r(g′w, s′w))dg′wds′w
The denominator serves as the normalizing constant, ensuring that q⋆(gw, sw | g, s) sums to one over its domain, thereby satisfying the properties of a probability distribution.
This concludes the proof.

heading: A.3 Proof of proposition 4.3
text: For the sake of clarity, we define ω as the tuple (gw, sw) and use q
⋆(ω) as shorthand for q⋆(gw, sw | g, s0). To rigorously proof this proposition, we define a series of functions and transition operators:
Definition 1. We introduce fj(·) for j ∈ {0, . . . , η} as a weighted blend of fη(·) and p(· | s, g; f) , given by fj(ω) = fη(ω)1−βjp(ω | s, g; f)βj . The sequence of weights βj satisfies 1 = β0 > β1 > . . . > βη = 0. Specifically, fη(ω) satisfies
fη(ω) Zf = q⋆(ω) where Zf is the normalizing constant.
Definition 2. Let Tj(ω, ω′) for j ∈ {1, . . . , η − 1} denote a transition operator, formulated as
Tj(ω, ω ′) = p(ω′ | ω, s, g;h)min ( 1, fj(ω ′)p(ω | ω′, s, g;h)
fj(ω)p(ω′ | ω, s, g;h)
) .
Then the process of sequentially sampling subgoals is defined as follows:
Definition 3. Let the process start with the sampling of ω1 from f0(·). Sequentially, ω2 is derived from ω1 via the transition operator T1, perpetuating this mechanism until ωη is obtained from ωη−1 through Tη−1. The joint distribution probability is articulated as
g(ω1,...,ωη) Zg
, wherein g(ω1, . . . , ωη) = f0(ω1)T1(ω1, ω2) . . . Tη−1(ωη−1, ωη) and Zg is the normalization constant.
Finally, the weight α for each sequence is given by α = ∏η
j=1 fj(wj) fj−1(ωj) .
To establish the validity of the proposition, we begin by proving the essential lemmas:
Lemma 1. Let fj(ω) and Tj(ω, ω′) be as specified in Definition 2. Define pj(ω) as
pj(ω) = fj(ω)∫
fj(ω′) dω′ .
Then, the following detailed balance condition holds:
pj(ω)Tj(ω, ω ′) = pj(ω′)Tj(ω′, ω).
Proof. The proof can be divided into two cases: Case 1: pj(ω′)p(ω | ω′, s, g;h) > pj(ω)p(ω′ | ω, s, g;h) Starting with pj(ω′)Tj(ω′, ω), we have:
pj(ω ′)Tj(ω′, ω) = pj(ω ′)((((
((((p(ω | ω′, s, g;h) pj(ω)p(ω ′ | ω, s, g;h)
pj(ω′)(((( ((((p(ω | ω′, s, g;h) = pj(ω)p(ω
′ | ω, s, g;h) = pj(ω)Tj(ω, ω ′).
Case 2: pj(ω′)p(ω | ω′, s, g;h) ≤ pj(ω)p(ω′ | ω, s, g;h) Starting with pj(ω)Tj(ω, ω′), we have:
pj(ω)Tj(ω, ω ′) = pj(ω)(((( ((((p(ω′ | ω, s, g;h)pj(ω
′)p(ω | ω′, s, g;h) pj(ω)(((( ((((p(ω′ | ω, s, g;h)
= pj(ω ′)p(ω | ω′, s, g;h) = pj(ω ′)Tj(ω′, ω).
In both cases, we find that pj(ω)Tj(ω, ω′) = pj(ω′)Tj(ω′, ω), thereby proving the lemma.
Lemma 2. Let fj(ω) and Tj(ω, ω′) be as defined in Definition 2. Define the normalized distribution pj(ω) as
pj(ω) = fj(ω)∫
fj(ω′) dω′ .
Then, Tj(ω, ω′) preserves the invariance of pj(ω), formally defined as ∫
Tj(ω ′, ω)pj(ω′) dω′ = pj(ω).
Proof. We proceed by leveraging the results from Lemma 1. Specifically, we have:
∫ Tj(ω ′, ω)pj(ω′) dω′ = ∫ Tj(ω, ω ′)pj(ω) dω′
= pj(ω) ∫ Tj(ω, ω ′) dω′
Given that ∫ Tj(ω, ω ′) dω′ = 1, we have ∫ Tj(ω
′, ω)pj(ω′) dω′ = pj(ω). This confirms that Tj(ω, ω′) preserves the invariance of pj(ω), thereby proving Lemma 2.
Now we give the proof of Proposition 4.3.
Proof. We first define the function f as follows:
f(ω1, . . . , ωη) = fη(ωη)
fη−1(ωη) Tη−1(ωη−1, ωη) . . .
f2(ω2) f1(ω2) T1(ω1, ω2)f1(ω1)
Given the definition of Zf , we have
Zf = ∫ fη(ω) dω
By Lemma 2, we have: ∫
Tj(ωj , ωj+1)fj(ωj) dωj = fj(ωj+1)
Thus, we can write: ∫
f(ω1, · · · , ωη) Zf dω1 · · · dωη
=
∫ fη(ωη)
Zf dωη
∫ Tη−1(ωη−1, ωη)fη−1(ωη−1)
fη−1(ωη) dωη−1 · · ·
∫ T1(ω1, ω2)f1(ω1)
f1(ω2) dω1
=
∫ fη(ωη)
Zf dωη
=1
This implies that Zf is also the normalizing constant of f(ω1, . . . , ωη). Since f0(·) is a distribution, it is evident that Zg = 1. We have:
Eg(ω1,...,ωη) [ 1
N
∑ α ] = Eg(ω1,...,ωη) [ f(ω1, . . . , ωη)
g(ω1, . . . , ωη)
]
= Zf
[∫ f(ω1, . . . , ωη)
Zf dω1 · · · dωη
]
= Zf
This concludes the proof of Proposition 4.3.

heading: B More Implementation Details for Each Module
text: The framework of SEGO is composed of five components, each serving a distinct purpose to enhance the system’s overall efficacy.

heading: B.1 Policy Network
text: The policy network π(a | s, g) takes as input the current state and intended goal and returns an action. Since the goal and the state can both be expressed as token sequences, we first concatenate these sequences before feeding them into the policy network. This network is tasked with predicting the subsequent action, also framed as a token sequence, utilizing standard decoding techniques like greedy search or top-k sampling (Holtzman et al., 2019).
The training of the policy network is conducted through instruction finetuning, utilizing the following instruction template:
Construct a Python script to address the given problem: {problem}
### Response: {solution}
In this template, problem and solution represent the goal g and the trajectory respectively. The base model for this process is CodeLLaMA, and it undergoes full parameter finetuning to optimize its performance. As the sequential subgoal optimization process progresses, the model is further trained by utilizing self-generated successful trajectories. This prompt template is also employed to generate the trajectory dataset using gpt-3.5-turbo-0613.

heading: B.2 Subgoal Generator
text: The subgoal generator, represented as f , aims to decompose a complex task into two more manageable sub-tasks. It works by taking the current state s and goal g, and outputting a pair consisting of a subgoal and its corresponding state: sw, gw = f(s, g). This approach ensures that both the journey from the current state to the subgoal and from the subgoal state to the intended goal become more tractable sub-tasks. Crucially, the subgoal state sw is a valid solution of the subgoal gw, adhering to the premise that a state is an aggregation of actions, each representing a step in the solution process.
The subgoal generator is trained through instruction finetuning, utilizing data collected from gpt-3.5-turbo-0613. The instruction template is defined as:
Break down the given problem into a smaller task (a subproblem) and devise a method to solve it, considering a provided partial solution to the original problem as a starting point.
### Input: {problem}
{partial solution}
### Output: {subproblem}{solution}[EOS]
This module, fundamentally built on the architecture of CodeLLaMA (Rozière et al., 2023), leverages the capabilities of LoRA (Hu et al., 2021) for efficient finetuning. The primary objective is to accurately predict {subproblem}{solution}[EOS] from its preceding context, realized through a causal language modeling. This prompt template is also utilized to predict both the subgoal and the corresponding state using gpt-3.5-turbo-0613.

heading: B.3 Subgoal Optimizer
text: The subgoal optimizer, denoted as h, is designed to refine subgoal gw and its corresponding state sw. Its objective is to yield improved subgoal g′s and state s ′ w that more effectively contribute to decomposing the overall intended goal: s′w, g ′ w = h(sw, gw, s, g). This component incorporates both the current state s and the intended goal g as inputs, providing insights into the complexity of the intended goal and the current status in the problem-solving process.
The subgoal optimizer is also trained through instruction finetuning, drawing upon data from gpt-3.5-turbo-0613. The instruction template for this module is as follows:
Optimize the given subproblem to make it more manageable. Then, develop a method to solve it, considering a provided partial solution to the original problem as a starting point.
### Input: {problem}
{partial solution}
{subproblem}{solution}
### Output: {optimized subproblem}{optimized solution}[EOS]
This module, also built on CodeLLaMA, utilizes LoRA for efficient parameter finetuning. The aim here is to accurately predict {optimized subproblem}{optimized solution}[EOS] from the provided context, ensuring the outputs are coherent and contextually aligned.

heading: B.4 Reward Network
text: The reward network, formulated as r(s, g), accepts a state s and a goal g as inputs and produces a score to evaluate whether the goal has been achieved in the current state. In mathematical problem-solving, this essentially translates to determining if the state s—which is an aggregate of executed actions, with each action representing a step towards the solution—is a valid solution for the problem posed by g. Given that the actual reward function, R (described in §2.1), is applicable only to problems where a ground-truth answer is available, the reward network serves as a surrogate that is crucial for evaluating sub-problems encountered during the algorithm’s execution.
This model is built on the architecture of CodeLLaMA and employs LoRA to achieve efficient finetuning. The reward model is trained through instruction finetuning, utilizing the following instruction template:
Does the provided solution accurately address the given problem? {problem} {solution} {Y/N}.
B.5 Value Network The value network, represented as vπ(s, g), is a regression model that takes a state s and an intended goal g as inputs, and outputs a score representing the likelihood of successfully achieving the goal from the state under the policy network π.
This model is trained to approximate the estimated α̂, utilizing instruction finetuning. The instruction template is defined as:
Determine the probability of resolving the problem, starting from the partial solution: {problem} {partial solution}.
This model, built on the CodeLLaMA architecture, is finetuned using LoRA. It is noted that, during each iteration of the sequential subgoal optimization process, a unique set of LoRA parameters is used to avoid any potential discrepancies between iterations. This approach ensures that the value network accurately reflects the real-time capabilities of the policy network.

heading: C Details about Trajectory Dataset Creation
text: To construct the goal collection Dg in Alg. 1, we incorporate mathematical problems sourced from the training subsets of three distinct datasets: GSM8k, MATH, and AQuA. For the generation of solutions corresponding to each problem, we apply a prompt as follows:
### Instruction Construct a Python script to address the given problem: {problem}
### Response: {solution}
In this format, “solution” is completed by GPT-3.5-turbo. The solution is subsequently broken down into steps, with the i-th state in a trajectory comprising the first i steps, and the i-th action defined as the (i+ 1)-th step.
In a trajectory, all states except the s0 which includes essential imports and function definitions (e.g., “import math; def solve():”) are considered intermediate states.

heading: D Details about Experimental Setup

heading: D.1 Training Data
text: For training SEGO, we use GSM8K, MATH, and AQuA (Ling et al., 2017) datasets. After filtering for correct answers, the resulting training set includes 10, 374 samples from GSM8K, 10, 981 from MATH, and 35, 355 from AQuA. These problems form the goal collection Dg in Alg. 1.

heading: D.2 Baselines
text: Closed-Source Models. (1) GPT-4: A model that sets a standard in various academic domains, including those that require intricate mathematical reasoning (OpenAI, 2023). (2) PaLM-2: A model that excels at logical reasoning and multilingual tasks, demonstrating advanced capabilities in reasoning and solving complex mathematical problems in multiple languages (Anil et al., 2023). (3) Minerva: A model that specializes in quantitative reasoning, providing precise and comprehensive solutions to advanced mathematical, scientific, and engineering problems (Lewkowycz et al., 2022).
Open-Source Models. (1) LLaMA2: A model that is trained on 2 trillion tokens of publicly accessible data, exhibits outstanding capabilities in mathematical reasoning (Touvron et al., 2023a). (2) WizardMATH: A model that enhances the mathematical reasoning capabilities of LLaMA2 by curating more complex and diverse supervised finetuning data (Luo et al., 2023). (3) MetaMath: This model employs a question bootstrapping technique, facilitating the generation of questions through both forward and backward reasoning paths. It further enhances its capabilities by incorporating Large Language Models (LLMs) to refine the phrasing of the question text (Yu et al., 2023). (4) CodeLLaMA: A model that
excels in code-related tasks with implications in mathematical programming and algorithm synthesis, demonstrating superior infilling capabilities and support for extensive input contexts in programming tasks (Rozière et al., 2023).3 (5) MAmmoTH-Coder: This model leverages a training dataset that incorporates both chain-of-thought (CoT) and program-of-thought (PoT) rationales, thereby not only facilitating the utilization of various tools but also accommodating diverse thought processes for solving distinct mathematical problems (Yue et al., 2023).
D.3 Implementation Details We maintain model consistency by employing CodeLLaMA as the base model for both the policy network and auxiliary modules, including the subgoal generator, subgoal Optimizer, reward network, and value network. Efficient finetuning of the auxiliary modules is achieved through the utilization of LoRA (Hu et al., 2021), configured with parameters r = 16, lora_alpha = 32, and lora_dropout = 0.05, targeting the “q_proj” and “k_proj” modules. The learning rates are set at 1e − 5 and 1e − 4 for the policy and auxiliary modules, respectively, with a uniform batch size of 32. When collecting data from gpt-3.5-turbo-0613, we set temperature and top_p as 0.8 and 1.0 respectively. All models go through an initial training phase of 4, 800 steps. Subsequently, a sequential optimization process is conducted, with the number (N) and length (η) of sequences set as 2 and 3 respectively, and the temperature and top_p for the Subgoal GeneratorOptimizer and the policy network configured at 0.2 and 0.95 respectively. This optimization is performed three times, each lasting 1, 200 steps, and when η = 3, the parameters β1 and β2 are precisely set at 0.33 and 0.66 respectively. Rigorous contamination checking, as delineated by OpenAI (2023), is executed to verify the purity of our test sets for GSM8K and MATH. During the test phase, a greedy search strategy is employed.

heading: E The Annotation of Problem Hardness
text: We employ the following prompt to automatically annotate the difficulty with gpt-3.5-turbo-0613:
Please assign a score between 1 and 5 to the following question, indicating its level of difficulty and complexity. A higher score should be given to denote greater difficulty and complexity.
Please provide only the score, without any additional explanations or reasons.
### Input: {question}
### Output:

heading: F More Discussions about Related Works
text: Mathematical Reasoning with LLMs. The exploration of mathematical reasoning in Large Language Models (LLMs) has been significantly influenced by the development of datasets such as GSM8K (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021), serving as crucial benchmarks for assessing machine learning models in mathematical domains. GSM8K encompasses a variety of grade school math problems, while MATH compiles challenging competition mathematics problems. The introduction of extensive datasets (Koncel-Kedziorski et al., 2016; Ling et al., 2017; Talmor et al., 2018; Geva et al., 2021) and platforms like MWPToolkit (Lan et al., 2022) has enriched the field. This exploration is systematically categorized into two main domains: prompting strategies and learning with verifications. In the realm of prompting strategies, a variety of methods have been conceptualized to enhance the reasoning capabilities of LLMs. Techniques such as Chain-of-Thought Prompting (Wei et al., 2023; Wang et al., 2022), Progressive-Hint Prompting (Zheng et al., 2023), Least-to-Most Prompting (Zhou et al., 2022), and
3For CodeLLaMA, we ensure consistency with our models by employing identical decoding methods and prompts during implementation, while for the other models, we refer to the results reported in their respective papers.
bi-modal behavioral alignment (Zhao et al., 2024) have been instrumental in progressively guiding LLMs to accurate conclusions and facilitating the generation of intermediate reasoning steps. Moreover, methodologies like Complexity-Based Prompting (Fu et al., 2023) and Self-Consistency(Wang et al., 2022) exploit higher reasoning complexity and diverse reasoning paths, respectively, to realize significant advancements in multi-step reasoning tasks. Within learning with verifications, the emphasis is on optimizing the mathematical proficiencies of LLMs through the integration of verifiers. Strategies like outcome-based verifiers (Cobbe et al., 2021), step-aware verifiers (Li et al., 2023; Lightman et al., 2023), and learning from partially-correct solutions (Ni et al., 2023) have been deployed to bolster reliability and precision in mathematical reasoning. While the aforementioned domains have significantly advanced mathematical reasoning within LLMs, our approach is orthogonal to these categories. We concentrate on the formulation of adaptive curricula, emphasizing the incorporation of subgoals, to facilitate nuanced learning pathways and enhance the model’s mathematical reasoning capabilities. A parallel and notably concurrent work, MAmmoTH (Yue et al., 2023), investigates the impact of instruction finetuning to empower large language models with mathematical problem-solving capabilities. This can be considered as an implementation of the instruction finetuning stage within our framework.
Subgoal-based RL. Subgoal Search is a central component in reinforcement learning, essential for empowering AI systems to navigate through complex, extensive tasks effectively. This concept has played a vital role in uncovering important aspects such as the benefits of recognizing and rewarding subgoals (Zhai et al., 2022), the proper structuring of Markov decision processes for hierarchical reinforcement learning (Wen et al., 2020), the difficulties in selecting the most suitable options for planning (Jinnai et al., 2019a), and the incorporation of temporal abstraction in RL (Fruit et al., 2017). The practical research in this field mainly focuses on exploring and creating subgoals for planning and developing learning curricula for subgoals. Exploration is aimed at finding the best or most efficient strategies, using diverse approaches like reducing cover time (Jinnai et al., 2019b), understanding dynamical distances (Hartikainen et al., 2019), increasing entropy (Pitis et al., 2020), and applying asymmetric self-play (OpenAI et al., 2021). In the area of subgoal planning, a variety of algorithms have been developed to refine decision-making processes. For example, SoRB (Eysenbach et al., 2019) utilizes RL to develop a graph for subgoal sequences, DC-MCTS (Parascandolo et al., 2020) employs learned subgoal proposals to divide tasks, PAIR (Li et al., 2022) combines online RL with offline supervised learning, and (Moro et al., 2022) improve MCTS with Hindsight Experience Replay for goal-oriented planning. Moreover, the work by (Chane-Sane et al., 2021) provides concise insights into improving goal-conditioned reinforcement learning by conceptualizing imagined subgoals, adding a fresh viewpoint to the field. Research in curriculum learning has developed innovative methods to construct curricula that systematically escalate the complexity of subgoals, thereby improving the speed and quality of learning (Zhang et al., 2020, 2021). The exploration of subgoal learning in the realm of complex mathematical problem-solving represents a largely unexplored field. Our work delves into the inherent challenges of applying subgoal learning in mathematical contexts, specifically, the difficulty in identifying the optimal subgoal within expansive state spaces, and introduces a theoretical framework to navigate these challenges.

heading: G Details about Time-complexity
text: This section presents the analysis of time-complexity for the sequential subgoal optimization process (see §3.2). For each example, the frequency of module invocation is shown in Table 3.
We acknowledge that the primary computational cost in our method stems from the decoding process conducted by the subgoal generator or optimizer. This process indeed requires significantly more time compared to the computation involved in calculating scores. Notably, as shown in §6.2, our method outperforms other approaches that produce more subgoals without sequential optimization, while maintaining a comparable computational budget. This result indicates the effectiveness of SEGO in identifying vital subgoals.

heading: H Analysis on Whether GPT-3.5-turbo Serves as an Upper Bound
text: This study seeks to explore the hypothesis that GPT-3.5-turbo represents a performance ceiling for SEGO. Moreover, it examines the applicability of SEGO when paired with more advanced foundational models, specifically employing Mistral, a language model with 7 billion parameters noted for its exceptional performance and efficiency (Jiang et al., 2023). In our experimental setup, both SEGO and GPT-3.5-turbo leverage a program-of-thought (Chen et al., 2022) rationale to ensure a fair comparison. The results are presented in Table 4.
The results suggest that SEGO’s performance potential is not limited by the upper limits of GPT3.5-turbo. This point is particularly supported by the results in the MATH benchmark, where SEGO configurations utilizing Mistral-7b and CodeLLaMA-13b models significantly surpass the performance of GPT-3.5-turbo. This performance differential is predominantly attributed to the subgoal-based fine-tuning phase within SEGO, which enables the policy network to generate novel solutions that exceed the upper limits of GPT-3.5-turbo.

heading: I Performance of Various Components
text: This section delves into the performance evaluation of key components within the SEGO framework.
Reward Network. The efficacy of the reward network was gauged through its performance on a binary classification task, aimed at determining the feasibility of achieving a goal state from a given state, as inferred from the reward scores. The classification accuracy achieved by the reward network is 62.8%.
Value Network. The performance of the value network was evaluated based on the metric recall1@10, which reflects the network’s ability to accurately identify viable subgoals from a set of ten candidates. The criteria for subgoal validity are detailed in §6.3. The results of this evaluation are presented in Table 5, illustrating the value network’s performance improvements post after subgoal-based fine-tuning.

heading: J Case Study
text: In this section, we delve into a specific example to illustrate the efficacy of our model, depicted in Figure 5. In this figure, the elements labeled as the problem, sub-problem, and solution (of the sub-problem) correspond to the final goal, intermediate goal, and intermediate state, respectively. The sub-problem showcased is derived through the sequential subgoal optimization process. Additionally, we provide the full solution, which is derived from the solution of the sub-problem. This case study indicates the model’s capability to search for a suitable sub-problem that ultimately facilitates the derivation of the accurate solution to the final goal.

heading: Future Work
text: Exploring more complex problems, such as those at the undergraduate level, is a promising future direction. (3) In the current SEGO framework, only the policy network is retained during inference. An intriguing future direction involves integrating the subgoal generator/optimizer and the value network to recursively decompose complex problems into simpler subgoals