💻 computer science

Aryabhata 2: Scaling Reinforcement Learning for Advanced STEM Reasoning

이 논문은 고품질 커리큘럼과 점진적 롤아웃 탐색을 활용하여 JEE 및 NEET 와 같은 경쟁적인 STEM 시험에서 우수한 성능과 더 높은 토큰 효율성을 달성하는 GPT-OSS-20B 기반의 강화 학습 훈련 언어 모델인 Aryabhata 2 를 소개합니다.

원저자: Ritvik Rastogi, Vishal Singh, Tejas Chaudhari, Sandeep Varma

게시일 2026-05-29

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ritvik Rastogi, Vishal Singh, Tejas Chaudhari, Sandeep Varma

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게 GPT-OSS-20B라는 이름의 천재적이지만 약간 산만한 학생이 있다고 상상해 보세요. 이 학생은 방대한 도서관의 책을 읽었고 많은 사실을 알고 있지만, 인도에서 가장 어려운 대학 입학 시험과 같은 까다롭고 다단계 수학 또는 과학 문제를 풀라고 요청하면 때로는 장황하게 말하거나 혼란스러워하거나 정답을 적는 데 너무 많은 시간을 보냅니다. 그들은 정답을 맞출 수도 있지만, 그렇게 하기 위해 많은 '잉크'(컴퓨터 토큰)를 사용하는데, 이는 비싸고 느립니다.

이 논문의 저자들은 더 큰 뇌(더 큰 모델)를 구매하지 않고도 이 학생을 챔피언 문제 해결사로 만들고 싶어 했습니다. 그들은 **아리아바타 2(Aryabhata 2)**라는 새로운 학생을 만들었습니다.

그들이 어떻게 했는지 간단한 비유를 통해 설명해 드리겠습니다:

1. 훈련장: 엄격한 코치

표준 훈련이 하는 것처럼 단순히 학생에게 더 많은 책을 읽게 하는 대신, 연구원들은 **강화 학습(Reinforcement Learning)**을 사용하는 엄격한 코치처럼 행동했습니다.

보상 시스템: 퍼즐을 정확하게 풀고 명확하게 설명할 때만 점수를 얻는 비디오 게임을 상상해 보세요. 학생이 정답을 틀리면 점수는 0 점입니다. 정답은 맞췄지만 한 단락으로 설명하면 될 것을 10 페이지 분량의 에세이를 작성하면 점수가 낮아집니다.
'심판': 코치는 학생의 말만 믿지 않았습니다. 그들은 모든 단계를 확인하기 위해 초지능적인 '심판'(다른 AI) 을 사용했습니다. 수학이 맞지 않거나 논리에 결함이 있으면 정답은 즉시 거부되었습니다.

2. 커리큘럼: 쉬운 것부터 불가능한 것까지

연구원들은 학생을 즉시 깊은 물속으로 던지지 않았습니다. 그들은 3 단계 훈련 캠프를 구축했습니다:

1 단계 (포맷팅): 먼저, 학생에게 펜을 올바르게 잡는 법을 가르쳤습니다. 정답이 깔끔하게 보이고 특정 구조를 따르도록 하는 데 중점을 두었습니다.
2 단계 (연마): 그런 다음, 학생에게 수천 개의 연습 문제를 주었습니다. 학생이 나아질수록 코치는 문제를 더 어렵게 만들었습니다. 학생이 계속 70% 를 맞췄다면, 코치는 훨씬 더 까다로운 질문으로 교체했습니다.
3 단계 (와일드카드): 마지막으로, 학생에게 동일한 어려운 문제를 해결하는 여러 가지 다른 방법을 동시에 시도하게 했습니다. 마치 학생에게 미로를 해결하기 위해 128 개의 서로 다른 경로를 동시에 시도해 보라고 요청하는 것과 같습니다. 이는 그들이 이전에 보지 못했던 교묘한 단축경을 발견하는 데 도움이 되었습니다.

3. 비밀 소스: "확장된 탐색"

일반적으로 AI 를 훈련할 때 한 번 문제를 풀도록 요청할 수 있습니다. 하지만 아리아바타 2 는 모든 단일 질문에 대해 많은 다른 시도를 생성하도록 훈련되었습니다.

비유: 범죄를 해결하는 형사를 생각해 보세요. 하나의 감에만 의존하는 대신, 형사는 단서를 찾기 위해 128 개의 서로 다른 팀을 파견합니다. 단 한 팀이라도 올바른 단서를 찾으면 사건은 해결됩니다. 이 방법은 모델이 정답에 이르는 '황금 경로'를 훨씬 더 빠르고 신뢰성 있게 찾도록 도왔습니다.

4. 결과: 더 빠르고, 더 똑똑하며, 더 날렵함

그들이 아리아바타 2 를 JEE 와 NEET 와 같은 실제 시험과 AIME 와 같은 초고난도 수학 대회에서 테스트했을 때, 결과는 인상적이었습니다:

더 높은 정확도: 이는 '부모' 모델 (GPT-OSS-20B) 보다 더 많은 문제를 정확하게 해결했으며, 훨씬 더 크고 비싼 일부 모델까지도 능가했습니다.
'토큰' 절감: 이것이 가장 큰 승리입니다. 원래 모델은 종종 길고 구불구불한 설명을 작성했습니다. 아리아바타 2 는 간결함을 배우게 되었습니다. 동일한 (또는 더 나은) 결과를 얻기 위해 최대 64% 적은 단어(토큰) 를 사용했습니다.
- 비유: 이전 모델이 완벽한 샷을 찾기 위해 100 장의 사진을 찍는 관광객이었다면, 아리아바타 2 는 한 번의 클릭으로 완벽한 샷을 찍는 프로 사진작가와 같습니다.

결론

이 논문은 매우 구체적이고 고품질의 연습 문제 세트와 스마트한 다단계 훈련 방법을 사용하여 표준 200 억 파라미터 AI 를 전문화된 STEM 추론 전문가로 변모시켰다고 주장합니다. 똑똑해지기 위해 거대한 모델이 될 필요는 없었습니다. 올바른 종류의 연습과 좋은 행동을 어떻게 보상해야 하는지 정확히 아는 코치만 필요했습니다.

그들이 주장하지 않은 것:
이 논문은 이 AI 가 인간 교사를 대체하거나 의학적 상태를 진단하거나 일반적인 대화에 사용될 수 있다고 주장하지 않습니다. 이는 경쟁 시험에서 발견되는 과학, 수학, 공학 문제를 해결하도록 특별히 설계되었으며, 이러한 특정 과목에 대한 매우 효율적인 튜터 역할을 합니다.

기술 요약: Aryabhata 2 – 고급 STEM 추론을 위한 강화 학습 확장

1. 문제 제기

경쟁적인 STEM 시험, 특히 인도의 합동 입학 시험 (JEE) 과 국가 자격 및 입학 시험 (NEET) 은 대규모 언어 모델 (LLM) 에게 독특한 과제를 제시합니다. 이러한 시험은 물리, 화학, 수학 전반에 걸쳐 다단계 기호 추론, 정밀한 수치 계산, 그리고 깊은 개념적 이해를 요구합니다. 표준 벤치마크와 달리, 이러한 문제들은 엄격한 제약 하에서 개념적 깊이를 테스트하도록 신중하게 구성됩니다.

최근 LLM 은 일반 추론 벤치마크에서 강력한 성능을 보이지만, 실제 교육 환경 배포에는 상당한 장애물에 직면해 있습니다:

확장성: 수백만 명의 학생들의 의문을 처리하기 위해 최첨단 모델을 대규모로 배포하는 것은 높은 추론 비용과 토큰당 가격으로 인해 비용이 너무 많이 듭니다.
효율성: 오픈소스 모델은 종종 대규모 파라미터 수를 필요로 하며 지나치게 긴 사고 체인을 생성하여 높은 지연 시간과 토큰 소모를 초래합니다.
도메인 특이성: 일반 모델은 경쟁 시험 전략에 필요한 특정 커리큘럼 범위와 단계별 정렬이 종종 부족합니다.

저자들은 고성능이지만 비싼 모델과 학생들의 의문에 정확하고 간결하며 구조화된 설명을 제공할 수 있는 효율적인 도메인 특화 솔루션 사이의 격차를 확인했습니다.

2. 방법론

본 논문은 강화 학습 (RL) 사후 훈련을 통해 훈련된 추론 중심 언어 모델인 Aryabhata 2를 소개합니다. 이 방법론은 제한된 컴퓨팅 자원 하에서 추론 능력을 극대화하도록 설계된 통합된 RL 전용 파이프라인을 기반으로 합니다.

2.1 기본 모델 및 아키텍처

기본 모델: 이 시스템은 36 억 개의 활성 파라미터를 가진 200 억 파라미터 오픈소스 전문가 혼합 (MoE) 모델인 GPT-OSS-20B를 기반으로 구축되었습니다.
파라미터 효율적 미세 조정 (PEFT): 메모리 제약 (두 개의 NVIDIA H100 NVL GPU 활용) 을 관리하기 위해 저랭크 적응 (LoRA) 을 사용합니다. 어댑터는 어텐션 프로젝션 레이어 ( $q, k, v, o$ ) 및 토큰 임베딩 레이어에 삽입됩니다. 이로 인해 학습 가능한 파라미터는 약 0.15%(209 억 개 중 3180 만 개) 로 줄어들어 메모리 풋프린트를 크게 줄이면서도 적응 능력을 유지합니다.

2.2 데이터 큐레이션 및 검증

훈련 코퍼스는 PhysicsWallah 의 내부 질문 은행에서 파생되었으며, 초기에는 178 만 개의 질문이 포함되어 있었습니다.

정제 파이프라인: 결정론적 파이프라인이 HTML 아티팩트 (특히 <img> 태그) 를 제거하고, pdflatex 를 사용하여 LaTeX 구문을 검증하며, LLM 분류기를 통해 불완전한 질문을 필터링하고, 비 STEM 콘텐츠를 제거합니다. 이를 통해 데이터셋이 125 만 개의 고품질 질문으로 축소됩니다.
답변 검증: 신뢰할 수 있는 보상 신호를 보장하기 위해 정책 모델(GPT-OSS-120B) 과 심판 모델(Qwen/Qwen3-30B-A3B-Thinking) 을 사용하는 자동화된 검증 파이프라인이 사용됩니다.
- 다중 패스 샘플링: 질문은 1 회, 4 회, 또는 16 회 샘플링 단계로 검증됩니다. 생성된 솔루션 중 어떤 것이 심판에 의해 정답으로 판정되면 질문이 수락되어 커리큘럼에 유효한 추론 궤적이 포함되도록 합니다.
커리큘럼 구성: 질문은 기본 모델의 생성 성공률에 따라 난이도 (사소함, 학습 가능함, 도전적) 로 분류됩니다. 초기 성능 불균형을 해결하기 위해 화학 질문이 오버샘플링됩니다.

2.3 강화 학습 프레임워크

훈련은 하드웨어 및 작업 제약을 준수하도록 특정 수정이 가해진 그룹 상대 정책 최적화 (GRPO) 를 활용합니다:

KL 프리 훈련: GPU 메모리를 절약하기 위해 참조 모델은 생략되며 KL 정규화는 제거됩니다.
목적 함수: 비대칭 상단 클리핑 임계값을 가진 DAPO 스타일의 클리핑 목적 함수가 사용됩니다.
보상 함수: 곱셈 보상 $R = R_{accuracy} \times R_{format}$ $R = R_{a cc u r a cy} \times R_{f or ma t}$ 이 사용됩니다:
- 정확도 보상 ( $R_{accuracy}$ ): 문자열 매칭, 수치 허용 오차 매칭, 그리고 (math-verify 사용) 기호 등가성을 통한 캐스케이드로 계산됩니다. 단일 정답 객관식 문제에서 라벨이 틀리더라도 정답 옵션 텍스트가 맞으면 부분 점수 (0.5) 가 부여됩니다.
- 형식 보상 ( $R_{format}$ ): 충분히 상세한 최종 답변 ( $S_{len}$ ) 과 추론 및 솔루션 길이 간의 균형 잡힌 비율 ( $S_{ratio}$ ) 을 장려하여 지나치게 간결하거나 과도하게 긴 출력을 방지합니다.

2.4 훈련 단계

훈련은 안정성과 탐색을 균형 있게 맞추기 위해 세 단계의 순차적 단계로 진행됩니다:

1 단계 (형식 정렬): 그룹 크기 8 로 300 스텝. 사소하고 학습 가능한 질문의 혼합을 사용하여 모델이 원하는 출력 구조에 정렬되도록 집중합니다.
2 단계 (연장된 RL): 약 5,000 스텝. 그룹 크기가 8 에서 16 으로 증가합니다. 정확도가 향상됨에 따라 커리큘럼이 적응적으로 더 어려운 질문으로 이동합니다. 보상이 정체될 때 훈련을 안정화하기 위해 EMA 기반 체크포인트 병합이 사용됩니다.
3 단계 (확장된 RL): 약 700 스텝. 그룹 크기가 크게 확장 (64 에서 128) 되어 추론 궤적의 탐색을 넓히고 대체 솔루션 전략을 발견합니다.

3. 주요 기여

Aryabhata 2 모델: 경쟁적인 STEM 추론을 위해 특별히 최적화된 200 억 파라미터 모델로, 도메인 특화 커리큘럼에 대한 표적 강화 학습이 훨씬 더 큰 기본 모델보다 우수한 성과를 낼 수 있음을 보여줍니다.
통합 RL 파이프라인: 엄격한 데이터 검증, 난이도 인식 커리큘럼 구성, 그리고 제한된 하드웨어 (2x H100) 에서 효과적으로 작동하는 위상별 RL(연장된 최적화 + 확장된 탐색) 을 결합한 새로운 훈련 전략.
효율성 향상: 이 모델은 오픈소스 모델 간에 최첨단 성능을 달성하면서도 기본 모델보다 훨씬 적은 토큰 (최대 64% 감소) 을 생성하여 대규모 교육 배포의 비용 장벽을 해결합니다.
종합 평가: 분포 내 경쟁 시험 (JEE Main/Advanced, NEET) 과 분포 외 추론 벤치마크 (AIME, HMMT, MMLU-Pro, GPQA) 전반에 걸친 검증.

4. 결과

4.1 분포 내 성능

경쟁 시험 벤치마크 (JEE Main 2025/2026, JEE Advanced 2025, NEET 2025) 에서 Aryabhata 2 는 평균 **Pass@1 88.95%**를 달성하여 기본 GPT-OSS-20B(83.00%) 를 능가하고 GPT-OSS-120B(88.28%) 및 Qwen3-30B-A3B(88.55%) 와 같은 훨씬 더 큰 모델과 경쟁합니다.

토큰 효율성: Aryabhata 2 는 GPT-OSS-20B 대비 출력 토큰을 52~64% 감소시킵니다.
토큰당 정확도: GPT-OSS-20B(15.68) 및 GPT-OSS-120B(26.66) 보다 훨씬 높은 42.31 Acc./1K 토큰을 달성합니다.

4.2 분포 외 성능

일반 추론 벤치마크 (AIME, HMMT, GPQA, MMLU-Pro, MMLU-Redux 2.0) 에서 Aryabhata 2 는 평균 **Pass@1 87.64%**를 달성합니다.

GPT-OSS-20B(84.95%) 및 Nemotron 3 Nano(83.48%) 를 상회합니다.
HMMT(+1.54, 기본 대비) 및 GPQA(+4.35, 기본 대비) 에서 상당한 향상을 보입니다.
최첨단 독점 모델 (예: GPT-5 Mini, Gemini 2.5 Flash) 보다 낮지만, 200 억 파라미터 모델로서 강력한 일반화를 보여줍니다.

4.3 효율성 비교

모든 분포 외 (OOD) 벤치마크에서 Aryabhata 2 는 GPT-OSS-20B 대비 출력 토큰을 24~71% 감소시켜 39.58 Acc./1K 토큰의 OOD 정확도를 달성합니다.

5. 중요성 및 주장

본 논문은 도메인 특화 커리큘럼에 대한 표적 강화 학습이 실용적인 STEM 추론 시스템을 확장하는 효과적인 전략이라고 주장합니다. 이 작업의 중요성은 다음을 입증하는 데 있습니다:

규모보다 전문화: 작고 효율적으로 훈련된 모델이 특정 고위험 추론 작업에서 더 큰 기본 모델보다 우수한 성과를 낼 수 있습니다.
배포 실현 가능성: 엄격한 데이터 큐레이션과 위상별 RL 을 결합함으로써 제한된 하드웨어 (2x H100) 에서 고품질 추론 모델을 훈련할 수 있어 수백만 명의 사용자를 가진 대규모 교육 시스템을 위한 고급 AI 튜터링이 실현 가능합니다.
토큰 효율성: 이 접근 방식은 정확도를 향상시킬 뿐만 아니라 더 간결하고 구조화된 추론 체인을 생성하여 추론 비용을 극적으로 줄여 교육 분야에 LLM 을 배포하는 경제적 장벽을 직접적으로 해결합니다.

저자들은 컴퓨팅 제약 하에서도 안정적인 훈련을 가능하게 하면서 강력한 추론 품질을 유지하는 파이프라인을 개발하여 실제 교육 워크로드를 위한 실현 가능한 길을 제시했다고 결론지었습니다.