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.

翻译：现有基于大型语言模型的代码生成流水线通常在解码过程中使用束搜索或采样算法。尽管其生成的程序在基于令牌匹配的评分上表现优异，但往往无法通过编译或产生错误输出。主要原因在于传统Transformer解码算法可能并非代码生成的最佳选择。本研究提出一种新型Transformer解码算法——规划引导Transformer解码（PG-TD），该算法利用规划算法进行前瞻搜索，引导Transformer生成更优程序。具体而言，Transformer不再单纯优化生成序列的似然概率，而是借助规划器生成候选程序并在公开测试用例上进行验证。这使得Transformer能够做出更明智的决策，生成最终能产出更高质量程序的令牌。我们还设计了一种信息共享机制，在Transformer与规划器间交换信息以提升算法计算效率。以多种大型语言模型为骨干网络，我们在公开编程挑战基准上进行了实证评估，结果表明：1）该方法生成的程序在性能上持续优于对比基线方法；2）通过优化修改后的目标函数，该方法可实现可控代码生成，例如生成简洁代码或高注释代码。