[논문 리뷰] BayesFormer: Transformer with Uncertainty Estimation

Karthik Abinav Sankararaman, Sinong Wang외 1명|arXiv (Cornell University)|2022. 06. 02.

Model Reduction and Neural Networks인용 수 8

한 줄 요약

BayesFormer는 새로운 드롭아웃 메커니즘을 사용한 근사 변분 추론을 통해 트랜스포머 모델에 이론적으로 탄탄한 불확실성 추정 방법을 도입한다. 주어진 어텐션 및 피드포워드 레이어에 베이지안 정규화와 함께 드롭아웃을 적용함으로써 예측 불확실성의 校정을 향상시키고 과적합을 줄이며, 활성 학습 및 장문 맥락 모델링을 포함한 다양한 NLP 작업에서 성능을 향상시킨다.

ABSTRACT

Transformer has become ubiquitous due to its dominant performance in various NLP and image processing tasks. However, it lacks understanding of how to generate mathematically grounded uncertainty estimates for transformer architectures. Models equipped with such uncertainty estimates can typically improve predictive performance, make networks robust, avoid over-fitting and used as acquisition function in active learning. In this paper, we introduce BayesFormer, a Transformer model with dropouts designed by Bayesian theory. We proposed a new theoretical framework to extend the approximate variational inference-based dropout to Transformer-based architectures. Through extensive experiments, we validate the proposed architecture in four paradigms and show improvements across the board: language modeling and classification, long-sequence understanding, machine translation and acquisition function for active learning.

연구 동기 및 목표

• 작은 데이터 또는 고위험 시나리오에서 트랜스포머 모델의 수학적으로 탄탄한 불확실성 측정의 부족을 해결한다.
• 실제 배포 환경에서 예측 신뢰도 추정의 강인성과 과적합을 줄이고 신뢰성을 향상시킨다.
• 더 정확한 불확실성 추정을 통해 샘플 선택에 활용함으로써 더 나은 활성 학습 성능을 가능하게 한다.
• 트랜스포머 아키텍처에 특화된, 이론적으로 탄탄한 드롭아웃 기반의 불확실성 추정 프레임워크를 개발한다.
• 피니튜닝, 장문 맥락 이해, 기계 번역을 포함한 다양한 NLP 패러다임으로의 일반화를 입증한다.

제안 방법

• 가중치에 대한 사후 분포를 유도하기 위해 근사 변분 추론을 적용하고, 가중치에 대해 정규 prior를 사용한다.
• 어텐션 및 피드포워드 하위층 모두에 드롭아웃을 적용하는 새로운 드롭아웃 구조를 도입하며, 각각 별도의 드롭아웃 비율 0.05를 사용한다.
• 11회의 전방 계산을 수행하는 몬테카를로 드롭아웃(MCDropout)을 사용해 예측 불확실성과 지식 기반 불확실성(epistemic uncertainty)을 추정하고, BALD 할당 함수를 활용한다.
• 증거 하한 경계(ELBO) 목적함수를 사용해 BayesFormer를 훈련하며, 근사 사후와 진짜 사후 간의 KL 발산을 최소화한다.
• 기존 트랜스포머 아키텍처를 유지하면서, 기존 드롭아웃 비율을 변경하지 않은 채 어텐션 및 피드포워드 레이어에만 베이지안 드롭아웃을 추가한다.
• 일반화 및 불확실성 校정을 향상시키기 위해 레이블 스무딩과 LabelSmoothedCrossEntropyCriterion를 사용해 훈련한다.

실험 결과

연구 질문

• RQ1 근사 변분 추론이 트랜스포머 아키텍처로 효과적으로 확장되어 이론적으로 탄탄한 불확실성 추정을 가능하게 할 수 있는가?
• RQ2 제안된 BayesFormer 드롭아웃 메커니즘이 표준 트랜스포머에서의 표준 MCDropout에 비해 불확실성 校정을 얼마나 향상시키는가?
• RQ3 BayesFormer의 불확실성 추정이 낮은 레이블 예산 조건에서 활성 학습 성능을 향상시키는가?
• RQ4 표준 RoBERTa에 비해 BayesFormer이 작은 데이터셋에서 피니튜닝 시 과적합을 얼마나 줄이는가?
• RQ5 BayesFormer의 불확실성 추정은 시퀀스 모델링, 번역, 제로샷 전이를 포함한 다양한 NLP 작업으로 일반화되는가?

주요 결과

• CoLA 개발 세트에서 BayesFormer는 32.1 MCC를 기록하여 RoBERTa base(21.3)를 크게 앞서며 일반화 및 불확실성 校정 향상을 입증한다.
• GLUE 벤치마크에서 BayesFormer는 SST-2에서 91.3% 정확도, MRPC에서 90.9% F1, MNLI에서 83.3% 정확도를 기록하며 다수 작업에서 RoBERTa base를 초월한다.
• 활성 학습에서, 10%의 데이터만 사용할 경우 BayesFormer는 개발 세트 MCC 62.8을 기록하여 RoBERTa보다 4점 높고, 무작위 샘플링보다 3.6점 높다.
• 가장 우수한 모델 성능는 일부 레이블된 데이터를 기각했을 때 달성되었으며, 이는 BayesFormer의 불확실성 추정이 노이즈가 많거나 모호한 예시에 대한 과적합을 피하는 데 기여함을 시사한다.
• 기본 Transformer 대비 기계 번역에서 테스트 BLEU 점수를 0.5–1.0 BLEU 포인트 향상시키며, 검증 BLEU 점수도 높아져 더 나은 일반화 능력을 보였다.
• 추가 하이퍼파라미터 튜닝 없이도 장문 맥락 이해 및 희소 학습을 포함한 모든 평가된 작업에서 성능 향상을 유지하였다.