SPS: Steering Probability Squeezing for Better Exploration in Reinforcement Learning for Large Language Models

1. 문제 정의 (Problem)

• 배경: 최근 RL(특히 규칙 기반 보상 신호를 활용하는 RLVR) 은 LLM 의 추론 능력을 향상시키는 데 효과적이지만, 주로 **단일 샘플 성공률 (Pass@1)**을 높이는 데 초점이 맞춰져 있습니다.
• 핵심 문제: RL 학습은 다양한 추론 궤적 (reasoning trajectories) 에 대한 탐색 (Exploration) 을 제한합니다. 이는 다중 샘플 성능 (Pass@k) 향상에 부정적인 영향을 미칩니다.
• 원인 분석 (Squeezing Effect): 저자들은 RL 학습 과정에서 발생하는 '확률 밀도 수축 (Probability Squeezing)' 현상을 발견했습니다.
  • RL 의 대조적 (contrastive) 보상 최적화 과정에서 낮은 확률의 응답에 대해 음의 기울기 (negative gradient) 를 적용할 때, 해당 확률이 감소하는 대신 이미 우세한 (dominant) 응답의 확률이 비례적으로 더욱 집중되는 현상이 발생합니다.
  • 이는 Softmax 함수의 정규화 속성에서 기인하며, 결과적으로 출력 분포가 특정 정답 궤적으로 과도하게 수축 (sharpening) 되어 모델이 새로운 해결책을 탐색하는 능력을 상실하게 만듭니다.

2. 제안 방법: SPS (Steering Probability Squeezing)

이러한 수축 효과를 완화하고 탐색 능력을 강화하기 위해 SPS를 제안합니다. 이는 기존 RL 과 **역강화 학습 (IRL, Inverse Reinforcement Learning)**을 교차적으로 수행하는 학습 프레임워크입니다.

• 핵심 아이디어:
  • RL 단계: 기존 온-정책 (on-policy) RL 을 수행하여 데이터셋을 탐색하고 롤아웃 (rollouts) 을 생성합니다.
  • IRL 단계: RL 단계에서 생성된 롤아웃을 '전문가 시연 (demonstrations)'으로 간주하여 IRL 을 수행합니다. 이때 외부 감독이나 사전 지식을 도입하지 않고, 모델이 스스로 생성한 온-정책 데이터만을 사용합니다.
  • 확률 재분배: IRL 은 과도하게 우세한 응답으로 쏠린 확률 질량을 미탐색되었으나 잠재적으로 가치 있는 다른 궤적으로 재분배하도록 정책을 재구성 (reshape) 합니다.
• 구체적 기법:
  • Forward-KL 목적 함수: 생성된 롤아웃 분포와 현재 정책 분포 간의 KL 발산을 최소화하여 정책을 조정합니다.
  • 저확률 궤적 강조 (L2TE): 수축 효과가 주로 모델이 매우 낮은 확률을 가진 샘플에서 발생할 때 발생하므로, IRL 단계에서 상대적으로 낮은 모델 확률을 가진 롤아웃을 우선적으로 샘플링하여 학습 신호를 증폭시킵니다.
  • 반복 학습 전략: RL 과 IRL 단계를 반복적으로 교차 (interleaving) 하여 확률 분포가 조기에 수축되는 것을 방지하고 지속적인 탐색을 유도합니다.

3. 주요 기여 (Key Contributions)

1. RL 학습 역동성 분석 및 경험적 상한선 발견: RL 학습에서 '확률 수축 (Squeezing Effect)'이 발생하며, 이것이 Pass@k 성능의 경험적 상한선을 결정하는 근본적인 원인임을 이론적 및 실험적으로 규명했습니다.
2. SPS 알고리즘 제안: IRL 을 활용하여 RL 의 수축 효과를 명시적으로 제어하고, 확률 분포를 재분배하여 탐색 능력을 강화하는 새로운 학습 패러다임을 제시했습니다.
3. 광범위한 실험 검증: 5 개의 올림피아드 수준 수학 벤치마크 (AIME, BRUMO, HMMT 등) 에서 SPS 가 기존 RL 방법론 (GRPO, DAPO, GSPO 등) 보다 Pass@k를 획기적으로 개선함을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋 및 모델: OpenR1-Math 데이터셋 (3k, 5k, 10k 샘플) 을 사용하여 Qwen2.5-Math (1.5B, 7B) 모델로 실험 수행.
• 성능 향상:
  • Pass@128에서 일관된 개선을 보였습니다. 특히 Qwen2.5-Math-1.5B 모델은 BRUMO 벤치마크에서 Pass@128 점수 63.33을 기록하여, 기존 GRPO 대비 +10.00 포인트 향상되었습니다.
  • Pass@1 (단일 샘플) 성능은 유지하거나 소폭 개선하는 동시에, Pass@k (다중 샘플) 성능이 크게 향상되어 모델이 더 다양한 정답 궤적을 발견할 수 있음을 시사합니다.
• 모델 크기별 차이: 작은 모델 (1.5B) 의 경우 기존 GRPO 가 오히려 분포를 과도하게 수축시켜 탐색을 저해했으나, SPS 는 이를 효과적으로 보완했습니다.
• 비용 효율성: IRL 단계는 전체 학습 시간의 약 3% 만 차지하여 계산 비용 증가가 미미함에도 불구하고 큰 성능 향상을 가져왔습니다.

5. 의의 및 결론 (Significance)

• 탐색의 한계 극복: 기존 RL 기반 추론 모델이 직면한 '엔트로피 붕괴'와 '탐색 부족' 문제를 해결하는 새로운 접근법을 제시했습니다.
• 자기지도 학습의 확장: 외부 보상 모델이나 전문가 데이터 없이, 모델이 생성한 데이터만으로 IRL 을 수행하여 분포를 조절하는 방식은 비용 효율적이고 확장 가능한 솔루션입니다.
• 미래 방향: RL 과 IRL 의 교차 학습이 추론 중심 LLM 의 탐색 능력을 확장하는 효과적인 경로임을 보여주었으며, 이는 추론 및 코드 생성 등 다양한 복잡한 작업에 적용 가능한 중요한 통찰을 제공합니다.

요약하자면, 이 논문은 RL 학습 중 발생하는 **확률 수축 현상을 IRL 을 통해 역으로 조종 (Steering)**함으로써, LLM 이 더 다양한 추론 경로를 탐색하고 복잡한 문제 해결 능력을 극대화할 수 있음을 증명했습니다.