用强化学习训练小模型,比用大模型提示更高效解决机器学习工程任务。
Reinforcement Learning for Machine Learning Engineering Agents
- 小模型通过强化学习逐步优化,克服静态提示的局限。
- 引入时长感知梯度更新,奖励耗时长但效果好的动作。
- 通过程序运行日志获取部分奖励,让模型学会区分接近正确的代码。
现有机器学习工程代理依赖强大语言模型的提示,无法随经验提升。本文证明,基于较弱模型并通过强化学习改进的代理,可超越使用更大但静态模型的代理。针对该场景提出两大挑战:其一,动作执行时间不固定导致异步梯度更新偏好快速但次优解;为此提出在分布式异步强化学习中采用时长感知梯度更新,增强高成本高回报动作的权重。其二,仅以测试集表现作为奖励反馈不足,几乎正确与完全失败的程序无差别对待;为此引入环境仪器化,利用独立静态语言模型在程序中插入打印语句,记录实验进展,从中提取部分奖励信号用于学习。在MLEBench上的实验表明,对较小模型(Qwen2.5-3B)进行强化学习训练,其在12个Kaggle任务上平均性能优于调用大型模型(Claude-3.5-Sonnet)的提示代理,提升达22%。
原文摘要 · Abstract (English)
Existing agents for solving tasks such as ML engineering rely on prompting powerful language models. As a result, these agents do not improve with more experience. In this paper, we show that agents backed by weaker models that improve via reinforcement learning (RL) can outperform agents backed by much larger, but static models. We identify two major challenges with RL in this setting. First, actions can take a variable amount of time (e.g., executing code for different solutions), which leads to asynchronous policy gradient updates that favor faster but suboptimal solutions. To tackle variable-duration actions, we propose duration-aware gradient updates in a distributed asynchronous RL framework to amplify high-cost but high-reward actions. Second, using only test split performance as a reward provides limited feedback. A program that is nearly correct is treated the same as one that fails entirely. To address this, we propose environment instrumentation to offer partial credit, distinguishing almost-correct programs from those that fail early (e.g., during data loading). Environment instrumentation uses a separate static language model to insert print statement to an existing program to log the agent's experimental progress, from which partial credit can be extracted as reward signals for learning. Our experimental results on MLEBench suggest that performing gradient updates on a much smaller model (Qwen2.5-3B) trained with RL outperforms prompting a much larger model (Claude-3.5-Sonnet) with agent scaffolds, by an average of 22% across 12 Kaggle tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。