用强化学习让大模型自动优化代码,提升性能。
MaxCode: A Max-Reward Reinforcement Learning Framework for Automated Code Optimization
- 基于最大奖励的强化学习框架,通过执行反馈迭代优化代码。
- 在CUDA和C++基准上分别提升20.3%速度和10.1%排名表现。
- 融合自然语言分析与奖励重排序,适合系统级代码优化研究者。
大型语言模型在通用编程任务中表现强劲,但在代码优化方面面临两大挑战:一是编写高性能代码(如高效CUDA核函数和竞赛级CPU代码)需掌握系统、算法及特定语言知识;二是需理解时延、设备利用率等性能指标,而不仅是正确性。本文提出MaxCode,一种基于推理时搜索的强化学习框架,通过执行反馈引导模型迭代优化。该方法将现有搜索策略统一于最大奖励框架下,使观察与动作价值函数可模块化调整。为增强观察空间,引入自然语言批评模型,将原始执行反馈转化为错误诊断与性能瓶颈分析,并结合历史最佳奖励。同时,训练生成式奖励预估模型,利用回溯动作值对候选方案重排序,提升探索效率。在KernelBench(CUDA)与PIE(C++)优化基准上的测试表明,相比基线,MaxCode分别实现20.3%的绝对加速提升和10.1%的相对加速排名提升。
原文摘要 · Abstract (English)
Large Language Models (LLMs) demonstrate strong capabilities in general coding tasks but encounter two key challenges when optimizing code: (i) the complexity of writing optimized code (such as performant CUDA kernels and competition-level CPU code) requires expertise in systems, algorithms and specific languages and (ii) requires interpretation of performance metrics like timing and device utilization beyond binary correctness. In this work, we explore inference-time search algorithms that guide the LLM to discover better solutions through iterative refinement based on execution feedback. Our approach, called MaxCode unifies existing search methods under a max-reward reinforcement learning framework, making the observation and action-value functions modular for modification. To enhance the observation space, we integrate a natural language critique model that converts raw execution feedback into diagnostic insights about errors and performance bottlenecks, and the best-discounted reward seen so far. Together, these provide richer input to the code proposal function. To improve exploration during search, we train a generative reward-to-go model using action values from rollouts to rerank potential solutions. Testing on the KernelBench (CUDA) and PIE (C++) optimization benchmarks shows that MaxCode improves optimized code performance compared to baselines, achieving 20.3% and 10.1% relative improvements in absolute speedup value and relative speedup ranking, respectively.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。