GFlowNet Fine-tuning for Diverse Correct Solutions in Mathematical Reasoning Tasks

TL;DR

The results show that GFlowNet fine-tuning derives correct final answers from diverse intermediate reasoning steps, indicating the improvement of the capability of alternative solution generation.

Abstract

Mathematical reasoning problems are among the most challenging, as they typically require an understanding of fundamental laws to solve. The laws are universal, but the derivation of the final answer changes depending on how a problem is approached. When training large language models (LLMs), learning the capability of generating such multiple solutions is essential to accelerate their use in mathematical education. To this end, we train LLMs using generative flow network (GFlowNet). Different from reward-maximizing reinforcement learning (RL), GFlowNet fine-tuning seeks to find diverse solutions by training the LLM whose distribution is proportional to a reward function. In numerical experiments, we evaluate GFlowNet fine-tuning and reward-maximizing RL in terms of accuracy and diversity. The results show that GFlowNet fine-tuning derives correct final answers from diverse intermediate reasoning steps, indicating the improvement of the capability of alternative solution generation.

Figures (2)

• Figure 1: GFlowNet fine-tuning consists of three steps: (1) multiple solutions are sampled for a given problem, (2) each solution is evaluated by a reward model, and (3) the parameters of an LLM are updated so that $\pi(s)\propto R(s)$. The solutions are sampled from the distribution of $R(s)$, which allows us to prevent the LLM from only generating high-reward answers.
• Figure 2: The number of distinct correct solutions when sampling $k=8$ solutions. The star symbol indicates the method that has produced the highest number of distinct correct solutions.