Building Math Agents with Multi-Turn Iterative Preference Learning
Authors: Wei Xiong
Chengshuai Shi
Jiaming Shen
Aviv Rosenberg
Zhen Qin
Daniele Calandriello
Misha Khalman
Rishabh Joshi
Bilal Piot
Mohammad Saleh
Chi Jin
Tong Zhang
Tianqi Liu
Venue: ICLR 2025 Poster (2025)
Keywords: large language model
RLHF
math
Primary area: reinforcement learning

Abstract: Recent studies have shown that large language models' (LLMs) mathematical problem-solving capabilities can be enhanced by integrating external tools, such as code interpreters, and employing multi-turn Chain-of-Thought (CoT) reasoning. While current methods focus on synthetic data generation and Supervised Fine-Tuning (SFT), this paper studies the complementary direct preference learning approach to further improve model performance. However, existing direct preference learning algorithms are originally designed for the single-turn chat task, and do not fully address the complexities of multi-turn reasoning and external tool integration required for tool-integrated mathematical reasoning tasks. To fill in this gap, we introduce a multi-turn direct preference learning framework, tailored for this context, that leverages feedback from code interpreters and optimizes trajectory-level preferences. This framework includes multi-turn DPO and multi-turn KTO as specific implementations. The effectiveness of our framework is validated through training of various language models using an augmented prompt set from the GSM8K and MATH datasets. Our results demonstrate substantial improvements: a supervised fine-tuned Gemma-1.1-it-7B model's performance increased from 77.5% to 83.9% on GSM8K and from 46.1% to 51.2% on MATH. Similarly, a Gemma-2-it-9B model improved from 84.1% to 86.3% on GSM8K and from 51.0% to 54.5% on MATH.

Published as a conference paper at ICLR 2025
BUILDING MATH AGENTS WITH MULTI-TURN ITERA- TIVE PREFERENCE LEARNING
Wei Xiong ', Chengshuai Shi?, Jiaming Shen Aviv Rosenberg', Zhen Qin', Daniele Calandriello
Misha Khalman', Rishabh Joshi', Bilal Piot", Mohammad Saleh", Chi Jin', Tong Zhang', Tiangi Liu?
University of Illinois Ttbana-Champaign University of Virginia?
Google Deepmind
Google Research
Princeton University
(wx13, Eozhang)e11linoo edu; tiangiil
ABSTRACT
Recent studies have shown that large language models' (LLMs) mathematical problem-solving capabilities can be enhanced by integrating external tools, such code interpreters, and employing multi-tury Chain-of-Thoughh (CoT) reason- ing. While current methods focus on synthetic data generation and Supervised Fine-Tuning (SFT), this paper studies the complementary direct preference learn- ing approach to further improve model performance. However, existing direct preference learning algorithms are originally designed for the single-turn chat task, and do not fully address the complexities of multi-turn reasoning and ex- ternal tool integration required for tool-integrated mathematical reasoning tasks. To fill in this gap, we introduce a multi-turn direct preference learning frame- work, tailored for this context, that leverages feedback from code interpreters and optimizes trajectory-level preferences. This framework includes multi-turn DPO and multi-turn KTO as specific implementations. The effectiveness of our frame- work is validated through training of various language models using augmented prompt set from the GSM8K and MATH datasets. Our results demonstrate sub- stantial improvements: supervised fine-tuned Gemma-1. 1-1tt--BB model's perfor- mance increased from 77.5% to 83.9% 9% on GSM8K and from 46.1 1% to 51.2% on MATH. Similarly, Gemma-2-it-9B model improved from 84.1% 86.3% on GSM8K and from 51.0% to 54.5% on MATH.
INTRODUCTION
Large language models (LLMs) have demonstrated remarkable capacities across variety of lan- guage tasks. Notable models include ChatGPT (OpenAl (023), Claude (Anthropic) (003), and the most advanced closed- source LLMs still struggle le with complex reasoning tasks that require multi-turn decision making. In particular, for the representative task of mathematical problem solving, LLMs often fail with basic arithmetic and symbolic computations (Hendrycks et al. 2021 Zheng et al. (2021). To address this issue, recent studies recommend the integration of external tools (e.g., caleulators, computational Python libraries and symbolic solvers) to augment the LLMs' mathematical woblem--soliig ca- pabilities (Shao et al., 2022] Mishra al. 2022] Zhang et al. 2024a). Specifically, by integrating natural language reasoning with the use of these external tools, these enhanced LLMS can receive external messages from tool interactions and reason based on both previously generated tokens and external messages, which significantly improves their performance in mathematical tasks (Gou et al. 2023b Toshniwal et al. 2024; Shao et al. 2024).
Gemini (Gemini et al., 2023 However, despite these advances,
These successes of tool-integrated LLMs lead to natural research question: how can we better train LLMs to combine tool usage with intrinsic reasoning to tackle complex reasoning tasks? For math- ematical problem solving, existing works primarily focus synthetic data generation (by strong teacher models) and supervised fine-tuning (SFT), as seen in ToRA (Gou et al.) 20236). Meta-
MathQA at. 2023). MAmmoTH (Yue al. 2023 2024). and Open-Mathlnston Toshniwal

---

Published as a conference paper at ICLR 2025
et al. 2024). These synthetic datasets have yielded significant improvements in test accuracy on standard benchmarks like MATH (Hendrycks et al. 2021) and GSMSK (Cobbe et al. 2021a).
Built strong SFT models, Reinforcement Learning from Human Feedback (RLHF) has proven to be a key technique to elicit LLMs' knowledge during the post-training stage and has become standard in the LLM training pipeline (Ouyang et al. 2022; Gemini et al." 2023). Broadly speaking,
the RLHF learning paradigm, which was originally designed for aligning LLMs with human values and preferences, is distinct from SFT as it learns from relative feedback. has notably enhanced the capabilities of models like ChatGPT, Claude, and Gemini, enabling them to generate responses that are more helpful, harmless, and honest (Bai al., 2022). Inspired by RLHF's success general chat applications, in this paper, we explore RLHF for improving LLMs' mathematical problem- solving abilities when equipped with external tools. In particular, since deep RI, methods (e.g., the proximal policy optimization, PPO algorithm (Schulman et al., 2017)) are often sample inefficient and unstable (Choshen et al. 2019), our goal is to derive direct preference learning algorithms that directly learn from the preference dataset (Zhao et al. 2023 Rafailov et al. 2023).
Contribution. We begin by formulating the learning process as Markov decision process (MDP), distinct from the contextual bandit approach typically used in RLHF for making general chatbots without external environment interactions (Xiong al. 2024 Rafailov et al. 2023). Then, we de- rive the optimality condition of the planning with such an MDP and our findings indicate that when the external randomness is low, we can develop multi multi-turn turn direct alignment algorithms (M-DPO and M-KTO). where the primary modification is to mask out irrelevant tokens during training. Further- more, we extend our approach to online iterative variants, which recent works demonstrated be promising (Xiong et al. 2024 Guo et al., 2024b). Finally, we evaluate our approach through case studies using augmented training sets from MATH and GSM8K benchmarks, employing various base models such as Gemma (Team et al. 2024), CodeGemma (Team] 2024). and Mistral (Jiang ct al. 2023). For instance, the performance of a supervised fine-tuned Cemmma-1.1-it-1 model increased from 77.5% to 83.9% on GSM8K and from 46.1% % to 51.2% on MATH. Similarly, Gemma-2-it-9B model improved from 84.1% 1% to 86.3% on GSM8K and from 51.0% to 54.5% on MATH. These empirical results indicate significant improvement in performance over standard SFT models, demonstrating the potential of RLHF in complex reasoning task. We also provide a comprehensive recipe for the practical implementation of our online iterative multi-turn methods, and will make our models, datasets, and code publicly available for further research and develop- ment.
2
ALGORITHMS DEVELOPMENT
2.1 PROBLEM FORMULATION
We first formally formulate the tool-integrated reasoning task. At the first step, prompt E X' is sampled from some distribution do as the initial state $1 = z. Then, at each step h E [H].
the agent observes the current state 8h, which is the history of the first interactions with the external environment, and takes action as according to some policy (-(s)) E A(A).
Action:
Observation: in response to the agent's 's action, the environment then returns an observation oy (- an based on the history s) and current action an.-
Then, we transit to a new state, which the history up to the step h + sh+1 = (sh, an 0%.) (Ir, a1, .1,... ah. 0x), and new step begins. This process repeats for H rounds in total and even- tually, we collect a trajectory: tool-integrated reasoning in Figure Typically, the action is in the ReAct manner, which consist of a reasoning step fn and an execution step εn (e.g., writing python code) (Yao et al. 2022). We mention in passing that such an MDP formulation of preference learning Zhong et al. (2024): Rafailov et al. (2024): Xie et al. (2024a) but with a focus on the single-turn chat task and without explicilly considering the external messages.
(xx, ... 00H-1 0.)) We present an example of multi-turn
recently studied in
When there is no ambiguity, the abbreviation 8h+1 P7 (-18n, an) is also adopted.
2

---

Published as a conference paper at ICLR 2025
To connect the problem with RLHF that learns from relative feedback, we follow Ouyang et al. 2022 Bai et al. (2022) to assume that we can query the Bradley-Tery Terry model for preference signal.
Definition (Bradley-Temy model). We denote 7/x - y, where the prompt excluded from the tra- jectory. We assume that there exists utility function of the trajectory u such that given (z, y v², one response y' is preferred over another response y². denoted as y' >> y². with probability
Prob(y' (y' > y² vr,y, ,y²) = = (2,, v -- (r,
(1)
where a is the sigmoid function ((z) = 1/(1 cxp(==)pp Also, given (I.y' yy² we denote the indicating y² > y² while
sampled preference signal as with
indicating y² > y².
Here we only assume access to the trajectory-level preference, but not an action-level one. However, we remark that the utility function itself can be defined a step-wise manner. Examples of the utility function include the binary reward from checking final result, outcome-superve reward models (Cobbe et al. 202Tb), and process- -upperisd reward model (Lightman et al. 2023).
2.2 PLANNING WITH A MODEL: OPTIMALITY CONDITION AND PRACTICAL ALGORITHM
We develop the main algorithms in this section with the general MDP formulation. Following Rafailov et al. (2023), we first establish the connection between a model M = (S,A, H. do. u) and its associated optimal policy. In particular, are interested in the following KL-regularized planning problem with respect to a reference policy Trof:
max (IT; M, m.at) = Eg mi (-1x))
[u(x,)) DRL. (mn()))s,,, Toet (-Isa)
(2) In the single-turn case with H 1, the optimal solution with respect utility function is the Gibbs distribution (see Lemma 3 Moving toward multi-turn case, we first consider H = 2 to illustrate the idea. The idea is to take a backward iteration from H = 2 to 1. Specifically, when we fix 82 and consider only the step 2, it reduces to the single-turn case:
T.M2(-182) arg max (1g (u(22 02) 17- DKL (#2(+18)), (-s.((1)))) (ne (een2(1152) exp [((2.)
m2
Then, we can define the value function associated with T,, as
VM,2($2)) (ag [u 02) nDKL. (T..2)1)32) (.d))
QM, a1) (. (V222(s2)].
For step 1, since we have determined TM22 with the definition of QM, az), we have
(-11(- arg max (1)(1)) (M..((s, DKL ( =eaa(()))) TTet1(-1s1)) exp (M,((1),))
m
By construction, M is optimal as it maximizes the KL-regularized target. For general MDP, we can repeat the process for Hl times starting with VM,H+1 0 where we recursively define
QM, an)
=
u(sH as), Bo, (n [VM. (5 HESSH-
if H,
(3)
Here the optimal policy and the V-values are given by
T.,n(an(an) = 1
Zn() Tee (anlsn) exp (om((snaa))
(Gibbs distribution of QM,a)
VM.((S)): = Ea at (l [Q.. an) - n DKL (1Ma(((s)) med (-ss)))
(4)
= 1 log Br (aglen) exp (mio(snaah) = 1 log (sR)..
17
where (sk) = Eases (an (a) exp (maalnnaa)) is the normalization constant. The second equality in the definition of the V-value is from Lemma Then, by definition, [ is optimal. Essentially, we solve H Gibbs distributions in terms of the Q-values. We summarize the results into the following proposition.
3

---

Published as a conference paper at ICLR 2025
Proposition 1. We can recursively define the following optimal value functions and optimal policies for a KL-regularized MDP with horizon H and external observation On- For Q value, we have
QMA( aN) ={
u(SH, aH), if = H if S H - 1
(5)
Bos Bog-Pa(((sn,a)) VVa.aa/aaaaaaallM......a... (V
Also, for all h E [ we have:
VMa(ss)) = log Ban exp (Min(sh
(:2n(on)
Ma, (s) Tret.h 1h))
exp (Mm((sn,
(6)
(s.)))
We have a few interesting observations that may be of independent interests.
1. The optimal value function is characterized by the expectation with respect to the initial reference policy due the additional KI. constraint.
2. For a fixed step h and state-action pair (*), a2), we can treat the future as bandit (with only one step), then, we have QM.a (sp, ax)) = E, u(S,, 2), where 2 is a complction staring from (sh, a.)) One use the Monte-Carlo estimation estimate this value by multiple roll-outs. We notice that the non-regularized version of this process, is commonly referred to as the process-supervised reward (PRM) in the literature (Wang al. (023a).
In other words, the PRM constructed in Wang et al. (2023a) is essentially Q learning process.
We remark that the results are essentially from the entropy-regularized MDPs (Williams & Peng 1991
Ziebart 2010).
Multi-turn DPO.
According equation
solve the Q-values
Qm on) n log TM, (anlsa)))s)) (s.)(s)).
Tret.n (aae.n(aaaie (sa)
Furthermore, combining equation with the definition of Q-values Qm,. in equation 3 have
Eog (-len VM,++1 (Sa+1) = rj log (a,.(an)as) + VM,)(sh),
if H
Feet, (an (8))
(8)
u(sH,aH) = rj log HM. (an s8))
+ VM,n(SH))
Feet (ag (2))
Summing over h E [H), we have the following re- re-parameterization result:
log (MMn(anlsa)
H
u(sH, aH) 77
1=1 Weef, (an sn) + VMMI($1)
VM Bon Bog (M.tss
term (B)
tarm (4)
term (C)
Here, term (A) is similar to the single-turn case and term (B) will be cancelled for the reward difference of two samples with the same prompt 81. However, in practice, term (C) is typically not feasible directly compute term (C)) is related to the randomness of the external environment.
For the focus of this work, i.e., the tool-integrated mathematical reasoning, luckily the code execu- tion result is determined by the history (the codes written by the LLMs). This leads to term (C) = 0. Therefore, we can plug equation 9 into the maximum likelihood estimation of the utility function with dataset D consisting of (z, " T'). to get the following multi-tun turn DPO (M-DPO) loss:
(0) = E JDDD log
log
Tae (a (x
log
TR.A A (af, (shh) Tref. (1)1))))
(10)
Feef A (a (1
Similarly, we can implement M-KTO under deterministic transition. We refer interested readers to Appendix A for the loss function details.
4

---

Published as a conference paper at ICLR 2025
2.3 ONLINE ITERATIVE TRAINING
We now combine the planning algorithm M-DPO with the online iterative learning framework, as inspired by its great success in the single-turn case (Xiong al., 2024 Guo et al. 2024b).
Learning objective. For more comprehensive understanding of its statistical behavior, we will consider two different learning objectives. The first objective is KL-regularized one:
H
max Re Ban Pf (1ls (x( y) n Σ h=1
DKL (T(-ss)) ao)(()s)),((((sss))) no(-1ss))] (1I)
i.e. max, To) where M (,, do. u²) ) is the groundtruth environment and To the initial SFT policy. This target is widely adopted in RLHF and requires us to search for the optimal policy only at fixed KL ball centered at the SFT policy no- In contrast, the second one is the non-regularized target, i.c., directly optimizing the reward:
max Egnddy Ban (n( ( (-18,,0))
(u' (z. v)].
(12)
37
This target is the standard one in canonical RL studies (Sutton & Barto 2018). One motivation for this target that the reasoning task, the reward function is more interpretable (e.g. final result checking) compared to the chat task.
Algorithmic framework. We present general online iterative algorithmic framework in Algo- rithm [ This framework termed as Online Iterative Multi-turn Gibbs Sampling from Human Feedback (M-GSHF) because the optimal policy is a layer- wise Gibbs distribution that generalizes the result in Xiong et al. (2024). We now discuss some features of the framework as follows.
Reference model choice for controlling regular- ization level. We unify the two different learn- ing targets in equation [1 and equation [2] by taking the reference model choice as a hyper- parameter. First, if we fix the reference model as the initial policy, i.e., Tr.ret = To, Vt E [T], we always search the optimal policy within the KI. ball centered at wo, and thus optimize the KI.-regularized target. In contrast, inspired by the mirror descent (Nemirovskij & Yudin 1983), if we update the reference policy every iteration to be the policy learned in the last iter- ation, i.e., Tr.ref = T[- Vt E [T), the cumula- tive update can make the model to move away from the original Tu (while a constraint is made on the per-iteration update magnitude) and we thus optimize the non-regularized target in equation 12 See Figure forr illustration.
Figure I: Illustration of the difference between the two learning objectives. Left: the KL-regularized target as we do not update the reference model. Right: the non-regularized target.
trade-off. We update our behavior policies non-symmetric way. The first agent aims to extract the historical information we have gathered so far and runs the M-DPO or M-DKO presented in Section 2.2 However, it is widely recognized
Non-symmetric policy choice for
in RL studies (Sutton & Barto 2018, Auer al. 2002) that simply exploiting the historical data via following the empirically best model is not sufficient to obtain a good final policy, while it also required to explore the environment so that new information can be collected to facilitate subsequent learning, i.e., the exploration-exploitttio tradeoff. Therefore, the second agent will strategically incorporate the uncertainty of the future relative to Tt to choose n?, which is referred to as the exploration policy.
comprehensive theoretical analysis is derived for Algorithm I deferred to Appendix D due to space constraint, with a focus the KL-regularized target. Here we highlight the following in- formal result (see Theorem 2 for the complete version), emphasizing the efficiency of Algorithm guaranteed by sublinear regret. The other target of optimizing the rewards has been theoretically studied in Wang et al. (203bb) while the techniques of analyzing mirror- perror-descentsstyy style algorithm have been developed in Cai et al. (2020).
5

---

Published as a conference paper at ICLR 2025
Theorem (Informal). Under the realizability assumption, with the KL-regularized target, the the- oretical version of Algorithm leads to a regret (defined in equation 15) that sublinear horizon T for a broad class reward and transition models.
The main take-away message from the theorem is that if we choose suitable exploration policy, the online itcrative learning is provably efficient. We also remark that without explicit mechanism to encourage exploration, the randomness of the LLM itself is not sufficient to learn the optimal policy Zhang 2022 if we do not make additional assumption.
Moving toward practical algorithm designs, the exploration is generally interpreted increasing the diversity of the collected data by adopting inference- eiterence-timeeeee methods with the base DPO policy T For instance, one may tune the sampling temperature in Llama project (Touvron al., 2023) or
use best-of-n sampling (Xu et al., 2023 Hoang Tran [2024] Dong et al. 2024), where these methods outperform the vanilla on-policy sampling with considerable margin. In this work, we mainly enrich the generated data by various intermediate checkpoints, as done in the Claude project (Bai et al.) 2022]. We refer this approach as mixture sampling. also natural adopt reward-guided Monte Carlo tree search (MCTS) (Xie et al. 2024b), which we leave for future work.
Algorithm Online Iterative M-GSHF
1: Input: KL coefficient > O, horizon T > O, initial policy To, batch size m >> 0.
2: Initialize D <- 0 and Ti nं T1,ref To-
3: for 2, T do Sample m pairs (z, 7", 2) as Dr by do, 7- n ní. receive the m preference signals following the Bradley-Tene model from Definition and update the preference dataset D +- DUD;
5: > Extract the empirically optimal policy from historical data
6: Practical: Perform the planning algorithms on D to get me (e.g., using the M-DPO loss in equation 10
or the M-KTO loss equation 13)
7: Theoretical: Perform MLE on D to obtain model estimation M. -me, P.) as in equation 16 and
equation 17 call Oracle 2 with M, 1), re to get
8: Select the exploration ion policy to facilitate learning
Practical: Given mi, select of an exploration policy using heuristic methods (such as mixture pling. inference parameters tuning and west-of-n sampling.
10:
Theoretical: Given n.. choose कद as an exploration policy following equation 18
11: Choose the reference model to control regularization
level
12:
Update Te+1.ref <- कद when considering the non-regularized target; keep ret1.ref when consid-
ering the KI.rregularized target
13: end for
14: Output: the best model in कं. by a validation set.
EXPERIMENTS
3.1 EXPERIMENT SETUP
Task, datasets, and models. We use the test sets of MATH (Hendryeks et al.
2021 and GSM8K
(Cobbe et al., 2021a) to measure the model's ability to solve the mathematical problems. To con- struct the training prompt set, we use the prompts from MetaMathQA (Yu al. 2023) and MMIQC (Liu & Yao 2024), which is an augmented prompt set from the 7.5K training problems of MATH and 7.47K traiiing problems of GSM8K. We provide an example of the data sample Figure We train with range of base models, including -eemma-1.1-it--
(Team et al. 2024), CodeGemma- 1.1-it-7B (Team] 2024), Mistral-7B-v0. (Jiang al. 2023). and Gemma2-it-99 We first fine-tune the model using a subset of the Open-Mathlnstruun dataset. The details of the SFT process are provided in Appendix B
Implementation of Iterative M-DPO and M-KTO. We run the iterative training for epochs in total. For each iteration, we have prompt set of 20K questions and generate 20 responses per prompt with current DPO model and 10 responses per prompt with the model from last iteration.
We use the pre-trained version because the chat template of its instruct model from huggingface is not consistent with their own codehase.

---

Published as a conference paper at ICLR 2025
We check the final answer of these responses to determine their correctness. Then, for each prompt, we randomly sample two responses with correct and incorrect final answers and add them into the training samples. Then, train the model the collected samples using the M-DPO/M-KTO loss. We also include an ablation of reference model choice. To implement the M-DPO, we simply set the labels of all the user-turn tokens to be -100 and mask the log-probability in the subsequent loss computation. We train the model for epoch at most and tune the learning rate in {2e-7, 4c-7, 7e-7, | with the first iteration of iterative training. Eventually, the learning rate of 4c-7 used for Gemma-1. 1.1 models and 2e-7 is used for Gemma-2 model and Mistral model. The global batch size 32 with warm-up step of 40. We evaluate the model every 50 training steps by the split prompt set. The hyper-parameters are of M-KTO are mostly the same as the M-DPO, We also set X+ A_ following the original KTO paper (Ethayarajh et al. 2024).
Baselines. The existing literature mainly focuses on the synthetic data generation and SFT to teach the models to the external tool. We use the results from Tooshiiall et al.) (2024) as baselines because we use the same SFT dataset so the results are generally comparable. For the CoT baselines, use the Wizardmath models from I
(2023). We also include the reward ranked fine-tuning (RAFT) as baseline (Dong et al. 2023). which is also known as rejection sampling fine-tuning (Tovvron al. 2023). Another baseline the single-turn online iterative DPO and KTO (Rafailov et al. [023] Ethayarajh et al. 2024), which ignore the problem structure (i.e., the external messages)
and treat the trajectory as a whole. In implementation, means that we do not mask the tokens of external messages.
Table I: Main results of different methods on the test sets of GSM8K and MATH. t: the model serves as the starting checkpoint of other methods. The results of the CoT methods are borrowed from the technical reports (Toshniwal al., 2024, Gou al. 2023b). For iterative M-DPO/M- KTO, we update the reference model by default if not specified. The gains relative to the SFT starting checkpoint are marked by 1.
Model
Method
GSMSK
AVG
WizardMath-7B
SFT for CoT
54.9
10.7
32.8
WizariMath-133
SFT for CaT
63.9
14.0
39.0
WizardMatb-7008
SFT for CoT
X
81.6
22.7
52.2
MMd22IaMA-2-7B
SFT SPT
V
75.9
43.6
59.8
CodeLLaMA-22133
45.5
62.2
CodeLLaaM-2234BB
SFT SFT
80.7
48.3
64.5
2.
84.6
50.7
67.7
Gemma- 1.1-it-7B
SFT!
77.5
46.1
61.8
Gemma-1 1. 1-it=7B
RAFT
79.2
47.3
63.3
Gemma-1 1.1-ttBB
Itcrative Single-turn DPO
65.3
Gemma-1 1.1-it-B
Iterative Single-turn KTO
80.6
49.0
64.8
Gemma-
IPO
Gemma-1 I. 1-it-7B
M-DPO Iteration
81.5
49.1
65.3
Gemma- 1.1 it 7B
DPO Iteration 2
82.5
49.7
66.1
Gemma- 1.1-it-BB
M-DPO ltraatinn
16.4
15.1
67.6
1..8
Gemma- 1.1-it-7B
Iterative M-KTO
82.1
14.6
49.5
13.4
65.8
14.0
ColeGemma. 1.1-it-77
SFTT
77.3
46.4
61.9
CodeGemma- 1---it-BB
RAFT
78.8
48.4
63.6
CodeGemma- 1-it-7B
Iterative Single-tum DPO
79.1
48.9
64.0
CodeGemmma- 1.1-1-7B
Iterative Single-turn
64.4
CodeeGemma: I. 1-i-1-7B
Ierattiv M-DPO
81.5
14.2
50.1
13.7
65.8
14.0
CodeGenmma 1.1-it.7B
Iterative M.KTO
81.6
14.3
49.6
13.2
65.6
13.8
Mistrll-BB-v0.3
SET
-Bstral-7B-v0.3
RAFT
79.8
43.7
61.8
-istral-7B-v0.3
Iterative Single-turn DPO
79.8
45.1
62.5
Mistral-7B 7B-v0.3
lterative Single-turn KTO
81.3
46.3
63.8
Mistral v0.3
lterative M.DPO
82.3
14.5
47,5
14.8
64.9
14.7
Mistral-RR-v0..
Iterative M-KTO
81.7
13.9
46.7
14.0
64.2
14.0
Gemma- 2B
SFTt
84.1
51.0
67.6
Gemma-2-4t-98
RAFT
84.2
52.6
68.4
lterative Single-tur
Gemma. i 9B
Iterative Single turn KTO
85.4
52.9
69.2
Gemma- 2-i-9B
Iterative M-DPO
86.3
12.2
54.5
13.5
70.4 12.9
Gemma-2-it
Iterarive M-KTO
86.1
12.0
13.5
12.8

---

Published as a conference paper at ICLR 2025
3.2 MAIN RESULTS
We evaluate the models in the zero-shot setting and report the main results in Table T From the first two sections in Tables] we first observe that the tool-integrated LLMs significantly outperform their CoT counterparts with only SFT, demonstrating the benefits of leveraging external tools. In the subsequent discussions, we focus on the comparison within the scope of tool-integrated LLMs.
Iterative M-DPO and M-KTO considerably improve the SFT models. Across all four base models, iterative training with M-DPO or M-KTO consistently leads to notable improvements over the initial SFT checkpoint on both GSM8K and MATH. In particular, with M-DPO. the aligned Gemma-1.1-it- model attains accuracies of 83.9% and 51.2% on GSM8K and MATH, respec- tively, and is comparable (slightly worse on GSM8K but also slightly better on MATH). Moreover, the aligned Gemma-2- it-9B model achieves accuracies of 86.3% and 54.5% on GSM8K and MATH, surpassing all of the open-source models trained with Open-MathInstruct in the 7B to 70B range. Overall, our framework can robustly further boost the tool-integrated models' ability after SFT.
the open-source Open-Matiinstruut--nettree
LaMALLL
Iterative M-DPO and M-KTO surpass existing RLHF baselines. We also observe that the it- erative M-DPO and M-KTO surpass other existing RLHF baselines. First, they consistently and significantly outperform RAFT across all four base models. This is because RAFT only imitates the correct trajectories, while the DPO-based and KTO-based algorithms further use the negative signal from incorrect trajectories. We note that the SFT stage in our pipeline also be viewed as an application of RAFT. Consequently, our results should be interpreted to be that after the first stage of SFT, algorithms with negative signal are more sample efficient. Moreover, while the on- line iterative single-turn DPO (KTO) also gives a better performance, it is generally worse than the multi-turn version. This suggests that learning to predict the off-policy external messages returned by the code interpreter usually has a negative impact on the reasoning ability improvement. We also present a representative example we encounter in Figure where LLMs generate poorly constructed code resulting in anomalous and lengthy external messages. Forcing LLMs to learn to predict these messages can significantly hurt the model's reasoning abilities.
Iterative training and reference update lead to better performance. Using Gemma- 1. 1. with M-DPO as an example, we observe that online iterative training leads to better results. The GSM8K test accuracy increases from 77.5% (SFT) to 81.5% (iter 1) to 82.5% (iter2) to 83.9% (iter3), and the test accuracy of MATH improves from 46.1% (SFT) to 49.1% (iter 1) to 49.7% (iter2) to 51.2% (iter3). This aligns with our theoretical insight that iterative training helps models progressively explore and learn the optimal policy. Additionally, if the reference model remains fixed at the SFT policy, the final performance is notably worse compared to updating the refer- ence model at each iteration. This likely occurs because the algorithm, in this case, optimizes the non-regularized reward, and the rewards in mathematical reasoning tasks are more accurate than in general chat tasks, leading to better in-domain performance. A detailed ablation study on the impact of KL regularization is deferred the next section.
GSMBK Test Accuracy
MATH Test Accuracy
00
73-
93
Germma-78 SFT
Gemma-78 K-DD0 Band Refeeence
53
Gemma-78 SFT
Gemma-T8 M-pop + Riad betereten
Germm 78 M. DPO Update Reforance
Gemma 7BM DPO + Update Reforence
22
44
12
n: the rumber of candidates
n: the rumber of candicates
Figure 2: The pass @n rate with respect the number of candidates We evaluate the models using temperature 0.7 following the previous works Shao et al. (2024); Toshniwal et al. (2024). We notice that preference learning only improves the metric pass @n when is relatively small.
8

---

Published as a conference paper at ICLR 2025
Preference learning improves pass@n only when is relatively small. We plot the pass @n accuracy in terms of the number of candidate trajectories n in Figure 2 A question is solved if at least one of the sampled trajectories is correct. We find that preference learning improves pass @n accuracy only when n is small. For n > 16, all models perform similarly on GSM8K and MATH, indicating that iterative M-DPO does not introduce new knowledge but instead enhances the quality of top-n responses. This observation also aligns with the result of CoT reasoning (Shao et al. 2024
3.3 ABLATION STUDY AND DISCUSSION
Moderate KL regularization balances per- iteration improvement and exploration. The effectiveness of iterative DPO is highly dependent on the reference model and KL coefficient. In our ablation study, we first consider two different choices of the reference model: (1) using the fixed reference model To: (2) updating the reference model to the last iteration's model at each round, which can be viewed as a trade-off between the generation diversity and reward optimization. As shown in Table (33] models with an updated reference model outperform those with a fixed reference model. We hypothesize that in reasoning tasks, the correct reasoning paths are highly concentrated, making diversity less crucial so optimizing the non-regularized reward gives superior model performance.
Method
GSM8K
MATH
SFT
77.5
46.1
update reference + 0.01
81.7
50.1
update reference 1 0.1
83.9
51.2
update reference + = 0.5
82.8
49.7
fixed reference + 1j 0.1 79.9
48.0
Table 2: Ablation study of the impact of KL regu- larization on iterative M-DPO.
Previous work (Tunstall et al. 2023) on offline DPO suggests that lower KL coefficient (0.01) improves performance by allowing the model to deviate more from the SFT model To. In our ablation study, we search the KL coefficient E f0.01, 0.1, 0.5}. According to Table 3.3 we find that the strongest model is obtained by a moderate KL coefficient of 0.1, outperforming both 0.01 and 0.5. To explain this, we plot the GSM8K test accuracy (Figure during iterative training. In the first iteration, lower KL values show larger improvements, consistent with Tunstall et al. (2023) 's results. However, models trained with very low KL coefficients lose diversity quickly, reducing their ability to generate diverse trajectories for later training, leading to diminishing returns in subsequent iterations. Conversely, higher KI. coefficient of 0.5 imposes too much regularization, limiting improvement per iteration. In summary, for online iterative training, a balance between per-iteration improvement and exploration efficiency is key to optimizing overall performance, intuition that also applies to sampling strategies and other experimental techniques.
GSMBK Acouracy
MATH Test Accuracy
ActUret
Accunayy
Recation
Iteration
Figure 3: Left: the test accuracy GSM8K dataset with different levels of KL regularization. Right: the test accuracy on MATH dataset with different sampling strategies.
The impact of sampling strategy: data diversity and coverage are crucial. During iterative training of Gemma-1. 1.nima---.1-it--1. 1-i- we see an increase in correct trajectories from 47% in the first iter- ation to 76% in last iteration. Moreover, as the reference model updates at each step, trajectory diversity declines, which is critical for DPO/KTO training due to its contrastive nature. We follow Bai et al. 2002);; Dong et al. (2024) to explore two data collection strategies: (1) on-policy sam- pling (trajectories sampled from the current model) and (2) mixture sampling (20 trajectories from the current model and 10 from model of previous iteration). As shown in Table 6 mixture sampling significantly outperforms on-policy sampling, particularly in the third iteration where on-policy sam-
9

---

Published as a conference paper at ICLR 2025
pling fails to improve MATH test accuracy. This highlights the importance of diversity in iterative training and aligns with previous findings that advanced exploration strategies help prevent diver- sity collapse and improve preference learning (Bai al.) 2022] Touvron et al., 2023 Xiong al. 2024 Pace et al., 2024, Dong et al., 2024). It would also be interested to explore more advanced exploration strategy like MCTS in the future study.
To ensure both correct and incorrect reason- ing paths exist, we collected N trajectories per prompt. A larger N generally improves prompt coverage. as more samples are needed for diffi- cult problems. For example, in iteration 1, with N=30, 92.5% of the prompts are covered, com- pared to 83.0% for N=12 and 60% for N=6. See Figure 2 for an illustration of the relationship between pass @1 and N. However, increasing N also increases computational costs. our abla- tion study (Table 3.3), we find that increasing N from to 12 leads to a significant performance boost, reflecting better coverage for complex prob- lems. However, increasing N from 12 to 30 yields only minor improvements, suggesting that the benefits of larger N diminish quickly in vanilla rejection sampling. We expect that difficulty-aware sampling can lead better performance, while maintaining a moderate inference cost.
Method
GSM8K
MATH
SFT N=30 + Mixture
77.5
46.1
83.9
51.2
N=12 + Mixture
83.5
51.2
N=6 + Mixture
82.0
49.2
N=30 On-policy
83.1
49.5
Table 3: Ablation study of the sampling strategy with iterative M-DPO and Gemma- 1.i-t-B.
CONCLUSION, LIMITATION, AND FUTURE RESEARCH DIRECTION
In this paper, we demonstrate that preference learning, as an alternative to supervised fine-tuning, further enhances the performance of tool-integrated reasoning LLMs after SFT. We introduce an online iterative multi-turn direct preference optimization algorithm, validated through extensive periments across multiple base models. Results show significant improvements in pass@) over the SFT policy, particularly benchmarks like GSM8K and MATH. Ablation studies highlight the im- portance of balancing per-iteration improvement with exploration, which is achieved by moderate levels of KL regularization and strategic exploration choices.
Several avenues for improvement remain unexplored. Our current approach only uses final result checks preference signals, limiting the comparison between trajectories with correct incorrect answers. One may use step-wise reward signal (Lightman et al. 2023) in the data ranking stage. Meanwhile, the fine-grained reward signals could enable the use of advanced exploration strategies like west-of-n sampling (Pace et al. 2024), or MCTS (Xie et al. 20246) in our heuristic exploration implementation. Finally, while the direct preference learning algorithms show promising gains for the mathematical reasoning tasks with code interpreter, it is not directly applicable to the general agent learning with more complex and stochastic external environments or against dynamic oppo- nents. In particular, it requires to construct value network for involving an adaptive margin in the optimization target and take the randomness of the external environment into considcration. We leave the study of this more involved algorithm to the future work.
ACKNOWLEDGEMENTS
Wei Xiong and Tong Zhang are partially supported by an NSF IIS grant No. 2416897
REPRODUCIBILITY STATEMENT
We believe that making the result reproducible important. Following the author guidance of ICLR, we present reproducibility statement here to help the interested readers to reproduce our result. Most implementation details, including hyperparameters, are provided Section 3. and Appendix Additionally, we have open-sourced our training code along with step-by-step guide, using Gemma- -. as an example. We have also made the processed SFT dataset, prompt set, and the training data for the first iteration of M-DPO/M-KTO available for easy download (see supplemental materials for details). The RLHF experiments this paper are run with 8xA100 80G GPUs, where an additional machine with 8xA 100 40G GPUs is also used to accelerate data
10

---

Published as a conference paper at ICLR 2025
collection and model evaluation. The main experiment of this paper can be reproduced within 24 48 hours with this setup. To improve the readability of this paper, we provide notation table in Appendix [ The informal version of our main theoretical result is summarized in Theorem is re-stated Theorem 2 and its proof is provided in Appendix
REFERENCES
Qwen2 technical report. 2024.
Alekh Agarwal, Yujia Jin, and Tong Zhang. VOQL: Towards optimal regret in model-free rl with nonlinear function approximation. In The Thirty Sixth Annual Conference on Learning Theory, pp. 987-1063. PMLR, 2023.
Anthropic. Introducing claude.
2023.
URL https:/www..nttaroicccooo/////aa
Introducing-claudee
Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47:235-256, 2002.
Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. A general theoretical paradigm to understand learning from human preferences. arXiv preprint arXiv:2310.1200 2023.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint urXiv:2204.05800 2022.
Qi Cai, Zhuoran Yang, Chi Jin, and Zhaoran Wang. Provably efficient exploration in policy opti- mization. In International Conference Machine Learning, pp. 1283-1294. PMLR, 2020.
Shicong Cen, Jincheng Mei, Katayoon Goshvadi, Hanjun Dai, Tong Yang, Sherry Yang. Dale Schu- urmans, Yuejie Chi, and Bo Dai. Value-incentivized preference optimization: A unified approach to online and offline rlhf. arXiv preprint arXiv:2405. 19320, 2024.
Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Step-level value preference optimization for mathematical reasoning. arXiv preprint arXiv:2446.10883 2024a.
Wenhu Chen, Xueguang Ma, Xinyi Wang, and William W Cohen. Program thoughts prompt- ing: Disentangling computation from reasoning for numerical reasoning tasks. arXiv preprint arXiv:2211.12588 2022.
Zixiang Chen, Yihe Deng, Huizhuo Yuan, Kaixuan Ji, and Quanquan Gu. Self-play fine-tuning converts weak language models to strong language models. arXiv preprint arXiv:2401.01335, 2024b.
Leshem Choshen, Lior Fox, Zohar Aizenbud, and Omri Abend. On the weaknesses of reinforcement learning for neural machine translation. arXiv preprint arXiv: 2019.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichir Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 1rXiv:2110.1411 14168, 2021a.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 10..11166 14168, 2021b.
Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, KaShun SHUM, and Tong Zhang. RAFT: Reward ranked finetuning for generative foundation model alignment. Transactions on Machine Learning Research, 2023. ISSN 2835- 8856. U.I.tttps:: htttps://openrett ew wwwpt/ttttt/ieeeeeeeeeeem. dtmpmmmm2mmmmmmm0m7717171 ottp:///eeeeeoooommmmmom///////1

---

Published as a conference paper at ICLR 2025
Hanze Dong, Wei Xiong, Bo Pang, Haoxiang Wang, Han Zhao, Yingbo Zhou, Nan Jiang, Doyen Sahoo, Caiming Xiong, and Tong Zhang. Rlhf workflow: From reward modeling to online rlhf. arXiv preprint arXhv:2405.07866 2024.
Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep policy gradients: A case study on ppo and trpo. arXiv preprini arXiv:2005. 12729, 2020.
Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01300 2024.
Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. Pal: Program-aided language models. In International Conference on Machine Learning, pp. 10764-10799. PMLR, 2023.
Team Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Sorieut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.118800 2023.
Zhibin Gou, Zhihong Shao, Yeyun Gong, Yelong Shen, Yujiu Yang, Nan Duan, and Weizhu Chen. Critic: Large language models can self-correct with tool-interactive critiquing, arXiv preprint urXiv:2305.11738. 2023a.
Zhibin Gou, Zhihong Shao, Yeyun Gong, Yujiu Yang, Minlie Huang, Nan Duan, Weizhu Chen, et al. Tora: tool-integrated reasoning agent for mathematical problem solving, arXiv preprint arXiv:2309. 17452, 2023b.
Lin Gui, Cristina Gârbacea, and Victor Veitch. Bonbon alignment for large language models and the sweetness of best-of-n sampling. arXiv preprint arXiv:2406.0003 2024.
Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, et al. Reinforced self-training (rest) for language modeling. arXiv preprint arXiv:2308. 08998, 089998888. 2023.
Shangmin Guo, Wei Xiong, and Chaoqi Wang. "alignment guidebook. Notion Blog, 2024a.
Shangmin Guo, Biao Zhang, Tianlin Liu, Tianqi Liu, Misha Khalman, Felipe Llinares, Alexandre Rame, Thomas Mesnard, Yao Zhao, Bilal Piot, et al. Direct language model alignment from online ai feedback. arXiv preprint arXiv:2402.0479. 2024b.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03872 2021.
Braden Hancock Hoang Tran, Chris Glaze. Smorkel-miista--ppirm--do 2024.
https: URL htwp:/ggww........ooooo./a/////a -Mitrkell--sssrrlllaa-- PeairRM-DDR httpss//huggiggfcc.ccom
Jiwoo Hong, Noah Lee, and James Thome. Orpo: Monolithic preference optimization without
reference model. arXiv preprint arXiv:2403.07661, 2(4):5, 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Fangkai Jiao, Chengwei Qin, Zhengyuan Liu, Nancy F Chen, and Shafiq Joty. Learning planning- based reasoning by trajectories collection and process reward synthesizing. arXiv preprint arXiv:2402.006558 2024.
Xin Lai, Zhuotao Tian, Yukang Chen, Senqiao Yang, Xiangru Peng, and Jiaya Jia. Step-dpo: Step- wise preference optimization for long-chain reasoning of llms. arxiv preprint 2rX.v:2406.1886 2024.
12

---

Published as a conference paper at ICLR 2025
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint wrXxi::3005.20000 2023.
Yong Lin, Lu Tan, Hangyu Lin, Zeming Zheng, Renjie Pi, Jipeng Zhang, Shizhe Diao, Haoxiang Wang, Han Zhao, Yuan Yao, et al. Speciality vs generality: An empirical study on catastrophic forgetting in fine-tuning foundation models. arXiv preprint arXiv:2309.06256, 2023.
Haoxiong Liu and Andrew Chi-Chih Yao. Augmenting math word problems via iterative question composing, arXiv preprint wwxiv:2401009000 2024.
Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle: A generic model- based algorithm for partially observable sequential decision making. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing, pp. 363-376, 2023a.
Tianqi Liu, Yao Zhao, Rishabh Joshi, Misha Khalman, Mohammad Saleh, Peter Liu, and Jialu Liu. Statistical rejection sampling improves preference optimization. arXiv preprint arXiv:2309.0660 06657, 2023b.
Tianqi Liu, Zhen Qin, Junru Wu, Jiaming Shen, Misha Khalman, Rishabh Joshi, Yao Zhao, Moham- mad Saleh, Simon Baumgartner, Jialu Liu, et al. Lipo: Listwise preference optimization through learning-to- tooninggto-rank. arXiv preprint arXiv:2402.0188 018888888 2024a.
Zhihan Liu, Miao Lu, Shenao Zhang. Boyi Liu, Hongyi Guo, Yingxiang Yang, Jose Blanchet, and Zhaoran Wang. Provably mitigating overoptimization in rlhf: Your sft loss is implicitly an adver- sarial regularizer. arXiv preprint arXiv:2405. /6436, 2024b.
Zimu Lu, Aojun Zhou, Ke Wang, Houxing Ren, Weikang Shi, Junting Pan, and Mingjie Zhan. Step- controlled dpo: Leveraging stepwise error for enhanced mathematical reasoning. arXiv preprint wvXiv:2407.00782. 2024.
Haipeng Luo, Qingfeng Sun, Can Xu, Pu Zhao, Jianguang Lou, Chongyang Tao, Xiubo Geng, Qing- wei Lin, Shifeng Chen, and Dongmei Zhang. Wizardmath: Empowering mathematical reasoning for large language models via reinforced evol-instruct. arXiv preprint arXiv:2308.0955: 2023.
Yu Meng, Mengzhou Xia, and Danqi Chen. Simpo: Simple preference optimization with reference-free reward. arXiv preprint arXiv:2405. 14734, 2024.
Meta. Introducing meta llama 3: The most capable openly available Ilm to date. Meta Al Blog, 2024. https://ai. meta com wwweem/blog/oe aatooo/mmetaall1 -lama----laaa-311
Swaroop Mishra, Matthew Finlayson, Pan Lu, Leonard Tang. Sean Welleck, Chitta Baral, Tanmay Rajpurohit, Oyvind Tafjord, Ashish Sabharwal, Peter Clark, et al. Lila: unified benchmark for mathematical reasoning. arXiv preprint 22Xiv:2210.175/11 2022.
Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orca-math: Unlocking the potential of slms in grade school math. arXiv preprint wxXi:2240..14800 2024.
Rémi Munos. Michal Valko, Daniele Calandriello, Mohammad Gheshlaghi Azar, Mark Rowland, Zhaohan Daniel Guo, Yunhao Tang, Matthieu Geist, Thomas Mesnard, Andrea Michi, et al. Nash learning from human feedback. arXiv preprint arXiv:23 2. 00886, 2023.
Arkadij Semenovië Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
OpenAl. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
Long Ouyang. Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Alizée Pace, Jonathan Mallinson, Eric Malmi, Sebastian Krause, and Aliaksei Severyn. West-of-n: Synthetic preference generation for improved reward modeling, arXiv preprint arXiv:2401.12086, 2024.
13

---

Published as a conference paper at ICLR 2025
Richard Yuanzhe Pang. Weizhe Yuan, Kyunghyun Cho, He He, Sainbayar Sukhbaatar, and Jason Weston. Iterative reasoning preference optimization. arXiv preprint arXiv:2404.19733 2024.
Renjie Pi, Tianyang Han, Wei Xiong, Jipeng Zhang, Runtao Liu, Rui Pan, and Tong Zhang. Strengthening multimodal large language model with bootstrapped preference optimization. arXiv preprini arXiv:2403.087 2024.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. arXiv preprint arXiv:2305. 18290, 2023.
Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From to q*: Your language model is secretly a q-function. arXiv preprint arXiv:2404.12355 2024.
Corby Rosset, Ching-An Cheng, Arindam Mitra, Michael Santacroce, Ahmed Awadallah, and Tengyang Xie. Direct nash optimization: Teaching language models to self-improve with general preferences. arXiv preprint arXiv:2404.03715. 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
Lior Shani, Aviv Rosenberg, Asaf Cassel, Oran Lang, Daniele Calandriello, Avital Zipori, Hila Noga, Orgad Keller, Bilal Piot, Idan Szpektor, et al. Multi-turn reinforcement learning from preference human feedback. arXiv preprint arXiv:2405. 14655, 2024.
Zhihong Shao, Fei Huang, and Minlie Huang. Chaining simultaneous thoughts for numerical rea- soning, arXiv preprint arXiv:2211. 16482, 2022.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.0332 2024.
Avi Singh, John D Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Peter Liu, James Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, et al. Beyond human data: Scaling self-training for problem-solving with language models. arXiv preprint arXiv:2312.06585, 2023.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Gokul Swamy, Christoph Dann, Rahul Kidambi, Zhiwei Steven Wu, and Alekh Agarwal. A minimaximalist approach to reinforcement learning from human feedback. arXiv preprint arXiv:2401.0400 2024.
Fahim Tajwar, Anikait Singh, Archit Sharma, Rafael Rafailov, Jeff Schneider, Tengyang Xie, Ste- fano Ermon, Chelsea Finn, and Aviral Kumar. Preference fine-tuning of iims should leverage suboptimal, oo-pol data. arXiv preprint arXiv:2404.14367, 2024.
Yunhao Tang, Zhaohan Daniel Guo, Zeyu Zheng, Daniele Calandriello, Rémi Munos, Mark Row- land, Pierre Harvey Richemond, Michal Valko, Bernardo Ávila Pires, and Bilal Piot. Gen- eralized preference optimization: A unified approach to offline alignment. arXiv preprint arXiv:2402.05749, 2024.
CodeGemma Team. Codegemma: Open code models based on gemma. arXiv preprint vrXiv:2406.11409 2024.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint UrXiv:2403.08295 2024.
Yuxuan Tong, Xiwen Zhang, Rui Wang, Ruidong Wu, and Junxian He. Dart-math: Difficulty-aware rejection tuning for mathematical problem-solving, 2024.
Shubham Toshniwal, Ivan Moshkov, Sean Narenthiran, Daria Gitman, Fei Jia, and Igor Git- man. Opennathinstrutt-II: A 1.8 million math instruction tuning dataset. arXiv preprint arXiv:2402. 10/76, 2024.
14

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXtv:2307.009288 2023.
Lewis Tunstall, Edward Beeching. Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro von Werra, Clémentine Fourrier, Nathan Habib, et al. Zephyr: Direct distillation of alignment. arXiv preprint arXiv:2310. 16944, 2023.
Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process-and outcome- based feedback. arXiv preprint arXiv22211.1427.. 2022.
Peiyi Wang, Lei Li, Zhihong Shao, RX Xu, Damai Dai, Yifei Li, Deli Chen, Y Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. CoRR, abs/2312.08935,, 2023a.
Xingyao Wang, Zihan Wang. Jiateng Liu, Yangyi Chen, Lifan Yuan, Hao Peng, and Heng Ji. Mint: Multi-turn interactive evaluation for tool-augmented iims with language feedback. In Proc. The Twelfth International Conference on Learning Representations (ICLR2024). 2024.
Yuanhao Wang, Qinghua Liu, and Chi Jin. Is rlhf more difficult than standard ri? arXiv preprint
arXiv:2306. 14111, 2023b.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229 256, 1992.
Ronald J Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science, 3(3):241-268, 1991.
Tengyang Xie, Dylan J Foster, Yu Bai, Nan Jiang, and Sham M Kakade. The role of coverage in online reinforcement learning. arXiv preprint arXiv:2210.04157 2022.
Tengyang Xie, Dylan J Foster, Akshay Krishnamurthy, Corby Rosset, Ahmed Awadallah, and Alexander Rakhlin. Exploratory preference optimization: Harnessing implicit qe-approximation for sample-efficient rlhf. arXiv preprint 2rXiv:2405.21044 2024a.
Yuxi Xie, Anirudh Goyal, Wenyue Zheng, Min- Yen Kan, Timothy P Lillicrap, Kenji Kawaguchi, and Michael Shich. Monte carlo tree search boosts reasoning via iterative preference learning. arXiv preprint S.ti::240000455.1 2024b.
Wei Xiong, Hanze Dong, Chenlu Ye, Ziqi Wang, Han Zhong, Heng Ji, Nan Jiang, and Tong Zhang. Iterative preference learning from human feedback: Bridging theory and practice for rlhf under kl-constraint. In Forty-first International Conference on Machine Learning, 2024.
Jing Xu, Andrew Lee, Sainbayar Sukhbaatar, and Jason Weston. Some things more cringe than others: Preference optimization with the pairwise cringe loss. arXiv preprint arXiv:2312. 16682, 2023.
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03629, 2022.
Chenlu Ye, Wei Xiong, Yuheng Zhang, Nan Jiang, and Tong Zhang. A theoretical analysis of nash learning from human feedback under general kl-regularized preference. arXiv preprint arXiv:2402.07314, 2024.
Longhui Yu. Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309. /2284 2023.
15

---

Published as a conference paper at ICLR 2025
Lifan Yuan, Ganqu Cui, Hanbin Wang, Ning Ding, Xingyao Wang, Jia Deng, Boji Shan, Huimin Chen, Ruobing Xie, Yankai Lin, et al. Advancing Ilm reasoning generalists with preference trees. arXiv preprint XXX::2444.000777 2024.
Zheng Yuan, Hongyi Yuan, Chengpeng Li, Guanting Dong, Chuanqi Tan, and Chang Zhou. Scal- ing relationship learning mathematical reasoning with large language models. arXiv preprini arXiv:2308.01825 2023a.
Zheng Yuan, Hongyi Yuan, Chuanqi Tan, Wei Wang, Songfang Huang, and Fei Huang. Rrhf: Rank responses to align language models with human feedback without tears. arXiv preprint 00X1::2304.05302, 2023b.
Xiang Yue, Ge Zhang Xingwei Qu, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653, 2023.
Xiang Yue, Tuney Zheng, Ge Zhang, and Wenhu Chen. Mammoth2: Scaling instructions from the web. arXiv preprint arXiv:2405.03355 2024.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning, Advances in Neural Information Processing Systems, 35:15476-15488, 2022.
Wenhao Zhan, Masatoshi Uehara, Nathan Kallus, Jason D Lee, and Wen Sun. Provable offline reinforcement learning with human feedback. arXiv preprint arXiv:2305. 14816, 2023.
Beichen Zhang, Kun Zhou, Xilin Wei, Xin Zhao, Jing Sha, Shijin Wang, and Ji-Rong Wen. Evalu- ating and improving tool-augmented iomputation-intensiv math reasoning. Advances in Neural Information Processing Systems, 36, 2024a.
Shenao Zhang, Donghan Yu, Hiteshi Sharma, Ziyi Yang, Shuohang Wang, Hany Hassan, and Zhao- ran Wang. Self-exploring language models: Active preference elicitation for online alignment.
arXiv preprint arwiv:2405. 2024b.
Tong Zhang. Feel-good thompson sampling for contextual bandits and reinforcement learning. SIAM
Journal on Mathematics of Data Science, 4(2):834-857, 2022.
Tong Zhang, Mathematical analysis of machine learning algorithms. Cambridge University Press, 2023.
Yuheng Zhang, Dian Yu, Baolin Peng, Linfeng Song. Ye Tian, Mingyue Huo, Nan Jiang, Haitao Mi, and Dong Yu. Iterative nash policy optimization: Aligning Ilms with general preferences via no-regret learning, arXiv preprint arXiv:2407.00617 2024c.
Yao Zhao, Rishabh Joshi, Tianqi Liu, Misha Khalman, Mohammad Saleh, and Peter Liu. Slic-hf: Sequence likelihood calibration with human feedback. arXiv preprint arXiv:2305. 2023.
Chujie Zheng, Ziqi Wang, Heng Ji, Minlie Huang, and Nanyun Peng. Weak-to-strong extrapolation expedites alignment. arXiv preprint 2671v 2024.
Kunhao Zheng, Jesse Michael Han, and Stanislas Polu. Minif2f: cross-system benchmark for formal olympiad-level mathematics. arXiv preprint arXiv:2109.00 10, 2021.
Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang, Zhuoran Yang, and Tong Zhang- Gec: A unified framework for interactive decision making in mdp, pomdp, and beyond. arXiv preprint arX1v22211.01962, 2022.
Han Zhong, Guhao Feng, Wei Xiong, Li Zhao, Di He, Jiang Bian, and Liwei Wang. Dpo meets ppo: Reinforced token optimization for rlhf. arXiv preprint arXiv:2404.189) 2024.
Denny Zhou, Nathanael Schärli. Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoc Le, et al. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205. 10625, 2022.
16

---

Published as a conference paper at ICLR 2025
Xinyu Zhu, Junjie Wang, Lin Zhang, Yuxiang Zhang, Yongfeng Huang, Ruyi Gan, Jiaxing Zhang, and Yujiu Yang. Solving math word problems via cooperative reasoning induced language mod- cls. arXiv preprint arXiv:2210.16255 2022.
Brian D Ziebart. Modeling purposefull adaptive behavior with the principle of maximum causal entropy. Carnegie Mellon University, 2010.
17

---

Published as a conference paper at ICLR 2025
A NOTATION TABLE, RELATED WORK, AND MISSING DETAILS
A.1 NOTATION TABLE
Notation
Description
X de
The prompt and the prompt space.
The distribution of initial state (prompt).
8h S an E On
The state, action, and observation.
Episode length. e.g., the maximal number of tool calls.
Fe [
The true observation kernel.
T (2,
is trajectory and y is the completion part, ie.. we exclude æ from т. The true utility function associated with BT model defined in Definition The true model with observation kemee P' and utility function u*
u
M (S, A, H, ,dd,,,,,,r" x))
x.))
((2) 1/(1 + exp(-2)) is the sigmoid function.
2 f0
Preference signal.
R [TN]
The policy, which is parameterized by the LLM.
M = (S, H, do, u)
One arbitrary environment with observation kernel I and utility function
Ruol = [
One arbitrary reference policy.
J (T; TTar))
The KL-regularized target (equation 2 with environment M and reference Tyaf-
The coefficient of KL penalty, defined in equation 2
Q. [
The optimal Q-values associated with [[;;; M,, Tror). defined in equation
VM [MMA
The optimal V-values associated with [x;:M, Trar). defined in equation
TM A.M.
The optimal policy associated with ([;; mrar). defined equation
DPO (-)
M-DPO loss, defined in equation 10
CA.KTol-)
M-KTO loss, defined in equation
J(T)
The abbreviation JT;; MM" , mo), defined in equation 14
[m N=1
The optimal policy associated with J(T).
The main and exploration policy round
Reg(T)
Regret over horizon T, defined in equation 15
U, P
Known sets such that u² E l and IP" E e
B
Assuming u" (z,w) w) E [0, B). V(r.u)) MLE of u² and IP" at round t, defined in equation 16 and equation 17
in, Pr
Ur Pe
Confidences sets of u" and PP" at round t, defined equation 19 Absolute constants.
C,,02,
1/(2 + exp(-B3) + (x)(
du
Eluder coefficient from Definition
dp,
Generalized Fluder-type condition from Danmition Total variation distance between two distributions.
TV(;))
Table 4: The table of notations used in this paper.
A.2 RELATED WORK
LLMs for Mathematical Problem Solving. A line of works proposes to prompt LLMs to solve the complex reasoning task in step-by-step manner, known as the Saiin-of-Though (CoT) prompt- ing (Wei et al. 2022 Zhou et al. 2022 Zhhu al. 2022 Tong et al. 20224), which has been standard practice in reasoning task. However, LLMs often struggle with basic arithmetic and symbolic ma- nipulations when relying solely on internal knowledge and natural language reasoning, as measured by standard benchmarks (Cobbe et al., 2221aa] Hendrycks et al. 202T). To overcome these limita- tions, several studies have explored the use of external tools to enhance the LLMs' problem-solving
abilities. This includes calculators (Cobbe et al. 20216] Shao et al. 2022). symbolic solvers (Zhang 2023), and code interpreters (Mishra et al. 2022 OpenAl| 2023). A particularly effective approach is the Program- based method (PoTT), which performs CoT reasoning by writing code and using the output of the written code as the final answer (Gao et al. 2023 Chen et al., 2022). This method significantly outperforms traditional CoT-based techniques in mathematical problem solving. How- ever, PoT also faces challenges in planning and error handling, where natural language reasoning is more suitable (Gou et al.) (203aa). In view of this, tool-integrated reasoning is proposed to combine the natural-lam langulllnngaaaaaaaae based intrinsic reasoning with the external tools (Gou et al. (023b) and has
achieved great progresses in recent studies (Gou et al. 2023b Yue et al. 2023
Yu et al. 2023
18

---

Published as a conference paper at ICLR 2025
Toshniwal et al. 2024). While these efforts have primarily focused on synthetic data generation for tool-integrated reasoning, our work aims to further boost the performance of tool-integrated LLMs by RLHF.
Shao et al. 2024
RLHF and RLHF Algorithms. The predominant approach in RLHF is the deep RL method, Proximal Policy Optimization Algorithms (PPO) (Schulman et al. 2017). which leads to the great successes in Chat-GPT (OpenAI, 2023), Gemini (Gemini et al., 2023). and Claude (Anthropic. 2023). However, applying PPO requires extensive efforts and resources (Choshen et 2019 En- gstrom et al. (020), often beyond the scope of open- open-source capabilities. In view this, alternative approaches have been developed. The rejection sampling fine-tuning was first proposed with the name RAFT (reward ranked fine-tuning) in RLHF (Dong et al.) 2023) and was later extended to machine translation (Gulcehre et al. 2023) and mathematical problem solving (Yuan et al. 2023a) Its theoretical advantage was explored in Gui et al. (2024). Subsequently, another long line of works proposes direct preference learning algorithms, including SLiC (Zhao et al., 2023), DPO (Rafallov
et al. 2023). IPO (Azar et al. 2023), KTO (Ethayarajh et al. 2024), and GPO (Tang et al.) 2024).
These algorithms bypass the reward modeling step and optimize carefully designed loss objectives directly on the preference dataset, hence the name direct preference learning. There are also some
works focusing on more general preference structure Munos al. (2023); (Swamy al. (2024); Ye
ct al.
20 Rosset et al. (2024) beyond the reward-based framework or post-processing of the (Lin et al. 2023 Zheng et al. 2024).
model
The newly proposed direct preference learning algorithms have largely advanced the RLHF area, particularly the post-training of open-source models, with the Zephyr project notable example
(Tunstall et al. 2023). After this, a long line of work re.g..LLiu et al. 2023b Xiong et al. [024]] Guo
ct al. 2024b 20.44b Xu et al.) 2023, Tajwar et al., 2024, Xie et al. 2024a Zhang et al. 2024b) Liu et a
2024a bj Meng et al. (224) demonstrates the effectiveness of on-policy sampling (the samples are generated by the policy to be trained) and online exploration in enhancing direct preference learning.
In particular, the online iterative DPO (Xiong et al. 2024 Xu et al., 2023] Hoang Tran 2024) and
its variants (e.g.. Chen et al. 2024b, Rosset et al. 2024| Cen et al. 2024 Zhang et al. (024c) have made state-of-the-art open-sourge models (Dong al. 2024), even the industry models (qwe)
2024 Meta 2024). Despite these advancements, most algorithms are proposed and designed for single-turn interactions and chat. The scenarios beyond single-turn chat remain largely unexplored in the existing literature. One exception is the very recent work by Shani et al. (2024) which studies multi-turn chat task under general preferences. In contrast, in this paper, we aim to explore the use of RLHF in multi-turn tasks that incorporate interactions with external tools. Meanwhile, they derive a serror-descent-aane policy optimization algorithm, which also different from ours.
RLHF for Math Problem Solving. Algorithms traditionally used in general chatbot applications have been adapted to enhance the reasoning capabilities of LLMs in mathematical contexts. For
instance, RAFT (Reward-rAnked Fine-Tuning) (Dong et al., 2023 Yuan et al.) 20236 Touvron
et al. 2023) is extensively employed for synthetic data generation, whether through on-policy (self-
improving) (Yuan et al. 2023a) or off-policy Kknowledge distillation) methods (Gou et al. 2023b
Yu et al.] 2023; Toshniwal et al. 2024 Singh al., 2023 Tong et al., [022). The reward signal in these scenarios is typically derived from either final result checking or sircome-superviseee Reward
Models (ORMs) (Uesato et al.) 2022 Zelikman et al. 2022). A novel approach by Lightman et al.
(202) introduces Process-supervised Reward Models (PRMs), which provide feedback at each step of the Sfhain-of-Thougee demonstrating significant improvements over ORMs when combined with rejection sampling (Lightman et al. 2023 Wang et al. (023a).
In addition to the RAFT, the GRPO algorithm proposed in (hao et al. (024) studies multi-tun turti-turr math problem solving but focuses on the CoT format without external inputs and the resulting model achieves the tate-of-thee-ar performance in its class. The GRPO a variant of Reinforce (Williams) 1992) thus falling into the scope of deep RL methods.
Further advancements include adapting direct preference learning algorithms to mathematical prob- lem solving. For instance, Jiao et al.) (2024); Yuan et al. (2024) have applied the original DPO or KTO by taking the trajectory completion as a "meta" action. Xie et al. (2024b); Pang et al. (2024) further adapt the online iterative DPO originally designed for chat (Xiong al. [024]] Xu et al.
2023 Hoang Tran 2024) and achieve better performance for CoT reasoning. Inspired by the suc- cess of PRMs, recent studies have explored generating proxy step-wise labels for the intermediate
19

---

Published as a conference paper at ICLR 2025
steps of the reasoning trajectories. For instance, Xie et al. (024b): Chen al. (2024a) Lai et al. (2024) leverage Monte Carlo Tree Search (MCTS) and use the estimated Q value to generate the
proxy labels for the intermediate steps. Lai et al. (2244) proposes use Al feedback like GPT-4
(Lai et al. 2024) to find the first error step in the trajectory. Meanwhile, [u et al. (024) identifies a trajectory with the correct final answer and errors preferable, and prompts the SFT model with a high temperature, starting from some intermediate step to collect rejected trajectory with errors (Pi et al., 2024). Finally, very recent study by Chen al. (2024a) proposes to use MCTS with backward iteration from the final leaf node to compute the proxy unregularized value of each node. Preference pairs are then extracted from the tree by fixing the prefix and comparing the next single reasoning step. Then, they run the original DPO on these intermediate actions with the proxy labels from MCTS. To summarize, these works present different ways of preference data collection and apply the original DPO algorithm (with some additional marginal loss and regularization adapted from the literature), thereby differing from our work in both algorithmic concepts and application scope. In contrast, study preference learning in the context of trajectory-level comparison, where we derive the optimality condition and introduce a multi-turn DPO within an online iterative frame- work, specifically for tool-integrated mathematical problem solving. However, remark that while we focus on the trajectory-lever comparison, the preference signal itself can be generated in step- by-step supervision (see Section 2.] for the detailed examples). When preference signals for partial trajectories with shared prefixes are available, our method can also adapt to learn these step-level signals (see the optimality condition in equation 9h. In particular, the algorithmic design presented in this paper can be readily combined with the MCTS-based data collection strategy outlined in recent literature, which we leave for future work.
A.3 MISSING DETAILS
Multi-furn KTO. With equation 9 implying that with term (C) 0, the implicit reward given
by A = 1 " log
[ec,, [x.] multi-turn version of KTO (Ethayarajh et al. 2024), denoted as
M-KTO, can also be naturally derived:
LM,TT(()) = Br.yD (N v(a, v)].
(13)
where
H
Tu, (an
ug(x, 3) = 17
log
Trof, (aof,a((assss (an(sn))
H
= Eeg w
DEL (To(-1)))))) (hee(ll
and
2 (n 20))
if y ydesirable 2
o(x v)
A A - us(x, (x)))))
if y ~
Here A and x_ are two hyper-parameters. We notice that Mitra et al. (024) developed an online iterative version of KTO for the CoT format reasoning task. Here we extend it to build the tool- integrated reasoning agent.
B IMPLEMENTATION DETAILS
Supervised fine-tuning (SFT). We first fine-tune the model for the tool-integrated reasoning task (Gou et al. 2023b), using a subset of the Open-Mathlnstruct dataset, which was generated by the permissively licensed Mi Mixtrall rattral-8x 1 model through in-context learning. The problems are from the training sets of MATH and GSM8K datasets. We restrict the number of samples for each question to 50 and remove the nearly duplicate responses. Eventually we get 510K samples the SFT dataset. We train the models for 4 epochs at most with learning rate of 5e-6 for Gemma instruct models (Team et al.) 2024) and a learning rate of le-5 for Mistral-vO.3 model (Jiang et al. 2023). The learning rates are determined by searching (2e-6, 5e-6, le-5). We use cosine learning rate scheduler and set the warm-up steps as 100. The samples are packed into blocks with length 4096 to accelerate training and global batch size of 64 is used. We also mask all the user messages (i.e., the
20

---

Published as a conference paper at ICLR 2025
Multi-turn Reasoning with Python Interpreter for Mathematical Problems
Prompt:
User: Let
ax + 3,
if 2,
f(x)
x-5 if =
2x
<
Find a + b if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper)
Assistant: The function must continuous 2 -2. That means 5 ax 3
and x - 2x at = -2. Let's now salve for in each case. (/1)
from sympy import symbols, Eq, solve
2 #a define the variables
5 x, a, symbols('x b')
#w the function must be continuous
eql Eq(x-5, axx + 3) eql_sol selve x)
10 ## the function must be continuous at -2
11 eq2 Ex 2*x
12 sol_sol solve
## print the solutions print (rc2_secc__
(e1)
User: "output [-8/(a 1)] [b 51" (o:)
Assistant: Thus, x -8/(a when and x when We can now solve for a and b. (fz)
#M -8/(a-1) when z
2 eq Eq(-8/(a-1), 2)
a soiva print(a)
## b-5 when -2
eq4 (b -2)
b solve b)
print (b)
(c2)
User: "output [:3]
[3] (oz)
Assistant: Thus, a -3 and = 3. Hence, a b- 0 (f)
Figure 4: An example multi-turn mathematical reasoning with Python interpreter. The action is in ReAct style (Yao et al. (022) where it consists of a reasoning step fr. and an execution step en-
prompt and the messages returned by the Python interpreter) in the training. The checkpoint at the end of the third epoch used for Gemma and the checkpoint of the end of the second epoch is used for Mistral as the starting point for RLHE This is because these models outperform the fast-iteration one with considerable margin and is very close to the next one. An ablation study the SFT epochs also included.
Data format and generation. We format the data into a multi-turn chat where the user asks the LLMs question, and provide the messages returned by the Python interpreter in the subsequent user rounds of chat. For all the data generation process, we adopt the following constraints: (1) for each turn, the model the maximal number token for each trajectory is 2048. Following Gou et al. (2023b); Toshniwal al. 2024), the I.LM agent is allowed to call the python interpreter when decodes python code
generate up to 512 tokens; (2) the maximal number of steps is H=6; (3)
21

---

Published as a conference paper at ICLR 2025
starting with " "python and ending with >. For each step h. to generate the observation oh, leverage the python package IPython, and run all the codes in the history one by one and treat each code snippet as a Jupyter cell. We only return the standard output or the enror message from the last snippet. When there exists some bug in the code, we only return the error message which is typically less than 20 tokens as in Toshniwal et al. (2024). We notice that some works (e.g. [Shao ct al. (2024)) also returns the first and the last 50 tokens of the trackback information.
Evaluation Configuration. All the models are evaluated in the zero-shot setting. For all the data generation process, we adopt the following constraints: (1) for each turn, the model can generate up to 512 tokens; (2) the maximal number of steps is H=6; (3) the maximal number of generated token for each trajectory is 2048. When collecting new data for online iterative M-DPO, we set temperature to be 1.0 and decode without top-K or top-p sampling. For evaluation, greedy decoding is employed so that the results are generally comparable with previous works Gou et al. (2023b); Toshniwal et al.] (2024). For evaluating the models with pass aass@n rate, we follow Toshniwal et al. 2024) to adopt a temperature of 0.7.
Python Experiment Environment. We find that the evaluation can influenced by the python environment, the precision (especially for the Gemma-1. 1.1 models), and even the virtual machine we use. This does not affect the overall trend and conclusion because the magnitude of oscillation is relatively small compared to the overall improvement. For completeness, however, we specify some of the key package versions here. We use transformers 4.42.4, torch 2.3.0, sympy 1.2, antir4- wython3-runtime 4.10.0, IPython 8.26.0 for all models. We evaluate the models using use vllm 0.5.0.postl for most the experiments except for Gemma-2 where vllm 0.5. is required. The inconsistency of vllm version is because Gemma-2 model was not released when we performed the main experiments of this project. We fix the python environment and machine for our evaluation throughout the experiment. For SFT, we use the open-source axolotl project with version 0.4.1 and for online iterative preference learning and RAFT, we use the code base from RLHF Workflow DDng et al. 2024).
and
RAFT implementation. RAFT first collects N trajectories per prompt, filters the low-quality data (by reward function), and fine-tune the selected trajectories. The data generation step is similar to the online iterative M-DPO training, except that we only keep the trajectories with correct final answer. For each prompt, we sample at most & trajectories where search & {1, 3, 8} and use eventually because we do not see improvement by leveraging more data. We run the algorithm for three iterations in total. The training parameters batch size of 32 so that there are enough optimization steps. For Gemma models, we use a learning rate of 5e-6. For each training stage, we train the models for epochs in total according to our parameter search. For Mistral model, we find that a smaller learning rate of le-6 and training for epoch give us much better performance.
similar to the SFT stage, but use smaller
Prompt template. We do not tune the prompt though we do observe that the prompt engineering can further improve the performance. For all the experiments, we simply adopt the chat template of the models form multi-turn chat as in Figure
C ADDITIONAL EXPERIMENTAL RESULTS
We include additional ablation studies in this section for more comprehensive understanding of the proposed algorithm.
The best model is obtained with starting checkpoint fine-tuned with more than epochs. Tun-
stall et al. (2023) finds that if the SFT model is trained for more than one epoch, the subsequent DPO training will lead to performance regression with longer training in terms of instruction-following ability and benchmark for general chatbot. In other words, there exists trade-off between the SFT training epochs and the DPO training steps. Moreover, the best model is obtained by SFT for one epoch in their practice. We also conduct ablation study the impact of the SFT epoch and summarize the results in T Consistently across all tested scenarios, the subsequent iterative M-DPO training leads to considerable model improvement compared the SFT model. Meanwhile,
22

---

Published as a conference paper at ICLR 2025
An example of external messages returned by the Python interpreter
Prompt:
User: Convert the periodic decimal notation 0.42 fallowed by an infinire sequence of 42s into reduced fraction such that and have common factors other than 1. Find the sum and
import вympy sp
Define the number as sympy number with infénite repeating digits
digits '42'
num sp ((.. (242422444,, 1000000) # 1000000 digits are accurate
other 320 '0' File http://wwwou..oo///////oooooeeea line 69, handler_ raise pttppeee__cccprooootiooottttt exception, excep- tion message)
User:
"ourput httpa//s...........0//0//0000000
File
decorator/tumeoue http:///mmuuuucccooooorrppppp""" Cecorator.p"", line 45,
raise exceptior:
raise exception.
timeout decuratoctimeout www.oo/oddooooooooomeoooooorero Timed Out"
Figure 5:
An example of external messages returned by the Python interpreter. The model writes down a bad python code leading to an anomalous and lengthy error message.
we also observe a similar trade-off between SFT and RLHF training because with more SFT epochs, the gains from the RLHF stage decrease. However, in our case, the strongest model is obtained with three epochs of SFT, followed by fine-tuning through iterative M-DPO. which is different from the offline DPO training (Tunstall et al., 2023) the iterative DPO for general chatbot (Dong et al. 2024 with only one epoch of SFT.
Table 5: Ablation study of the impact of SFT epoch. Mixture sampling is adopted for the iterative M-DPO training and we run for three iterations in total. The gains relative to their starting SFT checkpoints marked by
Model
Method
GSMSK
MATH
Gemma-1.1-it-778
SFT epoch
75.1
41.1
Gemma-1.1-it-77
SFT epoch + Iterative M-DPO
80.6 15.5
46.7 15.6
Gem-a-l.1-it-7B
SFT epoch
75.3
44.0
Cenmma-1.-it--77
SIT 2 epoch + Iterative M-DPO
82.4 17.1
49.8
15.8
Demmma-1.1it--1
SFT epoch
77.5
46.1
Gemma- Cemma--.1-it-7B
SFT 3 epoch + Iterative M-DPO
83.9 16.4
51.2
15.1
NLL loss helps when the SFT model is substantially underfitting. The recent work Pang et al. (0244 has introduced iterative RPO, specifically aimed at enhancing Chain of Thought (CoT)
pabilities for solving mathematical problems. A key feature of this approach is the inclusion of an additional negative log-likelihood (NLL) loss for the preferred response. The main intuition for adding the NLL loss is that the original DPO algorithm (Rafailov et al. 2023) tends to reduce the likelihood of the preferred responses, and this believed hurt the reasoning ability (Wang al. 2024). Motivated by their results, we explored the applicability of this idea to our setup. We conduct an ablation study by adding the NL.. loss into the iterative M-DPO training and observe performance regression reported in Table We observe that the best model is obtained in the second iteration if we add the additional NLL loss even though we use the mixture sampling to increase the diversity of the collected data. With time-weighted exponential moving average for smoothing training record,
we observe that the log probability the chosen responses and rejected responses are (-126, -222) at the 200th step of the third iteration training when we add the NLL loss, as compared to (-166, -350) in the case without the NLL loss. This is consistent with the result of Pang et al.) (2024) where with the additional NLL loss, both the log probability of chosen responses and that of rejected responses
23

---

Published as a conference paper at ICLR 2025
increase. These evidences indicate that the NLL loss further contributes to the model distribution collapse and eventually hurt the overall performance of online iterative learning. Finally, we notice that the additional NLL loss be viewed as an implementation of the pessimistic principle (Liu et al. 2024b). This also explains its inferior in-domain performance though it may be helpful to stable the training, which requires more in-depth studies.
However, one distinct feature between our setup and Pang et al. (024) is whether we first fine-tune the initialized SFT model with in-domain data. To further understand the phenomena, we fine-tune the Gemma-l1.1-t--D with only 100 steps (so that the model knows to leverage Python code to solve the problem) the starting checkpoint of preference learning and conduct an ablation study with the NLL loss using this model. We observe when the SFT model is substantially underfitting, the addition of NLL loss actually enhances performance. This scenario mirrors the findings of Pang et al. (024)), who utilized a general ALaMMA2-70B-hhh model (Touvron et al., (023) without firstly fine-tuning on the in-domain data. Our observations align with prior research in the context of developing general chatbots (Lin et al. 2023), which suggests that RLHF is less effective without preliminary SFT.
Table 6: Other ablation studies. Mixture sampling is adopted for the iterative M-DPO training and we run for three iterations in total. The gains relative to the iterative M-DPO are marked by 1.
Model
Method
GSM8K
MATH
Gemma-1.1-pp
SFT epoch
77.5
46.1
Demma-1.1-it-11
SFT 3 epoch + Iterative M-DPO
83.9
51.2
Gemma-1.1-it-77B
Iterative M-DPO with NLL loss
81.7 J2.2
49.5 41.7
SFT 100 steps
50.8
23.7
L.
+ M-DPO Iteration
57.8
27.9
Gemma-1.1-it-7B
+ M-DPO and NLL loss Iteration
61.0 13.2
30.1
12.2
On-policy sampling and small learning rate mitigate the probability drops in preferred re- sponses. In the literature, the Direct Preference Optimization (DPO) algorithm is often reported to diminish reasoning capabilities by reducing the likelihood of preferred responses (Yuan et al. 2024 Hong et al. 2024 Meng et al. 2024). In our preliminary experiments, we also observe simi- lar phenomena with large learning rate (1e-6), where the model's reasoning ability collapses after only a few training steps, preventing convergence to good reasoning performance. In contrast, we find that using on-policy sampling within our online iterative training framework, coupled with smaller learning rate (2e-7 or 4e-7), the DPO algorithm enhances the model's 's reasoning abilities. To interpret our observation, we can first write down the gradient of the DPO as follows:
VeLDPo Tret) -17 o (ro(z,)') re(2,8"))
(x) VOTT(y"|2)
1
T VOTT(
where ro(x, y) 17 log (a.y] is the implicit reward and we the single-turn one for simplicity. In practice, the probability of the rejected responses typically decrease, and their gradient quickly
me(w.))
dominates when Ty (o('(<z)) << To (u(")>>) and the optimization becomes unlearning of the rejected responses. In this case, the probability of the chosen responses cannot increase. This phenomenon was also discussed in the blog Guo et al. (2024a). When we adopt on-policy sampling, it leads to relatively large probability for both rejected and chosen responses at the initial stage, ensuring that both gradients remain valid and effective. Moreover, small learning rate prevents the model from deviating too significantly, maintaining the effectiveness of both gradients. We also notice that for the KTO algorithm, the preferred responses and the rejected responses do not appear in pairs. We suspect that the probability of the preferred response increases because the gradients of the rejected response do not dominate in every mini-batee of data. A more comprehensive understanding of the training dynamic of the direct preference learning algorithms remains largely open and we leave a more detailed study of this phenomena future study.
24

---

Published as a conference paper at ICLR 2025
D THEORETICAL PROOFS
D.I THEORETICAL RESULTS
In this following, we show that the multi-turn RLHF problem be solved in statistically efficient manner under standard assumptions in learning theory literature. In particular, for generality, we tar- get the most challenging scenario with stochastic and unknown transitions, while as aforementioned, multi-turn mathematical reasoning with external tools falls into a relatively easier regime with de- terministic transitions. As mentioned in the main paper, we mostly study the KL-regularized target due to the lack of theoretical research on it. The other target of optimizing the rewards has been theoretically studied in Wang al. (2023b) while the techniques of analyzing algorithm and corresponding guarantees have also been developed in Cai et al. (2020), which can be migrated to considering preference feedhacks. Also, to ease the presentation, we consider the scenario with batch size
surror-descent-syy
while the results can be easily generalized to large batches.
First, to measure the online learning process, we define the optimal policy as π == arg max (Tx) == J (T;;;; MM To))
(14)
and introduce the standard notion of regret as
Reg(T) = Σ J(TT) - J(T)), LE[2
(15)
which represents the cumulative performance loss over T steps comparing the learned policies [m]]?"" against the optimal policy '². In addition, we consider that bounded u (x,z) [0, B] for all (z, y) to maintain reasonable utillity regime. Also, it is assumed that we have accesses the following policy improvement oracle, that is analogue to the one considered in Xiong et al. (2024).
Definition (Policy Improvement Oracle), For any model M (S,A, H.P. P. do u) and a reference function Tref, we can compute the optimal policy associated with the model [TMMik== as in equation 4
IH iteratively
The overall algorithm, i.e., the theoretical version of online iterative M-GSHF, is also summarized in Algorithm At each round t, with D -t=1 Dj as the aggregated dataset, it starts with performing a maximum likelihood estimation (MLE) of the reward function u² over a set U1 whose elements are bounded in [0, B). as
ût arg max Lr (ū) Nell
Σ
2 (og( ú()))) + (1 - log(o((((7)) -(r?))))(r)))))
(16)
(2,T), AA
and also an MLE of the transition kernel IP* over a set P as
P..
arg max (I(()
Σ log p (T),
(17)
PEP
(TTT)EUf)))) D
where (r) denotes the probability of trajectory under policy and transition kernel P. With the obtained model M, (ū,, P.... the Oracle defined in Definitionn called with the reference policy Tref set as the initial policy To. whose output is adopted as the main policy m!-
Then, we specify how to choose theoretically sound exploration policy m?- The previous work of Xiong et al. (2024) on single-turn RLHF has demonstrated the intuition that the exploration policy should be in charge of collecting information of the uncertain parts of the environment M, which is thus often selected to maximize one uncertainty measurement. In the multi-turn RLHF setup considered in this work, the following proposition serves as the cornerstone to find a suitable uncertainty measurement to decide the exploration policy. In particular, we can observe that the optimal policy is parameterized by the optimal Q-function. If a different set Q-function adopted for policy parameterization, we can bound its performance as follows.
Proposition (Value Decomposition Lemma for KL-regularized MDP). If considering a set of ( and a reference policy Eree with the induced policy as
Q-functions
fr (an x Wrof (a)) exp (Q (Qn, un)/n).
25

---

Published as a conference paper at ICLR 2025
and the corresponding set of V -functions M
(ssss) (en (Qn(sh an)] (fn((( aroot(-)ss))),,
VH+I(SH+1)+ 0,
for any comparator policy it holds that
(T) -(+)
-.d (u" (sH, a (ur (SH, aH)l
Σ P. (Vet((0+12) (1k and
Vn+i(Sh+1) (hh an)
bE[H]
hE(H]
hE[H]
where the expectation E. p. is with respect to the prompt and response (i.e., the trajectory) gen- erated following do, P and n.
Based on Proposition the exploration policy oद is selected
π arg max max PEP.
A. [ü(SH, an)] E d (ū(H az)
E d. (Ri(sH,a)) - Ag ((e(HH
(18)
E. (ra+i(8n+1) (PanPest (sh an)
where Uf and P. are two confidence sets defined as
ü, = {u E 4 L(()) 2 L((()) - log((4(T/6)} P. [IP P (PP) L,((P) log( (PT/()))
(19)
with denoting absolute constant here. Note that for the theoretical convenience, we have as- sumed U and are finite here, which can be extended the infinite case using standard discretiza- tion techniques. It can be observed that πテ is selected maximize combination of uncertainty from estimations both rewards and transitions. If considering known transitions (i.e., without need to estimate IP), the uncertainty from the estimation of transitions dimimishes, which leads to similar uncertainty measurement adopted in Xiong et al. (2024).
The following theorem establishes rigorous guarantec for the regret incurred.
Theorem 2. Assuming ll and I P, with probability at least we have that
Reg(T") SK" -1 B
+ (/TTT)
Σ
Σ hE(H)
Edo. p [PK (t
1) *E
where k = 1/ exp(- + B)), e2 is absolute constant, du the Eluder coefficient defined in Desinition while dp and E(.) are from the generalized Eluder-type condition defined in Definition 4
We note that the Eluder coefficient and the generalized Eluder-type condition are standard and well- adopted conditions in the theoretical studies RL (Zhang 2023 Zhong al.) 2022 Liu et al.
2023a Xie et al.) 2202 Aggawaa et al. 2023) and also RLHF (Zhn et al. 2023 Wang et al,
b023b Ye et al. 2024). Moreover, for board class of RL problems (see Zhang (023) Liu et al.
(2023a) for more details), the Eluder coefficient du is small and the condition is satisfied with
E(dp, T, eg log([P HT/5)) < /d PopTlog([P)) /8), which implics that the regret theoretical
version of Algorithm is sublinear in T. further evidencing its statistical efficiency.
26

---

Published as a conference paper at ICLR 2025
D.2 PROOF OF PROPOSITION
2
Proof of Proposition For one policy x, starting with VM. = O, we recursively define its V- value and Q-value functions on one model M = (S, H, P. do. u) and the reference policy Fre
QMa(shaan) an):= {
aH) h HI
if H 1,
(S)) == (. (Qn, (sh, an) - 17 DRL (Th (-1s)) Tref & (-1s)))]-
is noted that with the optimal policy TM. QM. Q,A,, and MM, M.h. In the following discussions, we exclusively focus on the model M = (S, A, HPP do, u*) with abbreviations Q = Tu+.n and VÆ -VM-.
For any comparator policy it holds that
((T) J(+) = Edd (VT(s)) 11n)--i1 ((i))-((s1))
Eed [ - Ri(s1))
For any h E [H], we can obtain that
(sn) - (x()
Pii (1x) - ((()
(a) 1
En [0%(s an)l
DDL (T (-((-15)), Tzet (-15)))
-Edy Edg. Pia
Exn [ an)
17 DKL (fn (((s)) T (-1s))
-Edg - Pin
B*, [QA (sh, [m)]
1y DKL (fn ( (-18)) T (-1sh))
+ -ra
E*n, (2h,, an))
1y DK (#t (8h); Traf (r1s)
-Edo,rss.PPn-
(T an) Qu(ssn an)
@ (s,: an) -
+ +Edg 1.Pa--
Exn [2 an)
Exn
[ an)
17 Eeg [DKL (TT((199), [-
+ 17 Eio [DKL (#n(-199)) ( (-1ss))]
(b) Edgg 1 (T,, and - (n an) - Edo.ora.Pinn
( an) - (s an)
17 Eedg [DKL (-1x)) An (-1sn))]-
In the above derivation, equation (a) is from the definitions of Q" and VT, and the relationship between Q and V. The equation (b) is because
(term I) = Er,, [Qn(sh [n
EFA [ an)
= =) Ern log An (an (sn)
nn. (an (sn)
- 17 TEn,
log
Tree (an (sn)
Tref (an
= DKL (TM (-1sn), Tref (-
DKL (TT( (2
77- DKL (Ita (-1s)), Fref, is (-1sp))
where the second equation from the relationship that
#x (()a))))) log (xn)- #eof. (s))a))
(1h,, an) = 1) log
Furthermore, if h H. obtain that
Edo 1. PPPi-
V (sm) - Vr((sn) Edg 1PPin-1
27

---

Published as a conference paper at ICLR 2025
u (sH, (HH, an)] Edg u u'( QH
Edo Pim--
- 17. -doomm-- [DLL (TH (-)) #H (+1sH))]
-Bdo,aua, Pin-i [u" ($ an) Edg f.. [u" (S an)]
Edg Pi: VH++(SHH+) Q ( an)
Edg ( Qn(ss
17 -d [DLL (- AH (11s))],
where the second equality leverages that VH+1 (sH+1) O; otherwise, for all h < H holds that
B - (TT(s)) Vh(( Edo [VT( Vh(sn))
( an) (sn, and] Edo. (sn, (sh an)
-1 -1. -.do [DRL (I (1)
Rdo (n+i(SA+1) Qn an)
-
- -17 - [DRL (TM (-1sa))
Edo Pin (1+z1)5++1) Vh+ ( Eeg
[ (i+t2 Vn+ ((h+ii)
The proposition can obtained by iteratively using the above relationship for E [H].
D.3 PROOF OF THEOREM 2
First, with the assumption u" li and T P. the following lemma demonstrates that it and P. are valid confidence sets.
Lemma (Proposition B.1 from Liu et al. (2023a)) There exists an absolute constant ei such that for any ò E (0, 1]. with probability at least for all t [T], à E U. and
P, holds that
(. Le(w") S log( (T)()))))))
L.( (IP") I eı log( ()
which implies that u² E Ul and P E P.
Then, we provide an additional lemma demonstrating the in-sample error of the MLE and optimistic estimators.
Lemma 2. There exists an absolute constant 02 such that for any 8 E (0, 1}, with probability at least
1 - 8, for all E [T], have
( aZ.z) 0z(5), at (w ab.u) u² (s) at x)) log(((
set
Σ o (Tols)
n a Te n a(,n)) (u ag u² (st. at, )) (0g(((
ice
for all E [T], [H] we have
Σ Σ
TV ( T.. (Pi P... Phtil {do, क.. .P.H I 02 Log(P/HT//))
SE(1.2) RE[H] ist
Σ ΣEN
TV (Ido 7.. (f (F P., Hh (do, ml Pt.... log( ((HT)8))
SEIL,2) hE[H] ist
where TV n, P), (do, ,P)) denotes the TV distance between the probability distributions over the trajectories induced by do, n. and da, n²,
Proof of Lemma 2 First, for ūt absolute constant such that for all [T'].
obtain that with probability at least 1 - 6, there exists an
Σ o (7, ai.) ūe Te a)+))) (u" (s a..)) n" (st,,,a.,)))
28

---

Published as a conference paper at ICLR 2025
c. Σ log 2, a.0 (u'(u" (s,, cl.) 4" (" oi,H)) + (1 x)) (u (s7, aï,u) u² (s) at,h)) 2:.0 (ūi(s),H::), - ū. (., 4,,)) + (1 - z.) (57, 4 (( ūo ū. (x, H.
+c=log((((T88)
=[(((u")) Li (ū+) + (2((66))))
< <<(L((u)) Le (ūt) + log(R((7/8) + ()((())))))
S 1o)(
where the first inequality is from Proposition B.2 from Liu et al. (2023a) and the second inequality uses Lemmal
The result for ún can be similarly established.
Then, following similar steps, for Ft. we can obtain that with probability at least 1 - 8, there exists an absolute constant с such that for all t [T].
EEM
TV (fdo, o/, [P I. Pr,. 1h {do, xÿ. Fi.H
jE(1.2] kE(H) K<E
P P L.a))
Σ Σ
log
P.n 1 18 (g.h)) + 1og((Pa)//T/0))
jE(1,2) hE(H
ict
Σ
log
2 2 Log([P|HT/)) (T!)
(E(122) ict
(L((P)) - L, (P.) + 2 log (P)(()))/))
Se. (Ix(PT) - (P.)) + c1 log( PPT//) + log( (PP HT
(PHHT (8)
S [og]o (H))))).
The result for Prr can also be similarly established.
Proof of Theorem 2 In the following proofs, we omit the KL term in the decomposition to ease the presentation. Then, with probability at least 1 - 6, for all e [T', can obtain that
((T)) -
-Edo,='FF F. [u (sH, - Edo p- [u"
(ERoi [Lte (#H1, and) - Ed [ur(S [
Σ Edo (Vanta(ent1) | Venuaa (on, an) hE[H]
Σ
I. Be (onn-((9411) [Pei Ventta (on, (on.an)
NE(H]
$ Ed [ūe(s,, ax)) Bagi ,e, [uu ax]] (Bdiing.E, (ie(ss, ax)]
Edoint [ur (u(ss
term (D)
F Fe ((ntt(sne2)) - Pentis (sh, and +
Σ Edo.mt p. (E(M
[PenPeats]] (sn, an) - (eN+a(9n+1)
AE[H]
serm (II),
where the inequality is from the definition of गł and the fact that (u,,P*) E Ut x P. from Lemma |
We define the following terms:
term (A), := E dg p- [ū(HH,
Edg.r p- [ū (SH, an)] (B I- [u (sH,wH)) Ed pe (u' (SH, un)l) rr ((n( Ed [ūr( un)))
term (B), := Edg [u" (sH,
Edp.rt.r [u' (SH,, (sHr,
term (C), :=
E do
TVV (Pen an), IT (-1sn (-1sh:
(E NEIH
term (D), = Σ
E d. [TV (Pen(-1sh, an), 17 (- an))
(E hei
29

---

Published as a conference paper at ICLR 2025
For term (Dz, we have that
term (I) == d. (S1,, an)
= [su an)]
(E [0r(s [x)]s
= [sn an)l
2 [uc(s (x))
E of [ azr)]
E 1 [4] (sm amr)]
=t [u] (s11. an
of [0 (s21 an)
of [u] (str an)]
(d [ue(s an)) }]
dd [ (22 4
+ B 3
(SHr, an)
B do F
[e am)
1 [e(s az )
E o F [ (
(SH, as))
a (sH, ax))
[ 3 [s au )]
Bd [ue(s
I 7
[M ax)l
of [ as))]
E 2 [sH [#)] - Ed [M] ($ an))
= (sH, if [u] (sH aR)l
(Ede [ur(SH aH)]
Edo. [ (sH, aH)
4B TV {do. n. Pc], {do, के 11)) + 4B TV (do ne, Prh. (do, xł, 11
F [u((s,,aH) of [ue(s 1 (u (SH, aH)] Ed [u] (SH, an))
term (A),
B (s2, ax]) -. ( (su an)) - (Fut..n [Me(81s a11)l B.gi [1 an))
(B),
4B
Σ Σ
E
TV
JE(1,2) bE M]
term (C),
For term (II),, we have that
(e.++1(54+1) - [Pan Pereta] (sh, and
term (II),
do.
hE(H]
Σ
Ed Tf
[Pr, Pa (s)) - (8n411)
kE(H]
Σ
E 2
(1N++1(5+11) - [Pe. Vi,bot (sk, an)
bE(H]
Σ
E do
Venti
(sn,
ke(m]
Σ
Bdo.
(en+z(0n41) -
[Pen Pentt]
(shs (onsan)
KE[H]
+ Σ B.
az) - Ven+a(5ntz))
he(m]
< 2B
Σ
Σ
E
B
(V(Pe,a(-lse,,an)) PP (P un))
JE(1,2) KE[H]
+2BH - TV m
P.). (do, Ti,
2B Σ
Σ
E d
[TV(P (7 on)
jE(1.2) kE(H)
(D),
2BH
Σ
Σ
E Bo
JE(1.2 hE[H]
term (C),
30

---

Published as a conference paper at ICLR 2025
In the above derivations, have repeatedly used similar relationships follows:
TVd Pt), (do; x², -
Σ
E =
[TV (Prra((/sh an). Ph Ph. an))].
he|H]
which can be derived as TV((do. n², Pt). (do, m. 1
Σ
TV (Ido i Pr,. (do. ni, Piste Pentint)
kE[H]
Σ
TV (Pr, 1P). an) (
Then, we can obtain that
Σ 1E|T]
- - Σ
term (A), +
Σ
term (B);
(E(T]
tE[T]
+ (4B + 2BH)
term (C), + 2B
Σ
term (D);
(E[T]
AE(T]
Then, we control the sum of each individual term in the following. First, for term (A), with proba- bility at least 1 6, we have that
term (A);
Σ
2 [T(ss - Edo. [ūe(81, azs)] - (Rdom? (s an)] - Baoo p | (SH an).
AET
$ Σ
C (i de (s ah,u) (u' (s aZ u (s...a).H)) CBBOOBBBBO ((BV/TTog(L/))
du
1+
(T.( (. (s) u af)
1=2
O(B (lg((1/))
[-1
1
Σ oo (ū.(s) ai, im (ur oi,)) (s =
-
2m2
OBg/Tlo(()
$ K BB (og(((/))))
where the first inequality is from the Hoeffding inequality, the second inequality uses the Eluder coefficient du == EC 1, - u, T) from Definition 3 the third inequality leverages the mean value theorem with к = 1/(2 + exp(-B) exp(B) representing the minimum derivative of a-)) in the regime of [0, B). and the last inequality incorporates Lemma A similar result can obtained for term (B),-
For term (C),, we have that
Σ term (C); = Σ []
Σ
Σ
E do
TV (Pen ax), PP), (-1sh an))
JE(1,2) (E(T] hE(H)
Σ Σ
Σ
TV (Ido,
Th. [PPi. Pr.,, Ph+1::. (do, oi. Pita))
JE(1,2) (E(T] kE(m]
2H E(dp, e2 log( (O(((TT)88)
where the step is from the generalized Fluder-type condition in Definition g and Lemma 2 A similar result can be obtained for term (D),
Finally, we obtain that
BB Reg(T) -Bs-BBBar de log (T) Bd HE T2 (TT((P)HTT8))
31

---

Published as a conference paper at ICLR 2025
B. (TR(())())))) mi, (+1s))]]
77.
(E(T] [M(M]
which concludes the proof.
E TECHNICAL LEMMAS
Lemma 3 (Solution of KL-regularized Optimization (Proposition 7.16 and Theorem 15.3 of Zhang 20. Given a loss functional with respect to p(-[z), written
Bw-xp() (w) + ry PKL (p()),((())))))))))) (p()),
=nDKL (p()).P(( exp (1 (v()) 17 log Ew-no() exp
1 Uww))
7
Cr
where the minimizer the the loss functional is (w) CoPo(w) exp 11 distribution.
(LU(W) also known Gibbs
Definition 3 (Eluder Coefficient, Definition 17.17 in Zhang )023)). Given function class F, its Eluder coefficient EC(X, F. T) defined the smallest number d so that for any sequence {os :t E [T]} and {fe :tt [ F.F
t-1 X+
Σ -
d
Definition (Generalized Eluder-type Condition, Condition 3.1 in Liu at. (023a)), There exists
real number dr RR and function & such that for any (T, A) E NxRT transitions (TV, 1 (T])
and policies (Te & E [7]} we have
Vt [2],
Σ
TVdo, (d) m.)) <A
Kt
PV((do,P%.ss) (do P. p. S E(dp, T, A).
TE(T]
32


Reviews:
[{"review_id": "90SbmQFCiL", "reviewer": "Reviewer_Roam", "summary": "The paper investigates enhancing mathematical reasoning in large language models (LLMs) through a multi-turn preference learning approach with tool integration. The authors propose a multi-turn direct preference optimization framework that accommodates complex, multi-turn reasoning and external tool integration, crucial for mathematical tasks. By employing two specific implementations—multi-turn Direct Preference Optimization (DPO) and multi-turn Kernelized Trajectory Optimization (KTO)—the framework improves model accuracy on math benchmarks, GSM8K and MATH. Empirical results with multiple base models show notable gains, particularly for models fine-tuned with iterative multi-turn training, outperforming standard SFT and single-turn baselines. The study also highlights the importance of balancing model exploration and accuracy through KL regularization and mixture sampling to maintain performance in iterative training.", "questions": "For the example shown in figure 1, the human inputs (o1 and o2) do not provide additional information, why multi-turn conversation is needed in this case?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "a.\tThe paper proposed an iterative multi-turn DPO and KTO training framework, which effectively improves the models ability on multiturn reasoning for math problems.\nb.\tThe experiments are solid, and they conduct experiments with multiple base models.", "weaknesses": "a.\tOverall, the motivation makes sense, as for some complex math problems that require tool integration, human inputs might be needed. However, only GSM8K and MATH have been used as the benchmarks in this paper, which are relatively easy problems, and may not necessarily need human input as multi-turn conversation. It is more convincing if you can conduct experiments on more complex tasks, like formulating an optimization problem and solving it from a task description. \nb.\tThe presentation of the results is not clear enough. For example, table 1 is too long, and different models use different settings, e.g. Wizardmath, codellama and Gemma-1.1-it-7B have different settings from other models. It is more readable if the table is splitted, and each sub table has its own conclusion or findings.\nc.\tTypos, e.g. line 270 DKO, should be KTO.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "76VygBq9I7", "reviewer": "Reviewer_4eni", "summary": "This paper addresses the optimization of multi-turn MDP planning under conditions involving interactions with external tool integration. Through rigorous theoretical derivations, it models the optimal conditions for multi-turn MDP planning and demonstrates that when\nthe external randomness is low, a reward form using DPO can feasibly compute the utility function in multi-turn MDP planning. Furthermore, the authors propose multi-turn DPO and multi-turn KTO algorithms tailored for optimizing multi-turn MDP planning and validate the effectiveness of combining multi-turn MDP planning optimization with an online iterative approach, showing that performance loss relative to the optimal policy is sublinear, thus highlighting the statistical efficiency of the proposed algorithm. Extensive experiments conducted on various model series across GSM8K and MATH datasets affirm the efficacy of these methods in multi-turn MDP planning.", "questions": "None.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper proposes a solution to the language model’s multi-turn planning problem involving external environment interaction, offering well-grounded theoretical proofs and a straightforward formulation.\n\n2. The proposed algorithms are empirically validated to be effective in addressing mathematical problem-solving tasks with external tool integration.", "weaknesses": "1. The utility function form in Equation (7) needs a value network and is computable only under the form of DPO rewards and deterministic external conditions, which limits the practical applicability of this theoretical framework. Currently, the theoretic framework cannot utilize methods like PPO to address multi-turn MDP planning problems with external interactions.\n\n2. The experimental section’s baseline explanation for single-turn iterative DPO lacks clarity regarding the KL regularization strategy, especially in terms of whether a fixed reference is employed, which might contribute to the poorer performance observed in single-turn iterative DPO.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8FbwoBpfLj", "reviewer": "Reviewer_erZb", "summary": "This paper introduces a multi-turn direct preference learning framework, tailored for this context, that leverages feedback from code interpreters and optimizes trajectory-level preferences. This framework includes multi-turn DPO and multi-turn KTO as specific implementations. The effectiveness of our framework is validated through training of various language models using an augmented\nprompt set from the GSM8K and MATH datasets.", "questions": "Please see the weaknesses section. I also hope to see further details on the ORM, such as the base model and training method. \n\nOverall, the paper does not have any obvious weaknesses, and I personally find it quite impressive. I believe it meets the acceptance criteria for ICLR. If you address my concerns, I would consider raising the score.\n\n**Missing reference:**\n\n[1] DotaMath: Decomposition of Thought with Code Assistance and Self-correction for Mathematical Reasoning", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The integration of external tool augmented methods and the use of multi-turn preference learning is both interesting and reasonable, effectively addressing the shortcomings of natural language in mathematical reasoning. The exploration of tool-augmented methods in the mathematical domain is well-motivated.\n\n2. The methods presented in the paper are strongly supported by theoretical foundations, which enhances their credibility, and the overall writing is clear and easy to understand.", "weaknesses": "1. WizardMath, being a state-of-the-art math model from a year ago, seems somewhat outdated as a baseline for performance. I hope the authors can include results from recent stronger models in mathematics that incorporate MULTI-TURN ITERATIVE PREFERENCE LEARNING, such as Qwen 2.5 Math and DeepSeek Math.\n\n2. I would indeed like to see results of this method in the context of a process-based reward model (PRM) and combined with MCTS sampling strategies (even preliminary results would be helpful). Although the authors mention this limitation in the limitations section and in line 297, relying solely on ORM for reasoning verification is not sufficiently detailed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4RDzeqTh5D", "reviewer": "Reviewer_h8KG", "summary": "This paper proposes a new method to improve the mathematical reasoning abilities of large language models. Existing direct preference learning algorithms are mainly applied to single-turn dialogue tasks and fail to address the complexities of multi-turn reasoning and external tool integration. To tackle this, the paper introduces a multi-turn direct preference learning framework and specifically implements two algorithms (M-DPO and M-KTO). Experiments were conducted by training the model on the GSM8K and MATH datasets, and the results demonstrate that this new framework significantly enhances the model’s performance on multiple benchmark tasks.", "questions": "1.How is the definition and selection of preference signals specifically designed in this approach? Could different preference signals significantly impact the final optimization results?\n2.Is it possible to further test the method in other types of reasoning tasks or application domains to verify its generalizability?\n3.In terms of external tool integration, is the current method the optimal solution?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper improves the performance of large language models (LLMs) in multi-turn complex reasoning tasks. The study effectively validates the proposed approach using two datasets.", "weaknesses": "Although the proposed method shows outstanding performance, it is relatively focused on specific task types. Currently, the experiments are limited to the GSM8K and MATH datasets, and it would be better to include more datasets in the testing. Additionally, the framework's design is relatively complex, involving multi-turn preference learning and trajectory-level optimization, which may require high computational resources in practical applications. A comparative experiment to demonstrate the time cost and computational consumption compared to previous methods would be helpful here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)