マルチターン反復選好学習による数学エージェントの構築

Building Math Agents with Multi-Turn Iterative Preference Learning

September 4, 2024

著者: 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

cs.AI

要旨

近年の研究では、大規模言語モデル（LLM）の数学問題解決能力が、コードインタプリタなどの外部ツールの統合や、マルチターン連鎖思考（CoT）推論を活用することで強化できることが示されている。現在の手法は合成データ生成と教師ありファインチューニング（SFT）に焦点を当てているが、本論文ではモデル性能をさらに向上させるための補完的な直接選好学習アプローチを検討する。しかし、既存の直接選好学習アルゴリズムは単一ターンの対話タスク向けに設計されたものであり、ツール統合型数学推論タスクに必要なマルチターン推論と外部ツール統合の複雑性を十分に扱えていない。この課題を解決するため、我々はコードインタプリタからのフィードバックを活用し、軌跡レベルの選好を最適化する、この文脈に特化したマルチターン直接選好学習フレームワークを提案する。このフレームワークには、マルチターンDPOとマルチターンKTOを具体的な実装として含む。GSM8KおよびMATHデータセットから拡張したプロンプトセットを用いた各種言語モデルの学習を通じて、本フレームワークの有効性を検証した。結果として、教師ありファインチューニング済みのGemma-1.1-it-7Bモデルでは、GSM8Kで77.5%から83.9%へ、MATHで46.1%から51.2%へ、またGemma-2-it-9BモデルではGSM8Kで84.1%から86.3%へ、MATHで51.0%から54.5%へと、大幅な性能向上が実証された。

English

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.