AutoResearch-RL: 자율적 신경망 아키텍처 탐색을 위한 영구적 자기 평가 강화학습 에이전트

기존의 AI 모델 개발은 연구자가 가설을 세우고 코드를 수정하는 반복적인 수동 작업에 의존했다. 이 논문은 LLM 기반 에이전트가 스스로 코드를 편집하고 실험 결과를 바탕으로 전략을 개선하는 강화학습 프레임워크를 제안하여 연구 자동화의 새로운 가능성을 제시한다.

전통적인 신경망 아키텍처 탐색(NAS)은 미리 정의된 탐색 공간 안에서 최적의 조합을 찾는 방식이었다. 하지만 실제 연구는 손실 함수 변경, 최적화 도구 수정 등 코드 전반을 건드리는 개방형 작업이며, 이는 고정된 문법 기반의 탐색으로는 한계가 명확했다. AutoResearch-RL은 LLM을 강화학습의 에이전트로 활용하여 train.py 파일 자체를 행동 공간(Action Space)으로 취급한다. 에이전트는 코드를 직접 수정하고, 5분이라는 제한된 시간 동안 모델을 학습시킨 뒤 얻은 검증 성능(val-bpb)을 보상으로 받아 자신의 정책을 업데이트한다. 특히 과거의 실험 이력을 슬라이딩 윈도우 방식으로 기억에 저장하여 어떤 수정이 효과적이었는지 맥락을 파악한다. 이를 통해 에이전트는 단순한 파라미터 조정을 넘어 Muon 최적화 도구 스케일링이나 QK-norm 삽입과 같은 고차원적인 연구 휴리스틱을 스스로 터득하게 된다.

연구 과정을 MDP(S, A, T, R, γ)로 모델링한다. 상태(State)는 현재 소스 코드, 실험 이력, 시스템 진단 정보를 포함하며, 행동(Action)은 코드에 적용되는 구조화된 수정(insert/replace/delete)이다. 보상 함수 R은 검증 비트당 비트(val-bpb)의 개선량과 계산 효율성 보너스의 합으로 정의된다. [이전 bpb와 현재 bpb의 차이를 입력으로] → [개선된 만큼 양의 값을 부여하고 효율성 계수를 곱해 연산하여] → [최종 보상 숫자를 얻고] → [이 값이 클수록 에이전트가 해당 코드 수정을 긍정적으로 평가하게 된다]. 자기 평가(Self-Evaluation) 모듈은 파워 로우(Power-law) 모델을 사용하여 학습 곡선을 예측한다. [30초마다 관측된 손실 값을 입력으로] → [L(t) = a·t⁻ᵇ + c 수식에 피팅하여 최종 성능을 예측하고] → [예측값이 기존 최고 성능보다 낮으면 중단 신호를 출력하여] → [불필요한 연산 자원 낭비를 방지한다].

nanochat 프리트레이닝 벤치마크에서 AutoResearch-RL은 약 8시간(101회 실험) 만에 val-bpb 2.681을 달성했다. 이는 사람이 튜닝한 베이스라인(2.847)과 강화학습 없이 LLM만 사용한 경우(2.734)보다 우수한 수치이다. 에이전트는 Muon 최적화 도구의 학습률을 2.8×10⁻³으로 조정하고 AdamW 가중치 감쇠를 0.04로 줄이는 등 정교한 최적화 전략을 발견했다. 또한 쿼리와 키에 ℓ₂ 정규화를 적용하는 QK-norm 기법을 스스로 도입하여 안정성을 높였다. 자기 평가 모듈을 적용했을 때 유망하지 않은 실험의 54.3%를 조기에 중단시켰다. 이를 통해 시간당 실험 횟수가 1.35배 증가했으며, 정책 개선 효과가 누적되면서 최종적으로는 2.4배의 샘플 효율성 향상을 기록했다.

에이전트 정책은 Claude 3.5 Sonnet 모델을 LoRA로 파인튜닝하여 구현되었다. PPO 알고리즘을 사용하여 코드 수정 분포를 최적화하며, 상태 표현을 위해 64,000 토큰의 긴 컨텍스트 윈도우를 활용한다. 실험 이력 관리를 위해 최근 32개의 실험 결과와 역대 최고 성능(Best-ever) 요약을 포함하는 슬라이딩 윈도우 메커니즘을 사용한다. 이는 컨텍스트 길이를 제한하면서도 장기적인 성능 개선 흐름을 유지하게 한다. 탐색과 활용(Exploration vs Exploitation)의 균형을 위해 엔트로피 정규화와 함께 ϵ-novelty 보상을 도입했다. 기존에 시도된 코드 수정과 편집 거리(Edit Distance)가 먼 새로운 시도에 추가 보상을 부여하여 탐색 범위를 넓힌다. 안전성을 위해 수정 가능한 범위를 train.py 파일 하나로 제한하고 네트워크 접근을 차단하며 엄격한 실행 시간 제한을 두어 무한 루프나 자원 고갈을 방지한다.

현재 시스템은 단일 GPU와 고정된 데이터셋 환경으로 제한된다. 멀티 GPU 환경으로 확장하기 위해서는 노드 간 실험 실행 및 평가를 조율하는 복잡한 엔지니어링이 필요하며, 현재는 어휘집(Vocabulary)이나 데이터 파이프라인은 수정하지 못하는 한계가 있다.