Balanced Aggregation: Understanding and Fixing Aggregation Bias in GRPO

본 논문은 검증 가능한 보상을 사용하는 GRPO 스타일의 강화학습에서 발생하는 집계 편향을 식별하고 해결하기 위해, 긍정 및 부정 응답에 대한 토큰 수준의 그래디언트를 재결합하기 전에 각각 별도로 평균화하는 균형 집계(Balanced Aggregation) 방법을 제안함으로써 추론 및 코딩 벤치마크 전반에 걸쳐 훈련 안정성과 성능을 향상시킵니다.

핵심

"Balanced Aggregation: Understanding and Fixing Aggregation Bias in GRPO" 논문에 대한 설명을 간단한 언어와 창의적인 비유로 풀어냅니다.

큰 그림: AI 에게 퍼즐을 풀게 하기

수학 문제를 풀거나 코드를 작성하도록 로봇을 훈련한다고 상상해 보세요. 로봇에게 프롬프트를 주면, 로봇은 답을 생성해 보려고 합니다. 이를 가르치기 위해 **검증 가능한 보상을 활용한 강화 학습 (Reinforcement Learning with Verifiable Rewards, RLVR)**이라는 방법을 사용합니다.

이를 텔레비전 퀴즈 쇼라고 생각해 보세요. 로봇 (AI) 은 단일 질문에 대해 여러 가지 다른 답 (응답) 을 생성합니다. 심판 (간단한 컴퓨터 프로그램) 이 이를 확인합니다.

• 답이 맞으면 로봇은 "엄지척" (긍정적 보상) 을 받습니다.
• 답이 틀리면 로봇은 "엄지내림" (부정적 보상) 을 받습니다.

목표는 로봇이 "엄지척" 답을 더 많이, "엄지내림" 답은 더 적게 생성하도록 가르치는 것입니다. 이 논문은 간단하고 효과적이라는 이유로 인기 있는 GRPO라는 특정 훈련 방법에 초점을 맞춥니다.

문제: 표를 어떻게 세울 것인가

논문의 핵심 쟁점은 미묘하지만 결정적인 질문입니다: 로봇이 답의 집합을 생성할 때, 무엇을 배울지 "평균적인 교훈"을 어떻게 계산할까요?

로봇은 한 번에 16 개의 답을 생성할 수 있습니다. 어떤 것은 짧고 (5 단어), 어떤 것은 길며 (500 단어), 어떤 것은 맞고 어떤 것은 틀립니다. 훈련 알고리즘은 이러한 개별 단어들을 하나의 큰 "업데이트"로 결합하여 로봇의 두뇌를 개선해야 합니다.

사람들이 이를 수행하는 두 가지 주요 방식이 있으며, 논문은 둘 다 숨겨진 결함이 있다고 주장합니다.

1. "단어 수" 방식 (토큰 집계)

• 작동 원리: 모든 답에서 모든 단어 (토큰) 를 세어 모두 평균냅니다.
• 결함 (장황한 악당): 시험을 치르는 학생 그룹을 상상해 보세요.
  • 학생 A는 정답을 맞췄지만 매우 짧고 간결한 설명 (10 단어) 을 썼습니다.
  • 학생 B는 오답을 냈지만 방대하고 산만한 에세이 (500 단어) 를 썼습니다.
  • 단순히 단어 수만 세면, 학생 B 의 오답은 학생 A 의 정답보다 평균에서 50 배 더 많은 "가중치"를 갖게 됩니다.
  • 결과: AI 는 혼란을 겪습니다. 공간만 더 차지한다는 이유로 길고 틀린 답이 더 중요하다고 생각하게 됩니다. 이를 **"신호 - 길이 결합 (Sign-Length Coupling)"**이라고 합니다. 답의 길이가 우연히 교훈의 부호 (긍정 또는 부정) 를 바꿔버리는 것입니다.

2. "개인별" 방식 (시퀀스 집계)

• 작동 원리: 먼저 각 답마다 개별적으로 평균 교훈을 계산한 다음, 그 답들을 평균냅니다.
• 결함 (게으른 유권자): 같은 학생 예시를 사용해 보세요.
  • 학생 A(짧고, 정답) 는 1 표를 받습니다.
  • 학생 B(길고, 오답) 는 1 표를 받습니다.
  • 결과: 이는 "장황한 악당" 문제를 해결합니다. 하지만 이제 10 단어 답과 500 단어 답을 정확히 동일하게 취급합니다. AI 가 길고 상세한 설명에서 많이 배운다면, 이 방식은 그 추가 노력을 무시합니다. 긴 응답을 "가중치를 낮춰" 짧은 응답과 똑같이 단순한 것으로 취급합니다.

해결책: "균형 잡힌 집계 (Balanced Aggregation, BA)"

저자들은 **균형 잡힌 집계 (Balanced Aggregation, BA)**라는 새로운 방법을 제안합니다. 이는 이전 두 방법의 결함을 모두 수정하는 똑똑한 심판과 같습니다.

작동 원리:

1. 답을 분류: 먼저 심판이 답을 두 더미로 나눕니다. "좋음" 더미 (엄지척) 와 "나쁨" 더미 (엄지내림) 입니다.
2. 더미 내부 단어 세기: "좋음" 더미 안에서는 모든 단어를 세어 평균냅니다. "나쁨" 더미 안에서도 모든 단어를 세어 평균냅니다.
3. 더미 균형 맞추기: 마지막으로 두 더미를 합칩니다. 하지만 여기서 핵심은 무작위로 섞지 않는다는 점입니다. 각 더미에 몇 단어가 있든 상관없이 "좋음" 더미와 "나쁨" 더미가 최종 결정에 동등한 영향을 미치도록 합니다.

비유:
새로운 공원을 건립하는지 투표하는 마을 의회를 상상해 보세요.

• 옛 방식 1 (단어 수): 틀렸더라도 가장 오래 말하는 사람이 가장 많은 표를 얻습니다.
• 옛 방식 2 (개인별): 한 사람이 50 페이지 보고서를 작성하고 다른 사람이 단순히 "예"라고 말하더라도 모든 사람이 1 표를 얻습니다.
• 균형 잡힌 집계: 의회는 "공원 찬성" 그룹과 "공원 반대" 그룹으로 나뉩니다. 각 그룹 내부의 논쟁을 평균낸 다음, "찬성" 그룹과 "반대" 그룹에게 최종 결정에서 동등한 가중치를 부여합니다. 이렇게 하면 논쟁의 길이가 결과를 왜곡하지 않도록 보장합니다.

무엇을 발견했는가?

연구진은 수학 및 코딩 데이터셋을 사용하여 두 가지 다른 AI 모델 (Qwen2.5-Math-7B 및 Qwen3-1.7B) 에서 이 새로운 방법을 테스트했습니다.

1. 안정성이 핵심: 이전 방법들은 초기에는 잘 작동하다가 훈련 후반부에 충돌하거나 불안정해지는 경우가 많았습니다. 특히 AI 가 매우 길고 틀린 답을 쓰기 시작했을 때 "단어 수" 방식이 특히 불안정했습니다.
2. 더 나은 결과: 균형 잡힌 집계 방식은 일관되게 더 나은 최종 점수를 기록했습니다. 더 안정적이었으므로 AI 는 성능의 극심한 변동 없이 꾸준히 학습했습니다.
3. 왜 중요한가: 이 논문은 AI 훈련의 "최고"의 방법이 답의 길이 변동 정도에 달려 있음을 보여줍니다.
   • 답의 길이가 극단적으로 다양하면 "단어 수" 방식은 위험할 수 있습니다.
   • "좋음"과 "나쁨" 답의 길이 차이가 크다면 "개인별" 방식은 불공평할 수 있습니다.
   • 균형 잡힌 집계는 각 방법의 특정 편향을 수정하므로 두 상황 모두에서 잘 작동합니다.

결론

이 논문은 AI 훈련에서 "재료를 섞는" 방식 (데이터 집계) 이 단순한 사소한 기술적 세부 사항이 아니라, AI 가 효과적으로 학습할지 아니면 혼란을 겪을지를 결정하는 주요 설계 선택 사항이라고 결론지었습니다. 단순히 평균을 내기 전에 "좋은" 예시와 "나쁜" 예시를 분리함으로써, 저자들은 AI 에게 추론과 코딩을 가르치는 데 더 견고하고 안정적이며 효과적인 방법을 만들었습니다.

기술 요약

기술 요약: 균형 잡힌 집계 (Balanced Aggregation): GRPO 내 집계 편향의 이해 및 해결

문제 제기

검증 가능한 보상을 활용한 강화 학습 (RLVR) 은 대규모 언어 모델 (LLM) 의 추론 및 코드 생성 능력을 향상시키는 표준 패러다임이 되었으며, Group Relative Policy Optimization(GRPO) 은 그 단순성과 별도의 크리틱 (critic) 부재로 인해 널리 채택된 방법론입니다. 그러나 GRPO 내의 중요한 설계 선택 중 하나인 샘플링된 그룹 내 토큰 수준 정책 기울기 항에 대한 집계 규칙은 아직 충분히 탐구되지 않았습니다.

현재의 관행은 일반적으로 두 가지 범주로 나뉩니다:

1. 시퀀스 집계 (Sequence Aggregation): 표준 GRPO 의 기본값으로, 각 응답 내 토큰 기여도를 먼저 평균화한 후 응답 간에 평균화합니다. 이는 각 시퀀스가 토큰 수와 관계없이 동등하게 기여하므로, 더 긴 응답에 대해 암묵적으로 가중치를 낮추는 결과를 초래합니다.
2. 토큰 집계 (Token Aggregation): DAPO 및 Dr.GRPO 와 같은 최근 연구에서 주장하는 방식으로, 샘플링된 그룹 내 모든 토큰에 대해 잘린 목적 함수 (clipped objective) 를 직접 평균화합니다.

본 논문은 이러한 두 규칙이 체계적으로 다른 최적화 편향을 유발한다고 규명합니다:

• 토큰 집계는 **부호 - 길이 결합 편향 (sign-length coupling bias)**을 도입합니다. 양의 편차 (advantage > 0) 와 음의 편차 (advantage < 0) 샘플의 상대적 기여도는 정규화된 편차뿐만 아니라 평균 응답 길이에도 의존합니다. 양의 응답과 음의 응답의 길이 분포가 다르다면, 토큰 집계는 업데이트의 한쪽을 체계적으로 증폭시켜 불안정한 학습 역학을 초래할 수 있습니다.
• 시퀀스 집계는 각 응답에 동등한 가중치를 부여함으로써 부호 - 길이 결합을 제거합니다. 그러나 손실을 토큰 단위가 아닌 시퀀스 단위로 평균화하므로, 더 긴 응답이 암묵적으로 하향 평가되는 **시퀀스 동등 가중치 편향 (sequence equal-weighting bias)**을 도입합니다.

어떤 접근법도 보편적으로 최적이지 않습니다. 각 방법의 효과성은 응답 길이의 분산과 양/음성 샘플 간의 길이 차이에 따라 달라집니다.

방법론: 균형 잡힌 집계 (Balanced Aggregation, BA)

이러한 편향 간의 긴장 관계를 해결하기 위해 저자들은 GRPO 스타일 RLVR 의 집계 단계를 대체할 간단한 대안인 **균형 잡힌 집계 (Balanced Aggregation, BA)**를 제안합니다.

BA 의 핵심 메커니즘은 세 단계 프로세스로 구성됩니다:

1. 분할 (Partitioning): 정규화된 편차의 부호에 따라 샘플링된 응답 그룹을 두 개의 부분 집합으로 나눕니다. 양의 부분 집합 ($S_+$) 과 음의 부분 집합 ($S_-$) 입니다.
2. 부분 집합 내 평균화 (Intra-Subset Averaging): 각 부분 집합 내에서 토큰 수준의 평균을 별도로 계산합니다. 이는 부호 그룹 내에서 토큰 수준 평균화 속성을 유지하면서, 표준 시퀀스 집계에서 발생하는 강한 시퀀스 단위 동등 가중치를 피합니다.
3. 부분 집합 간 결합 (Inter-Subset Combination): 두 부분 집합의 손실을 각 부분 집합 내 시퀀스 수에 비례하는 가중치로 결합합니다 (양의 경우 $k/G$, 음의 경우 $(G-k)/G$. 여기서 $k$는 양의 시퀀스 수입니다).

이론적 정당성:
이진 보상 GRPO 설정에서 이 특정 가중치 체계는 BA 가 시퀀스 집계와 동일한 부호 간 균형 조정 인자 ($\sqrt{k(G-k)}/G$) 를 유도함을 보장합니다. 결과적으로 BA 는 긴 응답을 처벌하는 강한 시퀀스 동등 가중치 효과를 피하면서도 시퀀스 집계의 부호 균형 속성 (부호 - 길이 결합 제거) 을 유지합니다. 논문은 또한 가중치가 시퀀스 수가 아닌 편차 질량 (advantage mass) 에 의해 결정되는 비이진 보상에 대한 일반화된 공식을 제시합니다.

주요 기여

1. 집계 편향의 통합 분석: 본 논문은 GRPO 의 손실 집계가 단순한 구현 세부 사항이 아님을 보여주는 공식적 분석을 제공합니다. 토큰 집계의 특정 "부호 - 길이 결합" 편향과 시퀀스 집계의 "시퀀스 동등 가중치" 편향을 규명합니다.
2. 균형 잡힌 집계 (BA): 부호와 길이 편향을 분리하는 간단하고 즉시 적용 가능한 대안으로 BA 를 제안합니다. 이는 부호 그룹 내에서 토큰 수준 평균화를 수행하지만, 시퀀스 수에 기반하여 그룹 간 균형을 맞춥니다.
3. 실증적 검증 및 진단 기준: 응답 길이 분산과 양 - 음성 길이 차이가 토큰 대 시퀀스 집계의 상대적 효과를 지배함을 보여주는 광범위한 실험을 제시합니다. 또한 BA 가 다양한 모델과 데이터셋에서 두 기준선 (baseline) 을 일관되게 능가함을 입증합니다.

실험 결과

저자들은 두 가지 학습 데이터셋 (DAPO-17k 및 Polaris) 에서 Qwen2.5-Math-7B와 Qwen3-1.7B를 사용하여 BA 를 평가했습니다. 성능은 Math-500, AIME 2024, AIME 2025, OlympicBench, Minerva-MATH, LiveCodeBench 등 여섯 가지 벤치마크를 통해 측정되었습니다.

주요 발견:

• 학습 안정성: 토큰 집계는 종종 학습 후기 단계에서 심각한 성능 저하 (최고점 대비 마지막 단계 하락이 큼) 를 초래하는 반면, BA 는 견고한 마지막 단계 정확도를 유지합니다.
• 모델 의존적 역학:
  • Qwen2.5-Math-7B(더 큰 응답 길이 변이를 보임) 에서는 토큰 집계가 초기에 시퀀스 집계보다 우수했으나, BA 가 최고점 및 마지막 단계 성능 모두에서 두 방법을 능가했습니다.
  • Qwen3-1.7B(더 큰 양 - 음성 길이 차이를 보임) 에서는 시퀀스 집계가 토큰 집계보다 더 안정적이었으나, BA 가 다시 최고점 및 마지막 단계 지표에서 가장 높은 성과를 거두었습니다.
• 손실 역학: 정책 기울기 손실 궤적 분석 결과, 토큰 집계는 부호 - 길이 결합으로 인해 0 에서부터 대규모 편차를 유발하는 반면, BA 와 시퀀스 집계는 0 근처에서 안정적으로 유지되는 것으로 나타났습니다.
• 전체적 성능: BA 는 모든 테스트된 환경에서 표준 토큰 및 시퀀스 집계보다 일관되게 더 강력한 최종 성능과 더 나은 학습 안정성을 제공했습니다.

중요성 및 주장

본 논문은 GRPO 스타일 RLVR 에서 집계가 사소한 구현 세부 사항이 아닌 **1 급 설계 선택 (first-class design choice)**이라고 주장합니다. 이 연구의 중요성은 다음과 같습니다:

• 안정성: BA 는 종종 토큰 집계에서 후기 단계에 관찰되는 학습 붕괴를 방지하는 더 견고한 최적화 신호를 제공합니다.
• 보편성: 특정 길이 분포 조건 하에서만 잘 작동하는 토큰 또는 시퀀스 집계와 달리, BA 는 다양한 모델 크기와 데이터셋 전반에 걸쳐 견고합니다.
• 설계 원칙: 이 연구는 효과적인 RLVR 을 위해서는 편향을 방지하기 위한 부호 간 가중치 균형을 유지하면서도 (편향 방지), 긴 응답의 신호를 보존하기 위해 부호 내 토큰 정보를 폐기하지 않는 균형이 필요함을 강조합니다.

저자들은 균형 잡힌 집계 (Balanced Aggregation) 가 GRPO 의 고유한 트레이드오프에 대한 간단하면서도 효과적인 해결책을 제공하여, 추론 및 코딩 작업에서 더 안정적인 최적화와 향상된 최종 모델 성능으로 이어진다고 결론지었습니다.