목차
개요
AIMO2(AI Mathematical Olympiad 2) Kaggle 대회가 종료했다. 이번 대회는 LLM(대규모 언어 모델)을 활용한 수학 문제 해결 능력을 겨루는 장으로, 많은 참가자들이 L4 GPU 지원 하에 다양한 규모의 모델과 프롬프트 엔지니어링, 샘플링 파라미터 최적화 등을 시도했다.
이 글에서는 저의 대회 참가 경험과 함께, 압도적인 성능으로 1위를 차지한 NVIDIA NemoSkills 팀의 솔루션을 TIR(Tool-Integrated Reasoning), GenSelect(Generative Solution Selection), OpenMathReasoning 데이터셋 구축 중심으로 심층 분석하고자 한다.
대회 종료 후기
이전의 대회와 다르게 L4*4의 하드웨어를 지원하며, 더욱 큰 규모의 모델과 더 다양한 실험을 권장했다. 따라서, 실질적으로 효과가 있는 7B이상의 모델을 주로 사용하며, 다양한 방식을 시도해본 것 같다.
시도 및 성과
- 다양한 모델과 Prompt 변경 실험, Inference에 필요한 Sampling Parameter 변경 실험 등을 수행했다.
- HuggingFace에서 발표했던, PRM과 LLM을 활용한 Test Time Scailing을 적용한 기법의 노트북을 작성하여, 현재 기준 13개의 UpVote를 받았다.(첫 브론즈 메달 흭득..ㅎㅎ)
- Test에 필요한 AIME 데이터셋(~AIME2025)을 가공해서 업로드했다.
느낀점
- LLM의 성능(80%)
- Prompt 조작(15%)
- Sampling Parameter 등 생성 파라미터(5%)
가장 중요한건 역시 LLM의 성능이다. 기본 LLM의 성능이 좋아야지, 좋은 성과가 나온다. 이건 다른 도메인에 비해서 이 수학 도메인의 문제 해결에서 더 크게 차이가 나는 것 같다. 기본 LLM의 성능에 포함되는 부분은 CoT, TIR, Reasoning 등 모델의 능력에 영향을 줄 수 있는 다양한 방식을 모두 포함한다.
Qwen2.5
대회가 진행되면서, 사람들이 주로 사용하는 모델이 어느 정도 정해져있었다. 가장 처음에 사람들이 많이 사용한건 'Qwen2.5' 시리즈의 모델들이었다. 7B ~ 72B까지 다양한 사이즈의 모델을 사용했지만, 이 때까지만 해도 최대 Score가 10을 넘기기 힘들었다.
이 시기에 가장 효과적인 방법은 TIR 방법론을 활용해서, 모델이 문제 해결을 위한 코드를 생성하고, 생성한 코드를 이용해서 정답을 생성하는 것이었다.
QwQ-32B
대회가 진행되던 도중 Qwen에서 갑자기 출시한 정체를 알 수 없는 모델이었다. 하지만 AIME 데이터셋에 대한 성능이 엄청났고, 사람들은 바로 이 모델을 이용해서 제출하기 시작했다. 가장 먼저 제출했던, 참가자의 결과를 확인했는데 아무 것도 하지 않았는데도 Score가 18점이 나왔다. 이는 이전의 점수를 훨씬 넘어선 점수였기에 다들 이를 이용해서 출제하려고 시도하는 것 같았다.
이후 얼마 지나지 않아, 다양한 Prompt와 Sampling Parameter를 사용해서, 20점을 받은 사람이 나옴으로써 Early Sharing Prize를 받게되었다.
QwQ는 RL을 이용한 Long CoT를 검증한 모델이라고 생각한다. 생성을 진행하며, 계속 해서 고민하는 듯한 표현과 진행을 볼 수 있었고, 실제로 이렇게 생성을 진행했을 때 좋은 결과가 나왔다. 하지만 정답을 생성하기 위해서, 기존보다 훨씬 많은 Token을 사용해서 시간이 오래 걸리는 것, 종종 무한 자기 회귀에 빠지는 경우, Code Mixing 등 불안정한 모습을 보일 때도 있었다. 그럼에도 불구하고, 기존 모델들보다 아주 좋은 성능을 보였음으로 Long CoT를 활용한 Reasoning 방식에 다들 많은 관심을 보이기 시작했다.
DeepSeek-R1 Family
대중에게도 많이 알려진 DeepSeek 모델이다. RL을 활용해서 엄청난 Reasoning 성능을 이끌어냈다. 그냥 R1은 규모가 너무 커서 사용할 수는 없고, R1의 추론 결과를 Distill한 더 작은 규모의 모델을 이용했다. 1.5B(Math), 7B(Math), 14B, 32B 등 다양한 규모의 모델이 있지만, 다른 사람들의 제출결과들을 보면 14B정도의 큰 모델을 사용하는 것보다는 7B에서 더 많은 갯수의 정답을 생성해서 Majority를 정답으로 사용하는 것이 더욱 효율적이면서, 20점 ~ 25점 정도로 좋은 Score가 나왔다.
이후로 다들 새로운 모델들보다는 Prompt를 통한 개선을 시도한 것 같다. 나도 서류 작성을 하며, 이후로는 새로운 시도를 하지 못했다.
AIMO2 1위 솔루션: NVIDIA NemoSkills 전략 분석
1등은 NVIDIA의 팀원들이 Private Score 34점으로 마무리하며 차지했다. 결과적으로 이 문제를 해결할 때 어떤 인사이트를 얻을 수 있었는지, 학습할 필요가 있다. [https://arxiv.org/pdf/2504.16891]
NVIDIA의 팀원들은 다음과 같은 순서로 문제를 해결했다고 한다.
- 540K의 고품질(Olympiad수준 포함) 수학문제를 수집하고, 'DeepSeek-R1'과 'QwQ-32B'를 프롬프팅하여 3.2M의 Long-Reasoning 솔루션을 생성했다.
- 반복 훈련, 생성, 품질 필터링을 통해 Long-Reasoning 모델과 Code Excution을 통합하는 새로운 방법을 개발하여, 1.7M 고품질의 TIR(Tool-Integrated Reasoning) 솔루션을 얻었다.
- 다수의 후보에서 가장 유망한 솔루션을 선택하는 모델을 훈련시키는 파이프라인을 만들고, 이러한 Generative Solution Selection(GenSelect)이 Majority Voting의 베이스라인보다 눈에 띄게 향상시킬 수 있었다.
학습된 모델의 AIME, HMMT에 대한 성능비교를 보자. DeepSeek-Distill 모델과 동일한 규모에 대해서 성능을 비교한 결과, CoT maj@64의 결과는 모두 압도적으로 나왔다. 더해서 TIR, TIR+GenSelect를 적용하면, 비교하는게 우스울 정도로 높은 성능이다.
하지만 QwQ-32B의 CoT maj@64의 결과는 오히려 더 낮은 성능이 나온 것을 볼 수 있다. 이건 이후에 보도록 하자.
이제 데이터셋 준비 과정에 대해서 자세히 살펴보자.
문제 준비
Art of Problem Solving(AoPS)커뮤니티 포럼에서 수학문제들을 수집했다. 전체 포럼 토론에서 너무 기초적인 수학 문제(Middle School Math)는 제외하고, 수집했다. 전체 파이프라인에서 'Qwen2.5-32B-Instruct' 모델을 활용해서 처리를 진행했다.
문제 준비 파이프라인은 다음 순서로 진행된다.
- 문제 추출 : LLM에 Prompt를 이용해서, 초기의 포럼 게시글에서 모든 문제를 식별하고 추출했다. 대부분의 게시글에는 한 문제는 포함되어 있었지만, 여러 문제 or 한문제도 포함되지 않은 경우도 있었다.
- 문제 분류 : 각각의 추출된 문제들은 'Proof or not(증명)', 'Multiple-Choice or not(다지선다)', 'Binary or not(이지선다)', 'Valid or not(풀 수 있는지 여부)' 카테고리로 Prompt를 이용해서 분류한다. 최종 데이터셋에서는 Proof, Multiple-Choice, Invalid 문제는 모두 제거했다.
- 문제 변형 : Proof 문제에 대해서는 유사한 문제 해결 기법을 필요로 하도록하는 답변 문제로 Prompt를 이용해서 변환한다.
- 정답 추출 : non-Proof 문제에 대해서는 포럼 토론에서 최종 답안을 추출하도록 시도했다.
- 벤치마크 오염 제거(Decontamination) : LLM-based로 유명한 수학 문제 벤치마크와 유사한 질문들은 제거했다.
위에서 사용한 모든 프롬프트와 스크립트는 [NeMo-Skills]에 공개되어있다.
이렇게 처리한 전체 데이터셋의 갯수는 540K이고, 각각 [Proof 문제의 변형(260K)], [추출된 정답 문제(190K)], [추출되지 않은 정답 문제(90K)]으로 구성되었다. 저자는 이 데이터셋의 이름을 OpenMathReasoning 이라고 정의했다.
이후, 모델의 평가를 위한 Validation Set를 구성했다. 데이터 오염의 최소화를 위해 2024 ~ 2025년 사이 대회의 문제들로 제한하고, AIMO-2와 강력한 alignment를 갖고있는 AIME와 HMMT 문제를 Validation Set으로 선택했다. 결과적으로 'Comp-Math-24-25' 라고 하는 AIME, HMMT 문제들을 벤치마크로 구성했다.
핵심 기술 1: 합성 CoT 솔루션 생성
CoT 솔루션 생성을 위해서, DeepSeek-R1과 QwQ-32B을 사용해서 각 문제당 최대 32개의 솔루션 후보를 생성했다. 'Qwen2.5-72B-Math-Instruct' 모델을 이용해서 32개의 생성동안 평균 pass-rate를 난이도로 측정하여, 더 어려운 문제일수록 더 많은 솔루션을 생성했다.
최종 필터링 단계로, 기대했던 정답에 도달하지 못한 솔루션들은 제거하고, 예측한 정답과 기대한 정답이 동일한지는 'Qwen2.5-32B-Instruct'를 이용해서 판단했다. 정답을 추출하지 못한 문제(혹은, Proof 문제를 변환한 문제)의 경우 모든 솔루션 후보에서 다수로 생성된 정답을 Ground Truth로 판단하기로 결정했다.
핵심 기술 2: 도구 통합 추론 (Tool-Integrated Reasoning, TIR)
TIR 방법은 자연어 추론과 Python 코드 실행을 합쳐서, 수학 문제의 정답률을 올려주는 알려진 방법이다. 하지만 DeepSeek-R1, QwQ-32B에서 Prompting과 Few-Shot Examples를 사용해서 TIR을 시도해봤지만, 성공적이지 못했다고 한다.
이전의 연구결과(LIMO: Less is More for Reasoning)에 따르면, non-Reasoning 모델을 제한된 양의 Reasoning 데이터로 훈련을 시키면, 좋은 지시 수행 능력을 유지하는 경향을 보였다고 한다. NVIDIA팀은 이 직관에 의해 모델을 발전시켜나가기 시작했다.
먼저, LIMO-Qwen-32B를 Prompting 해서 TIR 솔루션을 생성해내는 것을 성공했다. 하지만 생성된 코드는 관련이 없거나, 단순히 CoT 단계를 검증하는데 사용되는 것들이었다.
이를 극복하기 위해서, Code 실행이 실질적인 추론에 이점을 제공하는 고품질의 예제만을 유지하기 위한 필터링 단계를 개발했다. 이렇게 필터링한 데이터셋으로, Reasoning 모델을 학습했더니 CoT only 모델보다 크게 성능이 향상되었다. 따라서, [TIR 모델 훈련 -> TIR 예제 생성 -> 필터링 -> 학습]의 반복적인 모델 개선 접근 방식을 채택했다. 이제 각 단계의 자세한 내용들에 대해서 보자.
Instruction-following reasoning Model
선행 연구들에 따르면, 1000개 정도의 적은 샘플의 훈련만으로도 LLM을 Long-CoT 솔루션을 생성하는데 충분하다. 따라서, 이렇게 훈련된 Instruct 모델이 Instruction-Following과 Long-Reasoning 능력을 유지할 수 있다고 가정하고, 이를 테스트해봤다.
'LIMO-Qwen-32B'모델을 이용하여, 복잡한 계산을 요구하는 문제를 Python Code로 해결하도록 Prompt를 제공했다.
결과적으로 문제의 절반정도에 대해 모델은 하나 이상의 Python 코드 블록을 포함하는 솔루션을 생성했고, OpenMathReasoning 문제들에 대해서 Temperature=0.7, Top-P=0.95, Max Generation Length=16,384로 1.2M개의 솔루션을 합성했다.(8개 이상의 코드 실행이 포함되면 생성을 멈추도록 제한했다.)
Filtering TIR Data
생성된 솔루션을 검사한 결과, 코드 실행이 문제 해결에 도움이 되지 않고 간단한 CoT 단계로 대체될 수 있는 경우가 많은 것을 확인했다. 따라서, 코드 사용을 필터링 하기 위해서, Qwen2.5-32B-Instruct를 사용하여, 각 코드 블록을 두가지 기준으로 분류했다.
- [새로운 계산 / 검증] : 코드 실행이 새로운 결과를 이끌어내는지 혹은 이전 단계를 그저 검증하는지 여부
- [중요한 / 중간 / 사소한] : 코드의 적용이 솔루션에서 중요한 부분인지 혹은 여러 단계의 CoT로 쉽게 대체될 수 있는지 여부
솔루션 선정 기준은
- 최소 하나의 독창적이고 중요한 코드 블록을 포함
- 절반 이상이 독창적이고 보통 수준의 코드 블록을 포함
또한 규칙 기반 필터링도 적용했는데,
- 최종 답이 잘못된 솔루션이나 코드 실행이 없는 솔루션은 제거
- 두 개 이상의 코드 블록이 포함된 솔루션도 제거
또한 코드 블록의 시작과 끝을 표시하는 태그를 교체하는 전처리 과정을 거쳤다.
- 초기 생성 단계에서 모델이 생성하기 쉽도록 코드 위치를 마크다운 스타일로 ‘‘‘python과 ‘‘‘\n으로 지정하도록 지시하며, 이를 <tool_call> 및 </tool_call> 태그로 교체하여 코드 블록을 쉽게 추출할 수 있도록 한다.
결과적으로 이러한 필터링 단계를 통해 최종적으로 15,000개의 TIR 샘플로 구성된 데이터 세트를 생성했다. 이 데이터셋을 Stage-0 TIR 데이터라고 부른다.
Iterative Data Generation
다음 단계의 TIR 솔루션을 생성하기 위해서, QwQ-32B 모델을 활용한다. Stage-0 데이터로 훈련을 진행하고, OpenMathReasoning 문제에 대한 솔루션을 생성한다. 700K개의 샘플을 생성하고, 부정확한 솔루션과 코드를 사용하지 않은 솔루션을 제거함으로써, 260K개로 필터링했다.
여기서 결과를 더욱 개선하기 위해서, 이 과정을 OpenMathReasoning 데이터의 CoT데이터로 훈련된 중간 크기(14B)의 모델을 이용해서 한번 더 반복했다. 14B모델을 QwQ-32B 솔루션에 대해 학습시킨 뒤, 데이터 생성과 필터링을 진행하여, 최종적으로 1.7M개의 TIR 데이터셋을 생성했다.
코드 블록의 갯수를 제어
모든 데이터 생성 단계에서, 남은 코드 실행 횟수에 대한 경고를 추가하여 코드 블록의 수를 제어하도록 유도했다. 결과적으로 모델의 Reasoning 과정에서 이를 참조하여, 제한에 도달하면 더 이상 코드 사용을 자제하는 것을 발견할 수 있었다.
각 문제에 대해 허용하는 코드 실행 횟수는 1 ~ 8 사이에서 무작위로 선택되었고, 이 정보를 Prompt에 제공했다. 훈련과정에서 정확한 행동을 강화하기 위해, 요청한 코드 실행 횟수보다 많은 코드 실행을 하려고 한다면, 해당 생성본은 제거했다.
핵심 기술 3: 생성형 솔루션 선택(Generative Solution Selection, GenSelect)
pass@k 와 maj@k 메트릭간의 큰 성능 차이가 있다는 것을 보고, 이는 이론적으로 maj@k보다도 더 좋은 결과를 낼 수 있음을 시사한다. 이를 위해 모델이 "솔루션 후보의 요약(Summary)이 주어졌을 때, 가장 유망한 솔루션을 선택"하도록 한다.
이전 연구(Generative Verifiers)에 따라, 모델의 구조 변경없이 제공된 솔루션 중에 하나를 선택하기 전에, 자연어로 추론하도록 한다.
Summary 생성하기
추론 모델(e.g. DeepSeek-R1)에 의해 생성된 결과는 thinking 부분과 summary 부분이 이어져 나온다. 하지만 summary 부분이 너무 요약적인 경우가 있어서, Qwen2.5-32B-Instruct 모델을 활용해서, 기존의 summary를 대체하는 새로운 summary를 생성했다. 각 솔루션에 대해 4개의 summary 후보를 2048 토큰으로 제한해서 생성했고, summary의 정확성을 위해서, 기존의 정답과 다른 것은 필터링했다. 만약 유효한 summary가 없다면, 해당 샘플은 제거했고 기존의 summary를 대체하기 위해서 가장 긴 길이의 summary를 사용했다.
OpenMathReasoning 전체에 대해서 summary를 재생성하고, 모델이 직접 생성할 수 있도록 훈련을 진행했다.
Selection 후보 생성하기
OpenMathReasoning 의 각 문제에 대해서 2~16개의 솔루션 요약 후보를 랜덤하게 선택했다. 적어도 한개의 옳은 샘플과 옳지 않은 샘플을 가지고 있도록 보장하고, 이 과정은 8개의 서로 다른 비교 그룹을 얻을 떄까지 반복한다.
GenSelect Prompt를 QwQ-32B에 이용해서, 각 그룹에 대해 가장 유망한 솔루션을 선택하도록 지시를 내렸다. 이 과정을 통해서 1M개의 선택지를 생성했고, 잘못된 해법이 선택된 경우를 제거하여 565K개까지 필터링했다.
OpenMath-Nemotron 모델
Training
최종 모델을 만들기 위해서 다음과 같은 크기의 모델을 SFT했다.
- 1.5B : Qwen2.5-Math-1.5B
- 7B : Qwen2.5-Math-7B
- 14B : Qwen2.5-Base-14B
- 32B : Qwen2.5-Base-32B
하지만 Qwen2.5-Math 모델은 컨텍스트 윈도우의 크기가 4096토큰으로 제한되어있어, Long-Reasoning 생성에 부적합했다. 이를 극복하기 위해서, NTK-Aware Scaled RoPE를 참고하여, RoPE의 base($\theta$)를 500K로 변경했다.
전체 모델은 6 Epoch 동안 CoT Solution, TIR Solution, GenSelect의 조합에 대해 훈련을 진행했다. 각 작업은 추론 시에 다른 생성 모드 간에 전환하는데 사용할 수 있는, 별개의 Prompt로 정의되었다.
또한, 모든 Task를 혼합해서 학습하는 것은 각 Task를 순차적(CoT -> TIR -> GenSelect)으로 학습하는 것과 비슷한 정확도를 보였다.
- Optimizer : AdamW - weight decay=0.01, cosine learning rate decay scheduler=10% linear warm-up
- Learning Rate : 1.5B=3e-4, 7B=2e-4, 14B,32B=1e-4
- Batch Size : 1024
처음 SFT를 진행하고, 이후에 한번 더 어려운 문제들에 대해 Olympiad Math 토론 포럼에 있는 문제에 대해 SFT를 진행했다.
Qwen2.5-Math-72B-Instruct 모델의 TIR의 32개의 생성 중 0.3보다 큰 통과율을 갖는 문제는 버리고, 솔루션의 길이가 5000토큰 미만인 것도 필터링했다. 또한, 4 Epoch 동안만 진행하고, 성능 저하가 관찰된 32B 모델은 제외하고 진행했다.
결과를 보면, CoT의 결과는 향상되는 반면에, TIR 결과는 안정적 or 저하 되는것을 볼 수 있었다.
Kaggle Submission
NVIDIA측에서 Kaggle에 제출해서 우승한 레시피는 다음과 같다.
Qwen2.5-14B-Base 모델을 Proof 문제를 변환한 문제를 제외한 CoT 솔루션에 대해서 8 Epoch동안 훈련했다. 또한, DeepSeek-R1에 의해서 생성된 솔루션만을 사용했다. 그 후 15K의 Stage-0 TIR 샘플들에 대해서 400 Step동안 1e-5의 LR로 학습을 진행했다. GenSelect는 Kaggle 제출에는 사용하지 않았다.
또한 CoT 체크포인트와 TIR의 체크포인트를 선형 결합하여 모델의 정확도, 생성 길이, 코드 사용 갯수 등을 최적화할 수 있었다.
실험을 통해 모델들에 최종 벤치마크 성능은 다음과 같다.
그 외 TensorRT 최적화, Speculative Decoding, Model Serving 등 추론 최적화와 관련된 내용이 더 있지만 이 정도까지만 하고 마무리해도 괜찮을 것 같다.
정리 및 느낀점
NVIDIA의 AIMO2 Kaggle 우승 솔루션을 한번 톺아봤다. 해당 솔루션을 보면서 느낀 점은 다음과 같다.
- 데이터 조작에 필요한 대부분의 작업을 LLM(Qwen2.5-32B-Instruct)을 활용한다. 정말 99% LLM의 Prompt변경을 통해서 작업을 수행했다.
- 역시 Reasoning(Long-CoT or Generation) 과정이 어려운 작업을 해결할 때 좋은 성능을 낸다.
- CoT, TIR 모델을 Merging할 생각을 했다니 예상치 못했다. 또 최적의 결과가 생각보다 간단하게(선형결합) 나왔다.
나도 최근 한국어 수학문제 해결에 적합한 모델을 만들고자, 먼저 평가를 위한 수능 수학 데이터셋을 구축하면서 Gemini-2.5-Flash를 이용하고 있었는데 역시 비용 상관 없으면 LLM이 중간의 복잡한 과정을 효율적으로 줄여줄 수 있는 수단이라는 것을 깨닫고 있었다. OCR부터, 구조화까지 너무 잘해줘서 감탄을 하면서 진행했지만... NVIDIA는 혹은 다른 더 유능한 사람들은 이것보다 더 나아가서 더욱 다양하고 많은 작업을 Prompt로 구체화해서 어시스턴트처럼 사용하고 있다는 것에 놀랐다.
이를 발판삼아 나도 작업을 Prompt로 구체화해서 지시하기 위한 고민과 현재 시대에서의 LLM 활용에 대해 다시 한번 곰곰히 생각해봐야할 것 같다.
다양한 사람, 다양한 장소에서 배움을 지향합니다.
최근 당신이 흥미롭게 보고계시는 주제를 저도 알고싶습니다!
또, 제가 작성한 글이 도움이 되셨거나, 의견과 느낀점을 편하게 댓글로 남겨주세요! 🤗