AI Papers Reader

Personalized digests of latest AI research

View on GitHub

Beyond Correctness: How Reinforcement Learning Teaches AI to Write Blazingly Fast Code

Artificial intelligence models have gotten remarkably good at writing software that works. However, code that merely works isn’t necessarily efficient; an AI might generate a solution that passes basic unit tests but hogs memory or runs agonizingly slowly.

Attempts to fix this using Reinforcement Learning (RL)—by simply giving AI models higher rewards for shorter execution times—have historically backfired. Small timing fluctuations, weak test suites, and noisy server environments often confuse the AI, prompting it to produce code that is either barely faster or outright broken.

Now, researchers at Meta FAIR, Inria, and Université Paris Dauphine have unraveled this bottleneck. In a new paper, the team presents a comprehensive framework that successfully trains large language models (LLMs) to write code that is both strictly correct and optimized for runtime speed.

The core challenge lies in making execution time a reliable learning signal. Standard RL for coding treats unit tests as binary pass/fail gates. Adding raw runtime metrics introduces significant noise: a momentary background process on a server can easily make a brilliant program look slow.

To overcome this, the researchers built DMC-Optim, a curated dataset of competitive programming problems equipped with heavy-input “stress tests” specifically designed to highlight performance differences. They isolated code execution on dedicated, calibrated remote CPU clusters rather than local training nodes, ensuring timing comparisons remained precise.

Crucially, the authors built an offline simulator to screen thousands of reward configurations before launching expensive GPU training. Instead of blindly rewarding lower execution times, they gated speed rewards behind strict correctness requirements and ranked model outputs against human reference solutions. They also adapted the underlying training algorithm, Group Relative Policy Optimization (GRPO), to remain stable under sparser, noisier timing feedback.

The results are striking. When applied to models like CWM 32B and Qwen 2.5, the optimization-RL framework dramatically boosted speed performance without degrading baseline correctness. On strict efficiency benchmarks requiring solutions to rank in the top 30% of human speed, CWM 32B achieved a 125% relative improvement (jumping from a 13.7% to a 30.9% pass rate) over standard correctness-focused RL. On external evaluations like LiveCodeBench, the optimized model beat standard models in head-to-head speed comparisons up to 83% of the time.

So, what does an optimization-trained model actually learn to do? The AI learns both practical system tweaks and high-level algorithmic leaps:

  • Algorithmic Redesign: In a complex vector-space problem (Euclid’s Nightmare), the baseline model wrote a standard Gaussian elimination algorithm. The optimization-trained model recognized an underlying graph structure and replaced the heavy matrix operations with a near-linear “Disjoint-Set Union” algorithm, slashing runtime from 46.4 seconds down to 28.5 seconds.
  • Mathematical Shortcuts: In a card game strategy problem (AtCoder ABS), the baseline implemented a heavy, recursive search algorithm. The optimized AI deduced a mathematical property of the game, collapsing the entire decision tree into a simple, constant-time mathematical formula that executed in 9.2 seconds—a six-fold speedup.

By demonstrating that efficiency can be systematically taught to AI models, this research marks an important step toward automated software engineers that don’t just solve problems, but build high-performance infrastructure.