4B 세계 모델이 ML 실험 결과를 예측해 에이전트 탐색을 늘린다

무엇인가

AI 연구 에이전트는 프로그램을 쓰고 실행하고 터미널 출력을 읽은 뒤 고치는 과정을 반복한다. 문제는 각 후보 프로그램을 실행할 때마다 모델 학습과 평가가 따라붙는다는 점이다. 그래서 제한된 시간 안에 시험해 볼 수 있는 아이디어 수가 실행 시간에 의해 결정된다. 기존 연구는 유망하지 않은 가지를 잘라내거나, 실행 전에 후보 순위를 매기거나, 성능을 예측하는 방식으로 비용을 줄이려 했다. 그러나 실행 1회당 비용 자체는 그대로이고, 오류 트레이스나 학습 로그처럼 이후 디버깅과 수정에 필요한 터미널 출력을 제공하지 못한다는 한계가 있다.

어떻게 동작하나

이 논문이 제안하는 MLE-World는 Qwen3.5-4B를 지도 파인튜닝한 실행 세계 모델이다. 입력은 프로그램 본문, 설치된 패키지와 버전, AIDE가 만드는 데이터 개요(파일 트리, CSV 차원, 컬럼 요약과 예시 값), 시간 제한으로 구성된다. 출력은 네 개 필드의 구조화된 레코드인데, 예외 타입(82개 클래스 중 하나로 첫 토큰에서 디코딩), 터미널 출력, 검증 메트릭, 제출 플래그 순서로 생성된다. 터미널 출력만 예측하면 자유 텍스트에서 크래시 여부를 추론해야 하고 AIDE는 별도 피드백 모델로 점수를 뽑아야 하므로, 예외 타입과 메트릭을 독립 필드로 분리했다. 그 결과 터미널 출력 일치율이 59.4%에서 62.8%로, 메트릭 재현율이 84.3%에서 89.5%로 올랐다. 학습 데이터는 4B에서 122B까지 다섯 개 에이전트 모델로 17개 MLE-bench 태스크에서 AIDE를 돌려 모은 497,056개 실행 로그이며, 실패한 에피소드도 버리지 않고 포함했다.

무엇과 다른가

신뢰도 점수는 첫 번째로 디코딩되는 예외 클래스 토큰의 확률 분포에서 나온다. 82개 클래스에 대해 재정규화한 뒤 최댓값을 p_max, 나머지 확률 질량을 불확실성 u=1-p_max, 정상 실행 확률을 p_clean으로 정의한다. 이 예측을 이용해 다섯 가지 테스트 타임 스케일링 정책을 만든다. 한 스텝에서 후보 N개를 뽑아 예측 출력을 비교하고 그중 하나만 실제로 실행하는 방식으로, Best-of-N은 예측 검증 점수가 가장 높은 것을 고르고, Sequential revision은 직전 후보의 예측 오류·로그·점수를 읽어 다음 후보를 수정하며, Pairwise는 예측 터미널 출력을 붙인 두 후보를 LLM 심판이 비교하고, Fusion은 후보 풀 전체와 예측 출력에서 새 프로그램을 합성한다. UCB는 예측 품질에 불확실성 보너스를 더해 v(s)+βu(x_s)를 최대화하는데, 탐욕적 선택이 과소평가된 유망 후보를 버리는 문제를 추가 추론 비용 없이 완화하려는 것이다. 여기에 신뢰도 게이트를 붙여 p_max가 임계값보다 낮을 때만 실제 실행을 요청한다. 성공 예측에는 γ_s=0.50, 실패 예측에는 정책에 따라 γ_f=0.56 또는 0.83을 쓰고, 타임아웃으로 예측된 프로그램은 아예 실행하지 않는다. 예측 검증 점수가 최고 기록을 갱신하는 후보는 신뢰도와 무관하게 항상 실제로 돌린다.

어떻게 쓰나

실험은 22개 중 17개 MLE-bench Lite 태스크에서 이뤄졌다. 태스크당 100개씩 총 1,700개 held-out 프로그램에서 MLE-World는 터미널 출력의 77.9%를 실제 실행과 일치시켰고, 가장 강한 프롬프트 베이스라인인 DeepSeek-V4-Flash는 41.1%였다. 예외 클래스 86.4%, 검증 메트릭 78.1%로 다른 필드에서도 앞섰으며, GPU 추론 시간은 프롬프트 모델보다 13~1,061배 적었다. 검색 성능에서는 Qwen3.5-9B가 코드를 쓰고 태스크당 3개 시드로 돌린 결과, UCB가 평균 리더보드 percentile 52.1로 세계 모델이 없는 Base AIDE의 45.4를 넘었고 다섯 정책 모두 관측 평균이 더 높았다. 최대 22.1% 향상이라는 수치는 여기서 나온다. 같은 150개 시작 상태에서 비교한 실험에서는 예측 터미널 출력을 붙였을 때 Fusion이 14.0에서 20.6으로, Sequential이 10.4에서 15.3으로 올랐다. 게이트를 켜면 실제 실행 시간이 30~40% 줄면서 성능은 대체로 유지됐고, 무작위로 실행을 건너뛰는 대조군보다 19% 더 절감했다.

전제와 한계

개발자 입장에서 이 논문의 핵심은 실행 전 스크리닝을 언어 모델의 다음 토큰 예측 문제로 바꿨다는 것이다. 학습·평가 파이프라인이 느리고 실행 로그가 이미 쌓여 있는 팀이라면, 후보 프로그램을 전부 돌리는 대신 예측으로 순위를 매기고 신뢰도가 낮은 것만 실제로 돌리는 구조를 시도할 수 있다. 다만 예측이 틀릴 수 있다는 전제가 설계 전체에 깔려 있다. held-out 기준으로 터미널 출력의 22%, 예외 클래스의 14%가 실제와 어긋났고, 저자들은 예측 오류가 검색을 잘못된 방향으로 이끌 수 있다고 명시한다. 따라서 게이트 임계값과 실패·성공 예측에 대한 비대칭 비용 설정을 자체 파이프라인 비용 구조에 맞춰 확인해야 한다.

저자가 밝힌 한계도 분명하다. 주 평가는 파인튜닝에서 본 17개 태스크에 대한 분포 내 평가이며, 학습에 약 50만 건의 과거 실행 로그가 필요하다. 파인튜닝에서 제외한 3개 태스크에서는 추론 없이도 53.3%로 DeepSeek(추론 사용, 42.1%)와 Coder(21.1%)를 앞섰지만 격차는 줄어든다. 타임아웃은 가장 부정확하게 예측되는 결과인데, 타임아웃을 놓치면 실행 한 번을 통째로 날린다. 최종 점수는 결국 실제 실행이 결정하며, 학습은 NVIDIA H200 8대로 수행됐다.