Aviary: training language agents on challenging scientific tasks

Motivation

• Known: 기존 연구는 language agent를 CoALA와 같은 개념적 프레임워크나 POMDP로 이론화했으며, TextGrad, DSPy, GPTSwarm 등 다양한 프레임워크가 prompt·tool usage·LLM weight 등 agent 구성요소의 최적화를 시도해왔다. 또한 GSM8K, HOTPOTQA 등 범용 벤치마크와 DiscoveryBench, ChemBench 같은 과학 특화 벤치마크가 존재한다.
• Gap: 그러나 language agent를 학습 문제로서 통일적으로 정의하는 이론적 틀이 부재하고, 특히 multi-step reasoning과 tool usage가 요구되는 도전적인 과학 연구 과제(분자 클로닝, 문헌 QA, 단백질 설계)를 대상으로 agent를 학습·평가할 수 있는 표준화된 environment 및 optimizer 프레임워크가 없었다.
• Why: 과학 연구 과정 자체가 분석·도구 사용·실험의 다단계 사이클을 요구하기 때문에, 이를 자동화할 수 있는 language agent를 저비용의 open-source LLM으로 frontier 모델 수준까지 학습시킬 수 있다면 과학 연구의 자동화 및 비용 효율화에 실질적 기여를 할 수 있다.
• Approach: language agent 문제를 language decision process(LDP)로 공식화하고 이를 stochastic computation graph로 표현하는 LDP 프레임워크와, 이를 기반으로 다섯 가지 환경(그 중 세 가지 과학 환경 포함)을 구현한 Aviary gymnasium을 제안하여, expert iteration 기반 online training과 inference-time compute scaling으로 non-frontier LLM agent를 학습시킨다.

Achievement

Training language agents to solve (A) SeqQA tasks using the molecular cloning environment and (B) LitQA2

1. LDP 이론적 정식화: language agent를 language-grounded POMDP로 정의하고 stochastic computation graph로 표현하여 기존 agent 아키텍처(예: ReAct 등)들을 통합적으로 포괄할 수 있는 프레임워크를 제시했다.
2. Aviary gymnasium 구현: GSM8K, HOTPOTQA를 포함해 DNA construct 조작(분자 클로닝), 문헌 기반 QA(LitQA2 기반), 단백질 안정성 설계 등 5개 환경을 parametrizable tool 기반으로 구현했다.
3. 저비용 고성능 agent 달성: Llama-3.1-8B-Instruct와 같은 소형 open-source LLM을 expert iteration으로 학습하고 majority vote sampling을 적용해, DNA construct 및 문헌 QA 과제에서 frontier LLM과 인간 전문가 성능을 최대 100배 낮은 inference 비용으로 능가하거나 맞먹는 성과를 보였다.
4. 오픈소스 공개: Aviary 환경 프레임워크와 LDP 계산 그래프·학습 코드를 모두 공개하여 재현성과 확장성을 확보했다.

How

Training language agents to solve (A) SeqQA tasks using the molecular cloning environment and (B) LitQA2

• language agent 구성요소(LLM weight, prompt, memory, sampling temperature 등)를 노드로 하는 stochastic computation graph 프레임워크(LDP 패키지)를 설계
• 다섯 개의 gym 환경을 Aviary로 구현: GSM8K, HOTPOTQA를 환경으로 재구성하고, DNA construct 조작(분자 클로닝 질문 응답, SeqQA), 과학 문헌 QA(LitQA2 기반), 단백질 안정성 mutation 제안 환경을 신규 구축
• expert iteration(EI)을 통한 online training: agent가 생성한 trajectory 중 고보상 trajectory에 대해 supervised fine-tuning을 반복 수행
• inference-time compute scaling: majority vote sampling 등을 적용해 성능 향상
• 학습된 Llama-3.1-8B-Instruct 기반 agent 성능을 frontier LLM(예: GPT-4급) agent 및 인간 전문가 성능과 비교 평가

Originality

• language agent를 language decision process(LDP)라는 명시적 이론적 대상으로 정의하고, stochastic computation graph라는 통일된 관점에서 다양한 기존 agent 아키텍처를 포괄한 점
• 단순 QA/추론 벤치마크를 넘어 분자 클로닝, 단백질 엔지니어링 등 실제 생물학 연구에 밀접한 다단계 과학 task를 environment로 정식 구현한 점
• expert iteration과 inference-time scaling을 결합해 소형 open-source LLM만으로 frontier LLM 및 인간 전문가를 능가하는 비용 효율적 접근을 실증한 점
• 환경(Aviary)과 학습 프레임워크(LDP)를 분리·모듈화하여 재사용 가능한 오픈소스 소프트웨어로 공개한 점

Limitation & Further Study

• 발췌된 내용상 세 과학 환경(특히 단백질 안정성 설계)에 대한 정량적 성능 비교 및 실패 사례 분석이 충분히 제시되지 않아, 일반화 가능성 검증이 제한적일 수 있음
• expert iteration 기반 online training은 self-generated trajectory에 의존하므로, 초기 정책 품질이 낮은 경우 학습이 정체되거나 편향된 trajectory에 과적합될 위험이 있음
• 소형 LLM이 frontier 모델을 능가한 결과가 특정 task(DNA construct, 문헌 QA)에 국한되어 있어, 더 복잡하거나 open-ended한 과학적 추론 과제로의 확장성은 추가 검증이 필요
• 후속 연구로 다양한 도메인(화학, 물리 등)으로의 environment 확장과, process supervision 등 더 정교한 optimization 기법과의 비교가 요구됨

Evaluation

총평: language agent 학습을 위한 이론적 정식화와 실용적 오픈소스 프레임워크를 함께 제시하고, 실제 과학 연구 과제에서 저비용 소형 모델의 우수성을 실증했다는 점에서 학술적·실용적 가치가 높은 연구이다.