UAT‑LITE 사전학습 변환기의 추론 시 불확실성 인식 어텐션

초록

UAT‑LITE는 사전학습된 BERT와 같은 변환기 모델에 대해 추론 단계에서 Monte‑Carlo 드롭아웃을 이용해 토큰 수준의 에피스테믹 불확실성을 추정하고, 이를 어텐션 가중치에 직접 반영함으로써 내부 증거 집합 과정을 불확실성에 민감하게 만든다. 모델 파라미터를 수정하거나 추가 학습 없이 구현 가능하며, 레이어별 분산 분해를 통해 불확실성이 깊이마다 어떻게 축적되는지 진단한다. SQuAD 2.0, MNLI, SST‑2 등에서 기존 BERT‑base 대비 평균 20 % 수준의 Expected Calibration Error(ECE)를 감소시키면서 정확도는 유지한다.

상세 분석

UAT‑LITE는 기존의 사후 보정(예: Temperature Scaling)과 달리, 모델 내부의 어텐션 연산 자체에 불확실성 정보를 주입한다는 점에서 혁신적이다. 구체적으로, 추론 시에도 드롭아웃을 유지하고 M번(보통 3~10회) 전방 패스를 수행해 각 토큰의 임베딩 표본을 얻는다. 이 표본들의 표준편차를 토큰‑레벨 불확실성 U(x_j)로 정의하고, 이를 어텐션 로그잇 a_ij에 exp(−λ·u_ij) 형태의 감쇠 인자로 곱한다. 여기서 u_ij는 Q‑only, K‑only, QK‑sum 등 여러 변형 중 하나이며, 논문에서는 Q‑only가 가장 좋은 보정‑정확도 트레이드오프를 보였다. λ는 하이퍼파라미터로, 불확실성이 큰 토큰의 기여를 크게 억제한다.

불확실성‑가중 어텐션은 각 레이어와 헤드에 동일하게 적용되며, 토큰‑레벨 불확실성은 매 패스마다 지연된 추정값 ˆU^{(m‑1)}(x_j)를 사용해 순환적으로 업데이트된다. 이렇게 하면 별도의 임베딩‑전용 샘플링 단계가 필요 없으며, 전체 연산 비용은 M배 정도만 증가한다.

또한, 저자들은 법칙의 전체 분산(Law of Total Variance)을 이용해 레이어별 기여도를 정량화하는 “레이어‑와이즈 불확실성 분해”를 제안한다. 이는 각 레이어의 출력 분산을 재귀적으로 분해해 어느 단계에서 불확실성이 급증하는지를 시각화할 수 있게 해준다. 비록 인과관계를 증명하지는 않지만, 모델 디버깅 및 설계에 유용한 진단 도구로 활용 가능하다.

실험에서는 BERT‑base를 기본 모델로 사용하고, UAT‑LITE와 기존 방법(Temperature Scaling, Deep Ensembles, MC‑Dropout만 적용)들을 동일한 MC‑샘플 수(M=5) 하에 비교했다. 주요 평가지표는 Expected Calibration Error(ECE)와 선택적 예측(selective prediction)에서의 AUC‑RC(Area Under Risk‑Coverage)이다. SQuAD 2.0의 답변 가능성 판단, MNLI의 인페런스, SST‑2의 감성 분류 모두에서 UAT‑LITE는 평균 20 % 정도 ECE를 감소시키면서 정확도 손실은 거의 없었다. 특히 분포 이동이 있는 MNLI‑mismatched와 도메인 전이된 MedQA/ PubMedQA에서도 불확실성‑가중 어텐션이 과신을 억제해 높은 신뢰도 임계값에서의 정밀도를 크게 향상시켰다.

한계점으로는 MC‑드롭아웃을 위한 추가 연산 비용이 존재하고, λ와 M 같은 하이퍼파라미터 선택이 태스크마다 민감하게 작용한다는 점이다. 또한 현재는 인코더‑기반 모델에만 적용했으며, 디코더‑전용 생성 모델에 대한 확장은 향후 연구 과제로 남아 있다.