[논문 리뷰] Accelerated Linearized Laplace Approximation for Bayesian Deep Learning

Zhijie Deng, Feng Zhou외 1명|arXiv (Cornell University)|2022. 10. 23.

Gaussian Processes and Bayesian Inference인용 수 5

한 줄 요약

이 논문은 Hessian나 Jacobian 행렬을 명시적으로 계산하지 않기 위해 Neural Tangent Kernel(NTK)의 Nyström 근사법을 사용하는 Bayesian deep learning를 위한 가속화된 Linearized Laplace Approximation(이하 ELLA)을 제안한다. 전방 모드 자동 미분을 활용하여 ELLA는 확장성 향상과 함께 상태 수준의 예측 불확실성 정량화를 달성하며, Vision Transformers를 포함한 다양한 아키텍처에서도 이론적 보장과 함께 기존의 표준 LLA 근사법보다 ImageNet 및 CIFAR-10 벤치마크에서 뛰어난 성능을 보인다.

ABSTRACT

Laplace approximation (LA) and its linearized variant (LLA) enable effortless adaptation of pretrained deep neural networks to Bayesian neural networks. The generalized Gauss-Newton (GGN) approximation is typically introduced to improve their tractability. However, LA and LLA are still confronted with non-trivial inefficiency issues and should rely on Kronecker-factored, diagonal, or even last-layer approximate GGN matrices in practical use. These approximations are likely to harm the fidelity of learning outcomes. To tackle this issue, inspired by the connections between LLA and neural tangent kernels (NTKs), we develop a Nystrom approximation to NTKs to accelerate LLA. Our method benefits from the capability of popular deep learning libraries for forward mode automatic differentiation, and enjoys reassuring theoretical guarantees. Extensive studies reflect the merits of the proposed method in aspects of both scalability and performance. Our method can even scale up to architectures like vision transformers. We also offer valuable ablation studies to diagnose our method. Code is available at \url{https://github.com/thudzj/ELLA}.

연구 동기 및 목표

• Kronecker-분해 또는 대각 GGN 근사법에 의존하는 기존 Linearized Laplace Approximation(이하 LLA) 방법의 비효율성과 정밀도 손실 문제를 해결하기 위해.
• 큰 Hessian나 Jacobian 행렬을 명시적으로 계산하거나 저장하지 않고도 Bayesian deep learning에서 확장 가능하고 정확한 불확실성 정량화를 가능하게 하기 위해.
• LLA와 Neural Tangent Kernels(NTKs) 간의 관계를 활용하여 저랭크 커널 근사를 통해 추론 속도를 향상시키기 위해.
• ELLA와 원래 LLA 간의 근사 오차에 대한 이론적 보장을 제공하여 예측 불확실성의 신뢰성 확보를 위해.
• 현대적인 아키텍처인 Vision Transformers와 표준 ImageNet 벤치마크에서 ELLA의 확장성과 성능을 입증하기 위해.

제안 방법

• LLA 공분산 행렬의 효율적 계산을 위해 Neural Tangent Kernel(NTK)에 대한 Nyström 근사를 제안하여 Hessian나 Jacobian를 명시적으로 저장하지 않도록 한다.
• NTK 근사를 위해 필요한 Jacobian-벡터 곱(JVP)을 효율적으로 계산하기 위해 전방 모드 자동 미분(fwAD)을 활용한다.
• 데이터 포인트의 부분 집합을 통해 NTK에 저랭크 근사를 적용하여 계산 비용을 감소시키면서도 예측 정밀도를 유지한다.
• 원래 LLA의 핵심 구조를 유지하면서 전체 GGN 계산을 커널 기반 근사로 대체하여 대규모 모델로의 확장성을 가능하게 한다.
• 이론적 분석을 통해 NTK 근사가 정밀해질수록 ELLA와 원래 LLA 간의 예측 오차가 감소함을 보여준다.
• 표준 딥 러닝 라이브러리와의 호환성을 고려하여 구현하여 사전 학습된 모델과 즉시 통합 가능한 플러그 앤 플레이 아키텍처를 제공한다.

실험 결과

연구 질문

• RQ1 NTK의 Nyström 근사법이 예측 정확도를 손상시키지 않으면서도 Linearized Laplace Approximation(이하 LLA)의 계산을 효과적으로 가속화할 수 있는가?
• RQ2 KFAC, 대각 또는 마지막 레이어 근사법을 사용하는 표준 LLA 변종과 비교했을 때, ELLA의 불확실성 정량화 성능은 표준 벤치마크에서 어떻게 되는가?
• RQ3 Hessian 전체 계산이 불가능한 현대적인 아키텍처인 Vision Transformers에 ELLA는 확장 가능한가?
• RQ4 NTK 근사의 정밀도와 ELLA에서 도출된 예측 불확실성 간의 이론적 관계는 어떠한가?
• RQ5 ImageNet-C 오염과 같은 분포 이탈 상황에서도 ELLA는 신뢰할 수 있는 불확실성 캘리브레이션을 유지하는가?

주요 결과

• CIFAR-10에서 ELLA는 ECE와 NLL 측면에서 KFAC, 대각 또는 마지막 레이어 근사법을 사용한 LLA와 비교해도 최상의 불확실성 캘리브레이션 성능을 달성한다.
• ResNet-18, ResNet-34, ResNet-50를 사용한 ImageNet에서 ELLA는 ECE를 0.015–0.018로 낮추어 MAP의 0.026–0.037보다 더 나은 불확실성 캘리브레이션을 보였다.
• ImageNet에서 ViT-B에 대해 ELLA는 ECE를 0.022로 낮추어 MAP의 0.039보다 개선된 불확실성 캘리브레이션을 보였다. 이는 현대적 아키텍처에서의 성능 향상을 보여준다.
• ViT-B에서 ELLA는 81.6%의 경쟁력 있는 정확도와 0.695의 낮은 NLL을 기록하여 MAP(81.5% 정확도, 0.700 NLL)보다 정확도 손실 없이 더 나은 캘리브레이션을 달성했다.
• ImageNet-C 오염 상황에서 ELLA는 NLL과 ECE 모두에서 MAP를 뛰어넘는 성능을 보였으며, 분포 이탈에 대한 강건성을 입증했다.
• 대규모 모델에서 발생하는 메모리 부족 오류와 장시간 학습 문제를 피함으로써, 표준 LLA 구현에 비해 훨씬 뛰어난 확장성을 입증했다.