Reasoning Distillation and Structural Alignment for Improved Code Generation

2025-10-20

AI Summary

To address weak reasoning capability, shallow intent understanding, and poor generalization in small-code-generation models, this paper proposes a structure-aware reasoning distillation framework. Methodologically, it transcends conventional token-level knowledge distillation by introducing (i) problem-solution structural alignment loss and (ii) fine-grained supervision over reasoning paths, enabling algorithmic-logic-level knowledge transfer. Evaluation on MBPP, MBPP+, and HumanEval benchmarks demonstrates substantial improvements: +8.2% in pass@1, +12.4% in average data-flow accuracy, and enhanced syntax matching rate—all achieved with lightweight training and efficient deployment. The core contribution is the first integration of structured semantic alignment into code generation distillation, endowing small models with interpretable, algorithm-aware reasoning capabilities.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Structured Learning
• Reasoning under Uncertainty: Probabilistic Programming

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods, algorithms and applications for the development of semantic models, knowledge graphs and other forms of structured data models with machine-interpretable semantics
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Distilling reasoning capabilities from large to small models
• Establishing structural alignment between problems and solutions
• Improving code generation accuracy and structural correctness

Innovation

• Distilling VLLM reasoning into smaller models
• Learning correct solution pathways from VLLM
• Using structure-aware loss for solution alignment