Reasoning Distillation and Structural Alignment for Improved Code Generation

TL;DR

This work tackles code-generation with LLMs by addressing the need for problem-intent understanding and solution-structure reasoning, which large models can perform but at high cost. It proposes reasoning distillation into a smaller model using LoRA fine-tuning and a structure-aware loss that combines token-level predictions with embedding-based structural alignment, formalized via a bridging context mechanism and a joint probability objective $\mathcal{L}(\theta)$ alongside a KL-divergence term $D_{KL}(p \parallel p_{\theta})$. Using the Taco dataset ($\approx 18{,}360$ tasks) and zero-shot evaluations on MBPP, MBPP Plus, and HumanEval, the approach yields notable improvements in pass@1 (e.g., MBPP+ from $48.23$ to $56.86$, MBPP from $37.67$ to $42.42$, and HumanEval from $21.95$ to $35.86$ with structure-aware loss) and better dataflow/syntax alignment, while maintaining cost efficiency (approx. $50 for the pipeline). The findings demonstrate that high-quality CoTs generated by very large models can be distilled into smaller models to enhance first-attempt correctness, though some challenges in mathematical reasoning remain. Future work includes exploring multi-agent sequential context distillation and extending structure-aware objectives to capture deeper problem structures.

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.

Figures (2)

• Figure 1: The perplexity of our model in understanding the coding question compared to base model.
• Figure 2: The perplexity of our model while generating the code for a given MBPP question compared to the base model.