TREX: 에이전트 기반 트리 탐색을 통한 LLM 파인튜닝 자동화

LLM 학습은 데이터 준비부터 하이퍼파라미터 튜닝까지 복잡한 수동 작업이 필요하지만, TREX는 이를 자율 에이전트 시스템으로 자동화했습니다. Monte Carlo Tree Search를 활용해 최적의 학습 전략을 스스로 탐색하며, 인간 전문가 수준의 성능 향상을 달성할 수 있음을 입증했습니다.

기존의 AI 에이전트는 코드 수정이나 파라미터 최적화 같은 단일 작업에는 능숙하지만, LLM 학습처럼 방대한 데이터와 긴 실행 시간이 필요한 복잡한 워크플로를 다루기에는 한계가 있었다. 특히 학습 데이터의 분포나 하이퍼파라미터 조합이 결과에 결정적인 영향을 미치는데, 이를 에이전트의 제한된 컨텍스트 윈도우 내에서 모두 처리하는 것은 불가능에 가깝다.

TREX는 이 문제를 해결하기 위해 인간 연구자의 반복적인 실험 과정을 '트리 탐색' 구조로 변환했다. Monte Carlo Tree Search(MCTS)의 원리를 적용하여, 에이전트는 현재까지의 실험 결과(보상)를 바탕으로 어떤 방향의 실험이 가장 유망할지 판단한다. 이는 Gradient Descent가 손실 함수를 줄이는 방향으로 가중치를 갱신하듯, MCTS가 전체 실험 공간에서 최적의 '학습 레시피'를 찾아가는 과정과 유사하다.

결과적으로 TREX는 단순히 명령을 수행하는 도구를 넘어, 이전 실험의 실패 원인을 분석하고(Bad-case Analysis), 이를 바탕으로 다음 실험의 데이터 믹싱 비율이나 학습률을 스스로 조정한다. 이러한 자율적 피드백 루프를 통해 제한된 컴퓨팅 자원 내에서도 인간 전문가가 설계한 것과 대등하거나 이를 능가하는 파인튜닝 성과를 낼 수 있게 된다.

TREX는 Researcher와 Executor라는 두 가지 핵심 모듈로 구성된 듀얼 루프 워크플로를 채택했다. Researcher는 상위 수준의 실험 계획을 수립하고 결과를 진단하며, Executor는 GPU 클러스터와 통합되어 실제 데이터 처리 및 모델 학습 코드를 실행한다. 내부 루프에서는 두 에이전트가 통신하며 단일 라운드 실험을 완료하고, 외부 루프에서는 MCTS 정책에 따라 실험 노드를 확장하며 전체 탐색 트리를 구축한다.

전략 탐색에는 UCT(Upper Confidence Bound for Trees) 공식을 사용한다. UCT(v) = Qv/Nv + c * sqrt(ln Np(v) / Nv) 계산식을 통해, 현재까지 성적이 좋은 노드(Exploitation)와 아직 충분히 시도되지 않은 노드(Exploration) 사이의 균형을 맞춘다. [노드의 누적 보상과 방문 횟수를 입력으로] → [상수 c를 이용해 탐색 가중치를 조절하는 연산을 수행해] → [최종 UCT 점수를 얻고] → [이 점수가 가장 높은 노드를 다음 실험 대상으로 선택한다].

데이터 처리를 위해 도입된 AIDP 라이브러리는 HuggingFace Datasets 기반의 결정론적 연산자들을 제공한다. 에이전트는 'select_by_score'나 'generate_dataset_with_llm' 같은 함수를 호출하여 복잡한 데이터 파이프라인을 구성한다. 또한, 실험 진단 단계에서는 실패 사례 분석(Bad-case Analysis)을 수행하여 성능 변화의 결정적 요인을 식별하고, 이를 요약된 메모리 컨텍스트로 변환하여 다음 실험 설계에 반영한다.

FT-Bench의 10개 과제 전반에서 TREX는 베이스 모델(Qwen3-1.7B) 대비 일관된 성능 향상을 보였다. 특히 분자 생성(TOMG-Bench) 과제에서는 인간 전문가가 설계한 Llama3.1-8B 기반 모델의 성능 향상폭(+0.189)보다 훨씬 큰 +0.498의 성능 이득을 기록하며 전문가 수준을 능가하는 최적화 능력을 입증했다.

Ablation Study 결과, MCTS 전략은 Greedy Best-First Search나 순차적 확장 방식보다 실험 과정의 변동성이 적고 안정적인 성능 향상을 이끌어냈다. 또한 AIDP 라이브러리를 사용하지 않았을 때보다 사용했을 때 데이터 처리 오류가 줄어들어 실험 성공률이 비약적으로 상승했으며, Bad-case Analysis를 포함한 진단 과정이 최종 성능 도달에 핵심적인 역할을 했음이 확인됐다.

TREX의 아키텍처는 에이전트의 추론 능력과 실제 실행 환경 간의 간극을 메우는 데 집중한다. Researcher는 Gemini 3 Pro 또는 Qwen3-Next-80B-Thinking과 같은 강력한 LLM을 백엔드로 사용하여 고도의 전략을 수립하며, Executor는 Claude 4.5 Sonnet을 통해 안정적인 코드 생성을 보장한다. 실험 이력은 'Condense' 과정을 거쳐 메모리 컨텍스트로 관리되며, 이는 궤적 정보, 형제 노드 정보, 그리고 성능 향상이 뚜렷했던 핵심 노드 정보를 포함한다.

성능 평가를 위해 도입된 '정규화된 상대적 성능 이득(GT)' 지표는 베이스 모델과 강력한 참조 모델(Qwen3-235B) 사이의 간극을 기준으로 성능 향상을 측정한다. GT = (E_FT - E_Base) / (E_Ref - E_Base) 공식을 통해 [파인튜닝 전후 점수와 참조 모델 점수를 입력으로] → [차이값의 비율을 계산하여] → [0에서 1 사이의 정규화된 값을 얻고] → [서로 다른 난이도의 과제 간 성능 향상도를 공정하게 비교한다].

논문은 TREX가 여전히 상당한 계산 비용을 수반하며, 특히 데이터셋을 처음부터 구축해야 하는 과제(CS-Bench, GTA 등)에서는 더 많은 실험 라운드와 깊이 있는 연구가 필요하다는 점을 한계로 언급했다.