Reward Models Know

• Related Project: Private
• Category: Paper Review
• Date: 2025-04-12

Reasoning Models Know When They’re Right: Probing Hidden States for Self-Verification

• url: https://arxiv.org/abs/2504.05419
• pdf: https://arxiv.org/pdf/2504.05419
• html: https://arxiv.org/html/2504.05419v1
• abstract: Reasoning models have achieved remarkable performance on tasks like math and logical reasoning thanks to their ability to search during reasoning. However, they still suffer from overthinking, often performing unnecessary reasoning steps even after reaching the correct answer. This raises the question: can models evaluate the correctness of their intermediate answers during reasoning? In this work, we study whether reasoning models encode information about answer correctness through probing the model’s hidden states. The resulting probe can verify intermediate answers with high accuracy and produces highly calibrated scores. Additionally, we find models’ hidden states encode correctness of future answers, enabling early prediction of the correctness before the intermediate answer is fully formulated. We then use the probe as a verifier to decide whether to exit reasoning at intermediate answers during inference, reducing the number of inference tokens by 24\% without compromising performance. These findings confirm that reasoning models do encode a notion of correctness yet fail to exploit it, revealing substantial untapped potential to enhance their efficiency.

TL;DR

핵심 내용
이 연구는 추론 모델의 내부 은닉 상태에 정답 여부(correctness) 정보가 내재되어 있음을 보여줍니다.

• 프로빙(Probe) 기법을 통해 체인 오브 사고(Chain-of-Thought) 중간 단계에서 추론 과정의 정답 여부를 판단할 수 있으며, 이 정보는 높은 ROC-AUC와 낮은 예상 보정 오차(ECE)를 나타냅니다.
• 또한, 모델의 은닉 상태는 정답이 완전히 생성되기 전에도 “미리보기” 형태의 신호를 전달해, 모델이 언제 올바른 답변에 도달했는지 예측할 수 있습니다.
• 이를 활용하여 신뢰 기반 얼리 스타핑(early-exit) 전략을 도입하면, 불필요한 추론 단계를 줄여 토큰 생성량을 최대 24%까지 감소시키면서도 성능 저하 없이 효율적인 추론이 가능함을 보여줍니다.

목차

1. 서론 및 배경
   • 1.1. 추론 모델의 발전과 한계: 과도한 추론(Overthinking) 문제
   • 1.2. “자기 검증(self-verification)”의 필요성
2. 관련 연구
   • 2.1. LLM의 불확실성 평가: 블랙박스 vs. 화이트박스 접근
   • 2.2. 추론 효율 개선을 위한 인퍼런스 기법 및 얼리 스타핑 전략
   • 2.3. 외부 vs. 내부 검증자(verifier) 기법 비교
3. 프로빙을 통한 중간 답변 정답 여부 추출
   • 3.1. 데이터 수집
     • 체인 오브 사고(Chain-of-Thought)를 여러 “청크(chunk)”로 분리
     • 각 청크 내에서 중간 답변(extracted intermediate answers)을 Gemini 2.0 Flash 등 도구로 추출하고, 정답 여부 라벨 할당
   • 3.2. 프로빙 모델 훈련
     • 은닉 상태(마지막 토큰의 표현)를 입력으로 하는 간단한 2층 MLP 프로브
     • 불균형 데이터 문제 해결을 위해 가중치가 부여된 이진 교차 엔트로피 손실(weighted binary cross-entropy loss)
4. 실험 및 평가
   • 4.1. 실험 설정
     • 사용 데이터셋: 수학 문제(GSM8K, MATH, AIME)와 논리 문제(KnowLogic)
     • 다양한 모델(예: DeepSeek-R1-Distill 시리즈, QwQ-32B)을 통한 평가
   • 4.2. 인-도메인 평가 결과
     • ROC-AUC, ECE, Brier score 등 지표로 내부 표현에서 정답 여부를 정확하게 예측하는 성능 검증
     • 많은 경우 단순 선형 프로브(숨은 차원 d=0)로도 우수한 결과 도출
   • 4.3. 도메인 간 일반화
     • 수학 추론 데이터셋 간, 그리고 논리와 수학 간 전이 성능 분석
   • 4.4. Long CoT와의 관계
     • 길게 구성된 추론(Chain-of-Thought)에서 정답 여부 정보가 더욱 명확하게 인코딩됨
     • 비추론(instruct) 모델과의 비교를 통해 학습 과정에서 획득되는 자체 검증 능력 강조
   • 4.5. 정답 생성 전 ‘미리보기’ 신호
     • 중간 답변이 완전히 생성되기 전, 초기 문단 단위의 은닉 상태에서도 정답 여부의 예측 신호가 존재함을 분석
5. 프로브를 활용한 얼리 스타핑 전략
   • 5.1. 전략 개요
     • 프로브의 출력(예측된 정답 확률)을 신뢰도 척도로 사용, 일정 임계치(Tₕᵣ)를 넘으면 추가 추론 단계를 중단
     • 모델이 과도하게 토큰을 생성하는 ‘과잉 추론(overthinking)’을 방지
   • 5.2. 실험 결과
     • 신뢰 기반 얼리 스타핑를 도입 시, 정답률 유지와 함께 평균 토큰 생성량을 약 19~24% 절감
     • 동일 토큰 사용량 대비 고정 얼리 스타핑(static early-exit) 방식보다 최대 5% 정확도 향상
6. 논의 및 인사이트
   • 6.1. 자체 검증 능력의 내재성
     • 모델 내부 표현에서 정답 여부 정보가 쉽게 추출 가능한 점을 통해, LLM이 ‘스스로 맞는지를 안다’는 가능성 확인
     • 단, 모델은 이 정보를 효과적으로 활용하지 못해 과도한 추론을 지속함
   • 6.2. 내부 메커니즘과 온-폴리시(on-policy) 제어의 잠재력
     • 경량화된 내부 프로브를 활용한 인퍼런스 제어 가능성 및 효율성 개선 제안
7. 결론 및 미래 연구 방향
   • 정리: 추론 모델은 자기 검증 정보를 은닉 상태에 풍부하게 인코딩하고 있으나, 이를 인퍼런스 단계에서 충분히 활용하지 못함
   • 제안: 보다 효과적인 온-폴리시 제어 및 얼리 스타핑 전략 개발, 다양한 도메인 및 모델에 대한 자기 검증 능력 연구

1. 서론 및 배경

• 모델의 검색 기반 추론
  최신 추론 모델은 문제 해결 과정에서 여러 경로(Chain-of-Thought)를 탐색하는 등 강력한 문제 해결 능력을 보여줍니다. 그러나 올바른 답변에 도달한 후에도 여전히 추가적인 추론 단계를 수행하는 ‘과잉 추론(overthinking)’ 현상이 나타납니다.

• 자기 검증(self-verification)의 필요성
  이러한 문제를 해결하기 위해, 모델이 자신의 중간 결과에 대해 정답 여부를 자체적으로 평가(자기 검증)할 수 있는지, 그리고 내부 은닉 상태가 그러한 정보를 어느 정도 담고 있는지를 탐구하고자 합니다.

2. 관련 연구

• 불확실성 및 신뢰도 평가
  기존 연구들은 주로 LLM이 제공하는 자가 보고(self-reported) 신뢰도나 외부 검증자를 활용한 방식을 사용했지만, 내부 은닉 상태를 활용하는 ‘화이트박스’ 접근법이 더 우수함이 제안되었습니다.

• 인퍼런스 효율 및 얼리 스타핑
  모델이 과도하게 추론하지 않고 필요한 경우 조기에 종료할 수 있도록 하는 연구가 진행되어 왔으며, 본 논문은 이와 연계하여 내부 신호를 활용하는 방안을 제시합니다.

3. 프로빙을 통한 중간 답변 정답 여부 추출

• 데이터 수집
  • 체인 오브 사고 전체를 문단 단위로 분리한 후, 각 청크에서 중간 답변을 추출합니다.
  • Gemini 2.0 Flash와 같은 도구를 활용해 추출된 답변에 대해 정답 여부를 자동으로 라벨링합니다.
• 프로브 모델 훈련
  • 각 청크의 마지막 토큰 은닉 상태를 특징 벡터로 사용합니다.
  • 간단한 2층 MLP를 이용해, 은닉 상태에서 정답 여부를 예측하도록 학습시키며, 불균형 데이터 문제를 해결하기 위해 가중치 조정 방식을 적용합니다.

4. 실험 및 평가

• 실험 설정
  • 수학 문제와 논리 문제 데이터셋을 통해 프로브의 인-도메인 성능을 평가합니다.
  • 다양한 모델(DeepSeek-R1-Distill 시리즈, QwQ-32B 등)을 대상으로 평가해, 모델 규모나 도메인에 따른 내부 신호의 강도를 비교합니다.
• 결과 분석
  • ROC-AUC, ECE, Brier score 등 다양한 평가 지표를 통해, 프로브가 높은 정확도와 신뢰도로 중간 답변의 정답 여부를 예측함을 확인합니다.
  • 또한, 학습 데이터 분포에 따른 도메인 간 전이 능력과, 길게 구성된 추론(Chain-of-Thought) 과정에서 정답 여부 신호가 더욱 명확하게 인코딩된다는 점이 드러납니다.
  • 심지어 답변이 완전히 생성되기 전, 초기 문단 수준의 은닉 상태에서도 예측 신호가 존재함을 분석합니다.

5. 프로브를 활용한 얼리 스타핑 전략

• 전략 개요
  • 학습된 프로브를 인퍼런스 시 검증자로 활용합니다.
  • 각 중간 답변에 대해 프로브가 산출한 신뢰도(정답 예측 확률)가 미리 정해진 임계치(Tₕᵣ)를 넘으면 추가 추론을 중단하고, 해당 답변을 최종 답변으로 채택합니다.
• 실험 결과
  • 신뢰 기반 얼리 스타핑는 정답률(약 88%)을 유지하면서 토큰 생성량을 최대 24%까지 절감하는 효과를 보여줍니다.
  • 동일한 토큰 사용량 대비 고정된 얼리 스타핑 방식에 비해 정답률이 최대 5%까지 향상되어, 내부 정보를 효과적으로 활용할 수 있음을 시사합니다.

6. 논의 및 인사이트

• 내부 검증 정보의 잠재력
  • 본 연구 결과는 모델의 은닉 상태에 정답 여부와 관련된 신호가 풍부하게 인코딩되어 있음을 보여줍니다.
  • 모델은 이미 자기 검증 능력을 내재화하고 있음에도 불구하고, 인퍼런스 시 이를 충분히 활용하지 않아 과도한 추론을 수행합니다.
• 향후 연구 방향
  • 온-폴리시 제어(on-policy control)와 같은 기법을 통해, 모델이 내부 신호를 더 효과적으로 활용하도록 유도하는 방법을 연구할 필요가 있습니다.
  • 또한, 자기 검증 능력의 학습 기원과, 다양한 모델 및 도메인에 걸친 일반화를 심도 있게 분석할 필요가 있습니다.

7. 결론 및 미래 연구 방향

• 결론
  • 추론 모델은 내부 은닉 상태에 정답 여부 정보를 잘 인코딩하고 있으며, 이를 간단한 프로빙 모델로 쉽게 추출할 수 있습니다.
  • 이 정보를 활용한 얼리 스타핑 전략은 인퍼런스 효율을 크게 향상시키며, 모델이 불필요한 추론을 줄이는 데 기여합니다.
• 미래 연구
  • 내부 검증 신호의 활용을 극대화할 수 있는 보다 정교한 온-폴리시 제어 메커니즘 개발
  • 다양한 도메인, 모델 및 학습 환경에서 자기 검증 능력의 기원과 확산 과정 분석
  • 효율적이면서도 신뢰도 높은 추론 전략 개발을 통한 실제 응용 분야에서의 성능 향상

모델의 내부 은닉 상태에서 정답 여부 정보를 효과적으로 추출하는 프로빙 기법과, 이를 활용한 얼리 스타핑 전략을 통해 추론 효율을 향상시킬 수 있음을 체계적으로 보여줍니다.