Elamur Research⌕
☰
home / papers / 2023 / Planning with Large Language Models for Code Generation

Planning with Large Language Models for Code Generation

Paper · 2023 · Shun Zhang, Zhenfang Chen, Yikang Shen et al.

Summary

Code generators built on large language models usually decode with beam search or sampling. Their outputs may resemble reference code at the token level yet still fail to compile or give wrong answers. The common sampling + filtering approach, used by systems like AlphaCode, checks candidates against test cases only after full programs are produced. This wastes many samples. The authors propose Planning-Guided Transformer Decoding (PG-TD). It is a tree-search planner, based on Monte-Carlo tree search, that runs inside the decoding loop. The planner completes partial programs, executes them on public test cases, and uses the pass rate as a reward to steer which token comes next. The Transformer's next-token probabilities and beam search act as heuristics within the search. The authors also add mechanisms that let the planner and the model share information, which avoids redundant computation. PG-TD needs no finetuning and works with any backbone model. On competitive programming benchmarks it beats beam search and sampling + filtering, which shows that test-aware search can strengthen frozen code models.

Key points

  • PG-TD treats code generation as sequential decision-making. Test cases are split into public ones, which the planner uses for reward, and private ones, which are held out to measure generalization.
  • The method is model-agnostic. It was demonstrated with GPT-2 (1.5B), GPT-Neo (2.7B) and Codex (code-davinci-002) as backbones, without retraining any of them.
  • Changing the planner's reward enables controllable generation, such as shorter programs or more heavily commented code.
  • Programs found by PG-TD can also serve as extra training data to finetune a code-generation Transformer.
  • In one illustrative example, PG-TD reached a pass rate of 1.00, compared with 0.125 for beam search and 0.625 for sampling + filtering. The project page is codeaimcts.github.io.

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

🎓Massachusetts Institute of TechnologyCambridge · United States243 works · 76,506 citations🏢IBMArmonk · United States154 works · 72,606 citations🔬MIT Computer Science and Artificial Intelligence LaboratoryCambridge · United States76 works · 32,207 citations🎓University of Hong KongPok Fu Lam · China118 works · 13,496 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksInformation Systems · 1,730 worksComputer Science Applications · 131 works

🧭 Topics 1Research subjects and languages the work is about.Code & programming languages · 168 works

Cite

Cite as Shun Zhang, Zhenfang Chen, Yikang Shen, Mingyu Ding, Joshua B. Tenenbaum and Chuang Gan. 2023. Planning with Large Language Models for Code Generation. arXiv:2303.05510.

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

@misc{zhang2023planning,
  title = {{Planning with Large Language Models for Code Generation}},
  author = {Shun Zhang, Zhenfang Chen, Yikang Shen, Mingyu Ding, Joshua B. Tenenbaum and Chuang Gan},
  year = {2023},
  doi = {10.48550/arxiv.2303.05510},
  eprint = {2303.05510},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2303.05510}
}

References

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

  1. Wasi Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. Unified pre-training for program understanding and generation. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2655–2668, Online, J…
  2. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and Charles Sutton. Program Synthesis with Large Language Models. arXiv:2108.07732 [cs], August 2021. URL http://arxiv.org/abs/2108.07732. arXiv: 2108.07732. 2 ↗
  3. Rudy Bunel, Matthew Hausknecht, Jacob Devlin, Rishabh Singh, and Pushmeet Kohli. Leveraging Grammar and Reinforcement Learning for Neural Program Synthesis. ICLR, May 2018. URL http://arxiv.org/abs/1805.04276. arXiv: 1805.04276. 3, 4 ↗
  4. Antoine Chaffin, Vincent Claveau, and Ewa Kijak. Ppl-mcts: Constrained textual generation through discriminator-guided mcts decoding. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2953–2967, 2022… find
  5. Guillaume MJ-B Chaslot, Mark HM Winands, and HJVD Herik. Parallel monte-carlo tree search. In International Conference on Computers and Games, pp. 60–71. Springer, 2008. 9 find
  6. Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian-Guang Lou, and Weizhu Chen. CodeT: Code Generation with Generated Tests, July 2022. URL http://arxiv.org/abs/2207.10397. arXiv:2207.10397 [cs]. 9, 17 ↗
  7. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhai… — Evaluating Large Language Models Trained on Code
  8. Xinyun Chen, Chang Liu, and Dawn Song. Execution-guided neural program synthesis. In International Conference on Learning Representations, 2018. 3, 9 find
  9. Xinyun Chen, Dawn Song, and Yuandong Tian. Latent execution for neural program synthesis beyond domain-specific languages. Advances in Neural Information Processing Systems, 2021b. 3 find
  10. Raj Dabre and Atsushi Fujita. Softmax tempering for training neural machine translation models. arXiv, 2020. 1, 4 find
  11. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volum…
  12. Kevin Ellis, Catherine Wong, Maxwell Nye, Mathias Sablé-Meyer, Lucas Morales, Luke Hewitt, Luc Cary, Armando Solar-Lezama, and Joshua B Tenenbaum. Dreamcoder: Bootstrapping inductive program synthesis with wake-sleep library learning. In Proceedings of the 42nd ACM SIGPLAN International Conference… find
  13. Kevin M. Ellis, Maxwell Nye, Yewen Pu, Felix Sosa, Joshua Tenenbaum, and Armando Solar-Lezama. Write, execute, assess: Program synthesis with a repl. Neural Information Processing Systems, 2019. 3, 4, 7 find
  14. Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 889–898, 2018. 1, 3, 4
  15. Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. Codebert: A pre-trained model for programming and natural languages. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 1536–1547, 2020. 2
  16. Sylvain Gelly and David Silver. Monte-carlo tree search and rapid action value estimation in computer go. Artificial Intelligence, 175(11):1856–1875, 2011. 3 find
  17. Alex Graves. Sequence transduction with recurrent neural networks. In International Conference on Machine Learning: Representation Learning Workshop, 2012. 3, 4
  18. Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, LIU Shujie, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, et al. Graphcodebert: Pre-training code representations with data flow. In International Conference on Learning Representations, 2020. 2
  19. Kavi Gupta, Peter Ebert Christensen, Xinyun Chen, and Dawn Song. Synthesize, execute and debug: Learning to repair for neural program synthesis. Advances in Neural Information Processing Systems, 2020. 3 find
  20. Peter E Hart, Nils J Nilsson, and Bertram Raphael. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Systems Science and Cybernetics, 4(2):100–107, 1968. 3 find
  21. Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, and Jacob Steinhardt. Measuring coding challenge competence with APPS. NeurIPS, 2021. 1, 4, 7, 8, 14, 17
  22. Aditya Kanade, Petros Maniatis, Gogul Balakrishnan, and Kensen Shi. Learning and evaluating contextual embedding of source code. In International Conference on Machine Learning, pp. 5110–5121. PMLR, 2020. 2 find
  23. Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European conference on machine learning, pp. 282–293. Springer, 2006. 21 find
  24. Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven C. H. Hoi. CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning. arXiv, July 2022. doi: 10.48550/arXiv.2207.01780. URL http://arxiv.org/abs/2207.01780. arXiv:2207.01780 [cs]. 3, 24 ↗
  25. Rémi Leblond, Jean-Baptiste Alayrac, Laurent Sifre, Miruna Pislar, Jean-Baptiste Lespiau, Ioannis Antonoglou, Karen Simonyan, and Oriol Vinyals. Machine translation decoding beyond beam search. arXiv preprint arXiv:2104.05336, 2021. 3 find
  26. Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, and Agustin Dal Lago. Competition-Level Code Generation with AlphaCode. arXiv preprint arXiv:2203.07814, 2022. 1, 2, 3, 4, 7, 8, 14, 17
  27. Jinsuk Lim and Shin Yoo. Field report: Applying monte carlo tree search for program synthesis. In International Symposium on Search Based Software Engineering, pp. 304–310. Springer, 2016. 3 find
  28. Ximing Lu, Sean Welleck, Peter West, Liwei Jiang, Jungo Kasai, Daniel Khashabi, Ronan Le Bras, Lianhui Qin, Youngjae Yu, Rowan Zellers, Noah A. Smith, and Yejin Choi. NeuroLogic A*esque Decoding: Constrained Text Generation with Lookahead Heuristics, December 2021. URL http://arxiv.org/abs/2112.0872… ↗
  29. Nadia Matulewicz. Inductive program synthesis through using monte carlo tree search guided by a heuristic-based loss function. 2022. 3 find
  30. Glenford J Myers. The art of software testing. 1979. 3 find
  31. Maxwell Nye, Armando Solar-Lezama, Josh Tenenbaum, and Brenden M Lake. Learning compositional rules via neural program synthesis. Advances in Neural Information Processing Systems, 33:10832–10842, 2020. 2 find
  32. Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. 2, 28 find
  33. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 2020. URL http://jmlr.org/papers/v21/20-074.html. 2 ↗
  34. Baptiste Roziere, Marie-Anne Lachaux, Lowik Chanussot, and Guillaume Lample. Unsupervised translation of programming languages. Advances in Neural Information Processing Systems, 33, 2020. 2
  35. Baptiste Roziere, Jie M Zhang, Francois Charton, Mark Harman, Gabriel Synnaeve, and Guillaume Lample. Leveraging automated unit tests for unsupervised code translation. In International Conference on Learning Representations, 2022. 1, 3 find
  36. Thomas Scialom, Paul-Alexis Dray, Jacopo Staiano, Sylvain Lamprier, and Benjamin Piwowarski. To beam or not to beam: That is a question of cooperation for language gans. Advances in neural information processing systems, 34:26585–26597, 2021. 3 find
  37. David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, and Marc Lanctot. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484, 2016. Publisher: Nature Publis… find
  38. David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, and Demis Hassabis. Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm…
  39. Riley Simmons-Edler, Anders Miltner, and Sebastian Seung. Program synthesis through reinforcement learning guided tree search. arXiv preprint arXiv:1806.02932, 2018. 3 — [arXiv:1806.02932](https://arxiv.org/abs/1806.02932) find
  40. Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. MIT Press, October 2018. ISBN 978-0-262-35270-3. 4 find
  41. Michele Tufano, Dawn Drain, Alexey Svyatkovskiy, Shao Kun Deng, and Neel Sundaresan. Unit test case generation with transformers and focal context. arXiv, 2020a. 3 find
  42. Michele Tufano, Dawn Drain, Alexey Svyatkovskiy, and Neel Sundaresan. Generating accurate assert statements for unit test cases using pretrained transformers. arXiv, 2020b. 3 find
  43. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 2017a. 1, 3
  44. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. arXiv:1706.03762 [cs], December 2017b. URL http://arxiv.org/abs/1706.03762. arXiv: 1706.03762. 2 ↗
  45. Yue Wang, Weishi Wang, Shafiq Joty, and Steven C.H. Hoi. CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 8696–8708, Online and Punta Cana, Dominic… find
  46. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, and Morgan Funtowicz. Huggingface’s transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771, 2019. 17
  47. Yifan Xu, Lu Dai, Udaikaran Singh, Kening Zhang, and Zhuowen Tu. Neural program synthesis by self-learning. arXiv preprint arXiv:1910.05865, 2019a. 3 — [arXiv:1910.05865](https://arxiv.org/abs/1910.05865) find
  48. Yifan Xu, Lu Dai, Udaikaran Singh, Kening Zhang, and Zhuowen Tu. Neural Program Synthesis By Self-Learning. October 2019b. doi: 10.48550/arXiv.1910.05865. URL https://arxiv.org/abs/1910.05865v1. 3 — [doi:10.48550/arxiv.1910.05865](https://doi.org/10.48550/arxiv.1910.05865) ↗
  49. A PPENDIX In this appendix, we supplement the main paper by providing more thorough empirical evaluations to back up our claims and more detailed descriptions of the algorithms to help readers better understand our paper. This appendix is organized as follows. • In Sec. A, we provide more comprehens… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2023-planning-large-language-models-code/