ReST^EM

Analysis

Why This Paper Matters

Large language models (LLMs) have shown remarkable problem-solving abilities, but their training often depends on vast amounts of human-generated data, which is costly and may not be optimal. ReST^EM addresses this by introducing a self-training loop that leverages the model's own generations, filtered by binary feedback, to iteratively improve performance. This approach is particularly relevant for AI practitioners seeking to scale model capabilities without proportional human effort. The paper demonstrates that self-generated data can surpass human data in effectiveness, especially for math and code tasks, and that careful iteration management preserves transfer to other benchmarks.

How It Works

ReST^EM is grounded in the expectation-maximization (EM) framework, treating the problem as a reinforcement learning task with binary rewards. The method consists of two alternating steps:

E-step (Generate): The current policy (language model) generates multiple output sequences for each input from the original dataset. These outputs are scored using a binary reward function (0 or 1), which in practice checks correctness (e.g., for math problems, whether the final answer matches; for code, whether the program passes tests). Only outputs with reward 1 are retained for the next step.

M-step (Improve): The base pretrained model is fine-tuned on the filtered dataset using a reward-weighted negative log-likelihood loss: (J(\theta) = \mathbb{E}{(x,y) \sim D_i} [r(x,y) \log p\theta(y|x)]). Crucially, each M-step starts from the original base model, not the model from the previous iteration. This prevents drift and preserves the model's general capabilities, leading to better transfer performance on held-out tasks.

Unlike the original ReST method, ReST^EM does not augment the dataset with human-generated outputs, as these may not be optimal or readily available. The binary reward simplifies the process: any threshold (\tau \in (0,1)) yields the same filtering, so only one Improve step per iteration is needed.

Results

The experiments use PaLM 2 models (S, S*, L) on MATH and APPS datasets, with evaluation on multiple benchmarks including GSM8K, HumanEval, and Big-Bench Hard.

Key findings:

• ReST^EM consistently outperforms fine-tuning on human-generated data for both MATH and APPS, with larger gains for larger models.
• Multiple iterations improve performance on MATH and transfer to GSM8K, but for APPS, most gains come from the first iteration; subsequent iterations cause regression due to overfitting.
• Pass@K and majority voting performance improve, indicating better solution diversity and correctness.
• A single iteration with 3x the data is less effective than multiple iterations, highlighting the benefit of iterative refinement.

• Model-generated data from ReST^EM is more effective for fine-tuning than human data, even when using a single solution per problem.
• This data can also be used to distill knowledge to smaller models effectively.

• ReST^EM shows better transfer to HumanEval compared to ReST, confirming the benefit of fine-tuning the base model each iteration.

• The method is data-efficient: significant gains are seen with as few as 1000 MATH questions.
• Improvements occur across all difficulty levels, with the largest gains on medium and hard questions.

• ReST^EM does not degrade general capabilities on Big-Bench Hard; in some cases, it even improves performance.

Significance

ReST^EM provides a practical, scalable self-training recipe for LLMs that reduces dependence on human annotations. By iteratively using the model's own correct generations, it achieves superior performance on math and code tasks while maintaining transferability. The finding that multiple iterations can overfit on code tasks underscores the need for task-specific tuning of iteration counts. For AI practitioners, this method offers a path to continuously improve models with minimal human intervention, potentially accelerating progress in problem-solving domains.