Reasoning Distillation and Structural Alignment for Improved Code Generation

TL;DR

This paper introduces a method to distill reasoning skills from large language models into smaller models for improved code generation, focusing on structural understanding and solution pathways.

Contribution

It proposes a novel structure-aware loss optimization technique to transfer reasoning capabilities from VLLMs to smaller models, enhancing code generation performance.

Findings

Significant improvements in pass@1, data flow, and syntax match metrics.

The method is simple, cost-effective, and outperforms baseline models.

Effective across multiple code generation benchmarks.

Abstract

Effective code generation with language models hinges on two critical factors: accurately understanding the intent of the prompt and generating code that applies algorithmic reasoning to produce correct solutions capable of passing diverse test cases while adhering to the syntax of the target programming language. Unlike other language tasks, code generation requires more than accurate token prediction; it demands comprehension of solution-level and structural relationships rather than merely generating the most likely tokens. very large language model (VLLM) are capable of generating detailed steps toward the correct solution of complex tasks where reasoning is crucial in solving the problem. Such reasoning capabilities may be absent in smaller language models. Therefore, in this work, we distill the reasoning capabilities of a VLLM into a smaller, more efficient model that is faster and cheaper to deploy. Our approach trains the model to emulate the reasoning and problem-solving abilities of the VLLM by learning to identify correct solution pathways and establishing a structural correspondence between problem definitions and potential solutions through a novel method of structure-aware loss optimization. This enables the model to transcend token-level generation and to deeply grasp the overarching structure of solutions for given problems. Experimental results show that our fine-tuned model, developed through a cheap and simple to implement process, significantly outperforms our baseline model in terms of pass@1, average data flow, and average syntax match metrics across the MBPP, MBPP Plus, and HumanEval benchmarks.