Self-Correcting Code Generation Using Small Language Models (2025)

Abstract: Self-correction has demonstrated potential in code generation by allowing language models to revise and improve their outputs through successive refinement. Recent studies have explored prompting-based strategies that incorporate verification or feedback loops using proprietary models, as well as training-based methods that leverage their strong reasoning capabilities. However, whether smaller models possess the capacity to effectively guide their outputs through self-reflection remains unexplored. Our findings reveal that smaller models struggle to exhibit reflective revision behavior across both self-correction paradigms. In response, we introduce CoCoS, an approach designed to enhance the ability of small language models for multi-turn code correction. Specifically, we propose an online reinforcement learning objective that trains the model to confidently maintain correct outputs while progressively correcting incorrect outputs as turns proceed. Our approach features an accumulated reward function that aggregates rewards across the entire trajectory and a fine-grained reward better suited to multi-turn correction scenarios. This facilitates the model in enhancing initial response quality while achieving substantial improvements through self-correction. With 1B-scale models, CoCoS achieves improvements of 35.8% on the MBPP and 27.7% on HumanEval compared to the baselines.

## introduction despite the progress in large language models (llms), generating the correct code snippets in a single attempt remains a challenging task across many programming scenarios. in response, recent studies have explored the use of external supervision, typically provided by a more capable teacher model, to guide the correction process @xcite . these approaches typically leverage teacher models, which either guide student models through iterative self-correction via feedback or perform both generation and correction in a self-directed manner without external models @xcite . given the dependence on powerful guidance, most of the successful results have been achieved using strong proprietary models such as gemini @xcite and gptseries @xcite , which are only available through external apis @xcite . considering the reliance on costly proprietary models in previous studies, it remains uncertain whether smaller, open-source language models (slms) can effectively guide their outputs through self-reflection @xcite . to investigate this, we explore whether slms can achieve meaningful selfcorrection without proprietary systems. our findings show that smaller models struggle to exhibit reflective revision behavior when they rely solely on prompting ( §4.3). building on these observations, we aim to explore training-based self-correction for code generation. however, we find that such methods are not directly applicable to slms, likely due to their assumption of sufficient self-correction capability ( §5.2). to tackle these limitations, we introduce co-cos (self-correcting code generation using small lms), a reinforcement learning (rl) framework designed to support intrinsic self-correction in small lms. our approach leverages an accumulated reward function that considers previous responses and outputs a scalar reward that captures the cumulative effect of multiple turns. this function encourages both successful self-correction and improved initial response quality by rewarding the entire response trajectory. additionally, our method adopts a progressive reward that evaluates incremental improvements in code quality by tailoring the setting to a multi-turn scenario, which offers a more fine-grained assessment than binary rewards based solely on code correctness. we validate our approach on 1b-scale models across various code generation tasks and observe robust generalization in unseen settings. our contributions are as follows: (1) we empirically demonstrate that previous prompting-based and training-based self-correction methods struggle to generalize to small language models. (2) we introduce a reinforcement learning reward function tailored for the multi-turn code generation setting, leveraging accumulated rewards with a discount factor and fine-grained assessments to encourage effective self-correction. (3) we validate our approach across diverse python code generation datasets and show consistent improvements over baselines, including in unseen settings. ## related work prompting-based correction guidance. several recent studies on self-correction have leveraged the advanced reasoning capabilities of llms to improve the accuracy of their initial responses @xcite . in coding tasks that require complex reasoning, selfcorrection often involves reviewing generated code snippets to identify and fix errors with the model's own judgment. however, prior works suggest that naive prompting for self-correction may reduce performance @xcite . as a result, sets @xcite combines sampling, self-verification, and self-correction into a unified framework that facilitates llms to improve their own outputs. self-refine @xcite conducts iterative evaluation and correction of the initial responses using few-shot prompting without any additional fine-tuning. however, these methods rely on the assumption that the model is already capable of revising its outputs. training-based correction guidance. in addition to the above methods, fine-tuning has been explored to endow models with intrinsic selfcorrection. self-corrector @xcite employs a separate correction model that shares the same backbone as the initial response generator. @xcite introduces a distinct verifier model to guide the revision process. @xcite propose a teacher-guided framework where a stronger model supervises both the intermediate reasoning steps and the final prediction of an slm. however, all of these methods depend on training separate corrector or verifier models, which require additional computational resources. self-correction without external guidance. re-vise @xcite implements internal selfcorrection by generating special tokens that determine whether to revise or terminate the current output. score @xcite adopts a two-stage rl framework that provides rewards to both the previous and current responses to guide learning. while both methods have demonstrated improved code generation, revise has limited generalization beyond the training turns due to its reliance on supervised fine-tuning (sft) @xcite . score is the first work to alleviate this limitation through online learning with self-generated data. while both score and our approach leverage online-rl, score remains constrained to proprietary models, whereas cocos demonstrates effectiveness even for slms that struggle with selfcorrection. ## problem description in this work, we focus on code generation tasks and study a small-scale llm to investigate its intrinsic ability to improve its outputs over subsequent trials. suppose @xmath0, y, u)} consists of data points, each represented as a tuple (x, y, u), where x is a problem, y is the canonical code snippet, and u is the unit test cases. we define the slm as π θ and aim to generate correct ŷt through the conditional distribution π θ (ŷ t |x, ŷ1:t-1 , p 1:t-1 ), where ŷ1:t-1 represents the history of slm's previous trials, and p 1:t-1 denotes auxiliary instructions for revising the previous responses. note that the subscripts indicate turn indices within the multi-turn setting. since we do not consider external feedback (e.g., compiler outcomes, other lms' feedback), p 1:t-1 is fixed for each turn. ## preliminary study: prompting-based self-correction prompting-based self-correction typically involves incorporating feedback into the prompt in the form of the auxiliary instruction sequence p 1:t-1 @xcite . this feedback can be either provided by a separate teacher model or generated by the model itself. in this section, we examine both settings to assess whether prompting-based selfcorrection is effective for small language models. ## proposed methodology based on the observations in §4, slms struggle to demonstrate reflective revision behavior when relying solely on prompting, without access to external llms. in this section, we introduce a trainingbased method designed to enable effective selfcorrection. ## experimental settings as prior studies are insufficient for enabling intrinsic self-correction ( §4), we compare cocos against fine-tuning approaches. to support this comparison, we introduce the experimental settings, including descriptions of the models ( §6.1), datasets ( §6.2), and baselines ( §6.3) used in our experiments. ## analyses in this section, we analyze the discount factor in the accumulated reward function ( §8.1), compare progressive and binary reward schemes ( §8.2), examine multi-turn correction beyond two turns ( §8.3), evaluate the model's robustness to varied correction instructions during inference ( §8.4), and conduct case studies on representative examples ( §8.5). ## conclusion in this paper, we investigated self-correction methods for enabling slms to revise their own generated code. our findings show that, despite their effectiveness in proprietary models, existing prompting-and training-based approaches fall short when applied to slms. to address these shortcomings, we propose cocos, which introduces an rl reward scheme tailored to the multiturn code generation setting. by incorporating fine-grained assessment and an accumulated reward function, cocos demonstrates both successful self-correction and improved initial response quality through trajectory-level optimization across diverse slms and previously unseen scenarios.