SAGE: Shaping Anchors for Guided Exploration in RLVR of LLMs Review

0. Introduction

Paper link

Code link

SAGE는 RLVR을 볼 때 자주 놓치는 지점을 정면으로 다루는 논문이다. 최근 reasoning post-training에서 RLVR은 pass@1을 꽤 안정적으로 올린다. 하지만 pass@k까지 같이 보면 이야기가 복잡해진다. 모델이 정말 새로운 reasoning mode를 배운 것인지, 아니면 이미 base model 안에 있던 몇몇 solution path를 더 잘 뽑게 된 것인지가 분리되지 않는다.

이 논문은 이 문제를 exploration의 관점에서 본다. 표준 RLVR은 reward가 맞는 trajectory를 강화하지만, reverse KL regularization이 reference model 주변에 policy를 묶는다. 이 anchor는 안정성에는 좋지만, reference model에서 낮은 확률로만 나오던 다른 풀이 branch를 계속 낮은 확률에 묶어둘 수 있다.

한 줄 요약: SAGE는 reverse KL을 제거하지 않고, reverse KL이 anchor하는 target distribution 자체를 guide function $q(x,y)$로 재형성해, 안정성은 유지하면서 underexplored reasoning mode를 더 잘 탐색하게 만드는 RLVR exploration framework다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같다.

• RLVR의 성능 향상을 pass@1 하나로만 보면 놓치는 support collapse 문제를 직접 다룬다.
• reverse KL을 단순한 exploration 억제 장치가 아니라, reshape 가능한 anchor로 다시 해석한다.
• reward shaping, sampling temperature, entropy bonus가 아니라 KL target 자체를 바꾸는 설계를 제안한다.
• token-level novelty와 branch-level reasoning diversity를 구분해, 어떤 exploration signal이 실제 reasoning mode에 가까운지 분석한다.
• SAGE가 GRPO뿐 아니라 BNPO, DAPO 같은 다른 RLVR 계열에도 붙을 수 있음을 보여준다.

이 논문을 단순히 “또 하나의 RL trick”으로 읽으면 핵심을 놓친다. SAGE의 핵심은 policy를 더 퍼뜨리는 것이 아니라, 어디로 퍼뜨릴지에 대한 anchor를 바꾸는 데 있다. 즉 exploration을 policy sampling 문제가 아니라 regularization geometry 문제로 보는 논문이다.

1. Problem Setting

1-1. Problem definition

이 논문이 겨냥하는 문제는 RLVR에서의 reasoning mode collapse다. RLVR은 수학이나 코드처럼 verifiable reward가 있는 영역에서 효과적이다. 정답이면 reward가 1이고, 아니면 0인 식으로 reward를 만들 수 있기 때문이다. 그래서 사람이 preference label을 많이 주지 않아도 모델을 강화할 수 있다.

하지만 reasoning task에서는 하나의 문제에 여러 풀이 경로가 존재한다. 어떤 풀이는 짧고 흔하게 나오고, 어떤 풀이는 base model에서 낮은 확률로만 나타나지만 여전히 정답이다. RLVR이 좋은 post-training이라면 단순히 자주 나오던 풀이를 더 자주 만들게 하는 것을 넘어, 낮은 확률에 숨어 있던 유효한 풀이 branch도 발견해야 한다.

논문은 이 차이를 pass@1과 pass@k의 차이로 읽는다.

• pass@1이 오르는 것은 가장 높은 확률의 답변이 좋아졌다는 의미다.
• pass@k가 같이 오르는 것은 여러 sample 안에 더 다양한 정답 trajectory가 포함될 가능성이 높아졌다는 의미다.
• pass@1은 오르는데 pass@k가 정체되거나 떨어지면, 모델이 더 다양한 풀이를 찾는 것이 아니라 특정 mode를 sharpen하고 있을 수 있다.

SAGE가 던지는 질문은 다음과 같다.

reverse KL의 안정성은 유지하면서, reference distribution의 낮은 밀도 영역에 있는 valid reasoning mode까지 empirical support를 넓힐 수 있는가?

1-2. Why previous approaches are insufficient

기존 접근은 크게 세 방향으로 나뉜다.

첫째, KL penalty를 줄이거나 없애는 방법이다. 직관적으로는 policy가 reference model에서 벗어나기 쉬워지니 exploration이 늘어날 것처럼 보인다. 하지만 논문은 이 방식이 reward hacking이나 train reward와 test performance의 불일치를 만들 수 있다고 본다. KL anchor를 버리면 sparse reward에 과적합되는 위험이 커진다.

둘째, reverse KL을 forward KL로 바꾸는 방법이다. Forward KL은 support covering 성향이 강하므로 coverage metric에는 좋아 보일 수 있다. 그러나 off-target region에 probability mass를 줄 수 있고, pass@1 측면에서는 손해가 날 수 있다. 즉 coverage는 늘지만 효율이 떨어지는 문제가 생긴다.

셋째, sampling이나 reward heuristic을 조정하는 방법이다. Temperature를 올리거나 entropy bonus를 주거나 advantage shaping을 하는 방식이다. 이런 방법은 diversity를 늘릴 수 있지만, policy가 anchor되는 reference distribution 자체는 그대로다. 그래서 exploration이 여전히 reference model의 high-density region 주변에 묶일 수 있다.

SAGE의 문제의식은 여기서 출발한다. Reverse KL은 버릴 대상이 아니라 안정성을 주는 구조다. 다만 그 anchor가 원래 reference distribution 그대로일 필요는 없다.

2. Core Idea

2-1. Main contribution

표준 reverse-KL RLVR objective를 단순화하면 다음처럼 볼 수 있다.

\[J(theta) = E_{x,y \sim pi_theta} [r(x,y) - beta D_{KL}(pi_theta || pi_{ref})]\]

이 구조에서 optimal policy는 reference probability에 reward exponential tilting이 곱해진 형태가 된다. 문제는 rare but valid trajectory도 처음부터 $pi_{ref}$에서 낮은 확률이면, reward를 받아도 empirical support 안으로 잘 들어오지 못한다는 점이다.

SAGE는 여기에 guide function $q(x,y)$를 넣는다.

\[J_{sage}(theta) = E_{x,y \sim pi_theta} [r(x,y) - beta D_{KL}(pi_theta || q(x,y) pi_{ref})]\]

핵심은 policy를 그냥 reference model에 묶는 것이 아니라, $q(x,y) pi_{ref}$라는 shaped anchor에 묶는 것이다. $q$가 특정 trajectory나 token prefix를 증폭하면, reverse KL은 여전히 안정화 장치로 남으면서도 policy가 조금 다른 영역을 탐색하도록 유도한다.

논문은 이를 empirical support expansion으로 설명한다. 이론적 support는 $pi_{ref}$의 support를 넘어서지 않는다. 하지만 실제 sampling과 training에서 관측 가능한 empirical support는 달라질 수 있다. SAGE는 이 empirical support를 reward-compatible region 쪽으로 넓히려 한다.

2-2. Design intuition

SAGE의 설계 직관은 다음 세 문장으로 정리할 수 있다.

1. Reverse KL은 버리면 안 된다. Sparse reward에서 학습을 안정화하는 anchor이기 때문이다.
2. 하지만 원래 reference model은 흔한 풀이 mode에 probability mass가 몰려 있을 수 있다.
3. 그러면 anchor 자체를 shaped anchor로 바꾸어, 낮은 밀도의 유효한 branch에 더 많은 기회를 줘야 한다.

이 관점에서 SAGE는 exploration을 random noise로 보지 않는다. Exploration은 “어디서 policy를 안정화할 것인가”를 정하는 문제다. 기존 RLVR이 $pi_{ref}$를 붙잡고 있다면, SAGE는 $q pi_{ref}$를 붙잡는다.

여기서 $q$는 reward label을 직접 쓰는 oracle이 아니다. 논문은 lightweight intrinsic signal을 쓴다. 대표적으로 token surprisal과 entropy가 있다. Token surprisal은 드문 lexical choice를 증폭하고, entropy는 여러 continuation이 경쟁하는 branch point를 포착하려 한다.

3. Architecture / Method

3-1. Overview

Item	Description
Goal	RLVR에서 pass@1과 pass@k를 함께 개선하는 guided exploration
Core mechanism	reverse KL anchor를 $q(x,y) pi_{ref}$로 reshape
Main signal	surprisal, entropy 같은 intrinsic reasoning signal
Main variant	branch-level exploration
Target tasks	mathematical reasoning benchmark
Key metric	pass@1, pass@k, 특히 pass@256 또는 pass@64
Difference from prior work	KL을 제거하지 않고 KL target distribution을 바꿈

3-2. Module breakdown

1) Shaped anchor

SAGE의 중심 구성요소는 shaped anchor다. 표준 reverse KL은 policy를 $pi_{ref}$로 끌어당긴다. SAGE는 여기에 $q(x,y)$를 곱한다. 이때 $q$는 normalized distribution일 필요가 없다. 그래서 논문은 pseudo-KL 형태로 objective를 정리한다.

실무적으로 보면 $q$는 trajectory나 token prefix별 exploration prior다. $q$가 크면 해당 trajectory는 reference model에서 낮은 확률이어도 더 탐색될 여지가 생긴다. 반대로 $q$가 1에 가까우면 표준 reverse KL과 거의 비슷하게 작동한다.

2) Random exploration

가장 단순한 variant는 random exploration이다. 이는 특정 reasoning branch를 겨냥하지 않고 probability mass를 더 넓게 흔드는 방식에 가깝다. 하지만 논문은 random exploration이 targeted mode discovery에는 제한적이라고 본다. 드문 token이나 branch를 구분하지 않기 때문이다.

3) Token-level exploration

Token-level exploration은 surprisal을 이용한다. 낮은 확률 token에 더 높은 exploration weight를 주는 방식이다. 이는 lexical novelty를 늘릴 수 있다. 하지만 token rarity가 reasoning branch rarity와 같지는 않다. 문장 표현만 바뀌어도 surprisal은 높아질 수 있고, 실제 풀이 구조는 그대로일 수 있다.

따라서 token-level exploration은 다양성을 늘리지만, 논문이 원하는 “새로운 reasoning mode”를 직접 겨냥한다고 보기는 어렵다.

4) Branch-level exploration

Branch-level exploration은 entropy를 사용한다. Entropy가 높다는 것은 현재 prefix에서 여러 plausible continuation이 경쟁하고 있다는 뜻이다. 이런 지점은 풀이 strategy가 갈라지는 branch point일 수 있다.

논문은 threshold와 amplification factor를 이용해 entropy가 충분히 높은 token 위치에서 exploration을 증폭한다.

\[q_{branch}(y_t | x,y_{<t}) = 1 + gamma max(H_t - tau, 0)\]

여기서 $H_t$는 token-level entropy이고, $tau$는 activation threshold, $gamma$는 증폭 강도다. 이 방식은 모든 rare token을 키우는 것이 아니라, reasoning branch가 갈라질 가능성이 있는 지점을 더 중요하게 본다.

5) Plug-and-play integration

SAGE는 특정 RL algorithm 하나에 묶이지 않는다. 논문은 GRPO 계열뿐 아니라 BNPO, DAPO에도 SAGE를 붙여 실험한다. 중요한 점은 SAGE가 reward function이나 sampling schedule을 대체하는 것이 아니라, KL anchor distribution을 바꾸는 모듈이라는 점이다.

4. Training / Data / Recipe

4-1. Data

논문은 AIME, AMC23, MATH-500 같은 수학 reasoning benchmark를 사용한다. 이 선택은 SAGE의 목적과 맞다. 수학 문제는 verifiable reward를 만들기 쉽고, 하나의 문제에 여러 reasoning path가 존재할 수 있다.

주요 실험은 Qwen2.5-Math-7B-Base를 중심으로 한다. 추가로 DeepSeek-R1-Distill-Qwen-7B에서도 model ablation을 수행한다. 이 부분은 SAGE가 짧은 reasoning model뿐 아니라 long-reasoning distill model에서도 coverage loss를 완화할 수 있는지를 보는 장치다.

4-2. Training strategy

전체 recipe는 다음처럼 볼 수 있다.

1. Non-instruction model에는 짧은 SFT stage를 적용해 output format과 reasoning structure를 맞춘다.
2. 이후 RLVR stage에서 GRPO, BNPO, DAPO 같은 algorithm을 사용한다.
3. 표준 algorithm과 SAGE-enhanced variant를 비교한다.
4. pass@1과 pass@k를 함께 보고, accuracy와 coverage trade-off를 확인한다.

논문 appendix 기준으로 SFT stage는 Qwen2.5-Math-7B-Base에 LoRA rank 8을 사용하고, RLVR stage는 Unsloth와 vLLM 기반으로 수행된다. 세부 hyperparameter는 원문 table에서 재확인하는 편이 안전하다.

4-3. Engineering notes

실무적으로 중요한 포인트는 세 가지다.

첫째, SAGE는 reward를 더 많이 만들지 않는다. Reward signal은 그대로 두고, KL term이 policy를 끌어당기는 방향을 바꾼다. 따라서 기존 RLVR codebase에 비교적 작은 변경으로 들어갈 수 있다.

둘째, $q$는 너무 강하면 위험하다. Low-density region을 키우는 것은 exploration에는 좋지만, off-target trajectory까지 키울 수 있다. 그래서 threshold, amplification factor, decay schedule 같은 안정화 장치가 중요하다.

셋째, pass@1만 보고 tuning하면 SAGE의 장점을 놓칠 수 있다. 이 논문은 pass@k coverage를 함께 보라고 말한다. 실제 reasoning model을 서비스에 쓸 때도 single sample quality와 multi-sample search capacity는 다른 능력이다.

5. Evaluation

5-1. Main results

논문의 main result는 GRPO에 SAGE를 붙였을 때 AIME, AMC23, MATH-500에서 pass@1과 pass@256이 함께 개선된다는 것이다. 특히 branch-level variant가 평균적으로 가장 강한 결과를 낸다.

중요한 관찰은 표준 RLVR이 pass@1을 올리면서 pass@256을 떨어뜨릴 수 있다는 점이다. 이는 모델이 더 좋은 answer mode를 sharpen하는 동시에, 다양한 정답 branch를 잃고 있다는 신호로 볼 수 있다. SAGE branch variant는 이 trade-off를 완화하려 한다.

논문은 또 KL을 제거한 variant도 비교한다. KL을 없애면 train reward는 높아질 수 있지만, test-time pass rate는 더 나빠질 수 있다. 이 결과는 reverse KL이 단순히 exploration을 막는 나쁜 항이 아니라, reward hacking을 막는 안정화 항이라는 논문 주장을 뒷받침한다.

5-2. What really matters in the experiments

이 논문의 실험에서 진짜 봐야 할 것은 headline accuracy보다 다음 세 가지다.

1. pass@1과 pass@k가 동시에 좋아지는가
   • pass@1만 오르면 mode sharpening일 수 있다.
   • pass@k도 같이 오르면 support coverage가 개선되었을 가능성이 있다.
2. Branch-level signal이 token-level signal보다 나은가
   • Token surprisal은 lexical novelty를 만들 수 있다.
   • Entropy-based branch signal은 reasoning continuation이 갈라지는 지점을 더 잘 포착할 수 있다.
3. Algorithm-agnostic하게 작동하는가
   • SAGE가 GRPO 전용 trick이라면 일반성이 약하다.
   • BNPO, DAPO에도 붙여 개선을 보인다는 점은 anchor shaping이라는 추상화가 유용할 수 있음을 시사한다.

이 논문의 가장 좋은 실험은 “KL 제거가 답이 아니다”를 보여주는 부분이다. Exploration이 부족하다고 해서 anchor를 끊어버리면 안정성도 같이 잃는다. SAGE는 이 문제를 anchor weakening이 아니라 anchor shaping으로 바꾼다.

6. Limitations

1. $q$의 설계는 여전히 heuristic 성격이 강하다.
   • Entropy가 branch point와 항상 일치하지는 않는다.
   • High entropy는 문제 재서술, 표현 선택, strategy 선택 등 여러 원인에서 생길 수 있다.
2. 평가가 수학 reasoning에 집중되어 있다.
   • Code RLVR, tool-use agent, long-horizon planning에서도 같은 방식이 안정적으로 통하는지는 추가 확인이 필요하다.
3. pass@k는 coverage proxy일 뿐이다.
   • pass@k가 올랐다고 해서 사람이 의미 있는 reasoning mode가 늘었다고 단정하기는 어렵다.
   • 논문은 behavior case study를 제공하지만, mode taxonomy 자체는 더 정교해질 필요가 있다.
4. Hyperparameter sensitivity가 있다.
   • $tau$, $gamma$, decay schedule에 따라 exploration intensity가 바뀐다.
   • 실무 적용에서는 reward hacking과 exploration gain 사이를 재조정해야 한다.
5. Reference model support 밖으로는 나가지 못한다.
   • SAGE는 theoretical support를 새로 만들지 않는다.
   • 이미 base model에 희미하게라도 존재하는 branch를 empirical support 안으로 끌어오는 방법에 가깝다.

7. My Take

7-1. Why this matters for my work

SAGE는 reasoning post-training을 볼 때 중요한 기준을 하나 더 준다. 단순히 benchmark score가 올랐는지보다, policy가 어떤 reasoning support를 유지하고 있는지를 봐야 한다. 특히 multi-sample inference, self-consistency, verifier search를 같이 쓰는 시스템에서는 pass@k coverage가 실제 성능과 직결될 수 있다.

7-2. Reuse potential

이 아이디어는 다음 영역에 재사용 가능성이 있다.

• RLVR math training에서 pass@1과 pass@k를 동시에 관리하는 reward design
• Code generation에서 여러 valid implementation strategy를 유지하는 exploration policy
• Tool-use agent에서 특정 tool path로만 collapse하지 않게 하는 anchor shaping
• Distillation에서 teacher trajectory의 high-density path만 복제하지 않게 하는 guide function
• Inference-time sampling에서 entropy branch point를 중심으로 candidate를 늘리는 adaptive sampling

실제 구현 관점에서는 $q$를 어디서 계산할지, 얼마나 자주 update할지, rollout batch 안에서 어떻게 normalize할지가 중요하다. 기존 GRPO loop에 붙일 때는 reward, KL, guide log term의 scale을 별도로 logging해야 한다.

7-3. Follow-up papers

• On the Origin of LLM Reasoning Mode Collapse
• Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning
• DAPO: An Open-Source LLM Reinforcement Learning System at Scale
• DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning

8. Summary

• SAGE는 RLVR의 exploration 문제를 reverse KL anchor의 구조적 문제로 재정의한다.
• KL을 없애거나 forward KL로 바꾸는 대신, $q(x,y) pi_{ref}$라는 shaped anchor를 만든다.
• Token-level surprisal보다 entropy 기반 branch-level exploration이 reasoning mode discovery에 더 적합하다고 본다.
• 실험은 pass@1과 pass@k를 함께 봐야 RLVR의 mode sharpening과 support expansion을 구분할 수 있음을 보여준다.
• 실무적으로는 reward hacking을 피하면서 exploration을 늘리는 KL-compatible recipe로 읽을 수 있다.