[논문 리뷰] Wide Bayesian neural networks have a simple weight posterior: theory and accelerated sampling

Jiri Hron, Roman Novak외 2명|arXiv (Cornell University)|2022. 06. 15.

Gaussian Processes and Bayesian Inference인용 수 4

한 줄 요약

이 논문은 넓은 베이지안 신경망(BNN)의 사후분포를 데이터에 따라 변형하는 repriorisation을 소개한다. 이 방법은 층의 너비가 증가함에 따라 표준 정규 사전분포와의 쿨백-라이블러(KL) 발산이 0으로 수렴하도록 사후분포를 재매arameter화한다. 이로 인해 마르코프 체인 몬테카를로(MCMC) 샘플링의 혼합 속도가 빨라지며, 특히 넓은 네트워크에서 효과적 샘플 크기(ESS)가 최대 50배 향상된다. 이는 완전히 연결된 네트워크와 리서벌 네트워크 모두에서 사후 추론 속도를 크게 향상시킨다.

ABSTRACT

We introduce repriorisation, a data-dependent reparameterisation which transforms a Bayesian neural network (BNN) posterior to a distribution whose KL divergence to the BNN prior vanishes as layer widths grow. The repriorisation map acts directly on parameters, and its analytic simplicity complements the known neural network Gaussian process (NNGP) behaviour of wide BNNs in function space. Exploiting the repriorisation, we develop a Markov chain Monte Carlo (MCMC) posterior sampling algorithm which mixes faster the wider the BNN. This contrasts with the typically poor performance of MCMC in high dimensions. We observe up to 50x higher effective sample size relative to no reparametrisation for both fully-connected and residual networks. Improvements are achieved at all widths, with the margin between reparametrised and standard BNNs growing with layer width.

연구 동기 및 목표

• 베이지안 신경망(BNN)의 이론적 이해 부족과 높은 계산 비용, 특히 매개변수 공간에서의 문제를 해결하기 위해.
• 기존의 함수 공간 결과(예: NNGP 근사)를 보완하기 위해, 넓은 BNN의 사후 행동을 가중치 공간에서 특성화하여 넓은 BNN에 대한 이해 격차를 메우기 위해.
• 너비에 따라 유리하게 스케일링되는 실용적이고 효율적인 MCMC 샘플링 방법을 개발하여, 일반적인 고차원 샘플링 과제를 극복하기 위해.
• 네트워크 너비가 증가함에 따라 사후분포가 표준 정규 분포에 점점 가까워지도록 변환함으로써, BNN에서 더 빠르고 효과적인 사후 추론을 가능하게 하기 위해.

제안 방법

• BNN 가중치 $ \theta $ 를 새로운 매개변수화 $ \theta = T(\theta) $ 로 변환하는 재매arameter화 맵 $ T(\theta) $ 를 도입하여, 층의 너비가 증가함에 따라 재매arameter화된 사후분포와 표준 정규 사전분포 사이의 KL 발산이 0으로 수렴하도록 한다.
• 이론적으로 (정리 2.1) $ \mathrm{KL}(\mathcal{N}(0,I_d) \| p(\phi|\mathcal{D})) \to 0 $ 이 되도록 증명하며, 이는 매개변수 공간에서 사후분포가 사전분포로 수렴함을 보장한다. 이는 최소 층 너비 $ d_{\min} \to \infty $ 일 때 성립한다.
• Cholesky 분해를 사용하여 동시에 재매arameter화와 그 자코비안 행렬식을 효율적으로 계산함으로써, MCMC 샘플링이 가능하도록 한다.
• 재매arameter화된 사후분포에 랭게빈 몬테카를로(LMC)를 적용하며, 이는 사후분포의 로그 밀도 기울기가 표준 정규 분포의 기울기 근처에 집중됨을 이용하여 혼합 속도를 향상시킨다.
• 이차 테일러 근사와 허치슨 추정기를 사용하여 자코비안의 로그행렬식을 효율적으로 근사함으로써 메모리 사용량을 줄이고 스케일링 가능한 기울기 계산을 가능하게 한다.
• JAX의 `vjp` 및 `jvp` 함수를 활용하여 기울기를 효율적으로 계산하며, 역전파 중 메모리 초과 오류를 방지하기 위한 커스텀 순방향 전파 전략을 도입한다.

실험 결과

연구 질문

• RQ1 네트워크의 너비가 증가함에 따라 넓은 베이지안 신경망의 사후분포는 매개변수 공간에서 어떻게 행동하는가?
• RQ2 무한한 너비 근사에서 재매arameter화된 사후분포와 표준 정규 사전분포 사이의 KL 발산이 0으로 수렴하도록 할 수 있는 재매arameter화를 설계할 수 있는가?
• RQ3 그러한 재매arameter화가 고차원 매개변수 공간에서 BNN의 MCMC 샘플링 혼합 속도를 빠르게 하는가?
• RQ4 이 방법은 완전히 연결된 네트워크와 리서벌 네트워크, 다양한 데이터셋 크기 등 다양한 아키텍처에서 효율적으로 구현되고 확장될 수 있는가?
• RQ5 특히 실용적인 유한한 너비 설정에서, repriorisation의 이점이 엄격한 NNGP 영역 외부에서도 유지되는가?

주요 결과

• Repriorisation는 무한한 너비 근사에서 재매arameter화된 BNN 사후분포와 표준 정규 사전분포 사이의 KL 발산을 0으로 줄여주며, 매개변수 공간에서 사후분포를 단순화하는 이론적 기반을 제공한다.
• 이 방법은 표준 BNN 대비 MCMC 샘플링에서 최대 50배 높은 효과적 샘플 크기(ESS)를 달성하며, 모든 너비에서 개선이 관찰되며 층 너비가 증가함에 따라 더욱 뚜렷해진다.
• CIFAR-10에서 3층의 완전히 연결된 네트워크(각 층 1024개 뉴런)를 사용한 실험에서, NNGP 근사에서 벗어난 상황에서도 ESS가 10배 향상된 것으로 관찰되었다.
• ResNet-20에서 ESS가 10배 향상된 것은 상단 층의 너비가 훈련 예제 수와 유사하거나 이를 초월할 때만 발생했으며, 이는 이 방법이 특정 영역에 의존함을 시사한다.
• Cholesky 기반 자코비안 계산과 허치슨 추정기, JAX의 `vjp`를 활용한 메모리 최적화된 기울기 추정을 통해 계산 효율성이 향상되었다.
• 비정규 분포의 가능도 함수에도 효과적이며, 더 작은 실험에서 다항 가능도를 사용한 결과, 재매arameter화된 경우 표준 매개변수화보다 더 높은 스텝 사이즈로 약 99%의 수용률을 달성하였다.