Enhancing Code Generation Performance of Smaller Models by Distilling the Reasoning Ability of LLMs

TL;DR

This paper introduces CodePLAN, a distillation framework that transfers the reasoning abilities of large language models to smaller models, significantly enhancing their code generation performance by over 130% on the APPS benchmark.

Contribution

The paper presents a novel multi-task distillation method using backward reasoning and plan sampling to improve small models' code generation by mimicking LLM reasoning.

Findings

Over 130% improvement in pass@1 metric on APPS benchmark

Effective transfer of reasoning abilities from LLMs to smaller models

Enhanced code generation performance with multi-task distillation

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.