[论文解读] Planning with Large Language Models for Code Generation

Shun Zhang, Zhenfang Chen+4 more|arXiv (Cornell University)|Mar 9, 2023

Software Testing and Debugging Techniques被引用 13

一句话总结

本文提出 Planning-Guided Transformer Decoding (PG-TD)，一种与预训练代码生成 Transformer 协作的基于规划的解码方法，通过规划者和测试用例奖励进行前瞻搜索来生成更高质量的程序，实现更高通过率并在不进行微调的情况下实现可控代码生成。

ABSTRACT

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 解码（束搜索/采样）在代码生成中的局限性，即使令牌匹配度很高也常产生无法编译或不正确的程序。
• 提出 PG-TD，即一个基于规划的解码框架，使用规划器在测试用例奖励的引导下搜索候选程序。
• 实现与任一标准代码生成 Transformer 的模型无关集成，无需微调或额外的样本收集。
• 展示 PG-TD 能带来更高的通过率，并通过奖励设计支持可控目标（如简洁性和注释）。

提出的方法

• 将代码生成形式化为一个 MDP，状态是将问题描述与部分程序拼接，行动是下一个 token。
• 应用一个受蒙特卡洛树搜索启发的规划器（带 P-UCB 选择），利用 Transformer 提供的下一个 token 概率作为引导来执行前瞻。
• 用 top-k 高概率的下一个 token 扩展节点，形成候选后续。
• 通过让 partial state 被 Transformer 的束搜索完成并在公开测试用例上测试以获得奖励来评估，并在树中反向传播。
• 引入缓存机制（树结构和生成序列缓存）以在 Transformer 与规划器之间共享信息并提高效率。
• 通过在 APPS 和 CodeContests 数据集上，用 GPT-2 和 GPT-Neo 主干进行测试来展示模型无关适用性。
• 通过修改规划器奖励来探索可控目标（如更短的代码、更多的注释），并展示通过 PG-TD 生成的样本进行微调的潜力。

实验结果

研究问题

• RQ1 将规划算法与预训练代码生成 Transformer 相结合，是否能产生比标准解码或基于采样的基线更好的程序？
• RQ2 PG-TD 是否在不同的 Transformer 主干上有效（模型无关）？
• RQ3 通过信息共享和缓存，基于规划的解码是否提高样本效率和计算时间？
• RQ4 PG-TD 是否在不对 Transformer 进行微调的情况下，支持可控的代码生成目标？
• RQ5 PG-TD 生成的样本是否可用于微调 Transformer 以进一步提升性能？

主要发现

• PG-TD 在 APPS 和 CodeContests 上的通过率方面，始终优于束搜索和采样+筛选基线。
• PG-TD 显示出应用于 GPT-2 和 GPT-Neo 主干时的模型无关改进。
• 基于规划的解码提供有效的前瞻性，优于 Sequential Monte-Carlo-Guided TD (SMCG-TD)，并在等效计算预算下达到或超过通过率。
• 缓存机制（树结构和序列缓存）显著减少计算时间和对 Transformer 的重复调用。
• PG-TD 通过改变规划器奖励实现可控代码生成（如更短的代码、更多的注释），生成的样本还可以微调 Transformer 以进一步提升性能。