Source-linked AI summary
CodeChain: Towards Modular Code Generation Through Chain of Self-revisions with Representative Sub-modules
Hung Le, Hailin Chen, Amrita Saha, Akash Gokul, Doyen Sahoo, Shafiq Joty
TL;DR
Complex coding tasks remain challenging for LLMs that often generate monolithic solutions rather than decomposing tasks into modules. CodeChain elicits modular code and revises it using representative sub-modules; it reports relative pass@1 gains on APPS and CodeContests.
Problem
LLMs remain less capable than skilled developers on highly complex coding tasks, and common generation approaches produce monolithic solutions instead of decomposing tasks into sub-tasks.
Source: §1 · p. 1Method
CodeChain prompts LLMs to generate modular solutions, clusters extracted sub-modules, and conditions successive self-revisions on representative modules.
Source: §1 · p. 2, §3.2 · p. 4, §3.3 · p. 5Results
More than 35% on APPS and 76% on CodeContests in relative average pass@1 improvement are reported, with consistent gains for OpenAI and WizardCoder models.
Source: §1 · p. 2Takeaways & Limitations
The experiments report CodeChain’s effectiveness with both OpenAI and open-source WizardCoder models on APPS and CodeContests.
Source: §5 · p. 9The authors suggest better embedding models that group output samples by functional correctness as well as programming semantics.
Source: §4.2 · p. 8
Abstract
from arXiv · showhide
Large Language Models (LLMs) have already become quite proficient at solving simpler programming tasks like those in HumanEval or MBPP benchmarks. However, solving more complex and competitive programming tasks is still quite challenging for these models - possibly due to their tendency to generate solutions as monolithic code blocks instead of decomposing them into logical sub-tasks and sub-modules. On the other hand, experienced programmers instinctively write modularized code with abstraction for solving complex tasks, often reusing previously developed modules. To address this gap, we propose CodeChain, a novel framework for inference that elicits modularized code generation through a chain of self-revisions, each being guided by some representative sub-modules generated in previous iterations. Concretely, CodeChain first instructs the LLM to generate modularized codes through chain-of-thought prompting. Then it applies a chain of self-revisions by iterating the two steps: 1) extracting and clustering the generated sub-modules and selecting the cluster representatives as the more generic and re-usable implementations, and 2) augmenting the original chain-of-thought prompt with these selected module-implementations and instructing the LLM to re-generate new modularized solutions. We find that by naturally encouraging the LLM to reuse the previously developed and verified sub-modules, CodeChain can significantly boost both modularity as well as correctness of the generated solutions, achieving relative pass@1 improvements of 35% on APPS and 76% on CodeContests. It is shown to be effective on both OpenAI LLMs as well as open-sourced LLMs like WizardCoder. We also conduct comprehensive ablation studies with different methods of prompting, number of clusters, model sizes, program qualities, etc., to provide useful insights that underpin CodeChain's success.
1 INTRODUCTION
Complex coding tasks remain challenging for LLMs, whose typical monolithic generation contrasts with developers’ modular, iterative reuse. CodeChain prompts modular solutions and revises them using representative modules from prior generations.
1 INTRODUCTION
- LLMs still struggle to match skilled developers on highly complex coding tasks, while many approaches generate monolithic code rather than decomposing tasks into sub-tasks.The introduction also describes generating many independent solutions in the hope that one passes private tests.Source: §1 · p. 1
- Prior self-revision methods use feedback from individual solutions, leaving collective insights across generation samples or their sub-components unused.Earlier public-test filtering approaches also assume selected programs pass private tests without revision or debugging.Source: §1 · p. 2
- CodeChain draws on developers’ practice of modularizing tasks and iteratively revising and reusing components from earlier implementations.The paper illustrates this iterative workflow in Figure 1.Source: §1 · p. 2, Figure 1 · p. 2
- CodeChain uses chain-of-thought prompting to generate modular solutions, then clusters extracted sub-modules and reuses representative modules in subsequent revisions.The process is intended to bring collective insights from previous generations into new solutions.Source: §1 · p. 2
- More than 35% on APPS and 76% on CodeContests in relative average pass@1 improvement are reported for CodeChain.The experiments also report improvements for OpenAI and open-source models, including WizardCoder.Source: §1 · p. 2
2 RELATED WORK
Related work covers advances in LLM code generation, output-feedback ranking, and iterative self-revision. CodeChain distinguishes itself by revising modular programs with reusable sub-modules rather than relying only on feedback from individual outputs.
2 RELATED WORK
- LLMs have been extended to learn from code and can generate syntactically correct programs and solve simpler coding problems with reasonable accuracy.The section situates this work alongside broader Transformer-based language-model research.Source: §2 · pp. 2–3
- Output-feedback methods rank or filter generated programs using public tests, execution behavior, or model-based ranking.Examples include public-test filtering, generated test cases, and execution-based ranking methods.Source: §2 · p. 3
- Prior self-revision work uses test outcomes, correctness prediction, generated explanations, or human feedback, whereas CodeChain revises modular programs using representative reusable sub-modules.The paper characterizes many earlier methods as relying on feedback associated with individual samples.Source: §2 · p. 3
3 CODECHAIN FRAMEWORK
CodeChain frames code generation as modular decomposition and improves it through revisions guided by representative sub-modules selected from earlier samples.
3.1 CODE GENERATION TASK
- The task maps a problem description to a program, which is checked for correctness against private tests; available public tests can also guide generation.Source: §3.1 · pp. 3–4
3.2 MODULAR CODE GENERATION WITH COT PROMPTING
- CoT prompting asks the model to outline high-level sub-modules before implementing them and combining them, though this modular style can reduce end-to-end correctness without self-revision.The paper reports that its self-revision technique benefits more from modularized generation.Source: §3.2 · p. 4
3.3 SELECT REPRESENTATIVE SUB-MODULES ACROSS MULTIPLE SAMPLES
- Because fully correct end-to-end candidates are rare on challenging tasks, CodeChain clusters modules from generated samples and selects representatives nearest each cluster’s embedding-space centroid.These centroid modules are intended to be semantically representative and reusable; public-test failures can be filtered before selection.Source: §3.3 · pp. 4–5
3.4 IMPROVE CODE GENERATION WITH CHAIN OF SELF-REVISIONS
- Each revision round conditions new generations on representative sub-modules clustered from all earlier samples, offering cross-sample feedback rather than feedback from an individual sample.The prompt is augmented with implementations of these modules to encourage reuse or adaptation.Source: §3.4 · p. 5
4 EXPERIMENTS
Experiments evaluate CodeChain on APPS and CodeContests, using several base models and pass@k metrics. Results cover benchmark performance, revision and clustering choices, difficulty trends, and generated-code quality.
4.1 EXPERIMENTAL SETUPS
- The experiments used APPS and CodeContests competitive-programming benchmarks, pass@1 as the main metric, up to five revision rounds, and GPT3.5, GPT4, and WizardCoder models.WizardCoder sizes ranged from 1B to 34B parameters; each generation or revision round used 20 samples per problem.Source: §4.1 · pp. 5–6
4.2 EXPERIMENTAL RESULTS
- 35.34% average pass@1 was achieved by filtered CodeChain+GPT3.5 on APPS; unfiltered results were 10.50% with WizardCoder and 30.24% with GPT3.5.The filtered result was reported across all problem difficulty levels.Source: Table 1 · p. 6, §4.2 · p. 7
- In one-round APPS ablations, chain-of-thought prompting alone hurt direct-generation performance, but modularized solutions performed better after self-revision; clustering worked best for representative selection, with public-test filtering also needed.The table evaluates prompting, public-test filtering, and revision sampling methods.Source: §4.2 · p. 7, Table 3 · p. 7
- On APPS, revision gains peaked at round 4 and dipped slightly at round 5, possibly because selected samples overfit the small public-test set; CodeContests’ more diverse official tests reduced this overfitting.The CodeContests result held even in the fifth revision round.Source: §4.2 · pp. 8–9
- CodeChain’s gains were larger on APPS competition- and interview-level problems than on introductory problems, with a similar trend across WizardCoder sizes; it also outperformed related self-revision methods using explanation or reflection feedback.The difficulty-level results are shown for GPT3.5 and WizardCoder models.Source: §4.2 · p. 8, Figure 5 · p. 8, Figure 7 · p. 8
- Decreasing the number of clusters across revision rounds performed best: larger early sets provided diverse representatives, while smaller later sets better matched increasingly similar sub-modules.The experiments used K = 5 as the selected cluster-count setting.Source: §4.2 · p. 8, Figure 6 · p. 8
- 13.75% pass@1 on the CodeContests test split was achieved by CodeChain+GPT3.5 with additional filtering, while CodeChain improved pass@1 and pass@5 over direct generation.On validation, performance continued improving through the final revision round.Source: §4.2 · p. 9
- GPT4 ratings of sampled APPS outputs found CodeChain more likely to produce highly modular and reusable programs, while direct generation produced score-0 outputs about 80% of the time.The evaluation used a 0–5 Likert scale, with most CodeChain outputs rated 3–5.Source: §4.2 · p. 9, Figure 9 · p. 9
5 CONCLUSION
CodeChain combines modularized code generation with iterative self-revisions guided by representative sub-modules. Experiments report improved performance with OpenAI and open-source models, including new state-of-the-art results on APPS and CodeContests.
5 CONCLUSION
- CodeChain achieves new state-of-the-art results on APPS and CodeContests with OpenAI and open-source WizardCoder models.The paper reports comprehensive ablations to examine factors behind these results.Source: §5 · p. 9
A COMPARISON WITH RELATED METHODS
Table 4 compares CodeChain and related approaches across four methodological dimensions: code execution, representative sampling, supervision-free operation, and iterative revision.
A COMPARISON WITH RELATED METHODS
- Code execution indicates whether a method uses outcomes from public or synthetic test cases.The passage defines comparison dimensions, not method-by-method outcomes.Source: Table 4 · p. 13
- Representative sampling indicates whether a method subsamples outputs for evaluation or revision.The table compares this alongside execution use, supervision requirements, and iterative revision.Source: Table 4 · p. 13
- The remaining dimensions capture specialized fine-tuning requirements and whether a method supports multiple self-revision rounds.The prose points readers to Table 4 for the systematic comparison.Source: Table 4 · p. 13, A COMPARISON WITH RELATED METHODS · p. 13
B DEMONSTRATION OF MODULARIZED CODE GENERATION
The example depicts modularized generation as outlining sub-modules before implementing and integrating them into a complete solution.
B DEMONSTRATION OF MODULARIZED CODE GENERATION
- The model first outlines the required sub-modules, each with a function header and descriptive docstring.These outline the intended use of each module.Source: Figure 10 · p. 13
- The model then implements each outlined module fully in code.Source: Figure 10 · p. 13
- Finally, the completed modules are integrated as parts of the complete solution.Figure 10 is cited as the illustration of this process.Source: Figure 10 · p. 13, B DEMONSTRATION OF MODULARIZED CODE… · p. 13
C ADDITIONAL DETAILS ON BENCHMARKS
The benchmark details describe an APPS validation sample, an APPS test subset, and the reported test-case statistics for APPS and CodeContests.
C ADDITIONAL DETAILS ON BENCHMARKS
- APPS validation uses 50 randomly selected training-split samples at each of three difficulty levels.The levels are Introductory, Interview, and Competition.Source: C ADDITIONAL DETAILS ON BENCHMARKS · pp. 13–14
- The APPS test subset contains 20 problems: 14 introductory, 3 interview, and 3 competition-level samples.The subset follows Olausson et al. (2023).Source: Table 5 · p. 14
- Table 6 summarizes APPS and CodeContests benchmark statistics.Its caption identifies the benchmark summary and marks statistics reported by Li et al. (2022).Source: Table 6 · p. 14
- APPS example input-output pairs were extracted by a rule-based method and treated as public test cases.The original APPS benchmark did not officially include public test cases; the paper also reports average public and private test-case counts for the benchmarks.Source: C ADDITIONAL DETAILS ON BENCHMARKS · p. 14
D ADDITIONAL EXPERIMENTAL RESULTS
Additional experiments examine revision rounds, filtering, embedding choices, stopping signals, and transfer to LeetCodeHardGym. They report gains across several settings while identifying sensitivity to test-case quality.
D ADDITIONAL EXPERIMENTAL RESULTS
- On APPS, CodeChain peaks at round 4, improving average performance over 1.6x/2x for GPT3.5/WizardCoder and over 2x/5x on competition-level problems.GPT3.5 drops slightly at round 5, while WizardCoder remains stable.Source: D ADDITIONAL EXPERIMENTAL RESULTS · p. 14, Figure 11 · p. 15, Figure 12 · p. 15
- 1.4x pass@1 improvement over direct generation on filtered outputs was observed with GPT3.5 across CodeChain revision rounds.The authors describe this as complementary to filtering or sub-sampling because models can iteratively revise selected outputs.Source: D ADDITIONAL EXPERIMENTAL RESULTS · p. 14, Figure 13 · p. 15
- Among synthetic-test filtering methods, CodeT sampling performed best, but synthetic-test results remained below public-test results and private-test filtering performed best.Complete-pass filtering and largest-cluster sampling reduced performance, possibly because synthetic tests produced noisy filtered outputs; CodeChain is sensitive to filtering correctness and relies on high-quality tests.Source: D ADDITIONAL EXPERIMENTAL RESULTS · pp. 14, 16, Table 8 · p. 16
- Among tested embedding methods, StarCoder achieved the best overall performance, while CodeT5+ outperformed CodeBert at every problem difficulty level.The authors note StarCoder's ability to encode long-context programs, which they found frequently in challenging tasks.Source: D ADDITIONAL EXPERIMENTAL RESULTS · p. 16, Table 9 · p. 17
- On APPS, public-test and private-test performance peaked at revision round 4, supporting public tests as a stopping signal when no validation set is available.The reported public/private performance correlation makes public-test results a proxy for choosing the revision step.Source: D ADDITIONAL EXPERIMENTAL RESULTS · p. 16, Figure 14 · p. 17
- On LeetCodeHardGym, GPT4+CodeChain improved over GPT4 alone and reached close to 7% pass@1 at round 3 on 39 of 40 tasks.The setup used available public tests from problem descriptions rather than the synthetic, AST-validated suite used in the original work.Source: D ADDITIONAL EXPERIMENTAL RESULTS · pp. 16–17
E EXAMPLE GENERATION SAMPLES
Example-generation analyses report increasingly modular and reusable outputs across revisions, alongside a five-round GPT4+CodeChain pass@1 evaluation.
E EXAMPLE GENERATION SAMPLES
- Figure 15 reports GPT4+CodeChain pass@1 results across five revision rounds on LeetCodeHardGym.The supplied figure caption gives the evaluation scope but no specific plotted values.Source: Figure 15 · p. 17
- Automatic evaluation found that generated samples became more modular across the revision chain and generally more reusable than direct-generation outputs.The comparison is against normal direct generation without CodeChain.Source: E EXAMPLE GENERATION SAMPLES · p. 17
- Manual inspection found CodeChain could produce better programs with appropriate modularization and explanatory function docstrings.Source: E EXAMPLE GENERATION SAMPLES · p. 17
E.1 EXAMPLE OF INTRODUCTORY-LEVEL PROBLEM AND CODECHAIN GENERATION SAMPLES
The introductory-level example contrasts direct generation with CodeChain revisions on an integer-divisibility task. Its evaluations show increased modularity, while reuse ratings distinguish generic input/output functions from task-specific computation.
E.1 EXAMPLE OF INTRODUCTORY-LEVEL PROBLEM AND CODECHAIN GENERATION SAMPLES
- Direct GPT4 generation received modularity 0, and its output was judged non-modular with reusability not applicable.Source: §E.1 · p. 18
- The round-0 chain-of-thought sample received modularity 2 and a reusability rating of 3 for its task-specific Nth_smallest_integer function.Its implementation increments N when N is 100 before multiplying by 100^D.Source: §E.1 · p. 18
- At revision round 3, modularity rose to 5, while evaluators rated get_input 3, happy_number 2, and print_output 4 for reusability.The evaluations describe happy_number as problem-specific and print_output as semi-generic but limited to integers.Source: §E.1 · p. 19
- At round 5, modularity remained 5; get_input and print_output each received reusability 5, while the task-specific computation function received 1.The computation function could be used in similar problems with slight modifications.Source: §E.1 · p. 20
E.2 EXAMPLE OF COMPETITION-LEVEL PROBLEM AND CODECHAIN GENERATION SAMPLES
This section presents CodeChain generation examples, including GPT-4 outputs for a Robber Language decoding task and prompts for modular Python code generation. It also covers synthetic test creation and evaluation of code modularity and reusability.
E.2 EXAMPLE OF COMPETITION-LEVEL PROBLEM AND CODECHAIN GENERATION SAMPLES
- The example task asks how many original passwords could yield an encrypted string when Edvin sometimes misses consonant transformations, with each miss treated independently.The Robber Language transformation doubles consonants with an “o” between them, while leaving other letters unchanged.Source: §E.2 · p. 20
- The evaluated output received modularity 4; reusability ratings were split_word 5, initialize_vowels_and_modulo 3, count_possible_passwords 2, and solve_task 2, with the latter modules described as problem-specific.The passage says functions setting constants can be reused, but initialize_vowels_and_modulo is task-specific; count_possible_passwords and solve_task are also described as problem-specific.Source: §E.2 · pp. 22–23
- In revision round 5, the example rated is_consonant’s reusability 5, split_word’s 3, and solve_task’s 1, reflecting their differing generality.The evaluation describes is_consonant as generic, split_word as suited to similar decoding tasks, and solve_task as specific to this problem.Source: §E.2 · p. 24
F PROMPTS WITH INSTRUCTION
- The prompt asks for a well-structured Python solution that considers edge cases, first outlining module headers and signatures, then implementing them.The guide is replaced with instructions to use either standard input streams or call-based functions.Source: F PROMPTS WITH INSTRUCTION · p. 24
- The code-generation prompt asks the model to outline modules, reuse or adapt supplied utility functions, and then implement a modular Python solution.Representative sub-modules selected by CodeChain are supplied as related functions; the prompt also asks for clear function names, signatures, and attention to edge cases.Source: F PROMPTS WITH INSTRUCTION · p. 24
- The synthetic-test prompt requests up to 20 new input-output test cases in a specified format, using available examples as a guide.The prompt replaces «example_test» with test cases from the problem description and asks the model to continue their format.Source: F PROMPTS WITH INSTRUCTION · pp. 24–25
- The evaluation rates solution modularity on a 0–5 Likert scale and rates each sub-module’s reusability on a 0–5 Likert scale, judging whether it is generic enough for reuse elsewhere.Modularity means decomposition into logical atomic sub-modules handling different subtasks, invoked from the main method.Source: F PROMPTS WITH INSTRUCTION · p. 25
G ONE-SHOT EXAMPLES
The examples use a binary-word reversal problem to illustrate ordinary code generation, code revision, and modularized generation. They contrast a one-shot solution with a modularized output organized around submodules.
G ONE-SHOT EXAMPLES
- Listings 15 and 16 show one-shot example inputs for normal or modularized generation and for code revision, respectively.Source: G ONE-SHOT EXAMPLES · pp. 25–26
- The task asks for the fewest word reversals needed to make all words orderable so each word starts with the previous word’s final character, while keeping the words unique.Source: G ONE-SHOT EXAMPLES · pp. 26–27
- The normal-generation output includes helper functions for reversing words and counting starting and ending characters, alongside a procedural solution.Source: G ONE-SHOT EXAMPLES · pp. 27–28
- The modularized output separates submodule generation from Python code generation, with functions to count endpoint characters, reverse selected words, and choose indices to reverse.Source: G ONE-SHOT EXAMPLES · pp. 28–29