기술 요약: 언어 모델 궤적에서 보정된 불확실성 읽기

1. 문제 제기

구조화된 출력을 가진 대규모 언어 모델 (LLM) 에 대한 불확실성 정량화 (UQ) 의 현재 접근 방식은 계산 효율성과 보정 (calibration) 사이에서 절충을 강요받습니다. 기본 방법인 최대 소프트맥스 확률 (MSP) 은 계산 비용이 거의 들지 않지만, 종종 잘못 보정되어 잘못된 예측에 높은 확신을 부여합니다. 반면, 내부 활성화 상태를 탐구하는 기존 방법들은 일반적으로 단일 계층에서 추출하거나 평균화한 숨겨진 상태를 정적 스냅샷으로 취급하여, 표현이 형성되는 계층별 궤적을 무시합니다. 이러한 정적 관점은 네트워크 깊이에 걸쳐 증거가 어떻게 누적되고, 강화되며, 혹은 반전되는지를 간과하여, 최종 출력 분포에는 나타나지 않지만 계산 경로에 존재할 수 있는 불확실성 신호를 가릴 수 있습니다.

2. 방법론

저자들은 모델의 활성화 상태를 정적 상태가 아닌 **표현 궤적 (representational trajectories)**으로 취급하는 방법을 제안합니다. 이 접근법은 세 가지 주요 구성 요소를 포함합니다:

2.1. 기질: MLP 쓰기 벡터 (MLP Write-Vectors)

이 방법은 원시 숨겨진 상태를 분석하는 대신, 각 트랜스포머 블록의 다층 퍼셉트론 (MLP) 이 잔류 스트림에 기여하는 **MLP 쓰기 벡터 ($m_\ell$)**에 초점을 맞춥니다. 이전의 기계적 해석 가능성 연구는 MLP 가 사실적 지식 저장의 핵심 장소임을 시사합니다.

• 이 방법은 단일 순전파 (forward pass) 동안 최종 프롬프트 위치 (읽기 위치) 에서 이러한 벡터들의 시퀀스를 포착합니다.
• 이 벡터들의 누적 합을 계산하여 표현 공간 내의 이산 궤적을 구성합니다: $s_\ell = \sum_{k \le \ell} m_k$.
• 궤적은 초기 상태 ($s_0 = 0$) 에서 최종 종점 ($s_L$) 까지 정의됩니다.

2.2. 기하학적 특징 추출

이 궤적은 네 가지 범주로 그룹화된 계층당 11 개의 척도 불변 기하학적 특징을 사용하여 요약됩니다:

1. 깊이 분포: 계산 노력 (업데이트 크기) 이 계층 간에 어떻게 분포되는지 측정합니다 (예: 앞쪽 집중 vs 뒤쪽 집중).
2. 국소 궤적 형태: 방향적 일관성 (연속 코사인), 급격한 변화 (곡률), 그리고 업데이트가 이전 상태를 강화하는지 아니면 모순하는지 (업데이트 - 상태 정렬) 를 포착합니다.
3. 종점과의 관계: 계층이 최종 방향과 어떻게 관련되는지 측정합니다. 여기에는 수렴 (최종까지의 방향), 개별 업데이트의 정렬 (최종까지의 업데이트), 그리고 종점에 적극적으로 반대하는 계층 식별 (모순적 지지) 이 포함됩니다.
4. 효율성: 순 변위와 총 경로 길이의 비율 (누적 일관성) 을 측정하여 궤적이 종점에 얼마나 직접적으로 도달하는지 나타냅니다.

2.3. 희소 선형 프로브 (Sparse Linear Probe)

이러한 특징들은 오류 확률을 예측하기 위해 **희소 선형 프로브 (희소 로지스틱 회귀)**에 입력됩니다.

• 입력: 계층당 11 개 특징의 벡터 ($\phi(x)$) 와 모델의 MSP 점수 ($p_{msp}$).
• 상호작용: 모델은 특징과 MSP 간의 상호작용 항을 포함하여, 프로브가 모델의 기준 확신도에 따라 기하학적 신호를 다르게 가중치할 수 있도록 합니다.
• 학습: 프로브는 오류 지시자에 대해 클래스 균형 이진 교차 엔트로피를 최소화하도록 학습되며, 희소성을 강제하기 위해 엘라스틱넷 페널티로 정규화됩니다. 이를 통해 모델은 별도의 선택 단계 없이 예측 특징을 직접 식별할 수 있습니다.

3. 주요 기여

1. 궤적 기반 UQ: 저자들은 네트워크 깊이에 걸친 표현의 진화를 설명하는 기하학적 특징의 컴팩트한 집합을 제안합니다. 이 특징들을 희소 선형 프로브에 입력하면, 선택적 거절 (selective abstention) 하에서 MSP 보다 우수한 성능을 보이며, 성능 향상은 기준 모델의 보정 오류가 클수록 커집니다 (최대 21 AURC 포인트 개선).
2. 종단 간 해석 가능성: 원시 활성화에 대한 불투명한 밀집 프로브와 달리, 모든 특징은 폐쇄형 기하학적 의미를 가집니다. 프로브의 계수는 오류가 어떻게 발생하는지 드러냅니다: 어떤 계층이 조기에 결정을 내리는지, 어떤 계층이 진행 중인 상태와 모순되는지, 그리고 궤적이 종점에서 어떻게 이탈하는지 식별합니다.
3. 포괄적인 경험적 평가: 이 방법은 9 개의 지시 튜닝 LLM (Qwen, Llama, DeepSeek 계열, 3B 에서 72B 파라미터까지) 에 대해 5 가지 다양한 NLP 작업 (MMLU, CosmosQA, HellaSwag, HaluDial, HaluSum) 에서 평가되었습니다.

4. 실험 결과

• 성능: 프로브는 45 개의 모델 - 데이터셋 쌍 중 41 개에서 MSP 보다 성능이 향상되었습니다. 향상은 MSP 가 가장 부진하게 작동하는 곳 (높은 보정 오류) 에서 가장 두드러졌으며, 가장 어려운 구성에서 AURC (위험 - 커버리지 곡선 아래 면적) 를 11.45 에서 21.83 포인트까지 감소시켰습니다.
• 기준 모델과의 비교: 프로브는 전체 숨겨진 상태로 학습된 고용량 "원시 활성화 프로브"와 비교하여 24 개 구성에서 일치하거나 우위를 점했으며, 6 개 다른 구성에서는 11 개의 스칼라 특징만 사용했음에도 2 AURC 포인트 이내로 유지되었습니다.
• 선택성: 프로브는 매끄럽고 단조로운 위험 - 커버리지 곡선을 생성하여 MSP 에서 흔히 보이는 불안정한 낮은 커버리지 스파이크를 피합니다. MSP 확신 점수가 동일할 때도 올바른 예측과 잘못된 예측을 효과적으로 분리합니다.
• 해석 가능성 발견:
  • 정답 vs 오답: 동일한 MSP 점수를 가진 예측은 뚜렷한 궤적 서명을 보입니다. 예를 들어, 잘못된 예측은 종종 초기 계층에서 최종 방향과 조기에 정렬된 후 "이중 후퇴"나 이탈을 보입니다.
  • 작업 특이성: 오류의 기하학적 서명은 작업에 따라 다릅니다. HaluSum에서는 오류가 조급한 결정을 followed by 늦은 상태 파괴 수정으로 나타났습니다. CosmosQA에서는 오류가 정답 방향에 수직인 중간 깊이 이탈로 나타났습니다.
  • 계수 분석: 프로브의 계수는 오류 메커니즘을 국소화합니다. 예를 들어, HaluSum에서는 오류 신호가 네트워크의 마지막 4 분의 1 에 집중된 반면, CosmosQA에서는 깊이에 걸쳐 분산되어 있었습니다.

5. 중요성과 주장

이 논문은 활성화 상태를 정적 스냅샷이 아닌 **운동 (motion)**으로 취급함으로써, 밀집 프로브의 불투명성과 MSP 의 제한된 신호 사이에서 "처리 가능한 중간 경로"를 제공한다고 주장합니다.

• 감사 가능한 불확실성: 특징들이 본질적으로 해석 가능하기 때문에, 결과적인 확신 신호는 전문가가 모델이 어디서 불확실한지뿐만 아니라 그 불확실성이 계산의 어디서 (예: 특정 계층 또는 기하학적 이탈) 발생하는지 검사할 수 있게 합니다.
• 효율성: 이 방법은 추론 시 거의 추가적인 계산 비용을 들이지 않으며, 단일 순전파와 간단한 기하학적 계산만 필요로 하여 샘플링 기반 또는 앙상블 접근법의 높은 비용을 피합니다.
• 한계: 저자들은 그들의 분석이 이산 선택 설정 (다중 선택 질문) 으로 제한되며, 프로브가 현재 (모델, 작업) 당 적합되어 있고, 모델 및 작업 간 전이 가능성은 여전히 열린 질문이라고 지적합니다.

요약하자면, 이 연구는 언어 모델이 예측을 구성하는 방식의 기하학이 최종 출력 분포에는 보이지 않는 풍부하고 해석 가능한 불확실성 신호를 포함하고 있음을 보여주며, 더 신뢰할 수 있고 감사 가능한 선택적 예측을 가능하게 합니다.