Source-linked AI summary

Afterburner: Reinforcement Learning Facilitates Self-Improving Code Efficiency Optimization

Mingzhe Du, Luu Anh Tuan, Yue Liu, Yuhao Qing, Dong Huang, Xinyi He, Qian Liu, Zejun Ma, See-kiong Ng

arXiv:2505.23387v3cs.SEcs.AI

TL;DR

LLM-generated code is often functionally correct but inefficient, a concern for real-world deployment. The paper uses a closed-loop framework with execution feedback to iteratively refine code, finding that GRPO sustains gains where SFT and DPO saturate, including higher PASS@1 and efficiency versus human submissions.

  • Problem

    LLM-generated code can be functionally correct yet computationally inefficient, a concern where latency and memory budgets matter.

    Source: §1 · p. 1
  • Method

    The Iterative Optimization Framework pairs Afterburner code generation with Monolith execution and empirical performance feedback in a repeated optimization loop.

    Source: §1 · p. 2, §3.2 · p. 4
  • Results

    GRPO raises PASS@1 from 47% to 62% and BEYOND-I from 31% to 45%, while SFT and DPO rapidly saturate in efficiency gains.

    Source: §1 · p. 2
  • Takeaways & Limitations

    Online RL with direct execution feedback achieved superior, sustained code-efficiency improvements across multiple iterations compared with SFT and DPO.

    Source: §1 · p. 2, §7 · p. 10

    The demonstrated efficiency optimization targets competition-level programming tasks; extension to larger real-world software projects remains for future investigation.

    Source: A Limitations · p. 14

Abstract

from arXiv · show

Large Language Models (LLMs) generate functionally correct solutions but often fall short in code efficiency, a critical bottleneck for real-world deployment. In this paper, we introduce a novel test-time iterative optimization framework to address this, employing a closed-loop system where LLMs iteratively refine code based on empirical performance feedback from an execution sandbox. We explore three training strategies: Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and Group Relative Policy Optimization (GRPO). Experiments on our Venus dataset and the APPS benchmark show that SFT and DPO rapidly saturate in efficiency gains. In contrast, GRPO, using reinforcement learning (RL) with execution feedback, continuously optimizes code performance, significantly boosting both pass@1 (from 47% to 62%) and the likelihood of outperforming human submissions in efficiency (from 31% to 45%). Our work demonstrates effective test-time code efficiency improvement and critically reveals the power of RL in teaching LLMs to truly self-improve code efficiency.

1 Introduction

LLM-generated code can be correct yet inefficient, limiting deployment where latency and memory matter. The paper introduces a closed-loop optimization framework and reports that GRPO sustains improvements where SFT and DPO saturate.

1 Introduction

  • Code that is functionally correct may still be inefficient, creating performance bottlenecks, higher computing costs, and system-wide latency in real-world systems.Prior efficiency-oriented attempts using prompting, self-optimization, or fine-tuning have had limited success, often failing to instill adaptive knowledge for robust improvement.Source: §1 · p. 1
  • The Iterative Optimization Framework uses Afterburner to generate improved code and Monolith to evaluate it and feed empirical performance back for another cycle.The authors describe the loop as resembling developers optimizing through trial and feedback.Source: §1 · p. 2
  • GRPO continuously improves code performance, raising PASS@1 from 47% to 62% and BEYOND-I from 31% to 45%, while SFT and DPO saturate.The experiments cover Venus and APPS; SFT gains initially, while DPO performs better than SFT but follows the same saturation trend.Source: §1 · p. 2

2 Related Work

Prior work measures code efficiency with benchmarks and has explored prompting, fine-tuning, and preference alignment, but efficiency remains an under-addressed optimization target.

2 Related Work

  • Efficiency benchmarks use different comparison strategies, including reference-solution ratios, system simulation, percentile ranks against human code, and categorical comparisons.Venus builds on this benchmark landscape with more tasks and solutions for efficiency assessment.Source: §2 · pp. 2–3
  • Prompting, self-optimization, and efficiency-focused fine-tuning have had limited success, while cross-entropy training may not directly instill nuanced efficiency preferences.The paper notes that execution-feedback methods have largely targeted functional correctness rather than computational efficiency.Source: §2 · p. 3
  • Existing execution-feedback approaches commonly reward correctness, so the paper introduces Monolith to provide real-time efficiency metrics for RL-based optimization.The related work describes execution feedback as advancing code correctness, while efficiency remains a less explored target.Source: §2 · p. 3

3 Iterative Optimization Framework

The Iterative Optimization Framework alternates code generation with empirical evaluation. Afterburner proposes refinements and Monolith measures their performance to inform subsequent iterations.

3 Iterative Optimization Framework

  • IOF closes the loop by having Afterburner generate improved code and Monolith execute it, returning empirical performance feedback for the next iteration.The framework is intended to improve code efficiency for resource-constrained or time-sensitive applications.Source: §3 · p. 3

3.1 Afterburner: Code Efficiency Optimization Models

  • Afterburner refines the current solution for a specified efficiency objective, using prior performance and selecting the best-performing code greedily for the next iteration.The objective may target time, memory, or an integral score; iteration continues for a predetermined number of rounds.Source: §3.1 · p. 3

3.2 Monolith: Code Execution Sandbox

  • Monolith executes generated code and supplies empirical performance feedback because complexity analysis alone misses practical effects such as constants, runtimes, libraries, compilers, and hardware.Its reported outputs include pass status, execution time, peak memory usage, and an integral score.Source: §3.2 · p. 4

4 Code Efficiency Optimization

Section 4 introduces Venus for code-efficiency assessment and describes three training strategies for Afterburner: supervised examples, offline preferences, and group-wise reinforcement learning. GRPO uses execution feedback and a reward combining format, correctness, and efficiency.

Data Preparation

  • Venus evaluates time, memory, and their combined impact using percentile ranks against diverse reference solutions; it averages 106.6 solutions per task, versus fewer than 20 in existing Python efficiency benchmarks.Its Python set contains 2,181 training and 300 test tasks.Source: §4 · p. 4

Supervised Fine-Tuning

  • SFT trains on pairs of functionally correct solutions, learning to generate the more efficient code from an inefficient-code prompt.The AfterburnerSFT objective minimizes cross-entropy for the preferred solution.Source: §4 · p. 5

Direct Preference Optimization

  • DPO aligns code with efficiency preferences offline by increasing the likelihood of preferred responses and decreasing that of dispreferred responses.Its reference model is the corresponding AfterburnerSFT model.Source: §4 · p. 5

Group Relative Policy Optimization

  • GRPO starts with format-alignment fine-tuning, then trains on execution-feedback prompts covering all 984 distinct Venus training tasks without ground-truth responses.For each prompt, GRPO compares multiple rollouts and learns their relative advantages.Source: §4 · pp. 5–6

Reward Functions

  • GRPO’s reward combines format control, functional correctness, and relative improvement in the instructed time, memory, or integral efficiency metric.Correctness reward distinguishes upgrades, maintained passing or failing status, and downgrades.Source: §4 · pp. 5–6

Objective

  • GRPO favors candidates with high intrinsic quality and better performance than peers generated for the same prompt.Its objective uses clipped policy ratios and reward-based relative advantages within each rollout group.Source: §4 · p. 6

5 Experiment Setup

The experiments use Venus as the primary code-efficiency dataset and APPS as an auxiliary benchmark. They assess functional correctness with PASS@1 and efficiency through task-level percentile ranks against reference solutions.

Dataset Recipe

  • Venus Python contains 2,181 algorithmic problems, validated test generators, and an average of 106.6 human solutions per problem; APPS is added as an auxiliary benchmark.APPS was originally designed around functional correctness, with 21.2 test cases and 23.4 solutions per problem.Source: §5 · p. 6

Functional Correctness

  • PASS@1 measures global functional correctness as the number of passed generations divided by the total number of test tasks.Source: §5 · p. 6

Computational Efficiency

  • Efficiency evaluation avoids absolute metrics because they are sensitive to hardware and operating systems, instead comparing execution time, peak memory, and integral memory use against reference distributions.Source: §5 · p. 7
  • Task-level efficiency percentile ranks are averaged globally, with higher scores indicating that generated code outperforms a larger fraction of reference solutions.Source: §5 · p. 7

Implementation Details

  • Models are trained on a single node with eight H100 GPUs; the paper uses Llama-Factory for SFT and DPO, and Verl for GRPO.Inference uses vLLM; further training-pipeline and hyperparameter details are provided in the appendix.Source: §5 · p. 7, Figure 4 · p. 7

6 Discussion and Key Takeaways

Vanilla LLMs can be functionally correct yet remain less efficient than human solutions, while Afterburner’s iterative strategies differ in how long they sustain improvement. GRPO’s online execution feedback supports continued refinement, though its exploration also produces some solutions less efficient than the human baseline.

6.1 How about the Code Efficiency Performance of Vanilla LLMs?

  • OpenAI o4 mini reaches 89.11% PASS@1 on Venus, but its code surpasses only 56.85% of human solutions on BEYOND-T there and 40.07% on APPS.The passage reports similar efficiency disparities for other leading models and metrics.Source: §6.1 · p. 7

6.2 Does Iterative Improvement Framework Work?

  • SFT’s early efficiency gains often exhaust its applicable patterns during iterative optimization.The passage attributes this to learning transformations from training data, with initial gains more likely when input code matches known inefficient patterns.Source: §6.2 · p. 7
  • DPO can improve more consistently than SFT, but tends to plateau once its learned preferences are fully exploited.The passage explains that DPO does not learn from its own generations without retraining, leaving exploration bounded by its initial preference dataset.Source: §6.2 · p. 8
  • GRPO uses online reinforcement learning: Monolith evaluates multiple generated candidates, and empirical feedback updates the policy toward more efficient code.The passage links this adaptive process to sustained performance improvements across iterations.Source: §6.2 · p. 8

6.3 Why GRPO Can Iteratively Enhance Code Efficiency?

  • Generation diversity and experience from the refinement loop support GRPO’s iterative capability; removing execution feedback or original-code context significantly diminishes its performance.The ablation passage describes results at iteration 4, with bracketed changes relative to baseline.Source: §6.3 · p. 9, Table 2 · p. 9

6.4 Can Afterburner Generate Code Surpassing Human Efficiency?

  • After 8 iterations, AfterburnerGRPO O achieves the highest B% scores on Venus across TIME (8.00%), MEMORY (7.00%), and INTEGRAL (5.33%), while also generating a larger fraction of solutions less efficient than the human baseline.B% denotes solutions better than all human solutions.Source: §6.4 · p. 9, Table 3 · p. 9

7 Conclusion

The paper introduces an iterative framework that uses real-time efficiency feedback to refine LLM-generated code. Across multiple iterations, GRPO with online reinforcement learning and execution feedback achieves superior, sustained efficiency improvements compared with SFT and DPO.

7 Conclusion

  • Online RL with direct execution feedback enables GRPO to achieve superior, sustained code-efficiency improvements across multiple iterations.The conclusion contrasts this with SFT learning superficial patterns and DPO internalizing efficiency preferences.Source: §7 · p. 10

A Limitations

The paper’s stated limitations concern the scope of evaluated tasks and the added inference time required by iterative optimization. The supplied material also identifies accompanying model and dataset-statistics tables.

A Limitations

  • Afterburner is demonstrated on competition-level programming tasks; extension to larger real-world software projects warrants further investigation.Such projects may involve more complex code context, efficiency criteria beyond algorithmic performance, and task decomposition strategies outside this work’s scope.Source: A Limitations · p. 14
  • Iterative optimization requires more inference time than single-pass code generation, so its cost-benefit trade-off needs evaluation for each application and deployment scenario.The authors argue that upfront optimization cost can be offset by cumulative runtime savings for frequently executed or performance-critical modules.Source: A Limitations · p. 14

B Model Details

  • Table 4 lists models with their size, reasoning ability, and model URL; a diamond marks reasoning models.Source: Table 4 · p. 14

C Dataset Curation and Statistics

  • Table 5 summarizes statistics for representative function-level code-generation datasets, including average solutions per problem and each dataset’s efficiency or correctness focus.The caption marks correctness-only datasets with ♣ and efficiency datasets with ♡.Source: Table 5 · p. 14

C.1 Venus Dataset

Venus combines curated problems with validated, language-agnostic test cases for efficiency benchmarking. The APPS subset is filtered for reference-solution coverage, while the iterative procedure evaluates revisions and retains the best-performing code.

C.1 Venus Dataset

  • 1,284 high-quality problems make up Venus, after filtering 3,535 LeetCode problems and splitting the retained set into 984 training and 300 held-out test problems.The pipeline also removes paid-only and non-algorithmic questions and excludes problems with fewer than 16 solutions passing all tests.Source: §C.1 · pp. 14–15
  • Each Venus problem gets a GPT-4O-generated test generator; candidate cases are accepted only when all canonical human solutions produce consistent outputs.The generators are designed to produce a virtually unbounded stream of diverse, valid inputs.Source: §C.1 · p. 15
  • Validated test cases use standard input and output, allowing the same cases to evaluate supported languages and support consistent cross-language efficiency comparisons.The paper lists Python, C++, Go, Java, and JavaScript among supported languages.Source: §C.1 · p. 16

C.2 APPS Dataset

  • 2,803 problems remain in the APPS subset after filtering its 5,000-problem official training split for sufficient accepted reference solutions.APPS contains 10,000 Python programming problems, averaging 21.2 test cases and 23.4 solutions per problem.Source: §C.2 · p. 16, Figure 8 · p. 17

D Iterative Efficiency Optimization Procedure

  • Given a problem, efficiency objective, test cases, and optional starting code, Afterburner proposes revisions that Monolith evaluates, retaining the better candidate and returning its performance metrics.If no initial code is provided, Afterburner generates it before the iterative evaluation-and-refinement loop.Source: D Iterative Efficiency Optimization… · p. 17

E Model Training Details

The training pipeline develops Afterburner with SFT, DPO, and GRPO, using different initialization and training data for each strategy. The GRPO setup combines format, correctness, and efficiency rewards, while prompts and empirical measurement procedures specify the optimization and evaluation conditions.

E.1 Training Pipeline.

  • SFT trains on its constructed dataset, DPO fine-tunes the corresponding SFT checkpoint, and GRPO uses cold-start format alignment before training on its GRPO dataset.The GRPO models are first fine-tuned on the cold-start dataset to adhere to the required response format.Source: §E.1 · p. 18

E.2 Details of Afterburner SFT

  • SFT fine-tunes Qwen2.5-3B-Instruct with LoRA for one epoch on its training dataset, using a learning rate of 3e-5 and an effective batch size of 64.Training uses BF16 precision and clips gradients at a norm of 1.0.Source: §E.2 · p. 18

E.3 Details of Afterburner DPO

  • DPO fine-tunes the SFT checkpoint for one epoch with LoRA, using a learning rate of 4e-5, an effective batch size of 16, and a beta of 0.1.The DPO setup uses a sigmoid loss and sets its SFT loss component to zero.Source: §E.3 · p. 18

E.4 Details of Afterburner Code Start

  • The GRPO code-start data is collected from Gemini responses passing a response-regex filter, and the base model is fully fine-tuned on that data for one epoch.The maximum input sequence length is 32,768 tokens.Source: §E.4 · p. 18

E.5 Details of Afterburner GRPO

  • GRPO starts from AfterburnerCS and trains for 20 epochs, generating 16 responses per prompt and parallelizing reward calculation because code execution and metric computation are time-consuming.The actor learning rate is 1e-6, with a PPO mini-batch size of 32.Source: §E.5 · p. 19
  • GRPO combines format compliance, correctness relative to a baseline, and a bounded tanh-scaled reward for relative efficiency gains.Format reward is +1 or −1; correctness rewards distinguish upgrades, downgrades, and unchanged pass/fail status, while efficiency gains are clipped before scaling.Source: §E.5 · pp. 19–20

Afterburner Prompt Template

  • The prompt supplies original pass, time, memory, and integral performance, and instructs the model to fix failing code or optimize the specified efficiency objective when the code passes.It also asks for a complete solution in one markdown code block and directs the model to generate initial code if no original solution is provided.Source: §F.3 · p. 21

G Uncertainty in Code Efficiency Measurement

  • Empirical runtime and memory measurements capture algorithm, code, compiler, and hardware effects, but the paper notes that system noise and run-to-run variability affect results.The authors choose bootstrapping to quantify uncertainty in these empirical measurements.Source: G Uncertainty in Code Efficiency… · p. 21

G.1 Details of Bootstrapping Evaluation

The evaluation uses task-level bootstrap sampling to estimate efficiency metrics and report uncertainty.

G.1 Details of Bootstrapping Evaluation

  • 128 bootstrap replicates sample 4 solutions per problem, with each generated code repeatedly evaluated 16 times, and report metric means with 95% confidence intervals.Task-level metrics are grouped by problem ID before sampling.Source: §G.1 · p. 21

H Monolith Implementation

Monolith standardizes code execution in isolated Docker containers and measures runtime and memory. The symbol list defines the paper’s model, reward, and efficiency-metric terminology.

H Monolith Implementation

  • Monolith runs 81 isolated Docker workers on a 96-vCPU, 96-GB GCP host, allocating each worker 1 vCPU and 1 GB of memory.Containers are recreated for each task; execution time and peak memory are measured with time -v, while VmRSS samples support integral-score calculation.Source: H Monolith Implementation · p. 22

I Symbol List

  • The symbol list defines Afterburner as the code-optimization model, Monolith as the performance-feedback sandbox, and IOF as the proposed Iterative Optimization Framework.It also names the SFT, DPO, and GRPO datasets and reward components.Source: I Symbol List · pp. 22–23
  • BEYOND-T, BEYOND-M, and BEYOND-I are average percentile ranks against human solutions for execution time, memory usage, and integral score, respectively; higher is better.The symbol list also defines B%, M%, and W% as generations better than, within, or worse than the range of human solutions.Source: I Symbol List · p. 23
Loading 2505.23387v3…