Balanced Aggregation: GRPO의 집계 편향 이해 및 해결

DeepSeek-R1 등에서 사용되는 GRPO 알고리즘의 토큰 집계 방식이 모델 성능과 학습 안정성에 미치는 영향을 분석했습니다. 기존 방식들이 가진 길이 편향 문제를 해결하는 Balanced Aggregation 기법을 통해 추론 및 코딩 작업에서 더 안정적이고 높은 성능을 달성할 수 있음을 입증했습니다.

GRPO는 별도의 Critic 모델 없이 그룹 내 상대적 보상을 통해 학습하는 강화학습 기법이다. 이때 각 토큰의 그래디언트를 합치는 방식이 중요한데, 기존의 Token Aggregation은 모든 토큰을 동일하게 취급한다. 이는 만약 틀린 답(부정 샘플)이 맞은 답(긍정 샘플)보다 훨씬 길 경우, 모델이 '틀린 것을 피하는' 방향으로만 과도하게 학습되어 최적화가 한쪽으로 쏠리는 문제를 낳는다.

반면 Sequence Aggregation은 모든 응답 시퀀스에 동일한 가중치를 부여하여 이 쏠림을 막는다. 하지만 이 방식은 긴 응답 속에 포함된 풍부한 정보 토큰들의 영향력을 인위적으로 낮추는 부작용이 있다. 즉, 정보량이 많은 긴 답변이 짧은 답변과 똑같은 취급을 받게 되어 학습 효율이 떨어진다.

Balanced Aggregation은 이 두 개념을 결합한다. 먼저 맞은 그룹과 틀린 그룹을 분리하여 각 그룹 내부에서는 토큰 단위로 평균을 내어 긴 답변의 정보를 보존한다. 그 다음, 두 그룹을 합칠 때는 시퀀스 개수 비율에 맞춰 가중치를 조절함으로써 특정 그룹의 길이가 길다고 해서 전체 학습 방향이 왜곡되지 않도록 균형을 잡는다. 이를 통해 학습의 안정성과 성능이라는 두 마리 토끼를 모두 잡았다.

GRPO-style RLVR 환경에서 토큰 수준의 PPO 목적 함수를 집계하는 세 가지 규칙을 정의하고 분석한다. Token Aggregation은 그룹 내 모든 토큰 수 N으로 전체 손실을 나누며, Sequence Aggregation은 각 시퀀스 i의 토큰 수 Ti로 먼저 나눈 뒤 시퀀스 개수 G로 평균을 낸다.

Balanced Aggregation (BA)은 응답을 이진 보상에 따라 긍정(S+)과 부정(S-) 서브셋으로 분리한다. 각 서브셋의 토큰 총합 N+, N-를 분모로 하여 서브셋별 평균 손실 L+, L-를 계산한다. [각 그룹의 손실 합 → 그룹 내 토큰 수로 나눔 → 그룹별 평균 손실 산출]

최종 목적 함수 J_BA는 긍정 시퀀스 수 k와 전체 시퀀스 수 G의 비율을 가중치로 사용한다. [L+에 k/G를 곱하고 L-에 (G-k)/G를 곱함 → 두 값을 합산 → 최종 손실값 도출]. 이 구조는 이진 보상 설정에서 Sequence Aggregation과 동일한 부호 간 균형(inter-sign balance)을 유지하면서도 그룹 내에서는 토큰 수준의 정보를 보존한다.

Qwen2.5-Math-7B와 Qwen3-1.7B 모델을 사용한 실험에서 BA는 일관되게 우수한 성능을 보였다. DAPO-17k 데이터셋에서 Qwen2.5-Math-7B의 경우 BA는 Peak Accuracy 36.34%를 기록하여 Token-agg(35.95%)와 Seq-agg(35.78%)를 모두 앞질렀다.

특히 학습 안정성 측면에서 BA의 효과가 두드러졌다. Token-agg는 학습 후반부에 성능이 급격히 하락하는 경향을 보였으나, BA는 Last-step Accuracy에서 Token-agg보다 높은 수치를 유지하며 성능 붕괴를 방지했다. Qwen3-1.7B 모델의 경우 Seq-agg가 Token-agg보다 유리한 고지를 점했음에도 불구하고, BA는 Seq-agg와 대등하거나 더 높은 Peak 성능을 달성했다.

분석 결과, Qwen2.5-Math-7B는 응답 길이의 변동성(Coefficient of Variation)이 커서 Token-agg 계열이 유리했고, Qwen3-1.7B는 긍정-부정 샘플 간 길이 차이(Length Gap)가 커서 Seq-agg 계열이 유리했다. BA는 이러한 모델별 특성에 상관없이 모든 시나리오에서 강건한 성능을 입증했다.

본 연구는 GRPO의 목적 함수가 토큰 수준에서 계산되지만 어드밴티지(Advantage)는 시퀀스 수준에서 공유된다는 점에 주목한다. Token Aggregation 식을 전개하면 긍정/부정 항이 각 그룹의 평균 길이(T_bar+, T_bar-)에 의해 가중치가 부여됨을 알 수 있으며, 이를 sign-length coupling bias로 정의한다.

BA는 이 편향을 제거하기 위해 긍정/부정 서브셋을 분리하여 정규화한다. 수학적으로 BA는 Sequence Aggregation과 동일한 inter-sign prefactor인 sqrt(k(G-k)/G)를 공유하지만, 서브셋 내부의 가중치 할당 방식(within-sign averaging)에서 차이가 난다. Seq-agg는 서브셋 내 모든 응답에 동일 가중치를 주지만, BA는 토큰 수에 비례한 가중치를 부여한다.

부록에서는 이를 비이진 보상(Non-binary rewards)으로 확장하는 일반화된 BA 수식을 제공한다. 일반화된 버전에서는 시퀀스 개수 대신 어드밴티지의 질량(Advantage mass)을 기준으로 가중치를 배분하여, 보상 값이 연속적인 경우에도 부호 간 균형을 유지할 수 있도록 설계되었다.

본 논문은 주로 이진 보상(Binary reward) 환경인 수학 및 코딩 작업에 집중하여 실험을 진행했다. 일반적인 텍스트 생성이나 다중 보상 체계에서의 효과는 이론적 확장은 제시되었으나 대규모 실험을 통한 검증은 향후 과제로 남겨두었다.