[논문 리뷰] Tractable Function-Space Variational Inference in Bayesian Neural Networks

Tim G. J. Rudner, Zonghao Chen외 2명|arXiv (Cornell University)|2023. 12. 28.

Adversarial Robustness in Machine Learning인용 수 6

한 줄 요약

이 논문은 베이지안 신경망을 위한 유의미한 사전 지식 통합과 상태별 예측 불확실성 및 정확도를 달성하기 위해 직접적으로 함수 공간에 대한 사후 분포를 모델링함으로써 확장 가능한 불확실성 추정을 가능하게 하는 유계 함수 공간 변분 추론(FSVI)을 제안한다. 함수 공간 Kullback-Leibler 발산에 대한 단순 추정기 도입으로, FSVI는 의미 있는 사전 분포 통합이 가능하며 다양한 작업, 특히 안전성이 중요한 의료 진단 벤치마크에서 최신 기술 수준의 예측 불확실성과 정확도를 달성한다.

ABSTRACT

Reliable predictive uncertainty estimation plays an important role in enabling the deployment of neural networks to safety-critical settings. A popular approach for estimating the predictive uncertainty of neural networks is to define a prior distribution over the network parameters, infer an approximate posterior distribution, and use it to make stochastic predictions. However, explicit inference over neural network parameters makes it difficult to incorporate meaningful prior information about the data-generating process into the model. In this paper, we pursue an alternative approach. Recognizing that the primary object of interest in most settings is the distribution over functions induced by the posterior distribution over neural network parameters, we frame Bayesian inference in neural networks explicitly as inferring a posterior distribution over functions and propose a scalable function-space variational inference method that allows incorporating prior information and results in reliable predictive uncertainty estimates. We show that the proposed method leads to state-of-the-art uncertainty estimation and predictive performance on a range of prediction tasks and demonstrate that it performs well on a challenging safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.

연구 동기 및 목표

• 베이지안 신경망에서 파rameter 공간 변분 추론의 한계를 해결하기 위해, 이는 종종 불확실성 정량화를 열악하게 만들고 비베이지안 방법보다 성능이 열 劣하다.
• 데이터 생성 과정에 대한 작업 특화 사전 지식을 베이지안 신경망 추론에 통합할 수 있도록 하기 위해.
• 고차원 및 과다 매개변수화된 설정에서 신뢰할 수 있는 예측 불확실성 추정을 지원하는 확장 가능하고 유계인 함수 공간 변분 추론 방법을 개발하기 위해.
• 표준 벤치마크와 실제 세계의 안전성이 중요한 의료 진단 작업에서 방법을 평가하여 분포 이탈 상황에서도 강건성을 입증하기 위해.

제안 방법

• 방법은 네트워크 파rameter에 의해 유도되는 함수 공간에 대한 사후 분포를 명시적으로 목표로 하는 베이지안 추론을 수립한다.
• 함수 공간 분포 간 Kullback-Leibler 발산에 대한 새로운 단순 추정기를 도입하여 확률적 변분 추론을 가능하게 한다.
• 각 경사 하강 단계에서 입력 포인트를 샘플링하기 위해 컨텍스트 분포 $p_{\mathbf{X}_{\mathcal{C}}}$ 를 사용하여 변분 목표에서 상한을 근사한다.
• 이미지 입력의 경우, 컨텍스트 분포는 훈련 데이터 픽셀에서 구성된 흑백 이미지 위에 균일하게 정의된다; 표본 데이터의 경우 경험적 입력 범위를 사용한다.
• 유연한 사전 설계를 지원하며, 외부 분포 영역에서 높은 불확실성을 유도하는 사전도 가능하다.
• Adam 최적화를 사용하고 표준 아키텍처인 ResNet과 MLP와 호환되는 표준 딥 러닝 프레임워크를 사용하여 구현된다.

실험 결과

연구 질문

• RQ1 고차원 및 과다 매개변수화된 신경망에 대해 함수 공간 변분 추론을 실현 가능하고 확장 가능하게 만들 수 있는가?
• RQ2 함수 공간 사전에 작업 특화 사전 지식을 통합하면 파rameter 공간 방법보다 예측 불확실성과 정확도가 향상되는가?
• RQ3 특히 안전성이 중요한 애플리케이션에서 분포 이탈 상황에서 제안된 방법의 성능은 어떠한가?
• RQ4 함수 공간 변분 추론이 베이지안 및 비베이지안 기준선을 모두 능가할 수 있는가, 특히 불확실성 추정과 예측 성능 측면에서?

주요 결과

• FSVI는 1D 회귀 및 Two Moons 데이터셋을 포함한 여러 회귀 및 분류 벤치마크에서 최신 기술 수준의 예측 성능과 불확실성 추정을 달성한다.
• Snelson 1D 회귀 작업에서 FSVI는 훈련 데이터에 잘 맞는 예측 분포를 생성하면서도 훈련 포인트에서 멀리 떨어진 영역에서는 높은 불확실성을 보이며 원하는 특성을 잘 구현한다.
• 당뇨병 망막병변 진단 작업에서 FSVI 기반의 전문가 참조 선택 예측은 베이스라인을 능가하며, 낮은 참조 비율에서도 높은 정확도를 유지한다.
• 방법은 분포 이탈에 대해 강건성을 보이며, OOD 샘플에 대해 높은 지식 불확실성(epistemic uncertainty)을 할당하여 효과적인 선택적 예측을 가능하게 한다.
• 작업 특화 사전을 사용한 FSVI는 표준 파rameter 공간 변분 추론 및 딥 앙상블 대비 불확실성 캘리브레이션과 예측 정확도를 크게 향상시킨다.
• 함수 공간 KL 발산에 대한 제안된 추정기는 안정적이고 확장 가능한 학습을 가능하게 하여 함수 공간 추론을 실세계 응용에 실현 가능하게 한다.