Elamur Research⌕
☰
home / papers / 2025 / Weight Ensembling Improves Reasoning in Language Models

Weight Ensembling Improves Reasoning in Language Models

Paper · 2025 · Xingyu Dang, Christina Baek, Kaiyue Wen et al.

Summary

When reasoning models are trained with supervised finetuning (SFT), the variety of their outputs shrinks. Pass@1 keeps rising, but Pass@k peaks early and then falls, which weakens test-time scaling methods such as majority voting or reranking with a reward model. The authors finetune Gemma-2-2B and Qwen-2.5-0.5B on rephrased GSM8k/MATH problems from MetaMath, then apply RL with GRPO. They test a simple fix, a variant of WiSE-FT: the most recent SFT checkpoint and an early one are merged by averaging their parameters. The merged model recovers nearly all of the lost Pass@k, raises Pass@1, scales better at test time, and makes a stronger starting point for RL while needing less data. The authors explain this by splitting a bound on expected Pass@k into a bias term and a variance term. High-temperature decoding and early stopping lower one term only by raising the other, whereas WiSE-FT lowers both. The result matters because preserving diversity during training, rather than only at decoding time, is presented as essential for getting the most out of reasoning models.

Key points

  • Training setup. SFT ran on a 30K subset of MetaMath40k for 1710 steps (10 epochs). GRPO followed on a separate 30K subset with a binary correctness reward. Evaluation used GSM8K and MATH500, with AIME also reported. Pass@1 was estimated from 100 sampled traces per question.
  • The merge. WiSE-FT here takes an equal-weight average of an early checkpoint and the current SFT checkpoint (½·w0 + ½·wt). It gave higher Pass@k than SFT at any decoding temperature tested (0.7–1.6).
  • Bias-variance view. Longer SFT makes per-question error bimodal, so problems end up either almost always solved or almost always missed. This lowers bias but raises variance.
  • Theory on RL. In a simplified bandit setting, the authors prove that REINFORCE and GRPO cannot restore diversity once it has collapsed.
  • Decoding gap. Standard decoders (nucleus, min-p, top-k) reach about 0.57/0.67 Pass@2/Pass@4 on GSM8k with Gemma. An optimal answer-level strategy would reach 0.76/0.83.

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

 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 2Research subjects and languages the work is about.Language models & pretraining · 4,145 worksReasoning · 1,948 works

Cite

Cite as Xingyu Dang, Christina Baek, Kaiyue Wen, Zico Kolter and Aditi Raghunathan. 2025. Weight Ensembling Improves Reasoning in Language Models. arXiv:2504.10478 · COLM 2025.

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

@misc{dang2025weight,
  title = {{Weight Ensembling Improves Reasoning in Language Models}},
  author = {Xingyu Dang, Christina Baek, Kaiyue Wen, Zico Kolter and Aditi Raghunathan},
  year = {2025},
  doi = {10.48550/arxiv.2504.10478},
  eprint = {2504.10478},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2504.10478}
}

References

35 references, 20 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. AlphaProof and AlphaGeometry teams. Ai achieves silver-medal standard solving international mathematical olympiad problems, jul 2024. URL https://deepmind.google/discover/blog/ ai-solves-imo-problems-at-silver-medal-level/. ↗
  2. Edward Beeching, Lewis Tunstall, and Sasha Rush. Scaling test-time compute with open models, 2024. URL https://huggingface.co/spaces/HuggingFaceH4/blogpost-scaling-test-time-compute. ↗
  3. Jeff Bilmes. Submodularity in machine learning and artificial intelligence. arXiv preprint arXiv:2202.00132, 2022. — [arXiv:2202.00132](https://arxiv.org/abs/2202.00132) find
  4. Feng Chen, Allan Raventos, Nan Cheng, Surya Ganguli, and Shaul Druckmann. Rethinking fine-tuning when scaling test-time compute: Limiting confidence improves mathematical reasoning. arXiv preprint arXiv:2502.07154, 2025.
  5. Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Alphamath almost zero: Process supervision without process, 2024. URL https://arxiv.org/abs/2405.03553. ↗
  6. Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Aleksandra Faust. Inference-aware fine-tuning for best-of-n sampling in large language models. arXiv preprint arXiv:2412.15287, 2024a. — [arXiv:2412.15287](https://arxiv.org/abs/2412.15287) find
  7. Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Aleksandra Faust. Inference-aware fine-tuning for best-of-n sampling in large language models, 2024b. URL https://arxiv.org/abs/2412.15287. — [arXiv:2412.15287](https://arxiv.org/abs/2412.15287) ↗
  8. Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shengbang Tong, Saining Xie, Dale Schuurmans, Quoc V. Le, Sergey Levine, and Yi Ma. Sft memorizes, rl generalizes: A comparative study of foundation model post-training, 2025. URL https://arxiv.org/abs/2501.17161. ↗
  9. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv.org/abs/2110.14168. ↗
  10. Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pp. 1321–1330. PMLR, 2017. find
  11. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
  12. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset, 2021. URL https://arxiv.org/abs/2103.03874. ↗
  13. Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration, 2020. URL https://arxiv.org/abs/1904.09751. ↗
  14. Audrey Huang, Adam Block, Dylan J Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz, Jordan T Ash, and Akshay Krishnamurthy. Self-improvement in language models: The sharpening mechanism. arXiv preprint arXiv:2412.01951, 2024.
  15. Yifei Li, Zeqi Lin, Shizhuo Zhang, Qiang Fu, Bei Chen, Jian-Guang Lou, and Weizhu Chen. Making large language models better reasoners with step-aware verifier, 2023. URL https://arxiv.org/abs/2206. 02336. ↗
  16. Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Zhi-Quan Luo, and Ruoyu Sun. Preserving diversity in supervised fine-tuning of large language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=NQEe7B7bSw. ↗
  17. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let’s verify step by step. arXiv preprint arXiv:2305.20050, 2023.
  18. Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. s1: Simple test-time scaling, 2025. URL https://arxiv.org/abs/2501.19393. ↗
  19. Minh Nguyen, Andrew Baker, Clement Neo, Allen Roush, Andreas Kirsch, and Ravid Shwartz-Ziv. Turning up the heat: Min-p sampling for creative and coherent llm outputs, 2024. URL https://arxiv.org/ abs/2407.01082. — [arXiv:2407.01082](https://arxiv.org/abs/2407.01082) ↗
  20. Theo X. Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar-Lezama. Is self-repair a silver bullet for code generation?, 2024. URL https://arxiv.org/abs/2306.09896. ↗
  21. Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amélie Héliou, Aliaksei Severyn, Matt Hoffman, Nikola Momc… — BOND: Aligning LLMs with Best-of-N Distillation
  22. Amrith Setlur, Chirag Nagpal, Adam Fisch, Xinyang Geng, Jacob Eisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. Rewarding progress: Scaling automated process verifiers for llm reasoning, 2024. URL https://arxiv.org/abs/2410.08146. ↗
  23. Louis Shao, Stephan Gouws, Denny Britz, Anna Goldie, Brian Strope, and Ray Kurzweil. Generating high-quality and informative conversation responses with sequence-to-sequence models. arXiv preprint arXiv:1701.03185, 2017.
  24. Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling llm test-time compute optimally can be more effective than scaling model parameters, 2024. URL https://arxiv.org/abs/2408.03314. ↗
  25. Yuda Song, Hanlin Zhang, Carson Eisenach, Sham Kakade, Dean Foster, and Udaya Ghai. Mind the gap: Examining the self-improvement capabilities of large language models. arXiv preprint arXiv:2412.02674, 2024. — [arXiv:2412.02674](https://arxiv.org/abs/2412.02674) find
  26. Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1–57, 2018. find
  27. Ashwin K Vijayakumar, Michael Cogswell, Ramprasath R. Selvaraju, Qing Sun, Stefan Lee, David Crandall, and Dhruv Batra. Diverse beam search: Decoding diverse solutions from neural sequence models, 2018. URL https://arxiv.org/abs/1610.02424. ↗
  28. Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models, 2023. URL https://arxiv.org/abs/2203.11171. ↗
  29. Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo-Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models, 2022. URL https://arxiv.org/abs/2109.01903. ↗
  30. Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. Inference scaling laws: An empirical analysis of compute-optimal inference for problem-solving with language models. arXiv preprint arXiv:2408.00724, 2024.
  31. Wei Xiong, Hanning Zhang, Nan Jiang, and Tong Zhang. An implementation of generative prm. https://github.com/RLHFlow/RLHF-Reward-Modeling, 2024. ↗
  32. Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long chain-of-thought reasoning in llms, 2025. URL https://arxiv.org/abs/2502.03373. ↗
  33. Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309.12284, 2023.
  34. Dun Zhang, Jiacheng Li, Ziyang Zeng, and Fulong Wang. Jasper and stella: distillation of sota embedding models. arXiv preprint arXiv:2412.19048, 2024a. — [arXiv:2412.19048](https://arxiv.org/abs/2412.19048) find
  35. Yiming Zhang, Avi Schwarzschild, Nicholas Carlini, Zico Kolter, and Daphne Ippolito. Forcing diffuse distributions out of language models, 2024b. URL https://arxiv.org/abs/2404.10859. — [arXiv:2404.10859](https://arxiv.org/abs/2404.10859) ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/dang-2025-weight-ensembling-improves-reasoning-language/