[Kaggle] AIMO2 : Test-Time Compute Scaling으로 LLM 성능 극대화하기: HuggingFace 연구 분석[5]

목차

• 1. 개요
• 2. Test-Time Compute Scaling 전략
  • 2.1 Self-Refinement (자기 정제)
  • 2.2 Search Against a Verifier (검증 기반 검색)
• 3. Search-Based Strategy 실험 설계
  • 3.1 실험 목표
  • 3.2 사용된 모델 및 데이터셋
  • 3.3 주요 실험 설정
• 4. 주요 Search Strategies
  • 4.1 Majority voting
  • 4.2 Best-of-N
  • 4.3 Beam Search with PRM
  • 4.4 Search Strategy결과 분석
  • 4.5 DVTS : Diversity로 성능을 높이자
• 5. AIMO에 적용하기
• 6. 결론 및 교훈
• 7. 참고 자료 및 코드

1. 개요

* 이 글은 HuggingFace의 연구결과를 적용한 방법입니다.

 

최근 몇 년간 대규모 언어 모델(LLM)의 성능은 모델 크기와 데이터셋 크기를 늘리는 Train-Time Compute Scaling 방식에 의존해 발전해왔습니다. 하지만 이 접근법은 비용이 크고, 자원이 제한된 환경에서 실행하기 어렵습니다. 이를 해결하기 위해 등장한 개념이 Test-Time Compute Scaling입니다.

 

Test-Time Compute Scaling은 모델이 추론 단계에서 더 "오래 생각(think longer)"하도록 설계된 방법론입니다. 대표적인 예로 OpenAI의 o1 모델이 있지만, closed-source 특성상 구체적인 구현 방법은 공개되지 않았습니다.

 

이 글에서는 Test-Time Compute Scaling의 주요 전략을 살펴보고, HuggingFace의 연구 결과를 바탕으로 한 실험 과정을 공유합니다.

2. Test-Time Compute Scaling 전략

Test-Time Compute Scaling에는 크게 두 가지 주요 접근법이 있습니다.

2.1 Self-Refinement (자기 정제)

• 모델이 자신의 출력을 반복적으로 검토하고, 오류를 수정하며 결과를 개선하는 방식입니다.
• 제한점: Self-Refinement는 빌트인 메커니즘을 요구하므로 모든 모델에 적용하기 어렵습니다.

2.2 Search Against a Verifier (검증 기반 검색)

• 여러 후보 답안을 생성하고, 검증기(verifier)를 통해 최적의 답안을 선택하는 방식입니다.
• 검증기는 하드코딩된 규칙부터 학습된 Reward Model까지 다양한 형태로 구현될 수 있습니다.
• 예: Best-of-N 샘플링, Beam Search, Tree Search 등.

이 글에서는 특히 Search-Based Strategy에 초점을 맞춰 실험을 진행했습니다.

Search Strategies in test-time compute scaling

• Best-of-N : 하나의 문제에 대해서 N개를 답을 생성하고, 각 후보 정답에 대해 스코어를 계산한다. 일반적으로 Reward Model을 사용해서 계산한다. 그리고 그 중 가장 높은 스코어를 갖는 정답을 선택한다.
• Beam Search : $N$개의 결과를 생성하고, $M$개의 Beam width 만큼 상위 점수의 결과를 유지한다. 이는 Process Reward Model(PRM)과 결합해서 각 과정에 대해 점수를 냄으로써, LLM을 사용한 검색 방식에 자연스럽게 적용될 수 있다.
• Diverse Verifier Tree Search : Beam Search의 확장으로, 초기의 Beam을 독립적인 서브트리로 나누고, 이후부터 PRM을 사용해서 Greedy하게 확장한다.

3. Search-Based Strategy 실험 설계

3.1 실험 목표

• 다양한 Search-Based Strategy(Best-of-N, Beam Search 등)의 성능을 비교합니다.
• Test-Time Compute Scaling의 효율성을 평가합니다.

3.2 사용된 모델 및 데이터셋

• 모델: meta-llama/Llama-3.2-1B-Instruct
• PRM (Process Reward Model): RLHFlow/Llama3.1-8B-PRM-Deepseek-Data
  • 동일 계열의 정책 모델보다 더 나은 성능을 제공.
• 데이터셋: MATH-500 subset

3.3 주요 실험 설정

Search-Based Strategy

1. 각 문제에 대해 $N$개의 후보 솔루션(중간 단계)을 생성.
2. PRM을 활용하여 각 단계별 점수를 계산.
3. Search Strategy에 따라 다음 단계 솔루션을 선택.
4. 최종적으로 PRM 점수를 기반으로 정답 순위를 매김.

4. 주요 Search Strategies

4.1 Majority voting

가장 기초적이고, 단순한 베이스라인입니다. $N$개의 후보 정답을 생성하고, 가장 많이 등장한 것을 선택하는 방법입니다.

$N=256$개의 후보를 생성했고, $T=0.8$, $Max Tokens=2048$로 각 문제에 대해 정답을 생성했습니다.

 

처음에 간단한 System Prompt를 이용해서 생성했을 땐, Meta의 Greedy Decoding실험 결과인 30.6%보다 낮게 측정됐습니다.

Please think step by step and put your final answer within \boxed{}.

 

그래서 Meta에서 발표했던 실험결과의 System Prompt로 변경하고 난 뒤에 결과를 다시 측정하여, 성능을 다시 이끌어낼 수 있었습니다.

Solve the following math problem efficiently and clearly:

    - For simple problems (2 steps or fewer):
    Provide a concise solution with minimal explanation.

    - For complex problems (3 steps or more):
    Use this step-by-step format:

    ## Step 1: [Concise description]
    [Brief explanation and calculations]

    ## Step 2: [Concise description]
    [Brief explanation and calculations]

    ...

    Regardless of the approach, always conclude with:

    Therefore, the final answer is: $\boxed{answer}$. I hope it is correct.

    Where [answer] is just the final number or expression that solves the problem.

 

Majority Voting Result

Greedy Decoding 베이스라인 보다는 확실한 점수 향상이 있습니다. 하지만, $N=64$ 이후로는 더이상 개선되지 않고, 유지되는 모습이 보입니다.

하지만, 추론이 필요한 문제나, 몇 세대에 걸치 오류가 있는 작업에선 어려움을 겪는 한계점이 있습니다. 또한, $N=1, 2$에서 0-shot CoT 베이스라인보다 더 낮은 이유가 궁금할 수 있는데, $T=0.8$을 사용하여 샘플링을 했기 때문에, 소수의 후보 중에선 정답이 나올 확률이 적어지기 때문입니다.

 

이제 Reward Model을 활용해서, 성능을 개선하는 것을 봅시다.

4.2 Best-of-N

Best-of-N은 간단하지만, Majority Voting을 효과적으로 개선하고, 괜찮은 정답을 결정합니다. BoN은 크게 두가지 변형이 있습니다.

• Vanilla Best-of-N : $N$개의 독립적인 결과를 생성하고, Reward Model 점수가 가장 높은 하나의 결과를 최종 정답으로 선택한다. 가장 확정적인 정답이 선택되는 것을 보장하지만, 답변간의 일관성을 보장하지는 않는다.
• Weighted Best-of-N : 모든 동일한 답변의 점수를 합산해서, 총 Reward Model 점수가 가장 높은 답변을 선택한다. 이 방법은 반복적으로 발생한 것들의 점수를 Boosting함으로써, 높은 퀄리티의 정답을 보장한다.
  $RM(p, s_i)$는 문제$p$에 대한, 솔루션$s_i$의 Reward Model에 의해 계산된 점수입니다. 솔루션 $s_i$는 모든 단계의 솔루션을 연결한 것이고, $a_i$는 일반적으로 \boxed{answer}에서 추출된 정답을 의미합니다.

보통, 단일의 솔루션 점수를 측정하기 위해서는 Outcome Reward Model(ORM)을 사용합니다. 하지만, 이후 검색전략과 공정하게 비교하기 위해서 Process Reward Model(PRM)을 사용해서, Best-of-N의 점수를 내겠습니다.

 

Process Reward Model을 이용하려면, 단일 솔루션 점수를 계산하기 위해 아래와 같이 각 단계의 점수를 이용해서 계산할 수 있습니다.

Score Reduction over the steps

• Min : 각 단계의 점수 중 최솟값을 사용
• Prod : 각 단계의 점수의 곱을 사용
• Last : 마지막 단계의 점수를 사용. 이 점수는 이전 단계의 누적 정보를 포함하기에, ORM처럼 사용 가능하다.

실험에서는 Last 점수를 활용합니다.

Best-of-N Result

실험 결과를 보면, weighted Best-of-N이 Majority Voting보다 좋은 결과를 보입니다. 위에서 언급한 장점처럼, 동일한 응답의 점수를 합산함으로써 빈도는 낮아도, 품질이 높은 답변에서의 우선순위를 통해 좋은 효과를 보입니다.

 

그럼에도 불구하고, Llama 8B 모델보다는 아직 조금 낮고, $N=256$이후로는 성능이 일정하게 유지됩니다.

 

4.3 Beam Search with PRM

Beam Search는 PRM과 결합해서, 문제 해결을 위한 과정에서 생성과 평가를 모두 최적화할 수 있습니다.

Beam Search의 과정은 다음과 같습니다.

1. 고정된 수의 "Beams"를 유지하면서, 후보 솔루션을 반복해서 생성한다.
2. 첫 반복에선, temperature$T$를 사용해서 LLM을 통해 $N$개의 독립된 스텝을 생성한다. 스텝은 '\n', '\n\n'과 같이 생성에서 스텝을 구분하는 문자들을 기준으로 stopping criterion으로 사용해서 생성한다.
3. 각가의 스텝을 PRM을 이용해서 점수를 내고 상위 $N/M$ 스텝들을 다음 생성을 위한 후보로 선택한다. 여기서 $M$은 beam_width로 유효한 경로의 갯수이다. 점수 계산에선 Best-of-N에서와 같이 "Last" reduction을 사용한다.
4. 솔루션의 $M$개의 다음 스텝을 샘플링해서 솔루션 스텝을 확장한다.
5. EOS 토큰에 도달하거나, 최대 search depth에 도달할 때까지 3단계와 4단계를 반복한다.

실험 과정에서 사용하는 하이퍼파라미터는 다음과 같습니다.

• $N$ beams 의 갯수 : 4, 16, 64, 256
• Beam Width $M$ : 4
• Temperature $T$ : 0.8
• Tree of maximum depth : 40

Beam Search Result

실험 결과를 보면, $N=4$일 때, Best-of-N의 $N=16$결과와 동일한 정확도를 보입니다. 유사하게, $N=16$일 때와 Best-of-N의 $N=256$의 결과과 비슷합니다. 게다가, Llama3.1-8B 모델의 성능을 $N=32$일 때 뛰어넘었습니다.

 

4.4 Search Strategy결과 분석

Beam Search가 다른 전략들에 비해서 좋은 성능을 내지만, 각 전략은 문제의 난이도와 Test-Time Compute 예산에 대해 트레이드오프가 있습니다. 문제의 난이도에 따라 $N=[4, 16, 64, 256]$의 Pass@1 점수에 해당하는 결과를 봅시다.

각 막대는 위의 각각의 $N$에 해당하는 결과이고, 각각의 값은 상대적인 정확도입니다.

예를 들어, Level 2의 결과를 보면,

• Majority Voting은 $N=256$을 제외한 전체 $N$에서 가장 안좋았다.
• Beam Search가 $N=[4, 16, 64]$에서 가장 좋았지만, Best-of-N이 $N=256$에서 가장 좋았다.

Level 3~5와 같은 중간~어려운 난이도의 문제에선 모든 $N$에 대해서 다른 전략들보다 Beam Search가 가장 좋은 결과를 보입니다. 하지만 쉬운 문제, 큰 $N=256$에선 Best-of-N, Majority Voting과 같은 전략들보다 오히려 더 안좋은 결과를 보여줍니다.

그래서 Beam Search방법으로 생성된 트리를 살펴봤을 때, 단일 단계에서 높은 Reward Score가 할당되면, 전체 트리가 그 것에 집중되어 다양성에 영향을 준다는 것을 관찰할 수 있습니다.

 

4.5 DVTS : Diversity로 성능을 높이자

위와 같이, Beam Search는 Best-of-N을 뛰어넘을 정도로 높은 성능을 보이지만, 쉬운 문제나 큰 Test-Time compute Budget($N$)에서는 성능이 떨어집니다. 이걸 다루기 위해, 큰 $N$에서 다양성을 극대화하기위해 Diverse Verifier Tree Search(DVTS)방법론을 제안합니다.

 

DVTS는 Beam Search와 거의 비슷하지만, 일부 다른점이 있습니다.

1. 주어진 $N$과 $M$에 대해서, $N/M$의 독립된 서브트리 Beams를 초기화한다.
2. 각각의 서브트리에 대해, 가장 높은 PRM 점수를 갖는 스텝을 선택한다.
3. 2번에서 선택된 노드로부터 $M$개의 새로운 스텝을 생성하고, 가장 높은 PRM 점수를 갖는 점수를 갖는 스텝을 선택한다.
4. 3번 단계를 EOS 토큰이 나오거나, maximum tree depth에 도달할 때까지 반복한다.

DVTS Result

실험 결과를 보면, DVTS는 Beam Search를 보완하는 전략으로, 작은 $N$에서 Beam Search가 훨씬 효율적으로 정답을 찾습니다. 그러나, $N$이 커질 수록, DVTS가 더 좋은 성능을 내는 것을 볼 수 있습니다.

 

실험 결과를 보면, Level 1~2단계의 쉬움~중간 정도의 난이도에서 큰 $N$에 대해 좋은 성능을 보입니다. 하지만 Beam Search는 전체적인 난이도의 작은 $N$일 때 가장 우수한 결과를 냈습니다.

 

5. AIMO에 적용하기

DVTS방법은 HuggingFace에서 Search-and-Learn이라는 라이브러리로 배포하고 있으니, 해당 [Repo]에서 Inference 코드를 확인할 수 있습니다.

AIMO2 대회에서는 주어진 하드웨어가 L4*4입니다. 메모리 제약이 24GB*4로, LLM과 Reward Model을 로드하고, 생성하는 길이를 고려할 필요가 있습니다. 조건을 조절해가며, 오류가 나지 않는한에서 저의 설정은 다음과 같습니다.

1. 사용된 모델:
   • LLM: Qwen2.5-Math-7B-Instruct
   • PRM: RLHFlow/Llama-8B-PRM-Deepseek-Data (4bit Quantization)
2. 설정:
   • $N=32$, $M=4$, Max Tokens=2048
   • Num Iteration(Search Depth)=15

모델 로드하는 것만으로는 메모리가 부족하지 않지만, AIMO의 문제 특성상 아주 어려운 난이도(위의 예시라면 Level5 정도에 해당)하기 때문에 문제 해결과정이 아주 길게 진행됩니다.

이는 Search Depth를 더 키움으로써, 문제 해결과정을 더 길게 이어나갈 수 있지만 문제 하나당 평균적으로 5분이내로 해결해야하기 때문에 시간도 부족하고, Depth가 깊어짐에 따라 컨텍스트 길이가 길어져서 채점과정에서 오류가 발생해서 실패했습니다. 결과적으로, 위와 같은 설정이 사용한 모델에 대한 제가 찾은 최적의 하이퍼파라미터 값이었습니다.

 

리더보드 최점결과는 4점으로 처참합니다... Num Iteration을 훨씬 키움으로써 해결해볼 수 있지않을까싶은데, 시간제약이 있기때문에 불가능했습니다.

 

6. 결론 및 교훈

Test-Time Compute Scaling은 추론 과정에서 더 많은 계산 자원을 활용하여 성능을 향상시키는 강력한 방법론입니다. 특히, DVTS와 같은 전략은 큰 $N$에서 뛰어난 성능을 보여주며, 제한된 하드웨어 환경에서도 효과적으로 작동할 수 있음을 입증했습니다.

그러나 시간 제약 조건이 있는 대회 환경에서는 적합하지 않을 수 있습니다. 이러한 경우 Train-Time Compute Scaling(더 많은 데이터와 학습 기법 활용)을 통해 성능 향상을 도모하는 것이 더 나은 선택일 수 있습니다.

7. 참고 자료 및 코드

실험에 사용된 코드는 아래 링크에서 확인할 수 있습니다:
GitHub Repository

Kaggle/AIMO2/notebooks/Version2-search-and-learn-qwen2-5-math.ipynb at AIMO2 · dbsrlskfdk/Kaggle

 

세줄 요약 by GPT-4o

1. HuggingFace 연구를 통해 Test-Time Compute Scaling 전략을 탐구했는데, 모델이 추론 중 더 "오래 생각"하도록 만들어 성능을 끌어올리는 방법이에요.
2. 다양한 검색 전략(Best-of-N, Beam Search, DVTS)을 실험해봤고, 상황에 따라 성능이 달라지는 재미있는 결과를 얻었어요.
3. 결론은? 추론 시간이 부족하면 훈련(Train-Time)에 더 투자하라는 교훈! 시간과 자원의 트레이드오프가 핵심이네요. 🚀

다양한 사람, 다양한 장소에서 배움을 지향합니다.
최근 당신이 흥미롭게 보고계시는 주제를 저도 알고싶습니다!
또, 제가 작성한 글이 도움이 되셨거나, 의견과 느낀점을 편하게 댓글로 남겨주세요! 🤗