Source-linked AI summary
Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models
Avi Singh, John D. Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Xavier Garcia, Peter J. Liu, James Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, Abhishek Kumar, Alex Alemi, Alex Rizkowsky, Azade Nova, Ben Adlam, Bernd Bohnet, Gamaleldin Elsayed, Hanie Sedghi, Igor Mordatch, Isabelle Simpson, Izzeddin Gur, Jasper Snoek, Jeffrey Pennington, Jiri Hron, Kathleen Kenealy, Kevin Swersky, Kshiteej Mahajan, Laura Culp, Lechao Xiao, Maxwell L. Bileschi, Noah Constant, Roman Novak, Rosanne Liu, Tris Warkentin, Yundi Qian, Yamini Bansal, Ethan Dyer, Behnam Neyshabur, Jascha Sohl-Dickstein, Noah Fiedel
TL;DR
Human-generated data limits language-model fine-tuning, especially for complex problem-solving tasks. The paper introduces ReST^EM, which filters model-generated samples with feedback and fine-tunes iteratively, and finds strong gains across math and coding tasks, suggesting reduced dependence on human data.
Problem
High-quality human-generated data is limited for fine-tuning language models on complex problem-solving tasks.
Method
ReST^EM alternates between generating model samples, filtering them with binary feedback, and fine-tuning on the accepted data.
Results
ReST^EM significantly improves mathematical reasoning and code generation across PaLM 2 scales and surpasses fine-tuning on human-generated solutions.
Takeaways & Limitations
Self-training with feedback can reduce dependence on human-generated data for problem-solving tasks.
Takeaways & Limitations
ReST^EM requires a moderately sized human-collected training set of problems or prompts and access to a manually designed or learned reward.
Abstract
from arXiv · showhide
Fine-tuning language models~(LMs) on human-generated data remains a prevalent practice. However, the performance of such models is often limited by the quantity and diversity of high-quality human data. In this paper, we explore whether we can go beyond human data on tasks where we have access to scalar feedback, for example, on math problems where one can verify correctness. To do so, we investigate a simple self-training method based on expectation-maximization, which we call ReST$^{EM}$, where we (1) generate samples from the model and filter them using binary feedback, (2) fine-tune the model on these samples, and (3) repeat this process a few times. Testing on advanced MATH reasoning and APPS coding benchmarks using PaLM-2 models, we find that ReST$^{EM}$ scales favorably with model size and significantly surpasses fine-tuning only on human data. Overall, our findings suggest self-training with feedback can substantially reduce dependence on human-generated data.
1. Introduction
The paper addresses the bottleneck of limited high-quality human data by training language models on filtered self-generated samples. ReST^EM uses expectation-maximization-style self-training and improves PaLM 2 performance across mathematical reasoning and code generation benchmarks.
- Motivation: High-quality human data is a costly bottleneck for complex problem-solving tasks, motivating scalable model-generated alternatives.Synthetic data is useful when its quality can be ensured.
- Method: ReST^EM alternates between generating multiple outputs, filtering them with binary reward, and fine-tuning on the accepted samples.The fine-tuned model supplies samples for the next generation step.
- Results: ReST^EM substantially improves PaLM 2 performance on the MATH and HumanEval benchmarks.Figure 1 compares PaLM 2 models with other models whose results are typically not comparable because of differing model scales.
- Motivation: Prior self-training work mainly used models up to 7B parameters, with limited scalability observed for larger models.
- Results: Self-generated synthetic data produces larger performance gains than human-written data, while ReST^EM also improves held-out-task performance and pass@k and majority-voting results.The reported transfer includes GSM8K, Hungarian HS finals, HumanEval, and Big-Bench Hard tasks.
- Takeaway: The paper presents ReST^EM as a way to reduce dependence on human data for problem-solving tasks.
2. Preliminaries
The preliminaries define autoregressive language-model generation, supervised fine-tuning, and a sequence-level reward objective. They also motivate an alternative to online reinforcement learning because continual sampling and fine-tuning become costly for very large models.
- Language-model generation: An autoregressive language model generates an output sequence token by token from a context using a parameterized conditional distribution.Each token is predicted using previously generated tokens.
- Language-model generation: The softmax temperature controls generation randomness: higher values increase randomness, while lower values favor more deterministic outputs.
- Supervised fine-tuning: Supervised fine-tuning trains the policy on human-generated input-output pairs by minimizing negative log likelihood.
- Reinforcement-learning objective: The reinforcement-learning formulation assumes access to a deterministic sequence-level reward for each input-output pair.
- Computational motivation: Direct online reinforcement learning requires repeated policy updates and sampling, making continual fine-tuning costly for models with tens or hundreds of billions of parameters.The paper therefore motivates an alternative to online methods.
3. Expectation-Maximization for Reinforced Self-Training
The paper formulates reinforced self-training as expectation-maximization: an E-step reweights model samples by reward, and an M-step updates the language model through weighted likelihood training. ReST^EM alternates these steps while separating data collection from optimization, enabling scaling to large language models.
- EM formulation: The EM formulation maximizes the likelihood of observing an optimality variable associated with high reward.The framework defines p(O = 1|x, y) from a non-decreasing function of reward and optimizes the corresponding objective.
- EM formulation: At each E-step, model outputs are weighted according to their likelihood of receiving high reward.The variational distribution is proportional to the current model distribution multiplied by the optimality likelihood.
- EM properties: Each EM iteration monotonically improves the evidence lower bound under the stated EM updates.The paper gives the inequality L(p_θt+1, q_t+1) ≥ L(p_θt, q_t+1) ≥ L(p_θt, q_t).
- Differences from standard RL: EM-based reinforcement learning decouples sampling from optimization by using a fixed previous-iteration policy to collect data.The paper identifies this decoupling as enabling easier scaling to large policy networks such as LLMs.
- ReST^EM algorithm: ReST^EM iteratively alternates Generate and Improve steps, filtering sampled outputs with binary rewards before fine-tuning.Inputs are resampled from the original dataset, outputs are generated from the current policy, and the resulting samples are scored before optimization.
- ReST^EM algorithm: At each M-step, the policy is updated by maximizing a reward-weighted negative log-likelihood objective.ReST^EM fine-tunes the base pretrained model on the generated dataset, weighting samples by their binary rewards.
- Differences from ReST: Unlike ReST, ReST^EM excludes human-generated outputs and fine-tunes the base model at every Improve step, improving held-out transfer performance.The paper reports comparable task-specific performance but much better transfer performance on held-out tasks.
4. Related work
The related-work discussion places ReST^EM among synthetic-data and EM-based approaches for language-model improvement. It distinguishes these methods by their search procedure, sampling regime, iteration structure, reward assumptions, and use of full datasets or mini-batches.
- Synthetic-data self-training: Expert Iteration alternates expert-sample generation through search with supervised policy distillation.Its expert-improvement step combines a base policy with a search procedure, whereas the cited setup assumes access to correctness rewards.
- Synthetic-data self-training: STaR uses greedy decoding and one model-generated solution per problem during data collection, unlike temperature-sampled multi-output collection.STaR also uses rationalization by providing the correct answer as part of the input for difficult problems.
- Synthetic-data self-training: RFT corresponds to one Generate and Improve iteration, while ReST^EM uses multiple iterations and reports larger gains on challenging APPS and MATH benchmarks as PaLM 2 capacity scales.The comparison contrasts RFT's limited scaling improvements on GSM8K with ReST^EM's reported results on more difficult tasks.
- EM-related methods: IML performs reward-weighted likelihood training on self-collected data, but applies its E- and M-steps to mini-batches rather than the entire training dataset.The related-work discussion also places IML applications in semantic parsing, machine translation, and simple math reasoning.
- EM-related methods: RWR is an EM instantiation using an exponential reward likelihood, allowing application to non-binary reward functions.Norouzi et al. build on RWR to propose a general IML variant for machine translation.
- EM-related methods: RAFT alternates E- and M-steps over mini-batches and selects the maximum-reward output for each input; with binary rewards, it is analogous to IML.The paper therefore treats RAFT as an instantiation of ReST^EM under binary rewards.
- Other related work: TRICE uses MCMC with a control variate to approximate the marginal-likelihood gradient, whereas ReST^EM samples and filters with a binary reward.TRICE treats the chain-of-thought rationale as a latent variable.
5. Experiments and analysis
Across math and code-generation experiments, ReST^EM improves problem-solving and transfer performance over human-data fine-tuning, though additional iterations can overfit task-specific data. Ablations show benefits from iterative sampling, model-generated solutions, and larger models, while broader evaluations find no major degradation.
- 5.1. ReST^EM on MATH and APPS: MATH benefits from multiple ReST^EM iterations, while APPS gains mostly occur in the first iteration and later iterations regress on APPS and HumanEval.The authors attribute the APPS regression to overfitting on its smaller training set.
- 5.1. ReST^EM on MATH and APPS: ReST^EM substantially outperforms fine-tuning on human-written solutions, especially for PaLM 2-L.The comparison is not entirely matched because ReST^EM can generate multiple correct solutions per problem, whereas human data typically provides one.
- 5.2. Impact on Pass@K and Majority-Voting Performance: ReST^EM improves Pass@K for every fixed K, with the largest typical performance gap at K=1.For MATH majority voting with 64 samples, ReST^EM reaches 48.82 test accuracy versus 44.02 for the base PaLM 2-L model.
- 5.3. Ablation Studies: Multiple ReST^EM iterations outperform a single iteration using three times as many sampled solutions: 41.9% versus 40.3% Pass@1 on MATH.The reported comparison uses PaLM 2-L and shows higher performance in the second and third iterations than in the larger single-iteration dataset.
- 5.3. Ablation Studies: Model-generated data remains more effective than human-written data in a restricted one-solution setting and can improve fine-tuning of smaller models through distillation.Distill* surpasses human-data fine-tuning despite using fewer training questions, while Distill (2-L) outperforms self-generated solutions for PaLM 2-S.
- 5.3. Ablation Studies: ReST^EM provides substantially better HumanEval transfer than ReST while achieving similar APPS performance on PaLM 2-S*.This comparison isolates the benefit of the modified training procedure for held-out-task transfer.
- 5.3. Ablation Studies: ReST^EM is sample efficient: using 1,000 MATH questions already produces significant gains, while gains generally increase with dataset size.The authors note a slight decrease at 4,000 versus 2,000 questions, suggesting variance in fine-tuning.
- 5.3. Ablation Studies: Fine-tuning on MATH or APPS causes no major degradation across BBH tasks, with chain-of-thought evaluation sometimes improving over the base model.Under direct prompting, the three evaluated models perform similarly.
6. Discussion
ReST^EM improves language-model problem-solving by training on model-generated data with reward feedback, while remaining data-efficient and relatively inexpensive. Its use is bounded by the need for task-specific training problems and rewards, vigilance against overfitting, and an incomplete pass@1–pass@K gap.
- ReST^EM combines model-generated data with a reward function to improve LLM performance on problem-solving tasks.
- The method is theoretically grounded in applying expectation-maximization to reinforcement learning and achieves gains at relatively low computational cost.
- ReST^EM is data-efficient, but repeated training requires vigilance to avoid over-fitting.
- The method requires a moderately sized training set of problems or prompts for each new task.
- ReST^EM requires a manually designed or learned reward function, ideally computable automatically.
- Although ReST^EM substantially improves pass@1, it may not close the gap to pass@K for sufficiently large K.