PaperlayerView exact revision on arXiv

Source-linked AI summary

Building Math Agents with Multi-Turn Iterative 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, Tianqi Liu

arXiv:2409.02392v2cs.LGstat.ML

TL;DR

Tool-integrated mathematical reasoning calls for preference-learning methods that handle multi-turn interactions, which existing single-turn chat algorithms do not fully address. The paper develops multi-turn, trajectory-level direct preference learning with online iterative training. Its experiments report pass@1 gains over SFT on GSM8K and MATH.

  • Problem

    Existing direct preference learning algorithms were designed mainly for single-turn chat and do not fully address multi-turn reasoning with external tools.

    Source: §1 · pp. 1–2

  • Method

    The paper formulates multi-turn interactions as an MDP and develops direct preference learning with external messages and online iterative training.

    Source: §1 · p. 2, §2.3 · p. 11

  • Results

    The method improves pass@1 over the SFT policy on GSM8K and MATH.

    Source: §4 · p. 25

  • Takeaways & Limitations

    The experiments validate online iterative multi-turn DPO for improving tool-integrated mathematical reasoning models over SFT.

    Source: §4 · p. 25

    The approach is not directly applicable to general agent learning in more complex, stochastic external environments or against dynamic opponents.

    Source: §4 · p. 25

Abstract

from arXiv · show

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.

1. Introduction

The introduction motivates preference learning for tool-integrated mathematical reasoning and frames the paper’s contribution as multi-turn, trajectory-level optimization. It formulates these interactions as an MDP and situates the approach among prior mathematical reasoning and preference-learning work.

1. Introduction

  • LLMs can struggle with arithmetic and symbolic computation, while existing mathematical problem-solving work has primarily emphasized synthetic data generation and supervised fine-tuning.Tool integration combines natural-language reasoning with external messages, which can improve mathematical-task performance.Source: §1 · p. 1
  • The paper derives multi-turn direct preference learning algorithms that incorporate external messages, mask irrelevant tokens, and extend to online iterative variants.The introduction distinguishes this approach from PPO and its sample inefficiency and instability.Source: §1 · p. 2
  • 77.5% to 83.9% on GSM8K and 46.1% to 51.2% on MATH for Gemma-1.1-it-7B; Gemma-2-it-9B also improved on both benchmarks.The Gemma-2-it-9B results were 84.1% to 86.3% on GSM8K and 51.0% to 54.5% on MATH.Source: §1 · p. 2

1.1. Problem Formulation

  • The formulation treats prompt-conditioned tool interactions as an MDP of agent actions and environment observations, with preferences comparing whole trajectories.The paper contrasts trajectory-level preferences with action-level supervision and describes result checking and reward models as possible utility signals.Source: §1.1 · pp. 2–3

We may also use a hard version

  • A hard process-reward utility can score a trajectory by its minimum step-wise reward.The passage identifies the step-wise reward function as r_θ.Source: We may also use a hard version · p. 4

1.2. Related Work

  • Prior work applies preference learning to math trajectories or intermediate steps; this paper derives trajectory-level comparison for tool-integrated mathematical problem solving.The approach can also adapt to step-level signals when preference pairs share prefixes.Source: §1.2 · p. 6

2. Algorithms Development

This section develops preference-learning algorithms for tool-integrated mathematical reasoning in a general MDP, including online iterative training and practical multi-turn preference losses based on deterministic external observations. It also presents a theoretical framework that motivates uncertainty-driven exploration and establishes a high-probability sublinear-in-horizon regret guarantee under stated complexity conditions.

2.1. Planning with a Fixed Model: Optimality Condition

  • For multi-turn KL-regularized planning, the optimal policy is a sequence of Gibbs distributions over Q-values computed recursively by dynamic programming.The recursion proceeds backward from the final step, with the optimal value characterized using the reference policy.Source: §2.1 · pp. 7–8

2.2. Planning with a Fixed Model: Practical Algorithm

  • The multi-turn framework derives M-DPO and M-KTO from trajectory-level preferences, with M-DPO valid for deterministic external observations.For stochastic observations, M-DPO is biased; the paper proposes estimating the additional value term with a value network and Monte Carlo methods for use as an adaptive margin, leaving this more complex setting for future work.Source: §2.2 · pp. 9–10

2.3. Learning with Online Iterative Training

  • Online iterative training collects trajectory pairs under two behavior policies, using one for exploitation and the other to explore uncertainty and diversify data.The framework updates the policies using accumulated preference data at each iteration.Source: §2.3 · p. 11
  • The framework supports KL-regularized and non-regularized objectives by fixing the reference policy at the initial policy or updating it to the previous learned policy, respectively.The paper frames the choice between these objectives as a tradeoff between out-of-domain and in-domain performance.Source: §2.3 · p. 12

2.4. Theoretical Result

  • The theoretical analysis targets statistically efficient online learning with stochastic, unknown transitions, focusing on the KL-regularized objective and batch size m=1.Its theoretical algorithm estimates reward and transition models from accumulated preference data and uses a policy-improvement oracle.Source: §2.4 · pp. 12–13
  • The value-decomposition lemma motivates exploration policies that maximize uncertainty from both reward and transition estimates.With known transitions, transition-estimation uncertainty disappears, yielding a measurement similar to that used in Xiong et al.Source: §2.4 · p. 14
  • Assuming u* ∈ U and P* ∈ P, the theorem guarantees its regret bound with probability at least 1 − δ.The bound uses Eluder-style complexity quantities and absolute constants; the supplied passage does not show its full expression.Source: §2.4 · pp. 14–15
  • Under the stated broad class of RL problems, the complexity condition implies sublinear regret in T and thus statistical efficiency.The analysis describes the Eluder coefficient and generalized Eluder-type condition as standard in RL and RLHF theory.Source: §2.4 · p. 15

3. Experiments

Across GSM8K and MATH, iterative multi-turn preference learning improves tool-integrated SFT models, with gains depending on the training and sampling choices. The experiments use augmented benchmark prompts and examine performance against baselines and across ablations.

3.1. Experiment Setup

  • The experiments use MATH and GSM8K test sets and 60K augmented training prompts, with preference trajectories filtered by final-answer correctness and heuristic checks.The augmented prompts include rephrased, backward, and bootstrapped questions; each iteration typically yields 15K–20K samples.Source: §3.1 · pp. 15–16

3.2. Main Results

  • 83.9% GSM8K and 51.2% MATH for Gemma-1.1-it-7B, and 86.3% GSM8K and 54.5% MATH for Gemma-2-it-9B, are reported after iterative M-DPO.The Gemma-2-it-9B results surpass the open-source Open-MathInstruct-trained models in the 7B–70B range.Source: §3.2 · p. 17, Table 1 · p. 18
  • Iterative M-DPO and M-KTO outperform RAFT across four base models, while single-turn DPO and KTO generally trail their multi-turn counterparts.The authors attribute the RAFT comparison to preference methods using negative as well as positive trajectory signals; single-turn methods also train on external-message tokens.Source: §3.2 · pp. 17, 19
  • Gemma-1.1-it-7B improves over three iterations from 77.5% to 83.9% GSM8K and from 46.1% to 51.2% MATH; updating the reference model also outperforms keeping the SFT reference fixed.Preference learning improves pass@n only at relatively small n; for n > 16, models perform similarly on both benchmarks.Source: §3.2 · pp. 19–20, Figure 2 · p. 19

3.3. Ablation Study and Discussion

  • A KL coefficient of 0.1 gives the strongest model, outperforming 0.01 and 0.5; low KL yields larger first-iteration gains but diminishing later gains.The discussion frames this as balancing per-iteration improvement against exploration efficiency.Source: §3.3 · p. 20, Figure 3 · p. 21, Table 2 · p. 21
  • Mixture sampling outperforms on-policy sampling, while increasing N from 12 to 30 yields only minor test-accuracy improvement.With N=30, iteration 1 covers 92.5% of prompts, versus 83.0% for N=12 and 60% for N=6; mixture sampling retains MATH gains in iteration three.Source: §3.3 · pp. 21–22, Table 3 · p. 22, Figure 4 · p. 22
  • Three SFT epochs produce the strongest starting checkpoint, although more SFT epochs reduce subsequent RLHF gains; adding NLL generally regresses performance.NLL helps when the starting SFT model is substantially underfitting.Source: §3.3 · pp. 23–24, Table 5 · p. 24
  • On-policy sampling with a smaller learning rate of 2e-7 or 4e-7 improves reasoning, whereas a 1e-6 learning rate caused collapse after a few steps in preliminary experiments.The authors suggest on-policy sampling keeps chosen and rejected response gradients effective, while a small learning rate limits policy deviation.Source: §3.3 · p. 24

4. Conclusion, Limitation, and Future Research Direction

The paper presents iterative multi-turn direct preference optimization as a way to further improve tool-integrated mathematical reasoning models, with performance depending on balancing improvement and exploration. It identifies final-answer-only preference signals and complex stochastic environments as boundaries for the approach.

4. Conclusion, Limitation, and Future Research Direction

  • Multi-turn direct preference optimization further improves tool-integrated reasoning models over SFT policies on GSM8K and MATH.The paper reports substantial pass@1 gains across various base models.Source: §4 · p. 25
  • Optimal performance requires balancing per-iteration improvement with exploration through moderate KL regularization and strategic exploration choices.The paper reports ablations supporting this balance.Source: §4 · p. 25
  • The experiments use only final-result checks as preference signals, limiting comparisons between trajectories that both end with correct or incorrect answers.The authors suggest step-wise AI feedback or a process-supervised reward model as future alternatives.Source: §4 · p. 25
  • The algorithms are not directly applicable to general agent learning in complex stochastic environments or against dynamic opponents.Such settings would require a value network and accounting for environmental randomness.Source: §4 · p. 25

A. Notation Table

The appendix includes a table of notations used in the paper.

A. Notation Table

  • The section consists of a table of notations used in the paper.The supplied passage identifies the table's purpose but does not specify its entries.Source: Table 6 · p. 33

B. Implementation Detail

The appendix describes the tool environment, data-generation constraints, annotation procedure, and implementation settings used for the experiments.

B. Implementation Detail

  • The Python interpreter runs code snippets as Jupyter cells and returns the final snippet's standard output or error message.Returned error messages are typically fewer than 20 tokens.Source: B. Implementation Detail · p. 33
  • Generation limits each turn to 512 tokens, each trajectory to 2048 tokens, and each trajectory to at most six steps.Online iterative M-DPO data collection uses temperature 1.0 without top-K or top-p sampling; evaluation uses greedy decoding.Source: B. Implementation Detail · pp. 33–34
  • Preference pairs are formed by sampling one correct-answer trajectory and one incorrect-answer trajectory after filtering buggy-answer wins and responses longer than 2048 tokens.Each iteration typically yields 15K–20K samples because some prompts have no correct answer.Source: B. Implementation Detail · p. 34
  • Evaluation can vary with the Python environment, numerical precision, and virtual machine, although the reported oscillation is small relative to the overall improvement.The experiments fix the environment and machine; Gemma-2 uses a different vLLM version because it was released after the main experiments.Source: B. Implementation Detail · p. 34
  • RAFT retains only correct-answer trajectories, tests k values of 1, 3, and 8, and ultimately uses k=1 because more data did not improve performance.The method runs for three iterations; Gemma models use a learning rate of 5e-6, while Mistral uses 1e-6 and one epoch.Source: B. Implementation Detail · p. 34
  • The experiments use each model's chat template without prompt tuning, although the authors observe that prompt engineering can further improve performance.Source: B. Implementation Detail · p. 34

C.1. Proof of Proposition 2

The proof establishes the proposition by iteratively applying relationships among the model's value and Q-value functions across the horizon.

C.1. Proof of Proposition 2

  • The proof focuses on the specified model M* and its associated Q-value functions.Source: §C.1 · p. 35
  • The displayed relationship expresses the estimated Q-value using the policy-to-reference-policy log ratio and a normalization term.Source: §C.1 · p. 35
  • The proposition follows by iteratively applying the relationship for each step in the horizon.The proof notes the terminal value condition V̂_{H+1}(s_{H+1})=0 when deriving the preceding relation.Source: §C.1 · p. 36

C.2. Proof of Theorem 1

The proof derives high-probability confidence sets and bounds theorem terms using estimation-error lemmas and complexity conditions.

C.2. Proof of Theorem 1

  • With probability at least 1 −δ, the true u* and P* belong to their confidence sets for every t, assuming u*∈U and P*∈P.The proof also identifies these as valid confidence sets.Source: §C.2 · p. 36
  • Lemma 2 bounds in-sample error for the MLE and optimistic estimators using total variation distance between induced trajectory distributions.Source: §C.2 · pp. 36–37
  • For term (A), the proof combines Hoeffding concentration, the Eluder coefficient, a lower bound on the sigmoid derivative, and Lemma 2.A similar result is stated for term (B).Source: §C.2 · pp. 39–40
  • The generalized Eluder-type condition and Lemma 2 support the bound for one term, with a similar result stated for term (D).Source: §C.2 · p. 40

D. Technical Lemmas

The technical lemmas identify a Gibbs-distribution optimizer and state two complexity conditions used in the analysis.

D. Technical Lemmas

  • The minimizer of the stated KL-regularized loss is identified as a Gibbs distribution.Source: D. Technical Lemmas · p. 40
  • The Eluder coefficient is defined as the smallest d satisfying a condition over sequences of inputs and functions from a function class.Source: D. Technical Lemmas · p. 40
  • The generalized Eluder-type condition posits a parameter dP and function ξ for transition and policy sequences.Source: D. Technical Lemmas · p. 41

Related papers

Loading 2409.02392v2…