SandMLE: 기계 학습 엔지니어링 에이전트 훈련을 위한 합성 샌드박스

이 논문은 대규모 언어 모델 (LLM) 에이전트가 소프트웨어 공학 (SWE) 을 넘어 기계 학습 엔지니어링 (MLE) 영역으로 확장될 때 발생하는 훈련의 병목 현상을 해결하기 위해 제안된 SandMLE 프레임워크에 대해 다룹니다.

1. 문제 정의 (Problem)

기존의 LLM 에이전트 훈련 방식은 기계 학습 엔지니어링 (MLE) 작업에 적용하기 위해 다음과 같은 치명적인 한계에 직면해 있습니다.

• 실행 지연 (Execution Latency): 소프트웨어 공학 (SWE) 작업은 단위 테스트 (unit test) 를 통해 수 초 내에 검증 가능하지만, MLE 작업은 데이터 전처리, 모델 훈련, 평가 지표 산출 등 전체 ML 파이프라인을 실행해야 합니다. 이는 표준 문제당 평균 200 초 이상의 실행 시간을 요구합니다.
• 온-폴리시 RL 의 비실용성: SWE 나 웹 검색과 같은 분야에서 성공적으로 적용된 **궤적 기반 온-폴리시 강화 학습 (Trajectory-wise On-policy RL, 예: GRPO)**은 각 스텝마다 환경 피드백이 필요합니다. MLE 의 긴 실행 시간으로 인해 대규모 롤아웃 (rollout) 을 수행하는 것이 계산 비용상 불가능해졌습니다.
• 대안들의 한계: 기존 연구들은 이 문제를 우회하기 위해 지도 미세 조정 (SFT) 이나 오프라인 프록시 보상을 사용하는 데 의존했습니다. 그러나 이는 에이전트의 탐색 (exploration) 능력과 일반화 (generalization) 능력을 희생하며, 에이전트의 내재적 추론 능력을 근본적으로 향상시키지 못합니다.

2. 방법론 (Methodology: SandMLE)

저자들은 MLE 작업의 실행 지연이 데이터셋의 크기에서 기인한다는 통찰을 바탕으로, SandMLE라는 새로운 프레임워크를 제안합니다. 이는 소수의 시드 (seed) 작업에서 다양한 합성 MLE 환경을 생성하여 훈련 속도를 획기적으로 높입니다.

2.1. 합성 샌드박스 생성 (Synthetic Sandbox Generation)

SandMLE 는 다중 에이전트 (Multi-agent) 협업 방식을 통해 실제 MLE 문제의 구조적 복잡성을 유지하면서 데이터셋 크기를 **마이크로 스케일 (50~200 샘플)**로 제한된 합성 환경을 자동 생성합니다.

• 데이터 전략가 (Data Strategist): 시드 작업의 구조적 DNA(모달리티, 레이블 분포 등) 를 추출하고, 도메인을 변경하거나 노이즈를 주입하여 새로운 시나리오를 설계합니다.
• MLE 개발자 (MLE Developer): 설계된 규칙에 따라 합성 데이터 (이미지, 테이블 등) 를 생성하고, 지상 진실 (Ground Truth) 라벨을 결정론적으로 매핑하는 숨겨진 수학적 규칙을 구현합니다.
• MLOps 엔지니어 (MLOps Engineer): 생성된 코드와 데이터를 기반으로 자동화된 평가 샌드박스를 구축하고, 성능 임계값 (Medal thresholds) 을 설정합니다.
• 기술 작가 (Technical Writer): 생성된 메타데이터를 바탕으로 에이전트가 이해할 수 있는 자연어 작업 지시문 (Prompt) 을 작성합니다.

이 과정을 통해 실제 MLE-bench 작업의 평균 실행 시간을 **196 초에서 14 초 (약 13 배 감소)**로 단축합니다.

2.2. 밀집 보상 기반 궤적 RL (Dense Reward-based Trajectory RL)

합성 환경의 빠른 실행 속도를 활용하여 GRPO (Group Relative Policy Optimization) 기반의 온-폴리시 RL 을 적용합니다.

• 밀집 보상 (Dense Reward): 최종 결과뿐만 아니라 중간 단계의 성공 (코드 실행 성공, 포맷 준수, 단계별 성능 임계값 달성 등) 에 대해 보상을 부여하는 마일스톤 기반 보상 함수를 설계합니다. 이는 희소 보상 (Sparse Reward) 으로 인한 학습 불안정성을 해결합니다.
• 선택적 마스킹 (Selective Masking): 역전파 시 에이전트가 생성한 행동 (Action) 토큰에만 손실을 계산하고, 환경의 관찰 (Observation) 은 마스킹하여 효율적인 학습을 유도합니다.

3. 주요 기여 (Key Contributions)

1. 확장 가능한 합성 MLE 샌드박스: 소수의 시드 작업에서 다중 에이전트 시스템을 통해 검증 가능하고 다양한 합성 MLE 태스크를 대량 생성하는 프레임워크를 제안했습니다.
2. MLE 를 위한 궤적 기반 RL: 합성 환경을 통해 실행 지연 문제를 해결함으로써, MLE 도메인에서 대규모 온-폴리시 RL 훈련을 현실적으로 가능하게 했습니다.
3. 프레임워크 무관한 일반화: 훈련된 모델이 학습 시 사용된 ReAct 프레임워크뿐만 아니라 AIDE, AIRA 등 다양한 에이전트 구조에서도 우수한 성능을 발휘함을 입증했습니다.

4. 실험 결과 (Results)

실험은 MLE-bench-lite와 MLE-Dojo 벤치마크를 사용하여 Qwen3 시리즈 (8B, 14B, 30B) 모델을 대상으로 수행되었습니다.

• 성능 향상: SandMLE 로 훈련된 모델은 SFT 베이스라인 대비 Any Medal(동메달 이상 획득) 비율에서 **20.3% ~ 66.9%**의 상대적 개선을 보였습니다.
  • 예: Qwen3-8B 모델은 SFT 대비 66.9% 향상, 30B 모델은 100% 이상 향상 (Base 대비).
• 모델 크기 확장성: 모델 크기가 커질수록 RL 의 효과가 더욱 뚜렷해졌습니다. 특히 30B 모델은 SFT 초기화 없이도 100% 의 유효 제출률 (Valid Submission) 을 달성했습니다.
• 일반화 능력: 훈련 시 사용되지 않은 에이전트 프레임워크 (AIDE, AIRA, MLE-Agent) 에서도 뛰어난 적응력을 보였습니다.
  • MLE-Dojo 벤치마크에서 Qwen3-30B-SandMLE 는 MLE-Agent 프레임워크와 결합 시 HumanRank 점수 38.56을 기록하며, 기존 SFT 모델 (7.34) 보다 월등히 높은 성능을 보였습니다.
• 보상 설계의 중요성: 밀집 보상 (Dense Reward) 을 사용한 경우, 희소 보상 (Sparse Reward) 을 사용한 경우보다 성능이 현저히 높았으며, 특히 30B 모델의 경우 희소 보상 시 동메달 획득률이 27.3% 에서 13.6% 로 급감했습니다.

5. 의의 및 결론 (Significance)

이 논문은 합성 마이크로 스케일 환경이 실제 세계의 MLE 작업을 위한 효과적인 대리자 (Proxy) 가 될 수 있음을 증명했습니다. SandMLE 는 다음과 같은 의의를 가집니다:

1. RL 의 부활: MLE 와 같은 고비용 작업에서도 온-폴리시 강화 학습을 가능하게 하여, 에이전트가 단순한 모방 (SFT) 을 넘어 환경과의 직접적인 상호작용을 통해 복잡한 문제 해결 전략을 학습할 수 있는 길을 열었습니다.
2. 효율성: 데이터셋 크기를 줄임으로써 실행 시간을 획기적으로 단축하면서도 문제의 구조적 복잡성과 수학적 규칙은 유지하여, 효율적이고 강력한 훈련을 가능하게 했습니다.
3. 미래 지향성: 이 접근법은 에이전트의 내재적 추론 능력을 향상시키고, 다양한 도구와 프레임워크에 적응할 수 있는 범용적인 MLE 에이전트 개발의 새로운 패러다임을 제시합니다.

요약하자면, SandMLE 는 MLE 에이전트 훈련의 가장 큰 장애물이었던 '실행 시간'을 해결함으로써, 강화 학습을 통한 에이전트 능력의 비약적인 발전을 이끌었습니다.