기술 요약: N-GRPO - 향상된 정책 최적화를 위한 임베딩 레벨의 이웃 혼합(Neighbor Mixing)

1. 문제 정의

수학적 추론에서 대규모 언어 모델(LLM)의 성공은 강화 학습(RL)의 롤아웃(rollout) 단계 동안 다양하고 유효한 솔루션 경로를 생성하는 데 크게 의존합니다. 그러나 현재의 탐색 기법들은 근본적인 트레이드오프에 직면해 있습니다:

• 토큰 레벨 샘플링(Token-Level Sampling): 표준 방식은 근본적인 추론 로직을 변경하지 않고 단순히 재구문화나 교환 가능한 순서 변경(예: "1+2"와 "2+1")만 발생하는 중복된 궤적을 생성하는 경우가 많습니다.
• 임베딩 레벨 노이즈(Embedding-Level Noise): 연속적인 임베딩에 무작위 노이즈(예: 가우시안 노이즈)를 주입하는 방식은 빈번하게 의미적 일관성을 해칩니다. 예비 실험 결과, 제약 없는 노이즈는 표현형을 의미적 매니폴드(semantic manifold)에서 벗어나게 하여, 유효한 토큰을 관련 없는 토큰으로 변질시키고 롤아웃 궤적을 탈선시키는 것으로 나타났습니다.

Soft Thinking이나 HRPO와 같은 기존의 시도들은 이 간극을 메우기 위해 이들의 무작위성을 위해 이산적(discrete) 토큰 샘플링에 의존하거나, 연속 공간을 섭동(perturb)할 때 의미적 일관성을 유지하는 데 실패합니다.

2. 방법론: N-GRPO

저자들은 N-GRPO를 제안합니다. 이는 그룹 상대 정책 최적화(GRPO) 프레임워크에 통합된 새로운 탐색 전략입니다. 핵심 혁신은 **의미적 이웃 혼합(Semantic Neighbor Mixing)**으로, 이는 로컬 의미 매니폴드를 엄격히 준수하면서 다양성을 주입하는 임베딩 레벨의 섭동 메커니즘입니다.

핵심 메커니즘: 의미적 이웃 혼합(Semantic Neighbor Mixing)

이 방식은 이산적 토큰을 샘플링하거나 제약 없는 노이즈를 추가하는 대신, 앵커 토큰의 임베딩과 그와 가장 가까운 의미적 이웃들을 혼합하여 입력 표현을 구성합니다.

1. 앵커 선택 (Anchor Selection): 각 생성 단계에서 모델은 온도 조절된 로짓(temperature-scaled logits)을 기반으로 가장 선호되는 토큰(의미적 앵커, $v^*_t$)을 식별합니다.
2. 이웃 검색 (Neighbor Retrieval): 앵커와의 코사인 유사도를 기반으로 임베딩 행렬에서 $k$개의 이웃 토큰($C_t$) 세트를 검색합니다. 이는 이웃들이 동일한 의미적 방향을 따라 클러스터링되도록 보장합니다.
3. 가중치 혼합 (Weighted Mixing): 이 이웃 세트에 대해 온도 조절된 로짓을 정규화하여 혼합 가중치($\alpha_t$)를 계산합니다.
4. 연속 임베딩 구축 (Continuous Embedding Construction): 이웃 임베딩의 가중 합으로서 연속적인 혼합 임베딩($\tilde{e}_{t+1}$)이 생성됩니다. 이는 다음 타임 스텝의 입력으로 사용됩니다.

GRPO와의 통합

N-GRPO는 확률적 게이팅 메커니즘을 통해 이 메커니즘을 GRPO 롤아웃 단계에 통합합니다:

• 혼합 마스크 (Mixing Mask): 고정된 확률 $\rho$(예: 0.1)를 가진 베르누이 마스크($m_{i,t}$)는 의미적 이웃 혼합을 활성화할지, 아니면 표준 이산 토큰 샘플링을 따를지를 결정합니다.
• 안정성 (Stability): 대부분의 단계를 이산 토큰으로 유지함으로써, 이 방식은 의미적 안정성을 보존하는 동시에 탐색에 필요한 간헐성을 도입합니다.
• 학습 (Training): 학습 단계 동안 후보 세트와 가중치를 기록하여 어드밴티지(advantage) 계산을 위한 재현 가능한 궤적 재구성을 보장합니다. 최종 보상은 혼합된 단계에 대응하는 이산 앵커 토큰을 기반으로 계산됩니다.

3. 주요 기여

• 제어 가능한 임베딩 레벨 탐색: 본 논문은 생성 과정 중에 의미적으로 유도된 섭동을 주입하는 방법을 도입하여, 이산적 토큰 레벨 추론과 연속적 임베딩 레벨 탐색 사이의 간극을 메웁니다.
• N-GRPO 프레임워크: 저자들은 Semantic Neighbor Mixing을 GRPO 프레임워크에 통합하여, 의미적 일관성을 희생하지 않으면서 탐색 공간을 확장하기 위해 로컬 의미 매니폴드를 활용하는 시스템을 구축했습니다.
• 실증적 검증: 다양한 모델 규모(1.5B ~ 7B) 및 아키텍처(DeepSeek-R1-Distill-Qwen, Llama-3.2, Qwen3)에 대한 광범적인 평가를 통해 강력한 베이스라인 대비 일관된 개선을 입증했습니다.

4. 실험 결과

본 방법은 도전적인 수학적 추론 벤치마크(AIME25, AMC23, MATH500) 및 분포 외(OOD) 작업(GPQA-Diamond)에 대해 평가되었습니다.

• 수학적 추론: N-GRPO는 표준 GRPO, Soft Thinking, 그리고 제약 없는 가우시안 노이즈를 사용하는 STHT 베이스라인을 포함한 베이스라인들을 일관되게 능가했습니다.
  • DeepSeek-R1-Distill-Qwen-1.5B 모델에서, N-GRPO는 평균 Pass@32를 77.41(GRPO)에서 79.17로 향상시켰습니다.
  • 7B 모델에서, 평균 Pass@32는 81.94(GRPO)에서 84.20으로 향상되었습니다.
  • 어려운 AIME25 벤치마크에서 상당한 이득이 관찰되었는데, N-GRPO는 1.5B에서 50.28, 7B에서 61.18의 Pass@32를 달성하여 GRPO와 STHT를 넘어섰습니다.
• 분포 외(OOD) 일반화: GPQA-Diamond 벤치마크에서 N-GRPO는 강력한 일반화 능력을 보여주었으며, 가장 높은 Pass@32 점수(1.5B의 경우 92.87, 7B의 경우 92.80)를 기록하여 이 방법이 수학 특화 기술뿐만 아니라 일반 과학적 추론 능력을 향상시킨다는 것을 나타냈습니다.
• 전이 가능성 (Transferability): 이 메커니즘은 GSPO 프레임워크(N-GSPO)로 성공적으로 전이되어 AIME25에서 +7.66의 이득을 보였으며, 이는 섭동 메커니즘이 GRPO의 특정 어드밴티지 정규화에 밀접하게 결합되어 있지 않음을 시사합니다.
• 절제 연구 (Ablation Studies):
  • 혼합 비율 (Mixing Rate): 성능은 다양한 혼합 비율(0.1 ~ 0.2)에 대해 안정적으로 유지되었으나, 작은 모델의 AIME25 작업에서 과도한 혼합(0.2)은 성능 저하를 일으켜 다양성과 안정성 사이의 균형이 필요함을 보여주었습니다.
  • 거리 측정 방식 (Distance Metric): 코사인 거리가 이웃 선택을 위한 L1 및 L2 메트릭보다 현저히 우수한 성능을 보였으며, 이는 고차원 임베딩 공간에서 의미 정보의 방향성 있는 특성 때문인 것으로 분석됩니다.
  • 추론 (Inference): 저자들은 (GRPO 필터링 메커니즘 없이) 추론 단계에서 Semantic Neighbor Mixing을 적용하면 성능이 저하된다고 언급하며, 이것이 디코딩 전략이 아닌 훈련 시의 탐색 도구임을 확인했습니다.

5. 의의 및 주장

본 논문은 N-GRPO가 임베딩 공간에서의 추론 역학을 탐색함으로써 LLM의 추론 능력을 발전시키는 유망한 경로를 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 트레이드오프 해결: 단순한 노이즈 주입으로 인한 의미적 표류(semantic drift)나 토큰 레벨 샘플링의 중복성 없이 롤아웃 과정에 다양성을 성공적으로 주입했습니다.
• 의미적 근거 (Semantic Grounding): 로컬 의미 매니폴드를 활용함으로써, 탐색이 유효한 추론 영역 내에 머물도록 보장합니다.
• 일반화: 이 접근 방식은 특정 모델 제품군이나 추론 증류 정책에 국한되지 않으며, 다양한 규모와 베이스 모델에 걸쳐 견고함을 보여줍니다.

저자들은 이 방법이 롤아웃 중에 계산 오버헤드를 발생시키고 현재 주로 수학 및 과학 작업에서 검증되었지만, 탐색과 의미적 일관성 사이의 균형을 맞추는 정책 최적화의 새로운 패러다임을 구축한다고 결론지었습니다.