Extrapolative Weight Averaging Reveals Correctness-Efficiency Frontiers in Code RL

2026-05-27Machine Learning

Machine LearningArtificial IntelligenceComputation and Language
AI summary

The authors studied how combining different versions of a trained program (called checkpoints) can help solve coding problems more effectively without extra training. They found that by mixing checkpoints trained on simpler tests with those trained on harder, more detailed tests, they could balance between just getting answers right and getting them efficiently. Furthermore, by going beyond just blending these versions and actually pushing beyond them (extrapolating), they could solve even more problems. This approach worked with different coding methods and model sizes, and when combined into groups (ensembles), it improved performance more than any single version alone.

linear interpolationextrapolative weight averagingreinforcement learningcompetitive programmingunit testsPareto frontiermodel checkpointsensemblespass@250
Authors
Kunhao Zheng, Pierre Chambon, Juliette Decugis, Jonas Gehring, Taco Cohen, Benjamin Negrevergne, Gabriel Synnaeve
Abstract
Linear interpolation between fine-tuned checkpoints has been shown to trace the Pareto front between competing objectives, but whether extrapolative weight averaging can extend such frontiers to new checkpoints useful at inference time, without additional RL training, remains unclear. We study this question in RL for competitive programming, where hidden unit tests under time and memory limits enforce both functional correctness and computational efficiency. Starting from a shared initialization, we train checkpoints under nested unit-test coverage: low-coverage rewards require passing smaller-input tests, while high-coverage rewards require passing progressively larger tests up to the full suite. This sweep reveals the emergence of a correctness-efficiency frontier: on hard problems, higher-coverage reward reduces optimization failures but increases correctness failures, leaving solve rate nearly unchanged. Interpolation between low- and high-coverage checkpoints recovers this frontier, while extrapolation extends it beyond the trained endpoints. Both the frontier and its extrapolative continuation appear across three inference settings, pure reasoning, tool use, and agentic coding, and across two model scales, 32B and 7B. At the problem level, moving along the frontier changes which problems are solved, making extrapolated checkpoints complementary policies in inference-time scaling. Ensembles with extrapolative weight averaging broaden coverage and improve pass@250 on LCB/hard by 3.3% over the best single checkpoint at matched sample budget. These results show that nested unit-test coverage in code RL induces a frontier that extrapolative weight averaging can navigate, extend, and exploit.