Agentic Discovery of Neural Architectures: AIRA-Compose and AIRA-Design

이 논문은 사람의 전문 지식과 직관에 기대던 언어 모델 구조 설계를, LLM agent가 primitive(모델 구조를 조립하는 기본 계산 block) 배치를 찾고 attention 코드와 학습 script를 직접 쓰는 task 12개로 다뤘다.

이 논문은 한마디로 LLM agent에게 언어 모델 구조 설계를 맡긴 실험 framework다. AIRA-Compose에서는 agent가 MLP(multi-layer perceptron, 층을 쌓은 기본 신경망 block), attention, Mamba(state space model 계열 block)를 16층으로 배치해 14개 구조를 찾았고, 이를 350M, 1B, 3B 크기로 늘려 Llama 3.2 같은 baseline과 비교했다. AIRA-Design에서는 agent가 LRA(Long Range Arena, 긴 입력의 먼 거리 의존성을 다루는 능력을 재는 benchmark)용 attention 코드와 Autoresearch(GPT 언어 모델 학습 script를 고쳐 5분 학습 예산 안에서 손실을 낮추는 benchmark)용 학습 script를 처음부터 썼다. 저자들은 LRA에서 agent가 만든 구조가 대부분 기존 아이디어를 다시 조합한 것이어서, 지금 agent는 공학 수준의 조합에 강하고 새 이론을 내는 데는 못 미친다고 보았다.

핵심 요약

• 문제: 지금 LLM은 대부분 attention과 MLP를 1:1로 번갈아 쌓은 Transformer이고, Mamba 같은 다른 primitive를 섞은 hybrid 구조로 가면 설계 공간이 크고 조합적으로 늘어난다. 저자들은 손으로만 탐색하면 성능 좋은 비직관적 배치나 새 primitive를 놓칠 수 있다고 본다. 그래서 이 탐색을 연구 agent에게 맡기고, 이를 recursive self-improvement(agent가 자기를 움직이는 신경망 구조를 스스로 찾고 고치는 것)로 가는 한 단계로 삼는다.
• 입력과 출력: 이 논문은 AIRS-Bench(agent의 독립적인 연구 능력을 재는 task 표준) 형식의 task 12개를 만들고 그 위에서 실험한 framework 논문이다. agent는 task 설명과 데이터를 받아 AIRA-Compose에서는 primitive 16개를 이은 구조 문자열(submission.csv)을, AIRA-Design에서는 LRA용 model.py나 Autoresearch용 train.py를 낸다. 따로 떨어진 평가 script가 이 결과물을 처음부터 학습시켜 agent가 보지 못한 test split이나 고정된 validation shard로 점수를 매긴다.
• 방법: agent는 LLM과 scaffold(탐색 방식)를 합친 것이고, AIRA-dojo(tree 탐색으로 코드 해답을 발전시키는 harness) 안에서 Draft, Debug, Improve, Analyze 네 operator로 해답을 고친다. greedy scaffold는 초기 해답 5개를 만든 뒤 validation 점수가 가장 높은 해답에 Improve를 거듭한다. AIRA-Compose는 Composer(Bayesian optimization으로 hybrid 구조를 찾는 기존 framework)의 탐색과 평가 단계를 agent task로 바꾸고, 뒤의 aggregation과 extrapolation 단계에서 찾은 16층 구조를 stretching(연속된 block을 비율대로 늘림)이나 stacking(구조 전체를 되풀이해 쌓음)으로 350M, 1B, 3B 크기까지 키운다.
• 구현: agent마다 H200 GPU 1장을 주고 run 하나는 최대 24시간 또는 500 step까지 실행하며, 데이터가 큰 BabiStories와 DCLM(DataComp-LM, 언어 모델 학습 dataset) 탐색은 run마다 60시간을 준다. AIRA-Compose는 6개 LLM을 쓰고, AIRA-Design의 LRA 실험은 Opus 4.6과 Gemini 3 Pro를 포함한 12개 LLM으로 agent 20개를, Autoresearch 실험은 그 가운데 7개 LLM으로 greedy agent를 만들었다. 찾은 구조를 큰 크기로 학습할 때 350M은 H200 8장, 1B와 3B는 H200 16장을 썼다.
• 결과: MLP와 attention 두 primitive로 만든 구조를 1B 크기에서 37.5B token으로 학습한 비교에서 AIRAformer-D(Stretched)가 validation loss 2.734와 평균 0-shot(예시 없이 푼) 정확도 59.7%로 Llama 3.2와 Composer가 찾은 모델을 앞섰고, Mamba를 더한 3-primitive 비교에서는 AIRAhybrid-D(Stretched)가 seed 하나로 validation loss 2.719를 냈다. LRA에서 agent가 설계한 모델의 최고 정확도는 사람이 설계한 SOTA(state of the art, 지금까지 보고된 최고 성능)보다 document matching에서 2.3pp(퍼센트포인트, SOTA 82%), text classification에서 2.6pp(SOTA 91%) 낮았다. Autoresearch에서는 문헌을 받은 Greedy Opus 4.5가 가장 좋은 seed에서 validation BPB(bits per byte, 낮을수록 좋은 언어 모델 손실 지표) 0.968을 내, 원 Autoresearch 최고값 0.9775보다 낮았다.
• 한계: 저자들은 AIRA-Compose가 작은 proxy 평가에 기대는데 이 평가가 큰 규모 성능과 어긋날 때가 있고, aggregation과 extrapolation은 아직 agent 바깥의 정해진 절차로 남아 있다고 밝혔다. LRA에서 agent가 만든 구조는 Performer, Longformer, Conformer 같은 기존 아이디어를 대부분 다시 조합한 것이었다. Autoresearch 실험은 원 실험(H100)과 달리 H200 GPU와 인터넷이 끊긴 sandbox에서 실행됐고, 저자들은 이 차이로 같은 train.py의 시작 BPB가 원 실험의 0.998보다 높은 1.012가 됐을 수 있다고 밝혔다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.