PaperlayerView exact revision on arXiv

Source-linked AI summary

Planning with Large Language Models for Code Generation

Shun Zhang, Zhenfang Chen, Yikang Shen, Mingyu Ding, Joshua B. Tenenbaum, Chuang Gan

arXiv:2303.05510v1cs.LGcs.AIcs.CLcs.PL

TL;DR

Conventional code-generation decoding can be sample-inefficient because it does not use public test cases while generating programs. PG-TD integrates Transformer-guided lookahead planning and information sharing, and it generates higher-quality programs across evaluated settings while supporting alternative objectives. Its main limitations are reliance on test cases and greater computational expense than pure beam search.

  • Problem

    Existing Transformer decoding methods use beam search or sampling without considering public test cases during generation, making code generation sample-inefficient.

  • Method

    PG-TD integrates a pretrained code-generation Transformer with planning-based lookahead search, public-test evaluation, and shared information for efficient decoding.

  • Results

    PG-TD consistently generates higher-quality programs than competing baselines across evaluated settings and can optimize objectives beyond pass rate.

  • Takeaways & Limitations

    PG-TD supports controllable code generation, including concise and highly commented solutions, without fine-tuning the Transformer.

  • Takeaways & Limitations

    PG-TD relies on test cases and is more computationally expensive than pure Transformer beam search because it runs beam search repeatedly within tree search.

Abstract

from arXiv · show

Existing large language model-based code generation pipelines typically use beam search or sampling algorithms during the decoding process. Although the programs they generate achieve high token-matching-based scores, they often fail to compile or generate incorrect outputs. The main reason is that conventional Transformer decoding algorithms may not be the best choice for code generation. In this work, we propose a novel Transformer decoding algorithm, Planning-Guided Transformer Decoding (PG-TD), that uses a planning algorithm to do lookahead search and guide the Transformer to generate better programs. Specifically, instead of simply optimizing the likelihood of the generated sequences, the Transformer makes use of a planner to generate candidate programs and test them on public test cases. The Transformer can therefore make more informed decisions and generate tokens that will eventually lead to higher-quality programs. We also design a mechanism that shares information between the Transformer and the planner to make our algorithm computationally efficient. We empirically evaluate our framework with several large language models as backbones on public coding challenge benchmarks, showing that 1) it can generate programs that consistently achieve higher performance compared with competing baseline methods; 2) it enables controllable code generation, such as concise codes and highly-commented codes by optimizing modified objective.

1 INTRODUCTION

Existing code-generation pipelines rely on beam search or sampling, but these methods ignore public test cases during generation and can be sample-inefficient. PG-TD adds lookahead planning, information sharing, and configurable objectives to guide Transformers toward higher-quality programs.

  • Motivation: Sampling + filtering evaluates programs only after generation, so it ignores test cases while choosing tokens and may require exhaustive sampling.Programs can fail because of a single incorrect token, making large-scale sampling necessary to find correct solutions.
  • Approach: PG-TD uses a planner during Transformer generation to perform lookahead search and favor tokens leading to programs with higher test-case pass rates.The planner uses generated-program quality to make Transformer decisions more informed.
  • Efficiency: PG-TD shares information between the Transformer and planner to reduce redundant computation from directly integrating planning with decoding.The Transformer supplies beam-search information and next-token probabilities as planner heuristics.
  • Properties: PG-TD is model-agnostic and does not require additional sample solutions or Transformer fine-tuning to improve performance.Any standard code-generation Transformer can serve as the backbone.
  • Controllability: Changing the planner reward enables controllable generation objectives, including concise programs and programs with more comments.The objective can be modified without fine-tuning the Transformer.

2 RELATED WORK

Prior work applies Transformers, test cases, reinforcement learning, and planning to program synthesis and generation. PG-TD combines a large language model with tree search for general-purpose programming-language generation while targeting computational efficiency.

  • Transformers for program synthesis: Transformer program-synthesis systems treat programming languages like natural languages and commonly use beam search or sampling during generation.AlphaCode combines large pretrained models with large-scale sampling for programming competitions.
  • Test cases for program synthesis: Related program-synthesis work uses unit tests for generating tests, constructing training data, or directly synthesizing domain-specific programs.These approaches motivate using test cases as signals in code generation.
  • Planning and reinforcement learning: Code generation has been formulated as sequential decision-making, enabling reinforcement-learning and planning algorithms during Transformer training or decoding.The paper situates PG-TD within this broader planning-oriented formulation.
  • Planning and reinforcement learning: MCTS becomes more applicable to large domains when combined with deep learning or a default policy, which motivates combining language models with tree search.PG-TD uses this recipe for competitive programming and adds efficiency mechanisms.
  • Planning in natural language generation: Prior NLP planning methods use pretrained discriminators or predefined metrics as rewards, whereas this work combines tree search with language models for programming generation.The paper positions this combination as novel for general-purpose programming-language generation.

3 METHOD

PG-TD formulates competitive-programming code generation as an MDP and integrates Transformer-guided tree search with beam-search evaluation. Caching and shared search information improve efficiency while public tests provide the optimization signal.

  • 3.1 OVERVIEW: The task gives an agent a natural-language problem description and test cases, requiring a program whose outputs exactly match expected strings.Public tests are available during generation, while private tests evaluate generalization.
  • 3.1 OVERVIEW: Beam search and sampling cannot easily optimize objectives different from reference-solution similarity, whereas planning can directly optimize pass rate or another programming objective.The code-generation problem is formulated as an MDP with token actions and deterministic concatenation transitions.
  • 3.1 OVERVIEW: PG-TD uses MCTS-inspired tree search to explore program states and find terminal states with high rewards.Tree nodes represent states, edges represent token actions, and the search tracks visits and values.
  • 3.2 PLANNING-GUIDED TRANSFORMER DECODING: P-UCB uses Transformer next-token probabilities to prioritize branches, while TOP K expands each selected node with the most likely tokens.TOP K avoids random expansions that could readily produce syntax errors.
  • 3.2 PLANNING-GUIDED TRANSFORMER DECODING: Beam search completes selected partial programs so the planner can evaluate their quality using generated complete programs.BEAM SEARCH(s, b) generates a sequence from prefix s with beam size b.
  • 3.2 PLANNING-GUIDED TRANSFORMER DECODING: Tree-structure caching reuses states visited by earlier Transformer beam searches, while sequence caching avoids regenerating complete programs.These mechanisms address redundant Transformer computation across PG-TD iterations.

4 EMPIRICAL EVALUATION

The evaluation compares PG-TD with beam search, sampling, and sequential Monte-Carlo baselines across coding benchmarks, budgets, caching choices, and controllable objectives. PG-TD generally achieves higher pass rates and uses planning and caching to improve efficiency, while also supporting concise or commented code generation.

  • Experimental setup: The experiments compare PG-TD with beam search, Sampling + Filtering, and SMCG-TD on APPS and CodeContests.The evaluation uses APPS and CodeContests test sets with pass rate and strict accuracy metrics.
  • Effectiveness: PG-TD consistently outperforms the baselines on pass rate across the evaluated datasets.On strict accuracy, Sampling + Filtering matches or exceeds PG-TD on the APPS introductory subset, where the gaps are small.
  • Efficiency: With the same computational budget, PG-TD achieves higher pass rates than Sampling + Filtering by considering test-case pass rates during generation.SMCG-TD performs worse because its sampling-based procedure lacks PG-TD’s multi-step lookahead search.
  • Efficiency: Caching reduces PG-TD computation time by avoiding regeneration of whole sequences and repeated Transformer calls for likely next tokens.Sequence caching has the larger effect, while removing tree-structure caching makes the algorithm slightly slower.
  • Controllable code generation: Alternative planner rewards generate more concise code or code with more comments while retaining reasonable pass rates on public test cases.The evaluated objectives are code length penalty and comment encouragement.
  • Automatically generated test cases: PG-TD retains higher strict accuracy than beam search when automatically generated test cases are used instead of provided test cases.This evaluates PG-TD in settings where human-specified test cases are unavailable.

5 DISCUSSION AND CONCLUSION

The paper presents PG-TD as a tree-search extension of pretrained Transformer decoding that improves program quality and computational efficiency across settings. It also supports alternative code-generation objectives without Transformer fine-tuning, while relying on test cases and costing more than pure beam search.

  • Conclusion: PG-TD combines a pretrained Transformer with Monte-Carlo-tree-search-inspired planning to generate higher-quality programs than competing baselines.The paper reports this across different evaluation settings.
  • Conclusion: Model-structure-specific caching mechanisms reduce the computational expenses of PG-TD.The framework also supports objectives beyond pass rate without fine-tuning the Transformer.
  • Limitations and future work: PG-TD relies on test cases and is more computationally expensive than pure Transformer beam search because tree search runs beam search multiple times.The paper identifies value-function estimation and parallel tree search as directions for reducing this cost.

APPENDIX

The appendix provides expanded empirical results, special-setting evaluations, algorithmic details, generated-code examples, and discussion of benefits and potential negative social impacts. It also examines automatically generated tests and objectives beyond pass rate.

  • Appendix organization: Section A adds comprehensive algorithm and baseline results, dataset license information, and further parameter studies.The appendix includes evaluations across datasets and variations of PG-TD settings.
  • Appendix organization: Section B evaluates PG-TD with automatically generated test cases when human-specified test cases are unavailable.This tests the method under a setting without provided test cases.
  • Appendix organization: Section C studies concise-code and highly-commented-code objectives in addition to pass rate.These objectives are implemented through alternative reward functions.
  • Appendix organization: Sections D and E provide component and baseline details alongside examples of codes generated by PG-TD and baseline algorithms.These sections support closer inspection of the methods and their outputs.
  • Appendix organization: Section F discusses the advantages and potential negative social impacts of automatic code generation.The discussion addresses both assistance for software engineering and possible malware development.

A EMPIRICAL EVALUATION

Experiments across APPS, CodeContests, Codex, and multiple evaluation settings show that PG-TD generally improves code-generation quality, while its benefits depend on search budget, submission budget, and evaluation design.

  • Evaluation settings: PG-TD is evaluated on the APPS dataset and CodeContests, including APPS introductory, interview, and competition problems.The APPS test dataset contains 1000 introductory, 3000 interview, and 1000 competition problems.
  • Search budget: With a Transformer-generation budget of 512, the performance gap between PG-TD and other methods becomes smaller, and Sampling + Filtering outperforms PG-TD on strict accuracy for some dataset subsets.These results are reported in Table 5.
  • Search configuration: Increasing beam size improves PG-TD pass rate, whereas increasing tree width can waste Transformer-generation and computation budgets by considering less-likely tokens.This comparison concerns the effects of beam size and the number of candidate tokens considered during rollouts.
  • Evaluation metrics: For smaller n or k values, PG-TD finds programs more likely to pass all private test cases, but Sampling + Filtering matches or outperforms it when more programs are generated or submitted.The comparison uses n@k and pass@k on APPS introductory problems.
  • Public test cases: Even with 1, 3, or 5 public test cases, PG-TD achieves higher pass rate and strict accuracy than beam search on the first 500 APPS introductory problems.The experiments use the remaining test cases as private test cases.
  • Backbone models: PG-TD helps Codex generate better code, although the Codex evaluation uses only a subset of data because of OpenAI API limits.A concrete example where PG-TD outperforms beam search is reported in Fig. 16.
  • Evaluation-step search: Beam search performs better than sampling in PG-TD’s evaluation step, while sampling multiple times could estimate node values more accurately at greater computational cost.The comparison is reported in Table 9.
  • Failure analysis: PG-TD dramatically reduces compilation and runtime error percentages by executing generated programs and selecting those with higher public-test pass rates.The error percentages are averaged over the APPS introductory dataset.

B USING AUTOMATICALLY-GENERATED TEST CASES

PG-TD can use automatically generated test cases when human-specified tests are unavailable, and these tests are sufficiently accurate to support improved program verification.

  • Evaluation procedure: The HumanEval evaluation uses Codex to generate solutions and automatically generated tests to assess the strict accuracy of those solutions.The prompt includes the natural-language problem description and avoids directly copying example input-output pairs.
  • Test-case quality: On HumanEval, automatically generated test cases achieve 72.56% strict accuracy for human-written sample solutions.This indicates that the generated tests are mostly correct, although they do not reach the 100% strict accuracy of ground-truth tests.
  • PG-TD results: PG-TD achieves higher strict accuracy than beam search on a HumanEval subset when it uses high-quality automatically generated test cases to verify programs.The comparison is performed with Codex as the backbone model.

C PLANNING FOR OTHER CODE GENERATION OBJECTIVES

PG-TD supports controllable code generation by changing the planner’s reward function, producing either shorter programs or programs with more comments while preserving test-case performance in the shown example.

  • PG-TD supports alternative objectives beyond pass rate, including code length penalty and comment encouragement.
  • Code length penalty: The code length penalty uses pass rate and code-string length, with λ=0.1 and t=20 controlling the penalty.
  • Code length penalty: Pass Rate: 1.00 accompanies a code-string length reduction from 187 to 78 with the code length penalty.
  • Comment encouragement: Comment encouragement rewards comment lines while adding a length penalty and capping the comment contribution to prevent repeated meaningless comments.
  • Comment encouragement: The designed comment reward generates solutions with more comment lines, as illustrated in Figure 10.
  • PG-TD’s search procedure selects nodes with a P-UCB criterion balancing program quality, Transformer likelihood, and under-explored successors.

D.2 EFFICIENT IMPLEMENTATION BY INFORMATION SHARING

PG-TD shares generated sequences and rewards across evaluations to reduce Transformer computation, while acknowledging that sequence caching can be suboptimal when beam width exceeds one.

  • Sampling + Filtering: Sampling + Filtering generates multiple Transformer samples, evaluates their rewards, and returns the sample with the largest reward.
  • Sequence caching reduces evaluation cost by caching complete programs and their rewards for reuse in later PG-TD iterations.
  • Sequence caching: When b > 1, cached sequences may be suboptimal because they were found using fewer beams than a later evaluation requires.
  • Sequence caching: Despite possible value underestimation for b > 1, the implementation retains sequence caching for computational savings.
  • SMCG-TD: SMCG-TD also uses Transformer evaluation and public test cases, but samples partial programs instead of using PG-TD’s tree search.
  • SMCG-TD: Sequence caching in SMCG-TD avoids repeated Transformer generation for identical prefixes, reducing calls and computation time.

D.4 FINETUNING

The section describes finetuning the Transformer with PG-TD-generated solutions and pass-rate-aware training, followed by beam-search generation. This procedure improves pass rate and strict accuracy over beam search with the original model.

  • Finetuning procedure: PG-TD-generated samples are used to further finetune the Transformer before generating solutions with beam search.The finetuned model is compared against beam search using the original model.
  • Finetuning procedure: The SMCG-TD algorithm initializes a population from the problem description, iteratively updates generations and fitness values, and returns the complete program with the largest reward.Its loop continues while the population is nonempty and the maximum number of steps has not been reached.
  • Results: Pass rate and strict accuracy both improve over beam search with the original Transformer model.The section reports improvement on both metrics after using PG-TD-generated samples for finetuning.
  • Pass-rate-aware training: Solutions with pass rates above 80% are collected for finetuning rather than retaining only 100%-pass solutions.This preserves more training samples while exposing the Transformer to differing solution quality.
  • Pass-rate-aware training: The training objective uses a program’s pass rate as its reward, with future work proposed to compare alternative loss functions.The passage defines W as a complete program, r(W) as its pass rate, and PD as the program description.

E ILLUSTRATIVE EXAMPLES

The illustrative examples compare PG-TD with baseline code-generation methods on competitive-programming problems. The examples include higher-pass-rate solutions and a Codex case where PG-TD succeeds while beam search fails.

  • Qualitative comparisons: PG-TD generates solutions with higher pass rates than the baseline methods in the qualitative examples.The examples are presented in Figures 14 and 15, with an additional comparison against beam search in Figure 16.
  • Backbone-model comparison: The project website provides a step-by-step illustration of PG-TD.The cited site is codeaimcts.github.io.
  • Qualitative comparisons: PG-TD achieves a pass rate of 1.00 on the illustrated travel-time problem.The task asks whether Takahashi can travel distance D within T minutes at speed S.
  • Qualitative comparisons: Beam search achieves a pass rate of 0.125 on the illustrated circle-area problem.The task asks for the area ratio between circles of radius r and 1 as an integer.
  • Qualitative comparisons: PG-TD achieves a pass rate of 1.00 on the illustrated array-instability problem.The problem requires removing exactly one array element to minimize the resulting instability.
  • Backbone-model comparison: With Codex as the backbone Transformer, beam search finds an incorrect solution while PG-TD finds a correct solution.The example is presented as evidence that PG-TD can work with different backbone Transformers.

F MORE DISCUSSIONS

The discussion highlights PG-TD’s flexibility and potential efficiency benefits without finetuning, while noting that automatic code generation could also facilitate malware development. It suggests screening harmful natural-language requests as a future safeguard.

  • Potential benefits: PG-TD can improve pass rates and optimize different code objectives without finetuning pretrained Transformer models.Different planner reward functions adapt the framework to different tasks.
  • Potential negative social impacts: Automatic code generation may make it easier to develop malware from malicious goals expressed in natural language.The discussion suggests a separate screening module to reject requests that could lead to harmful code.
Loading 2303.05510v1…