Reverse Probing: Supervised Token-level Uncertainty Quantification for Large Language Models in Clinical Text

기술 요약: 임상 텍스트를 위한 토큰 수준 불확실성 정량화를 위한 역탐사 (Reverse Probing)

문제 제기
임상 텍스트 요약에 대형 언어 모델 (LLM) 을 배포하는 것은 신뢰성과 관련된 중대한 안전 문제를 야기합니다. 불확실성 정량화 (UQ) 는 일반 도메인에서 잘 연구되어 왔으나, 임상 텍스트에 적용할 때는 고유한 도전 과제에 직면합니다: 임상 오류는 종종 전체적으로 정확한 문장 내의 단일 단어나 숫자에서 발생하는 희소하며, 오류의 비용은 매우 높습니다. 기존 UQ 방법은 주로 시퀀스 수준에서 작동하거나 불확실성을 추정하기 위해 여러 출력 후보를 생성 (샘플링) 하는 데 의존합니다. 이러한 접근 방식은 다음과 같은 이유로 임상 환경에 적합하지 않습니다:

1. 세분화 (Granularity): 긴 문서 내의 특정 토큰이나 구간에 대한 불확실성을 국소화하지 못합니다.
2. 효율성: 긴 임상 문서를 반복적으로 샘플링하는 것은 계산 비용이 매우 큽니다.
3. 데이터 부족: 전문가가 주석한 임상 텍스트는 희소하여 새로운 생성물을 신뢰성 있게 평가하기 어렵습니다.
4. 아키텍처: 전통적인 토큰 수준 탐사 (예: 마스킹 언어 모델링) 는 BERT 와 같은 양방향 모델에는 내재되어 있지만, 자기회귀 디코더 전용 LLM 에는 내재되어 있지 않습니다.

또한, 이 논문은 UQ 를 환각 (hallucination) 탐지와 구별합니다. UQ 는 기준 진실 (ground truth) 에 독립적으로 모델이 자신의 예측에 대해 가지는 신뢰도를 측정하는 자기 평가의 한 형태인 반면, 환각 탐지는 외부 사실 확인에 의존합니다.

방법론: 역탐사 (Reverse Probing)
저자들은 새로운 텍스트를 생성하지 않고 기존에 레이블이 지정된 요약에서 직접 토큰 수준의 불확실성을 추정하는 프레임워크인 역탐사 (Reverse Probing) 를 제안합니다. 핵심 혁신은 표준 UQ 패러다임을 "역전"시키는 것입니다: 모델에게 텍스트를 생성하게 하고 출력을 평가하는 대신, 이 방법은 기존 임상 텍스트를 동결된 (frozen) LLM 에 다시 입력하여 내부 상태를 탐사합니다.

방법론은 다음과 같이 진행됩니다:

1. 프롬프트 기반 마스킹: 디코더 전용 모델과의 아키텍처 격차를 해소하기 위해, 시스템은 임상 요약의 각 대상 토큰을 순차적으로 마스킹 자리표시자로 대체합니다. 그런 다음 모델은 주변 문맥을 기반으로 마스킹된 토큰을 예측하도록 프롬프트됩니다.
2. 문맥 교란 ("역" 메커니즘): 이 프레임워크는 환자별 증거의 원천으로 간단한 병원 과정 (Brief Hospital Course, BHC) 을 활용합니다. 각 토큰에 대해 두 번의 순방향 전달을 수행합니다:
   • 문맥 포함: BHC 와 함께 요약을 입력합니다.
   • 문맥 제외: BHC 없이 요약을 입력합니다.
     핵심 가설은 지지되는 토큰은 모델의 표현 공간에서 BHC 와 정렬되는 반면, 지지되지 않는 (환각된) 토큰은 이러한 근거가 부족하다는 것입니다. 정렬의 부재는 잠재적 불확실성 신호로 작용합니다.
3. 특징 추출: 시스템은 모델의 내부 활성화 및 출력 분포에서 네 가지 범주의 특징을 추출합니다:
   • 내부 신호: 숨겨진 상태 노름, 레이어 간 변화 (L2 거리, 코사인 유사도), 그리고 주의 메트릭 (엔트로피, BHC 로 향하는 가중치, 레이어 간 드리프트).
   • 불확실성 신호: 출력 분포 형태 메트릭 (섀넌 엔트로피, 마진, 지니 계수, 자유 에너지).
   • 의학적 지식 신호: scispaCy/UMLS 를 사용한 개체명 인식 (NER) 을 통한 도메인 특화 문맥, 의료 개체 밀도, 그리고 토큰 빈도/IDF.
   • 분포 신호: "BHC 포함"과 "BHC 제외" 전달 간의 확률 및 엔트로피 차이 (예: $\Delta E$, $\Delta H$, KL 발산).
4. 지도 분류: 이러한 특징을 기반으로 지도 분류기 (XGBoost 또는 CatBoost) 를 훈련시켜 각 토큰에 대한 연속적인 불확실성 점수를 예측합니다. 이 모델은 지지되지 않는 구간과 관련된 내부 활성화 패턴을 식별하도록 훈련됩니다.

주요 기여

1. 임상 요약을 위한 최초의 토큰 수준 UQ 프레임워크: 역탐사는 임상 요약 작업에서 토큰 수준 불확실성을 추정하도록 특별히 설계된 최초의 프레임워크입니다.
2. 계산 효율성: 다중 샘플링을 제거하고 기존 텍스트를 기반으로 작동함으로써, 이 방법은 추론 시간을 샘플링 기반 기준선에서 요구되는 몇 시간에서 몇 분으로 단축합니다.
3. 새로운 벤치마크: 저자들은 두 개의 전문가 주석 데이터셋 (Hallucinations-MIMIC-DI 및 Hallucinations-Generated-DI) 에서 여섯 개의 LLM(일반 및 생의학 모델을 포함) 에 대해 여덟 개의 적응된 기준선에 대해 평가함으로써 임상 요약용 UQ 에 대한 최초의 벤치마크를 확립했습니다.
4. 해석 가능한 통찰력: 이 연구는 불확실성과 상관관계가 있는 특정 내부 신호를 식별하여 임상 환경에서 설명 가능한 AI 로 가는 길을 제시합니다.

결과
평가 결과, 역탐사는 모든 적응된 기준선을 크게 능가하는 것으로 나타났습니다:

• 성능: MIMIC-DI 데이터셋 (인간 작성 요약) 에서 역탐사는 BioMistral-7B 를 사용하여 0.5478의 AUPRC 를 달성한 반면, 가장 강력한 기준선 (SE Probes) 은 0.1313이었습니다. 이는 약 4 배의 개선을 의미합니다.
• 강건성: 이 방법은 다양한 크기 (7B~70B) 와 사전 훈련 조건 (일반 대 생의학) 의 모델에서 일관된 우위를 유지하지만, 일반적으로 70B 모델보다 작은 모델 (7B/8B) 이 더 좋은 성능을 발휘했습니다. 저자들은 더 큰 모델이 과도한 확신을 가지며 불확실성 신호의 동적 범위를 압축할 수 있다고 가정합니다.
• 효율성: 역탐사는 단일 GPU 에서 100 개의 샘플을 10 분 미만으로 처리하는 반면, 샘플링 기반 방법 (예: Semantic Entropy) 은 동일한 작업에 약 7 시간이 소요됩니다.
• 특징 분석: 특징 중요도 분석에 따르면, 델타 에너지 (BHC 가 존재할 때의 자유 에너지 변화) 와 이웃 문맥 (슬라이딩 윈도우 통계) 이 모든 모델에서 불확실성의 가장 일관된 예측 변수임이 밝혀졌습니다. 이는 임상 텍스트의 불확실성이 종종 고립된 토큰이 아닌 지역적 구간의 속성임을 시사합니다.

의의 및 주장
이 논문은 새로운 텍스트를 생성하여 평가하는 것이 불가능한 고위험 임상 환경에서 불확실성 추정을 위한 실용적이고 확장 가능한 솔루션으로 역탐사를 제시합니다. 기존 텍스트를 탐사로 간주함으로써, 이 방법은 예측을 특정 내부 신호 (특히 임상 증거와의 정렬 부재) 로 거슬러 추적하여 불확실성 추정을 해석 가능한 과정으로 변환합니다.

저자들은 겸손하게 다음과 같은 한계를 지적합니다:

• 모델 범위: 이 프레임워크는 주로 Mistral 및 Llama 아키텍처에서 검증되었습니다.
• 대규모 모델 적응성: 70B 모델에서 성능이 저하되어 더 깊은 네트워크를 위한 적응형 특징 추출 전략이 필요함을 시사합니다.
• 데이터 부족: 전문가가 주석한 데이터셋에 대한 의존성은 평가 규모를 제한합니다.
• 작업 일반화: 이 연구는 퇴원 요약에 국한되어 있으며, 다른 임상 텍스트 유형 (예: 방사선 보고서) 으로 일반화되는 것은 아직 검증되지 않았습니다.

궁극적으로 이 연구는 모델이 지지되지 않는 내용에 반응하여 내부적으로 불확실성을 인코딩하며, 역탐사는 생성적 샘플링의 계산 오버헤드 없이 이러한 신호에 효율적으로 접근할 수 있는 메커니즘을 제공한다고 주장합니다.