Synthetic Sandbox for Training Machine Learning Engineering Agents

발행일
출처
arXiv
논문 번호
1269
분야
Research Agents
arXiv 번호
2604.04872

이 논문은 MLE(machine learning engineering) agent의 trajectory 단위 on-policy(지금 학습 중인 모델이 직접 만든 trajectory로 바로 학습하는) RL이 ML pipeline 실행 지연 때문에 너무 느린 문제를, 표본 50에서 200개짜리 합성 과제를 만드는 multi-agent framework로 실행 시간을 줄여 풀었다.

이 논문은 한마디로 작은 합성 데이터 환경으로 MLE agent를 on-policy RL로 학습시키는 방법 SandMLE를 낸 연구다. 네 역할의 LLM agent가 MLE-bench seed 과제 60개를 바탕으로 데이터, hidden rule(특징에서 정답을 정하는 숨긴 규칙), evaluator, 과제 설명을 갖춘 합성 과제 912개를 만들었다. 평균 코드 실행 시간이 196.17초에서 14.31초로 줄어 trajectory 단위 GRPO(그룹 안의 상대 reward로 policy를 고치는 RL 알고리즘) 학습이 가능해졌다. ReAct로 잰 Qwen3-30B의 MLE-Bench-Lite Any Medal 비율(Kaggle 동메달 threshold 이상을 낸 과제 비율)은 13.6%에서 27.3%로 올랐다.

핵심 요약

  • 문제: MLE agent는 rollout 단계마다 데이터 전처리, 모델 학습, 지표 평가까지 전체 ML pipeline을 실행해야 해서, 표준 문제에서 코드 한 번 실행에 평균 200초 가까이 걸린다. 이 지연 때문에 trajectory 단위 on-policy RL이 너무 느려서, 기존 연구는 SFT나 offline proxy reward를 쓰는 step 단위 RL로 물러났다.
  • 입력과 출력: 구현과 실험을 갖춘 학습 방법(method) 논문이다. 입력은 MLE-bench에서 고른 seed 과제 60개이고, 출력은 학습 데이터가 50에서 200개 표본뿐인 합성 MLE 과제(데이터, hidden rule, 평가 스크립트, 과제 설명)와 이 과제로 RL 학습한 Qwen3 policy 모델이다.
  • 방법: Data Strategist가 seed 과제에서 Structural DNA(도메인 맥락을 뺀 수학적 과제 구조)를 뽑아 새 도메인과 noise를 입히고 특징과 정답을 잇는 hidden rule을 정하며, ML Developer가 데이터를 만들고 baseline 방법을 학습, 평가해 단계별 milestone threshold를 정한다. MLOps(machine learning operations) Engineer는 숨긴 test set으로 채점하는 evaluator를, Technical Writer는 과제 설명을 쓰고, threshold가 정해진 순서를 지키는 과제만 학습에 남긴다. 학습은 ReAct로 rollout하는 trajectory 단위 GRPO이고, reward는 형식 준수, 실행 성공, median, bronze, silver, gold threshold 통과에 가중치를 나눈 dense reward다.
  • 구현: seed 60개에서 합성 과제 1200개로 시작해 실행 검증과 threshold 점검을 거친 912개를 남기고, 그중 848개를 학습에, 64개를 검증에 썼다. Qwen3-8B, Qwen3-14B, Qwen3-30B-A3B-2507을 RLLM(language agent의 post-training을 위한 framework)으로 학습했고, 학습 중 최대 20 turn과 단계당 90초 실행 제한을 두었으며 과제마다 NVIDIA H200 GPU 한 장을 썼다. reward 가중치는 형식 0.1, 실행 0.3, median 0.1, bronze 0.2, silver 0.2, gold 0.1이다.
  • 결과: Gemini-2.5-Flash가 쓴 코드 기준 평균 실행 시간은 원래 MLE-bench 과제의 196.17초에서 합성 과제의 14.31초로 줄었다. ReAct로 잰 MLE-Bench-Lite Any Medal 비율은 8B가 13.6%(Base, 그리고 Claude-4.5-Sonnet trajectory로 SFT한 Seed-SFT)에서 22.7%로, 30B가 Base 13.6%에서 27.3%로 올랐고, Claude-4.5-Sonnet은 31.8%였다. MLE-Dojo에서 30B SandMLE는 MLE-Agent scaffold로 HumanRank(사람 참가자 leaderboard 안의 순위로 매긴 점수) 38.56을 얻어 Base의 29.12보다 높았다.
  • 한계: 저자들은 최대 turn 수를 늘리는 test-time 실험에서 성능이 30 turn에서 가장 높았고, 그 뒤로는 context window를 자주 넘겨 떨어졌다고 밝혔다. 8B SandMLE의 Valid Submission 비율은 63.6%로 Base의 68.2%보다 낮았고, 30B를 AIDE(코드 공간을 탐색하는 agent scaffold)로 평가한 MLE-Bench-Lite에서는 Any Medal 비율이 Base와 같은 13.6%였다. 평가는 22문항의 MLE-Bench-Lite와 62개 과제의 MLE-Dojo에서 했다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.

원문 보기 (새 탭에서 열림)