PaperlayerView exact revision on arXiv

Source-linked AI summary

Physics of Language Models: Part 2.2, How to Learn From Mistakes on Grade-School Math Problems

Tian Ye, Zicheng Xu, Yuanzhi Li, Zeyuan Allen-Zhu

arXiv:2408.16293v1cs.CLcs.AIcs.LG

TL;DR

Language models still make reasoning mistakes, and whether immediate error-correction examples improve their reasoning remains an open question. The paper studies retry-data pretraining in controlled synthetic math experiments and reports higher accuracy than training on the same amount of error-free data. It also finds that fine-tuning retry data after error-free pretraining may require costly continued pretraining.

  • Problem

    Language models still make reasoning mistakes, and it is unclear whether training on errors followed by immediate corrections can improve reasoning accuracy.

    Source: §1 · pp. 2–3

  • Method

    The authors use controlled, program-generated iGSM math data to pretrain models on erroneous solution steps followed immediately by corrections.

    Source: §1 · p. 4

  • Results

    Models trained on retry data outperform models trained on the same amount of error-free data, and high error rates do not tend to make generations erroneous.

    Source: §1 · p. 4, §7 · p. 14

  • Takeaways & Limitations

    The findings support including retry data during pretraining because error correction is not easily adapted through parameter-efficient fine-tuning.

    Source: §7 · p. 15

    Full fine-tuning can learn error correction with sufficiently many retry samples, but its cost is no less than pretraining with retry data.

    Source: §5 · p. 12

Abstract

from arXiv · show

Language models have demonstrated remarkable performance in solving reasoning tasks; however, even the strongest models still occasionally make reasoning mistakes. Recently, there has been active research aimed at improving reasoning accuracy, particularly by using pretrained language models to "self-correct" their mistakes via multi-round prompting. In this paper, we follow this line of work but focus on understanding the usefulness of incorporating "error-correction" data directly into the pretraining stage. This data consists of erroneous solution steps immediately followed by their corrections. Using a synthetic math dataset, we show promising results: this type of pretrain data can help language models achieve higher reasoning accuracy directly (i.e., through simple auto-regression, without multi-round prompting) compared to pretraining on the same amount of error-free data. We also delve into many details, such as (1) how this approach differs from beam search, (2) how such data can be prepared, (3) whether masking is needed on the erroneous tokens, (4) the amount of error required, (5) whether such data can be deferred to the fine-tuning stage, and many others.

Meta FAIR

The acknowledgments thank colleagues and teams for conversations and experimental support.

Meta FAIR

  • The authors thank colleagues and Meta FAIR, Meta Cloud Foundation, and W&B contributors for conversations and experimental support.They state that the experiments would not have been possible without this support.Source: Meta FAIR · p. 1

1 Introduction

The introduction asks whether training on immediate error-and-correction data can improve language-model reasoning, then outlines controlled iGSM experiments and their findings. It also discusses practical ways to construct retry data and limits claims about transferring synthetic-data results to future LLMs.

1 Introduction

  • The study asks whether retry data can teach error correction and improve reasoning, amid concerns about mistaken-token training and distribution shift from error-free solutions.The introduction notes that immediate correction can save inference tokens and simplify inference by avoiding multiple verification rounds.Source: §1 · pp. 2–3
  • The authors use program-generated iGSM math problems to control error generation and correction rather than relying on LLM-generated examples.The setting is chosen to avoid uncertainty about whether GPT-4 can reliably generate errors and corrections.Source: §1 · pp. 3–4
  • Retry upon regret improves accuracy beyond beam search, but depends on repeated regeneration and highly accurate error detection.The method detects an error and regenerates from the end of the previous sentence; the introduction contrasts it with prompting-based self-correction.Source: §1 · pp. 2, 4
  • Retry-data pretraining improves reasoning accuracy over the same amount of error-free data and can teach automatic correction without additional prompting or multi-stage generation.The paper also reports higher accuracy than retry upon regret.Source: §1 · p. 4
  • Pretraining with retry data teaches models to detect and correct errors, unlike prompting-based self-correction, which requires additional prompting or multi-stage generation.The paper reports higher accuracy than retry upon regret.Source: §1 · p. 4
  • Even with error rates of p = 20% or p = 50%, models do not tend to generate erroneous steps, and higher p improves accuracy within a reasonable range.Label masking on errors is unnecessary; comparisons across p use the same number of training tokens.Source: §1 · p. 4
  • Retry-data fine-tuning after error-free pretraining does not significantly improve accuracy with PEFT, while full fine-tuning requires enough data and resembles continued pretraining.The authors conclude that error correction may require major weight changes and recommend including retry data during pretraining.Source: §1 · p. 5

Conclusion.

  • Fake retry data that inserts a future step as an apparent error and then restores the current step is nearly as effective as carefully designed retry data in iGSM.The conclusion frames controlled synthetic-data experiments as guiding principles and does not claim the synthetic data directly aids future commercial LLMs.Source: §1 · p. 5, Conclusion · p. 5

2 Synthetic Math Data From Prior Work

The paper uses controllable synthetic math datasets whose step-by-step solutions isolate logical reasoning and can be verified. Problems are organized by parameter dependencies and operation count, with out-of-distribution tests including reask variants.

2 Synthetic Math Data From Prior Work

  • The synthetic datasets simulate GSM8K while removing arithmetic and common-sense difficulties, leaving logic reasoning with fully verifiable solutions.They contain over 90 trillion solution templates.Source: §2 · p. 5
  • Each problem’s parameters form a dependency graph, and computations are split into binary operations so failures are not due to arithmetic difficulty.A parameter can be computed only after all its predecessors have been computed.Source: §2 · p. 5
  • iGSM-med trains on op ≤15 and tests OOD at op ∈{20, 21, 22, 23}, while iGSM-hard trains on op ≤21 and tests at op ∈{28, 29, 30, 31, 32}.Source: §2 · p. 6
  • Reask evaluation re-samples a parameter to query, substantially changing the math-data distribution and providing an OOD test.The problem’s operation count may change after reasking.Source: §2 · p. 6
  • GPT-4 and GPT-4o cannot solve these problems for op > 10, even with few-shot learning and efforts to remove English and arithmetic difficulties.The paper also illustrates an op = 21 example as requiring non-trivial reasoning for humans.Source: §2 · p. 6

3 Result 0-1: Language Models Can Retry Upon Regret

On synthetic math problems, models can often detect that a generated step is erroneous, but using detection to trigger retries improves accuracy only modestly over beam search. The retry process depends on highly accurate detection and adds inference complexity.

3 Result 0-1: Language Models Can Retry Upon Regret

  • The studied reasoning mistake is generating a parameter before its dependencies are ready; arithmetic and common-sense difficulties are removed to focus on reasoning.This error is also described as skipping steps.Source: §3 · p. 6

3.1 Result 0: Models Can Be “Regretful” After Making Mistakes

  • The can-next probe predicts whether a parameter is computable next with 99% accuracy, and on erroneous steps it identifies can next(A) = false about 60% of the time.The contrast between these cases is evidence of a detectable regretful pattern in the model’s internal states, even after training only on correct math solutions.Source: §3.1 · p. 7

3.2 Result 1: Let Models Retry Upon Regret

  • A model may first believe a parameter is computable, then recognize the error after naming it, when attention shifts to its computation.Source: §3.2 · p. 7
  • Retry upon regret uses the can-next probe after each solution sentence and regenerates from the previous sentence when it detects an error.The described procedure uses multinomial sampling with a maximum of 10 retries per solution.Source: §3.2 · p. 7
  • About 2% improvement over beam32 occurs for op = 23 on iGSM-med and op = 32 on iGSM-hard, while increasing detector accuracy from 99% to 100% can significantly improve final accuracy.Retry upon regret can outperform beam search, but its gains remain marginal even with a 99%-accurate detector.Source: §3.2 · pp. 8–9
  • Retry upon regret adds inference complexity because it requires an error detector and repeated checking, with possible multiple rounds of regeneration.Its accuracy gains depend strongly on detector accuracy.Source: §3.2 · p. 9

4 Result 2-6: Pretrain with Retry Data

Pretraining on retry data improves reasoning accuracy, including on hard problems, without requiring masking the mistakes. The model can still produce concise solutions, and the paper distinguishes this learned correction skill from beam search and inference-time retries.

Data preparation

  • Retry data inserts a wrong parameter followed by [BACK], with retry rate controlling the probability of insertion.The wrong parameter is selected from parameters not already appearing in the solution.Source: §4 · p. 9

Result 2-3

  • 78% to 94%: at retry rate = 0.5, accuracy on iGSM-medop=23qp rises from 78% to 94%.The authors describe this as a result within a reasonable range of retry rates.Source: §4 · p. 10
  • Masking the erroneous tokens is generally unnecessary, even at retry rate = 0.5.The comparison tests standard autoregressive pretraining against masking the loss on wrong parameter tokens.Source: §4 · p. 10

Result 4

  • At retry rate = 0.2, models average fewer than 0.3 retries per problem; at 0.5, they average 2–4, and label masking can mitigate retries.The low retry count at 0.2 holds even for problems with large op; the paper attributes it to the 0.8 chance of an error-free step.Source: §4 · p. 11

Result 5

  • Retry-pretrained models can still produce shortest math solutions, supporting the authors’ conclusion that this data can improve accuracy without changing pretraining or inference processes.The paper characterizes retry data as safe to include and reports no need to increase training time.Source: §4 · p. 11

Result 6

  • Retry-data pretraining raises accuracy even with beam=1; error-free-pretrained beam search shows no noticeable improvement, while in the reported iGSM-medop=23pq case retry data gives 78% ⇒95%, versus 78% ⇒80% for retry upon regret.Source: §4 · pp. 11–12

5 Result 7: Finetune with Retry Data

The experiments find that LoRA finetuning on retry data does not significantly improve an error-free-pretrained model, unlike direct pretraining with retry data. Full finetuning can learn correction, but its cost is essentially continued pretraining.

Setup

  • The experiments consider parameter-efficient LoRA updates to a frozen pretrained model, varying trainable low-rank matrices.The paper says LoRA fine-tunes low-rank matrix updates while keeping the original pretrained weights frozen.Source: §5 · p. 12

Result 7

  • LoRA finetuning on retry data does not significantly improve over the original model, while direct retry-data pretraining outperforms it.At small LoRA ranks, finetuning can even underperform the original error-free-pretrained model; the comparison includes retry rate = 0.2.Source: §5 · pp. 12–13
  • Full finetuning with enough retry samples can learn error correction, but its cost is no less than pretraining with retry data.The authors describe this full finetuning as essentially continued pretraining.Source: §5 · p. 12
  • The paper contrasts correction with error detection, which error-free-pretrained models can acquire almost for free, including through probing or LoRA.The authors conjecture that correction requires revising internal computations to find alternative solutions, but do not analyze those computations.Source: §5 · p. 12

6 Result 8: Pretrain with Fake Mistakes

The paper tests simple ways to add fake mistakes and retries to correct math solutions without semantic parsing. Retry weak data improves accuracy, while the slightly more complex retry miss data improves it little.

Result 8

  • Retry weak data significantly improves accuracy, especially on the hardest op = 23/32 tasks, while retry miss data improves accuracy little.The authors describe retry weak as easy to obtain in practice and say there is no need to use the slightly harder-to-acquire retry miss data.Source: §6 · p. 14

Retry weak

  • Retry weak data selects a future solution sentence and appends it with [BACK], encouraging the model not to skip steps.The inserted future step may be computable at the current position, so it can be a fake rather than a true mistake.Source: §6 · pp. 13–14

Retry miss

  • Retry miss data inserts a problem parameter not yet used in the solution, encouraging the model not to skip steps or compute unnecessary parameters.It is harder to obtain than retry weak data but easier than perfect retry data.Source: §6 · p. 14

Practical implications

  • The authors propose auxiliary models to rewrite math data with fake mistakes of the retry weak type, while noting their experiments use synthetic data.They do not claim the synthetic data itself can directly aid future LLM building.Source: §6 · pp. 13–14

7 Conclusion

The paper finds that pretraining with errors followed by immediate corrections outperforms pretraining on the same amount of error-free data. It argues that retry data supports correction during ordinary autoregressive generation, while identifying limits on the study’s scope and the data’s direct commercial use.

7 Conclusion

  • Pretraining on data with mistakes followed by immediate corrections outperforms pretraining on the same amount of error-free data.Source: §7 · p. 14
  • Retry data is reported as safe even at high error rates, and ordinary autoregressive training needs no error-token label masking.The authors say the model learns to correct errors if needed rather than being encouraged to make them.Source: §7 · pp. 14–15
  • Comparable correction skill does not come from beam search or retry-upon-regret without nearly perfect error detection, or from PEFT such as LoRA.The conclusion therefore argues for including retry data during pretraining.Source: §7 · p. 15
  • The study focuses on logical reasoning in grade-school math, one of the weaker aspects of GPT-4, rather than other difficulties such as arithmetic or common sense [29].Source: §7 · p. 15
  • The authors seek guiding principles rather than direct use of iGSM retry data in commercial-scale pretraining, and report that GPU limits prevented follow-up data-preparation studies.They suggest teaching models not to skip steps, using weak retry data or prompting an auxiliary model to rewrite math data.Source: §7 · p. 15

A A Harder Example with Retry

  • Figure 9 presents a harder iGSM-hardpq training example with op = 21 and a retry example with retry rate = 0.2.Source: Figure 9 · p. 16

B More Experiments for Finetune (Result 7)

Additional fine-tuning comparisons find rare cases where retry fine-tuning outperforms the original model, while doubling retry-pretraining tokens does not significantly improve accuracy over double-time retry pretraining.

B More Experiments for Finetune (Result 7)

C More Experiments for Fake Retry Data (Result 8)

With fake retry data, models retry more often than with perfect retry data at retry rates of at least 0.1, without a corresponding tendency to perform more unnecessary operations. The appendix also reports the model and data settings and notes that sample complexity was not studied.

C More Experiments for Fake Retry Data (Result 8)

  • For retry rates ≥0.1, retry-weak and retry-miss models retry more often than perfect-retry models, but do not tend to perform more unnecessary operations.Figure 11 reports retry and unnecessary-operation statistics; the models can still generate shortest solutions most of the time.Source: C More Experiments for Fake Retry Data… · p. 18, Figure 11 · p. 18

D Experiment Details and Parameters

  • The experiments primarily use GPT2-small with rotary positional embeddings, GPT2Tokenizer, and contexts of 768/1024 tokens for pretraining and 2048 for evaluation.The authors report no benefit from trying Llama, including gated MLP layers.Source: D Experiment Details and Parameters · p. 19
  • Training data was generated on the fly without a data-amount limit, and sample complexity was not explored.Source: D Experiment Details and Parameters · p. 19

D.1 Pretrain Experiment Details

The appendix specifies data construction, training and evaluation settings, and the V-probing procedures used in the experiments.

D.1 Pretrain Experiment Details

  • Pretraining follows the cited AdamW and cosine-decay setup; iGSM-med uses 100,000 steps, while iGSM-hard uses 200,000 steps.The med and hard settings use learning rate 0.002, with batch sizes 512 and 256 and context lengths 768 and 1024, respectively.Source: §D.1 · p. 19
  • Problems and solutions are randomly generated, concatenated, and right-truncated to fit the context window; examples longer than that window are discarded.Source: §D.1 · p. 19
  • Evaluation uses context length 2048, greedy decoding or beam-4 sampling, and a parser that checks intermediate steps; accuracy averages 4096 problems.Retry-count and unnecessary-operation statistics also use 4096 problems.Source: §D.1 · pp. 19–20
  • When comparing error-free and retry-pretrained models, evaluation uses the same test distribution; reported accuracies select the best of two seeds and two beam choices.Source: §D.1 · pp. 19–20
  • For the double-time retry experiments, iGSM-med training runs 200,000 steps with weight decay 0.03, and iGSM-hard runs 400,000 steps with weight decay 0.02.Source: §D.1 · p. 20

D.2 V-Probing Experiment Details

  • V-probing version 1 freezes the pretrained model and trains a linear head with an embedding update; version 2 adds trainable parameters but uses fewer samples.The paper also reports retry-upon-regret accuracies for both versions in Figure 3(b).Source: §D.2 · p. 20

D.3 Finetune Experiment Details

The appendix specifies full and LoRA finetuning configurations, rank choices, and evaluation procedures for the iGSM experiments. The LoRA experiments vary update rank, with the paper emphasizing practical sample complexity.

D.3 Finetune Experiment Details

  • The authors prioritize practical sample complexity, unlike [29], which aimed to demonstrate the existence of very small-rank updates.They cite probing to guide generation as an example of the practicality they seek.Source: §D.3 · p. 20
  • LoRA experiments varied query/value rank r from 4 to 256 and used rank-2r embedding updates; r = 256 is nearly full finetuning for GPT2-12-12.LoRA freezes pretrained weights and adds trainable low-rank updates; the model's hidden dimension is 768.Source: §D.3 · pp. 20–21
  • LoRA finetuning used learning rate 0.001 for both datasets, with 200,000 steps and dataset-specific weight decay, batch size, and context length.For iGSM-med: weight decay 0.05, batch size 256, and context length 768; for iGSM-hard: weight decay 0.05, batch size 128, and context length 1024. The token count matched pretraining for iGSM-med and was half as large for iGSM-hard, which the authors say sufficed for the curve to plateau.Source: §D.3 · p. 21
  • Figure 7 reports accuracy on 4,096 math problems per cell, selecting the best result across two pretrained models, two seeds, and beam sizes 1 and 4.The same evaluation and selection procedure is stated for both full and LoRA finetuning.Source: §D.3 · p. 21
  • The full finetuning runs used pretraining parameters except for the specified dataset settings, with the hard-dataset run using the same training-token count as pretraining.For iGSM-med: learning rate 0.001, weight decay 0.05, batch size 512, context length 768, and 100,000 steps; for iGSM-hard: learning rate 0.001, weight decay 0.03, batch size 256, context length 1024, and 200,000 steps.Source: §D.3 · p. 21

Related papers

Loading 2408.16293v1…