Flow-DPO: Improving LLM Mathematical Reasoning through Online Multi-Agent Learning

TL;DR

The quality of reasoning traces generated by the method are compared with those produced through direct model inference, demonstrating the effectiveness of the approach in improving LLM performance in mathematical reasoning tasks.

Abstract

Mathematical reasoning is a crucial capability for Large Language Models (LLMs), yet generating detailed and accurate reasoning traces remains a significant challenge. This paper introduces a novel approach to produce high-quality reasoning traces for LLM fine-tuning using online learning \textbf{Flows}. Our method employs an incremental output production Flow, where component LLMs collaboratively construct solutions through iterative communication. We train the Flow using online Direct Preference Optimization (DPO) learning with rollouts, generating DPO pairs for each training example and updating models in real-time. We directly compare the quality of reasoning traces generated by our method with those produced through direct model inference, demonstrating the effectiveness of our approach in improving LLM performance in mathematical reasoning tasks.

Figures (4)

• Figure 1: Illustration of the incremental production flow. The Answer LLM is designated to generate an answer chunk with a limited number of tokens. The Stop LLM determines if the current partial answer has reached a satisfying final answer.
• Figure 2: Illustration of the DPO training with rollouts. At each node of the initial generation, we do a random rollout that is different from the original node and continue generation to a final answer. A pair that leads to different answers (correct and incorrect) is considered a DPO training data.
• Figure 3: Progressive validation accuracy of Llama-3-Instruct on MetaMath.
• Figure 4: Progressive validation accuracy of Phi-3-Medium on MetaMath.