UAT-LITE: 사전 학습된 트랜스포머를 위한 추론 시점 불확실성 인식 어텐션

UAT-LITE: Inference-Time Uncertainty-Aware Attention for Pretrained Transformers

신경망 NLP 모델은 종종 잘못된 예측에 높은 확신을 부여하는 등 캘리브레이션이 제대로 되지 않아, 선택적 예측과 고위험 환경에서의 배포를 저해한다. 사후 보정 방법들은 출력 확률만 조정할 뿐 내부 연산은 변경하지 않으며, 앙상블 및 베이지안 접근법은 불확실성을 개선하지만 상당한 학습 또는 저장 비용이 소요된다. 본 논문에서는 사전 학습된 트랜스포머 분류기 내에서 몬테카를로 드롭아웃을 통한 근사 베이지안 추론을 활용하여, 셀프 어텐션이 불확실성을 인식하도록 만드는 추론 시점 프레임워크인 UAT-LITE를 제안한다. 확률적 순전파를 통해 토큰 수준의 인식적 불확실성을 추정하고, 이를 사용하여 사전 학습된 가중치나 학습 목표의 수정 없이 문맥화 과정에서 셀프 어텐션을 조절한다. 또한 예측 불확실성이 트랜스포머의 깊이에 따라 어떻게 누적되는지 진단하기 위해 층별 분산 분해를 도입한다. SQuAD 2.0 답변 가능성 판별, MNLI, SST-2 태스크 전반에 걸쳐 UAT-LITE는 파인 튜닝된 BERT-base 베이스라인 대비 태스크 정확도를 유지하면서 예상 캘리브레이션 오류(ECE)를 평균 약 20% 감소시켰으며, 분포 변화 상황에서의 선택적 예측 성능과 견고성을 향상시켰다.

Neural NLP models are often miscalibrated, assigning high confidence to incorrect predictions, which undermines selective prediction and high-stakes deployment. Post-hoc calibration methods adjust output probabilities but leave internal computation unchanged, while ensemble and Bayesian approaches improve uncertainty at substantial training or storage cost. We propose UAT-LITE, an inference-time framework that makes self-attention uncertainty-aware using approximate Bayesian inference via Monte Carlo dropout in pretrained transformer classifiers. Token-level epistemic uncertainty is estimated from stochastic forward passes and used to modulate self-attention during contextualization, without modifying pretrained weights or training objectives. We additionally introduce a layerwise variance decomposition to diagnose how predictive uncertainty accumulates across transformer depth. Across the SQuAD 2.0 answerability, MNLI, and SST-2, UAT-LITE reduces Expected Calibration Error by approximately 20% on average relative to a fine-tuned BERT-base baseline while preserving task accuracy, and improves selective prediction and robustness under distribution shift.