[论文解读] Enhancing Code Generation Performance of Smaller Models by Distilling the Reasoning Ability of LLMs

Zhihong Sun, Chen Lyu+5 more|arXiv (Cornell University)|Mar 20, 2024

Natural Language Processing Techniques被引用 5

一句话总结

本文提出 CodePLAN，一种多任务蒸馏框架，用来教导较小的代码生成模型生成并使用从大型语言模型中蒸馏出的高质量解题计划，在 APPS 的 pass@1 上显著提升超过 130%。

ABSTRACT

Large Language Models (LLMs) have recently made significant advances in code generation through the 'Chain-of-Thought' prompting technique. This technique empowers the model to autonomously devise "solution plans" to tackle intricate programming challenges, thereby improving its performance in code generation. Nevertheless, smaller models have been struggling to keep up with LLMs in deducing these plans, adversely affecting their code generation capabilities. Given the considerable size and associated deployment costs, along with concerns about data security, many teams opt for deploying smaller models for code generation. Consequently, there arises a compelling need for transferring LLMs' code generation reasoning abilities to the smaller models. In this paper, we propose the CodePLAN framework, which aims to transfer LLMs' reasoning capabilities to smaller models through distillation. We adopt a multi-task learning approach, jointly undertaking code generation and solution plan generation tasks, to enhance the code generation capabilities of the smaller model. To ensure the superior quality of the solution plans, we advocate for the utilization of backward reasoning and plan sampling strategies. Our experiments show that in comparison to the conventional fine-tuning approach, our approach improves the smaller model's code generation performance (measured in pass@1 metric) by over 130% on the challenging APPS benchmark.

研究动机与目标

• 促使更小的模型在不产生部署成本或数据安全问题的前提下，达到与 LLM 相似的代码生成性能。
• 开发一个多任务框架，通过蒸馏 LLM 的推理，联合训练代码生成与解题计划的生成。
• 通过反向推理和计划采样来提升计划质量，以指导训练和推理。
• 研究计划质量和采样对下游代码生成性能的影响。

提出的方法

• 提出 CodePLAN：一个多任务框架，使基础模型学习同时生成代码和解题计划。
• 使用 LLM 作为教师，提供高质量的解题计划，通过联合损失进行训练，优化代码生成和计划生成（L = (1-λ)L_code + λL_plan，λ=0.5）。
• 引入反向推理，从实际解答中提取程序员式的解题计划，相较基于正向问题描述的规划提升计划质量。
• 在推理阶段，应用计划采样生成多个候选计划，利用单元测试对其进行评估，并选择最佳计划来指导代码生成。
• 为基础模型增加一个计划头以输出计划，使得在代码任务和计划任务之间可以实现交替式多任务微调。
• 在 APPS 和 MBPP 数据集上使用 pass@k 指标进行评估，并与标准微调、CoT 和基于 RL 的基线进行比较。

实验结果

研究问题

• RQ1 将 LLM 的解题计划蒸馏到更小的模型中，是否能提升其代码生成性能？
• RQ2 从解答中进行的反向推理是否比基于问题描述的正向规划提供更高质量的规划信号？
• RQ3 推理阶段的计划采样如何影响生成代码的质量与可靠性？
• RQ4 CodePLAN 相对于标准微调在不同数据集和不同难度问题上的相对改进是多少？

主要发现

• 与标准微调相比，CodePLAN 在较小模型的代码生成性能（APPS 的 pass@1）提升超过 130%。
• 聚焦计划的蒸馏在 APPS 和 MBPP 上相比标准微调和基于 CoT 的方法，带来持续的提升。
• 反向推理产生的计划比直接来自问题描述的计划具有更高的质量，提升下游代码质量。
• 计划采样（用单元测试对多份计划进行测试）显著提高代码正确性，随着采样数量大于 N=1，收益提升。
• CodePLAN 与基于排序的后处理（CodeRanker）是互补的，在代码正确性方面产生了显著的组合增益。