[论文解读] Energy-Based Models for Code Generation under Compilability Constraints

Tomasz Korbak, Hady Elsahar+2 more|arXiv (Cornell University)|Jun 9, 2021

Topic Modeling参考文献 43被引用 5

一句话总结

该论文提出一种基于能量的模型（EBM），通过约束自回归模型仅生成语法正确的代码，以确保神经代码生成的可编译性。采用KL自适应分布策略梯度方法，微调生成模型以近似EBM，显著提升了可编译率，同时保持了多样性与复杂性，优于导致灾难性遗忘的标准强化学习微调方法。

ABSTRACT

Neural language models can be successfully trained on source code, leading to applications such as code completion. However, their versatile autoregressive self-supervision objective overlooks important global sequence-level features that are present in the data such as syntactic correctness or compilability. In this work, we pose the problem of learning to generate compilable code as constraint satisfaction. We define an Energy-Based Model (EBM) representing a pre-trained generative model with an imposed constraint of generating only compilable sequences. We then use the KL-Adaptive Distributional Policy Gradient algorithm (Khalifa et al., 2021) to train a generative model approximating the EBM. We conduct experiments showing that our proposed approach is able to improve compilability rates without sacrificing diversity and complexity of the generated samples.

研究动机与目标

• 解决自回归神经代码模型在生成语法正确、可编译代码时因视野受限和缺乏全局序列约束而导致的局限性。
• 将带可编译性约束的代码生成建模为约束满足问题，从而实现系统性的分布方法。
• 训练一个生成模型，使其近似受限的EBM分布，同时保持原始模型的多样性与复杂性。
• 克服强化学习微调中常见的灾难性遗忘问题，尤其针对序列级奖励。
• 提供一种方法，在不降低生成代码流畅性或结构复杂性的情况下提升可编译性。

提出的方法

• 定义一个基于能量的模型（EBM），仅对可编译的代码序列赋予低能量（高概率），将可编译性约束编码为全局特征。
• 使用KL自适应分布策略梯度（KL-DPG）算法，训练自回归生成模型以近似EBM分布。
• 利用原始预训练语言模型作为先验，最小化微调策略与原始模型之间的KL散度。
• 通过可微分代理在训练过程中应用EBM约束，实现与策略梯度方法的端到端优化。
• 使用编译反馈作为奖励信号，引导策略仅生成有效且可编译的代码序列。
• 将EBM集成到强化学习框架中，平衡约束满足与分布保真度。

实验结果

研究问题

• RQ1 基于能量的模型能否有效在代码生成中编码可编译性作为全局约束？
• RQ2 KL自适应分布策略梯度算法能否成功训练生成模型以近似EBM，同时保持多样性？
• RQ3 与标准强化学习微调相比，通过EBM强制可编译性是否能降低灾难性遗忘的风险？
• RQ4 所提方法在不降低生成代码复杂性或流畅性的情况下，能在多大程度上提升可编译率？
• RQ5 不同的微调策略（如RL使用$R(x)=b(x)$、$R(x)=P(x)$或KL-DPG）如何影响编译错误的分布？

主要发现

• 所提出的基于EBM的方法显著优于标准自回归模型和强化学习微调基线，大幅提升了可编译率。
• 该方法在生成代码中保持了高水平的多样性与复杂性，避免了标准强化学习微调中常见的严重退化。
• KL-DPG微调相比标准RL方法实现了与原始模型更好的分布对齐，从而减少了灾难性遗忘。
• 通过强制实施全局约束，该方法有效减少了长距离句法依赖中的编译错误。
• 定性分析表明，与基线相比，基于EBM的方法产生的编译错误更少且更集中。
• 该方法优于使用启发式奖励的标准RL以及仅最大化编译准确率的策略梯度方法。