Elamur Research⌕
☰
home / papers / 2024 / Physics of Language Models: Part 2.2, How to Learn From Mistakes on Grade-School Math Problems

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

Paper · 2024 · Tian Ye, Zicheng Xu, Yuanzhi Li et al.

Summary

Strong language models still make reasoning slips. They can often fix these when prompted to re-check their work, but not while they are generating. This paper asks whether the ability to catch and repair errors mid-solution can be built in during pretraining. To keep the experiments controlled, the authors use iGSM, a corpus of program-generated grade-school math problems. Because a program creates the problems, wrong steps and their exact corrections can be produced with full reliability, without relying on GPT-4.

The authors find three main things:

  • A GPT2 model pretrained only on clean solutions already "knows" when a step is wrong. A tiny probe can read this out, and regenerating flagged steps ("retry upon regret") beats beam search.
  • Pretraining on "retry data", where a faulty step is immediately followed by its fix, gives higher accuracy under plain autoregressive decoding than the same amount of error-free data.
  • They also examine masking, error rate, data preparation and moving such data into fine-tuning.

The work matters because it gives practical guidance on building pretraining corpora that teach models to correct themselves.

Key points

  • Setting: the iGSM dataset (iGSM-med and iGSM-hard splits), with GPT2-12-12 models (hidden size 768). Each configuration is evaluated on 4096 problems, using 2 seeds and beam size 1 or 4.
  • Error detection uses V-probing: a rank-8 update on the embedding layer plus a linear head. It reaches over 99% accuracy, far above what prompting-based self-verification achieves.
  • A second, cheaper probe uses rank-32 embedding updates plus rank-2 query/value updates. It needs only about 0.25% of the pretraining tokens.
  • Deferring correction data to fine-tuning was tested in two ways:
    • LoRA with ranks r ∈ {4, …, 256}
    • Full fine-tuning, treated as continued pretraining on the same number of tokens as pretraining
  • Motivating example: GPT-4o (2024-08-06) often insists a wrong solution is correct even when told which line contains the error.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢GoogleMountain View · United States533 works · 443,885 citations🎓Carnegie Mellon UniversityPittsburgh · United States633 works · 195,218 citations🏢MetaMenlo Park · United States224 works · 241,960 citations🎓Mohamed bin Zayed University of Artificial IntelligenceAbu Dhabi · United Arab Emirates145 works · 6,810 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Tian Ye, Zicheng Xu, Yuanzhi Li and Zeyuan Allen-Zhu. 2024. Physics of Language Models: Part 2.2, How to Learn From Mistakes on Grade-School Math Problems. arXiv:2408.16293.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2408.16293 · arXiv 2408.16293

@misc{ye2024physics,
  title = {{Physics of Language Models: Part 2.2, How to Learn From Mistakes on Grade-School Math Problems}},
  author = {Tian Ye, Zicheng Xu, Yuanzhi Li and Zeyuan Allen-Zhu},
  year = {2024},
  doi = {10.48550/arxiv.2408.16293},
  eprint = {2408.16293},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2408.16293}
}

References

3 references, 3 of them with a card of ours. ↗ leads to the work.

  1. Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
  2. Zeyuan Allen-Zhu and Yuanzhi Li. Physics of Language Models: Part 1, Learning Hierarchical Language Structures. ArXiv e-prints, abs/2305.13673, May 2023. Full version available at http://arxiv.org/abs/2305.13673. ↗
  3. Zeyuan Allen-Zhu and Yuanzhi Li. Physics of Language Models: Part 3.2, Knowledge Manipulation. ArXiv e-prints, abs/2309.14402, September 2023. Full version available at http://arxiv.org/abs/2309.14402. ↗

Updated

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ye-2024-physics-language-models-part-2-2/