Restoring Exploration after Post-Training: Latent Exploration Decoding for Large Reasoning Models

이 논문은 강화학습 후 학습을 거친 대형 추론 모델에서 발생하는 탐색 붕괴 문제를 해결하기 위해 추가 학습 없이 중간 계층의 엔트로피를 활용하는 '잠재적 탐색 디코딩 (LED)' 전략을 제안하여 추론 성능을 효과적으로 향상시킨다는 내용을 담고 있습니다.

핵심

🧠 "잠재된 호기심"을 깨우는 마법: 대형 추론 모델의 새로운 해법

이 논문은 최근 인공지능(AI)이 수학이나 코딩 같은 어려운 문제를 풀 때, **"왜 더 많이 시도해본다고 해서 정답을 찾을 확률이 오르지 않는지"**에 대한 놀라운 발견과 그 해결책을 제시합니다.

비유를 들어 쉽게 설명해 드릴게요.

---

1. 문제: "완벽한 학생"이 된 AI 의 함정

최근의 AI(대형 추론 모델) 들은 **강화 학습 (RL)**이라는 훈련을 통해 아주 똑똑해졌습니다. 마치 시험을 볼 때 "정답만 맞으면 점수를 주고, 틀리면 감점"하는 방식으로 훈련받은 거죠.

이 훈련 덕분에 AI 는 **한 번에 정답을 맞출 확률 (Pass@1)**이 매우 높아졌습니다. 하지만 여기서 뜻밖의 문제가 생겼습니다.

• 과거의 AI: "어떤 문제를 풀까?"라고 생각할 때, 다양한 가능성을 열어두고 여러 가지 시도를 했습니다. (온도 조절을 높이면 더 창의적으로 여러 답을 냈죠.)
• 현재의 AI: 훈련을 너무 잘 받아서, 자신의 첫 번째 생각이 100% 맞다고 확신하게 되었습니다.
  • 마치 "이 문제는 내가 이미 다 알고 있어! 다른 생각은 필요 없어!"라고 외치는 과신 (Overconfidence) 상태가 된 것입니다.
  • 그래서 "다른 답도 한번 찾아봐"라고 명령해도 (온도를 높여도), AI 는 여전히 같은 생각만 반복하며 실패합니다. 이를 논문에서는 **'탐험의 붕괴 (Exploration Collapse)'**라고 부릅니다.

2. 발견: "중간 단계"에 숨겨진 보물

연구진은 AI 의 두뇌 구조를 자세히 들여다보며 흥미로운 사실을 발견했습니다.

• 최종 단계 (마지막 층): AI 가 최종적으로 내뱉는 생각은 너무 확신에 차서 단조롭고 (엔트로피가 낮음), 다른 가능성을 전혀 고려하지 않습니다.
• 중간 단계 (중간 층들): 하지만 AI 가 생각의 과정을 거치는 중간 단계에서는 여전히 **"아직 확실하지 않아, 다른 방법도 있을까?"**라는 호기심과 불확실성이 살아있습니다.

비유:
AI 가 문제를 풀 때, 마지막 결론은 "정답은 A 야!"라고 단호하게 외치지만, 생각하는 과정 중반에는 "A 일 수도 있고, B 일 수도 있고, C 일 수도 있겠네?"라고 다양한 가능성을 고민하고 있습니다.
문제는 AI 가 이 중간의 호기심을 무시하고 마지막 결론만 믿고 있다는 점입니다.

3. 해결책: LED(잠재적 탐험 해독)

이 문제를 해결하기 위해 연구진은 **LED(Latent Exploration Decoding)**라는 새로운 방법을 만들었습니다.

LED 가 어떻게 작동하나요?

1. 중간 단계의 목소리를 듣기: AI 가 최종 결론을 내기 전, 중간 단계의 생각들을 모두 모아봅니다.
2. 가장 호기심 많은 순간 찾기: 중간 단계들 중에서 **"가장 다양한 가능성을 고민하고 있는 순간 (엔트로피가 가장 높은 순간)"**을 찾아냅니다.
3. 스마트한 선택:
   • AI 가 이미 확신하는 쉬운 문제라면, 그냥 최종 결론을 따릅니다 (효율성).
   • 하지만 AI 가 헷갈려 하거나 어려운 문제라면, 중간 단계에서 찾은 가장 다양한 가능성을 선택해 새로운 길로 탐험합니다 (창의성).

창의적인 비유:
이 방법은 마치 등산과 같습니다.

• 기존 AI: 정상에 거의 다 왔다고 믿고, 한 가지 길만 고집하며 넘어질 위험이 큽니다.
• LED: "아직 정상인가? 아니면 다른 길이 있을까?"라고 중간 지대에서 잠시 멈춰서 지도를 다시 확인하고, 가장 유망한 여러 길을 동시에 탐색한 후 가장 좋은 길을 선택합니다.

4. 결과: 더 적은 비용으로 더 많은 성공

이 방법을 적용한 결과, 놀라운 성과가 나왔습니다.

• 정답률 상승: 한 번에 맞출 확률도, 여러 번 시도했을 때 맞출 확률도 모두 향상되었습니다.
• 비용 절감: AI 를 다시 훈련시키거나 복잡한 장비를 추가할 필요가 없습니다. 단순히 생각하는 방식을 조금만 바꿔주면 됩니다.
• 실제 효과: 수학 문제나 코딩 과제에서 AI 가 더 유연하게 문제를 해결하게 되었습니다.

📝 한 줄 요약

**"AI 가 너무 자신만만해져서 다른 가능성을 잃어버렸을 때, AI 가 생각하던 '중간 과정'의 호기심을 다시 꺼내들여 더 똑똑하게 문제를 해결하게 해주는 방법"**입니다.

이 기술은 AI 가 단순히 정답만 외우는 기계가 아니라, 진짜처럼 고민하고 탐험하는 지능을 되찾는 데 중요한 열쇠가 될 것입니다.

기술 요약

1. 문제 정의: RL 후학습에 의한 탐색 붕괴

최근 LRMs (예: DeepSeek-R1, Qwen3 Thinking 등) 은 강화 학습 (특히 GRPO) 을 통해 수학, 코딩, 과학 문제 해결 능력에서 뛰어난 성능을 보였습니다. 그러나 저자들은 이러한 RL 기반 후학습이 의도치 않은 부작용을 초래한다고 지적합니다.

• 최종층 확률 분포의 붕괴: RL 학습은 정답 확률을 극대화하는 방향으로 모델을 최적화하여, 최종 레이어 (Final Layer) 의 확률 분포가 매우 높은 확신 (High Confidence) 을 가지게 되고 엔트로피가 급격히 낮아집니다.
• 온도 기반 샘플링의 실패: 기존 LLM 에서는 샘플링 온도 (Temperature) 를 높여 분포를 부드럽게 만들면 다양한 답을 시도할 수 있어 pass@n (n 번 시도 중 하나라도 정답일 확률) 성능이 향상되었습니다. 하지만 RL 학습을 거친 최신 LRM 에서는 온도를 높여도 pass@n 이 향상되지 않거나 오히려 성능이 저하됩니다. 이는 최종층의 확률 분포가 이미 너무 좁게 수렴 (Collapsed) 하여 탐색이 불가능해졌기 때문입니다.
• 잠재적 엔트로피의 존재: 흥미롭게도, 최종층의 엔트로피가 낮아진 반면, 중간 레이어 (Intermediate Layers) 의 잠재적 상태 (Latent States) 에서는 여전히 상당한 불확실성 (엔트로피) 이 유지되고 있습니다.

2. 방법론: 잠재적 탐색 해독 (LED)

저자들은 최종층이 아닌 중간 레이어의 잠재적 확률 분포 (Latent Posteriors) 를 활용하여 탐색 능력을 복원하는 LED를 제안합니다. 이 방법은 추가적인 학습이나 모델 파라미터 증가 없이 적용 가능합니다.

핵심 단계

1. 잠재적 확률 분포 추출 (Early Exit):
   • 모델의 중간 레이어 (Layer $L-d+1$부터 $L$까지) 의 은닉 상태 (Hidden States) 를 언어 모델 헤드 (LM-Head) 에 직접 입력하여 각 레이어별 확률 분포를 추출합니다.
2. Top-k 필터링:
   • 최종 레이어의 Top-k 토큰 후보를 기준으로 중간 레이어의 확률 분포를 필터링합니다. 이는 희귀하거나 의미 없는 토큰이 탐색 과정에 노이즈로 작용하는 것을 방지하기 위함입니다.
3. 누적 집계 및 엔트로피 기반 선택:
   • 최종 레이어부터 중간 레이어까지의 필터링된 확률 분포를 누적 합 (Cumulative Sum) 으로 집계합니다.
   • 각 깊이 (Depth) 별 집계된 분포의 엔트로피를 계산하여, 가장 높은 엔트로피를 가진 깊이를 '탐색 후보 (Exploration Candidate)'로 선택합니다. 이는 가장 풍부한 탐색 신호를 제공하는 레이어를 자동으로 찾습니다.
4. 탐색과 활용의 균형 (Balancing):
   • DeepThink 단계 (추론 중): 모델의 확신도가 낮을 때 (Top-1 확률이 낮을 때) LED 를 적용하여 잠재적 공간을 탐색합니다.
   • 답변 생성 단계: 확신도가 높거나 trivial 한 토큰일 때는 표준 디코딩 (Exploitation) 을 사용하여 효율성을 유지합니다.

3. 주요 기여 (Contributions)

1. 현상 규명: RL 기반 후학습이 최종층 확률 분포의 엔트로피 붕괴를 유발하며, 이로 인해 기존 온도 기반 탐색이 무효화됨을 실증적으로 분석하고, 중간 레이어에 '잠재적 엔트로피 저수지 (Latent Entropy Reservoir)'가 남아있음을 발견했습니다.
2. 새로운 디코딩 전략 제안: 추가 학습 없이 중간 레이어의 잠재적 표현을 활용하여 탐색을 복원하는 LED를 제안했습니다.
3. 광범위한 실험 검증: 5 개의 모델 (Qwen3, MiMo, DeepSeek 등) 과 6 개의 벤치마크 (수학, 과학, 코딩) 에서 일관된 성능 향상을 입증했습니다.

4. 실험 결과

• 성능 향상:
  • Pass@1: 평균 0.61% 향상 (기존 CoT 대비).
  • Pass@16: 평균 1.03% 향상. 이는 모델의 탐색 능력이 복원되었음을 의미합니다.
  • 다양한 모델 (Qwen3-4B-T, MiMo-7B-RL 등) 에서 DoLa, SoftThinking 등 기존 최첨단 베이스라인보다 우수한 성능을 기록했습니다.
• 온도 민감성 회복: LED 를 적용하면 RL 학습된 모델에서도 온도를 높였을 때 pass@n 성능이 향상되는 양상 (양의 기울기 $\alpha$) 이 다시 나타납니다.
• 효율성:
  • 추가 학습 비용이 없으며, 추론 오버헤드는 미미합니다 (생성 길이 증가 < 1%).
  • Case Study: AIME 2025 문제에서 LED 는 CoT 가 실패한 문제에서 "wait"이라는 토큰을 선택하여 사고의 전환점을 찾고 정답을 도출하는 것을 확인했습니다.

5. 의의 및 결론

이 논문은 RL 기반 후학습이 LRM 의 '생각하는 과정 (DeepThink)'에서 발생하는 다양성 (Exploration) 을 억압한다는 근본적인 문제를 지적하고, 이를 해결하기 위해 모델 내부의 중간 레이어 정보를 활용하는 새로운 패러다임을 제시했습니다.

• 기술적 의의: 단순히 출력층을 조작하는 것을 넘어, 모델의 내부 표현 (Latent Representation) 을 디코딩 과정에 통합하여 효율적인 탐색을 가능하게 합니다.
• 실용적 가치: 추가 학습 없이 기존 모델을 즉시 개선할 수 있어, 코드 생성, 수학 증명, 복잡한 추론이 필요한 실제 응용 분야에서 모델의 신뢰성과 정확도를 높이는 데 기여할 수 있습니다.

요약하자면, LED는 RL 학습으로 인해 '단단해져버린' 모델의 마지막 단추를 풀고, 그 안에 숨겨진 '유연한 사고 과정'을 다시 끌어내어 더 나은 답을 찾도록 돕는 혁신적인 디코딩 기술입니다.