Skip-Connected Policy Optimization for Implicit Advantage

1. 문제 제기 (Problem)

대규모 언어 모델 (LLM) 의 추론 능력을 향상시키기 위해 **검증 가능한 보상 (Verifiable Rewards) 을 통한 강화 학습 (RLVR)**이 주류 패러다임으로 자리 잡았습니다. 특히 **그룹 상대적 정책 최적화 (GRPO)**는 결과 기반의 보상과 그룹 내 상대적 이점 (Advantage) 계산을 통해 안정적인 학습을 보여주고 있습니다.

그러나 기존 연구들은 다음과 같은 구조적 딜레마에 직면해 있습니다:

• 미세한 밀집 보상 (Fine-grained Dense Rewards) 의 이론적 매력: 토큰 단위의 상세한 보상을 제공하면 추론 과정의 각 단계에서 더 정교한 학습이 가능해져 성능 상한선이 높아질 수 있습니다.
• 실제 샘플링 비용의 한계: 몬테카를로 (Monte Carlo, MC) 샘플링을 통해 토큰 단위의 보상을 추정하려면 많은 샘플이 필요합니다. 하지만 실제 제한된 샘플링 예산 (예: 그룹당 8 개) 하에서는 초기 추론 토큰들에 대해 이점 (Advantage) 추정의 분산이 매우 크고 부호 (Sign) 가 불일치하는 문제가 발생합니다.
• 역설적 결과: 이러한 높은 분산과 부호 오류로 인해, 미세한 보상을 적용한 방법이 단순한 결과 기반의 GRPO 보다 오히려 성능이 저하되는 역설이 발생합니다. 즉, 기존 방법들은 미세한 정보의 유효성과 계산 효율성 사이에서 타협을 강요받고 있습니다.

2. 제안 방법론: SKPO (Methodology)

이러한 딜레마를 해결하기 위해 저자들은 **Skip-Connected Policy Optimization (SKPO)**을 제안합니다. SKPO 는 추론 과정을 **상류 (Upstream)**와 하류 (Downstream) 두 단계로 분해하고, 스킵 연결 (Skip Connection) 구조를 통해 밀집 보상의 이점을 유지하면서 샘플링 비용을 GRPO 수준으로 낮춥니다.

핵심 아키텍처 및 프로세스

1. 상류 (Upstream) 단계:

   • 모델이 문제 $q$에 대해 초기 추론 세그먼트 $s$를 생성합니다 (조기 종료).
   • 이 단계에서는 **단일 스트림 최적화 (Single-stream Optimization, SPO)**를 사용합니다. 그룹 내 상대적 이점 계산이 불가능하므로, 과거 샘플을 기반으로 한 시간적 기준선 (Temporal Baseline) 을 사용하여 안정적인 학습 신호를 제공합니다.
   • 밀집 보상 획득: 상류 세그먼트 $s$에 대한 보상은 하류에서 생성된 여러 완성된 답변들을 통해 몬테카를로 방식으로 추정합니다. 즉, $s$가 올바른 해답으로 이어질 확률을 하류 샘플들의 평균 정확도로 평가합니다.

2. 하류 (Downstream) 단계:

   • 입력은 스킵 연결된 형태인 $[s, q]$ (상류 세그먼트 + 원본 문제) 입니다.
   • 이 입력으로부터 $G$개의 완성된 답변을 생성합니다.
   • 그룹 상대적 최적화 (Group Relative Optimization): 하류 단계에서는 충분한 샘플 ($G$개) 이 존재하므로 GRPO 와 유사한 그룹 상대적 이점 계산을 적용하여 안정적인 학습을 수행합니다.
   • 스킵 연결의 역할: $[s, q]$ 입력은 두 가지 기능을 합니다.
     1. 상류의 추론이 하류에 영향을 미치게 하여 (밀집 보상 학습).
     2. 동시에 원본 문제 $q$에 직접 접근할 수 있는 경로를 제공하여, 상류의 추론이 틀렸을 경우 모델이 이를 우회하거나 수정할 수 있는 **탐색의 자유 (Exploration Freedom)**를 보장합니다.

3. 공학적 최적화 (Engineering Optimization):

   • naive 한 구현은 상류 생성 후 하류 생성을 별도의 배치로 나누어 오버헤드를 발생시킵니다.
   • SKPO 는 **vLLM 엔진을 수정하여 KV 캐시 (Key-Value Cache) 를 중간에 재작성 (Rewriting)**하는 방식을 사용합니다. 하나의 GPU 배치 내에서 모든 샘플을 병렬로 생성하고, 중간 지점에서 선택된 상류 세그먼트의 KV 캐시 포인터를 공유함으로써 GRPO 와 동일한 계산 비용을 달성합니다.

비대칭 최적화 전략 (Asymmetric Optimization)

• 상류: 샘플이 적으므로 시간적 기준선 (Temporal Baseline) 을 사용하는 SPO 알고리즘 적용.
• 하류: 샘플이 충분하므로 그룹 상대적 이점을 사용하는 GRPO 변형 적용.

3. 주요 기여 (Key Contributions)

1. 문제 규명: 실제 샘플링 예산 하에서 미세한 몬테카를로 보상이 실패하는 이유 (높은 분산과 빈번한 부호 오류) 를 정량적으로 분석하고 증명했습니다.
2. SKPO 제안: 상류 - 하류 구조와 스킵 연결을 통해, 상류에는 밀집 보상을 제공하고 하류에는 그룹 상대적 최적화를 적용하는 새로운 아키텍처를 설계했습니다. 이는 추가적인 샘플링 비용 없이도 단계별 정보를 활용하게 합니다.
3. 암시적 이점 (Implicit Advantage) 발견: 최종 정답 여부가 동일한 경우에도, SKPO 로 학습된 모델은 중간 단계의 추론 품질이 더 높음을 발견했습니다. 이는 외부 평가자를 통한 오프 - 정책 (Off-policy) 분석을 통해 확인되었습니다.

4. 실험 결과 (Results)

• 벤치마크: Qwen2.5-Math-7B 와 Llama-3.2-3B 모델을 사용하여 수학 (In-Domain) 및 일반 추론/코드 생성 (Out-of-Domain) 태스크에서 평가했습니다.
• 성능 향상:
  • Qwen2.5-Math-7B: 평균 42.5% 점수를 기록하여 2 위인 CISPO(40.9%) 보다 1.6%p, GRPO(30.2%) 보다 12.3%p 높은 성능을 보였습니다.
  • Llama-3.2-3B: 평균 24.1% 점수로 2 위 DAPO(22.7%) 보다 1.4%p, GRPO(18.9%) 보다 5.2%p 우위를 점했습니다.
  • OOD 일반화: 수학뿐만 아니라 MMLU-Pro(일반 추론) 와 LiveCodeBench(코드 생성) 에서도 일관된 성능 향상을 보여주어 구조의 일반화 능력을 입증했습니다.
• 스킵 연결 검증: 스킵 연결이 없는 경우 (단순 접두사 사용) 하류의 탐색 다양성이 급격히 떨어지고 이점 신호가 0 이 되는 현상이 발생했으나, SKPO 는 이를 방지하여 탐색의 자유를 유지했습니다.

5. 의의 및 결론 (Significance)

이 논문은 RLVR 분야에서 미세한 보상 (Fine-grained Rewards) 과 계산 효율성 사이의 긴장 관계를 해결하는 새로운 패러다임을 제시합니다.

• 구조적 해결책: 단순히 샘플링 수를 늘리는 것이 아니라, 추론 과정을 분할하고 스킵 연결을 도입하여 비용 없이 단계별 학습 신호를 획득하는 방법을 제시했습니다.
• 내재적 추론 품질 향상: 외부 감독 없이도 모델이 더 질 높은 중간 추론 단계를 생성하도록 유도한다는 점은 (Implicit Advantage), 추론 모델의 신뢰성과 해석 가능성을 높이는 중요한 통찰입니다.
• 실용성: vLLM 기반의 KV 캐시 재작성 기법을 통해 기존 GRPO 와 동일한 학습 속도를 유지하므로, 실제 대규모 모델 학습에 즉시 적용 가능한 실용적인 솔루션입니다.

요약하자면, SKPO 는 제한된 샘플링 예산 하에서 발생하는 몬테카를로 추정의 불안정성을 구조적으로 우회하여, 효율적이면서도 고품질의 단계별 추론 학습을 가능하게 하는 혁신적인 강화 학습 알고리즘입니다.