Calibrating Overconfidence Without Sacrificing Confidence: Probe-Conditioned Head Intervention for LLMs

기술 요약: LLM을 위한 프로브 조건부 헤드 개입 (Probe-Conditioned Head Intervention)

문제 정의
대규모 언어 모델(LLM)은 종종 "언어적 과잉 확신(verbalized overconfidence)"을 보이며, 틀린 답에 대해서도 높은 확신을 표현하곤 한다. 온도 스케일링(temperature scaling)이나 사후 지시어 튜닝(post-hoc instruction tuning)과 같은 표준 교정 기법들은 일반적으로 출력 분포에 전역적인 조정을 적용한다. 이러한 방법들은 집계된 교정 지표를 개선할 수는 있지만, 결정적인 한계를 갖는다. 즉, 정답에 대한 정당한 확신을 동시에 훼손하지 않으면서 틀린 답에 대한 부당한 확신만을 선택적으로 억제할 수 없다는 점이다. 이는 잘못된 양성(틀렸으나 확신함)을 줄이는 것이 필연적으로 잘못된 음성(맞았으나 확신하지 못함)을 증가시키는 결과를 초래하여, 두 현상 사이의 긴장을 유발한다. 본 논문은 더 선택적인 접근 방식이 필요하다고 주장한다. 즉, 응답이 '틀렸지만 확신하는(wrong-but-confident)' 상태일 가능성이 높을 때만 개입함으로써, 부당한 확신의 억제와 정당한 확신의 손실을 분리해야 한다는 것이다.

방법론: 프로브 조건부 헤드 개입 (PCHI)
저자들은 PCHI를 제안한다. 이는 추론 시점에 모델 내부에서 작동하며, 과잉 확신의 증거에 따라 어텐션 헤드(attention-head) 출력을 선택적으로 재조정하는 추론 시점(inference-time) 방법이다. 이 방법은 세 가지 핵심 구성 요소로 이루어진다:

1. 구조화된 자기 평가(Structured Self-Evaluation): 모델은 reasoning(추론), answer(답변), 그리고 이진 형태의 is_confident(yes/no) 필드를 포함하는 구조화된 JSON 응로를 생성하도록 지시받는다. 개입은 "틀렸지만 확신하는(WY)" 그룹(틀린 답을 내놓으면서도 높은 확신을 보이는 경우)을 대상으로 하며, "맞았고 확신하는(CY)" 그룹은 보존하는 것을 목표로 한다.
2. 확신 템플릿 프로브(Confidence-Template Probes): 동결된 선형 프로브(frozen linear probe)가 WY 응답의 잠재적 증거를 탐지하도록 훈련된다. 이 프로브는 특정 "확신 템플릿" 토큰(답변과 최종 확신 값 사이의 토큰 구간)에서의 은닉 상태(hidden states)를 읽는다. 프로브는 확신 있는 응답이 실제로 틀렸을 확률을 예측한다. 프로브 레이어($\ell_p$)는 WY/CY의 구별 가능성이 강한 네트워크 중간 지점으로 선택되며, 이를 통해 하위 레이어들이 개입을 위해 사용될 수 있도록 보장한다.
3. 조건부 헤드 재조정(Conditional Head Rescaling): 추론 과정에서 프로브 점수는 학습된 다운스트림 어텐션 헤드 출력의 재조정치를 게이팅(gating)한다.
   • $\ell_p$의 하위 레이어에 있는 각 어텐션 헤드 $h$에 대해, 학습된 스칼라 계수 $g_h$(초기값 1)가 적용된다.
   • 개입 강도 $s_{i,t}$는 프로브 점수 $p_{i,t}$에 의해 결정된다. 점수가 임계값 $\tau$(예: 0.5)를 초과하면, 헤드 출력 $z_{i,t,h}$는 다음과 같이 재조정된다: $z'_{i,t,h} = (1 + (g_h - 1)s_{i,t}) z_{i,t,h}$.
   • 만약 프로브가 해당 응답을 '틀렸지만 확신하는' 것으로 플래그를 지정하지 않으면($s_{i,t} = 0$), 모델은 베이스라인과 동일하게 동작한다.
   • 선택적 사항으로, 프로브 이후에 어텐션 마스크를 적용하여 템플릿 쿼리가 참조하는 컨텍스트 영역(프롬프트, 추론, 답변)을 제한함으로써 개입 신호를 강화할 수 있다.

학습 목적 함수는 힌지 스타일(hinge-style) 손실을 사용하여, CY 사례를 저하시키지 않으면서 WY 사례를 확신 마진 아래로 밀어내도록 하며, 불필요한 헤드 계수를 항등원 근처로 유지하기 위해 L1 페널티로 규제한다.

주요 기여

• 선택적 내부 개입: 전역적인 사후 교정과 달리, PCHI는 내부 증거에 기반하여 조건부로 개입함으로써, 전역적으로 확신을 억제하지 않고도 특정 오류 유형(WY)을 교정할 수 있게 한다.
• 교정과 정확도의 분리: 이 방법은 부당한 확신을 줄이면서도 정당한 확신을 대부분 보존할 수 있음을 입증하여, 균일한 재조정 방식에 내재된 트레이드오프 문제를 해결한다.
• 메커니즘적 통찰: 이 접근 방식은 언어적 확신을 단순히 매핑되어야 할 출력 분포가 아니라, 특정 지점(템플ate 토큰)에서 제어 가능한 내부 계산 과정으로 취급한다.

실험 결과
방법론은 구조화된 이진 확신 필드를 가진 Qwen3-4B-Instruct 및 Gemma3-4B-it 모델을 사용하여 OpenMathInstruct 문제에 대해 평가되었다.

• 리드아웃 토큰 개입(Readout-Token Intervention): 최종 확신 토큰에 PCHI를 적용했을 때 유의미한 개선이 나타났다. Qwen3-4B의 경우, 기대 교정 오차(ECE)는 21.9%에서 11.7%로 감소했고, AUROC는 66.5에서 90.3으로 증가했다. 결정적으로, 원래의 WY 리드아웃 중 82.2%가 "no"로 전환된 반면, 원래의 CY 리드아웃 중에는 10.8%만이 손상되었다.
• 업스트림 개입(Upstream Intervention): 상위 확신 템플릿 토큰에서 개입했을 때는 가변적인 결과가 나타났다. Qwen3-4B의 경우, 토큰 1~5에 걸친 공동 개입은 ECE를 9.2%로 낮추었으며, 63.3%의 WY 교정과 5.1%의 CY 손상을 기록했다. Gemma3-4B의 경우, 업스트림 효과는 더 약했으며 어텐션 마스킹에 더 의존적이었다. 컨텍스트를 답변 필드로 제한했을 때 WY 교정 성능은 40.4%로 향상되었으나, CY 손상이 약간 증가했다.
• 절제 연구(Ablation): 대조 실험을 통해 이러한 이득이 단순히 어텐션 마스킹이나 일반적인 활성화 스티어링(activation steering) 때문이 아님을 확인했다. 프로브의 증거와 학습된 헤드별 재조정의 결합이 관찰된 선택성을 위해 필수적이다.

의의 및 주장
본 논문은 PCHI가 모델이 올바를 때 확신을 표현하는 능력을 희생하지 않으면서도, 신뢰할 수 있는 방식으로 불확실성을 표현할 수 있는 길을 제시한다고 주장한다. 교정을 전역적 매핑이 아닌 선택적 제어 문제로 다룸으로써, 저자들은 내부 개입이 부당한 확신의 억제와 정당한 확신의 손실을 부분적으로 분리할 수 있음을 보여주었다.

저자들은 연구의 범위를 겸허히 유지하며, 결과가 수학적 질문 답변 및 구조화된 JSON 출력과 이진 확신 필드에 특화되어 있음을 명시한다. 또한 이 방법이 모델 및 데이터 특정적인 프로브 훈련에 의존한다는 점을 인정하며, 이를 자유 형식의 확신 표현이나 개방형 작업으로 확장하는 것은 여전히 해결해야 할 과제로 남겨두었다. 이 연구는 언어적 확신에 대한 완전한 메커니즘적 설명을 제공하려는 것이 아니라, 그러한 증거가 디코딩 가능하며 선택적 교정을 위해 실행 가능하다는 증거를 제공하는 데 목적이 있다.