[논문 리뷰] Bayesian Neural Ordinary Differential Equations

Raj Dandekar, Chung, Karen외 5명|arXiv (Cornell University)|2020. 12. 14.

Model Reduction and Neural Networks참고 문헌 36인용 수 16

한 줄 요약

이 논문은 NUTS, SGHMC, SGLD를 통한 베이지안 추론을 통합하여 신경미분방정식(Neural ODEs)에서 지식 기반 불확실성(epistemic uncertainty)을 정량화하는 베이지안 신경미분방정식(Bayesian NODEs)을 소개한다. MNIST에서 98.5%의 테스트 정확도를 달성하고, 베이지안 유니버설 ODE를 활용해 부분적으로 기술된 동역학계에서 누락된 항을 확률적으로 탐지함으로써 불확실성 인식이 가능한 과학적 머신러닝을 실현한다.

ABSTRACT

Recently, Neural Ordinary Differential Equations has emerged as a powerful framework for modeling physical simulations without explicitly defining the ODEs governing the system, but instead learning them via machine learning. However, the question: "Can Bayesian learning frameworks be integrated with Neural ODE's to robustly quantify the uncertainty in the weights of a Neural ODE?" remains unanswered. In an effort to address this question, we primarily evaluate the following categories of inference methods: (a) The No-U-Turn MCMC sampler (NUTS), (b) Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) and (c) Stochastic Langevin Gradient Descent (SGLD). We demonstrate the successful integration of Neural ODEs with the above Bayesian inference frameworks on classical physical systems, as well as on standard machine learning datasets like MNIST, using GPU acceleration. On the MNIST dataset, we achieve a posterior sample accuracy of 98.5% on the test ensemble of 10,000 images. Subsequently, for the first time, we demonstrate the successful integration of variational inference with normalizing flows and Neural ODEs, leading to a powerful Bayesian Neural ODE object. Finally, considering a predator-prey model and an epidemiological system, we demonstrate the probabilistic identification of model specification in partially-described dynamical systems using universal ordinary differential equations. Together, this gives a scientific machine learning tool for probabilistic estimation of epistemic uncertainties.

연구 동기 및 목표

• 베이지안 학습 프레임워크가 신경미분방정식과 통합되어 모델 가중치의 불확실성을 견고하게 정량화할 수 있는지에 대한 열린 질문을 해결하기 위해.
• 물리계와 머신러닝 데이터셋 모두에서 세 가지 베이지안 추론 방법(NUTS, SGHMC, SGLD)의 성능을 평가하고 비교하기 위해.
• 변분 추론과 노멀라이징 플로우, 신경미분방정식을 융합하는 것을 선구적으로 시도하여 예측 성능과 불확실성 추정을 향상시키기 위해.
• 베이지안 유니버설 ODE(UDE)를 활용해 부분적으로 기술된 동역학계에서 누락된 항을 확률적으로 복원하는 것을 보여주기 위해.
• 물리 시뮬레이션과 대규모 딥러닝 작업 모두에서 지식 기반 불확실성을 정량화할 수 있는 과학적 머신러닝 프레임워크를 제공하기 위해.

제안 방법

• 차분 가능한 프로그래밍 스택인 Julia를 활용하여 DifferentialEquations.jl의 미분방정식 해법기와 확률적 프로그래밍 언어인 Turing.jl을 원활하게 조합한다.
• 세 가지 베이지안 추론 방법인 No-U-Turn Sampler(NUTS), Stochastic Gradient Hamiltonian Monte Carlo(SGHMC), Stochastic Gradient Langevin Descent(SGLD)를 적용하여 신경미분방정식의 가중치에 대한 사후분포를 추정한다.
• 새로운 아키텍처를 도입하여 컨볼루션 레이어와 신경미분방정식을 결합하고, SGHMC를 통해 훈련하여 이미지 분류 성능을 향상시킨다.
• 노멀라이징 플로우를 활용한 변분 추론을 베이지안 신경미분방정식에 도입하여 복잡한 사후분포를 근사함으로써 표현력과 예측 정확도를 향상시킨다.
• 예측된 SGLD(PSGLD)를 베이지안 유니버설 ODE(UDE)에 적용하여 동역학계에서 누락된 항을 탐지하며, 예를 들어 포식자-피식자 모델과 SIR 모델을 포함한다.
• GPU 가속과 자동 미분을 활용하여 대규모 데이터셋(MNIST)과 복잡한 물리계에까지 확장 가능한 베이지안 추론을 실현한다.

실험 결과

연구 질문

• RQ1 베이지안 추론 프레임워크는 신경미분방정식과 성공적으로 통합되어 모델 가중치의 불확실성을 정량화할 수 있는가?
• RQ2 NUTS, SGHMC, SGLD는 신경미분방정식에 적용되었을 때 예측 정확도, 수렴성, 편향 측면에서 어떻게 비교되는가?
• RQ3 노멀라이징 플로우를 활용한 변분 추론은 베이지안 신경미분방정식의 성능과 불확실성 추정을 향상시킬 수 있는가?
• RQ4 베이지안 UDE는 포식자-피식자 모델과 전염병학적 모델과 같은 부분적으로 기술된 동역학계에서 누락된 항을 견고하게 식별할 수 있는가?
• RQ5 제안된 베이지안 신경미분방정식 프레임워크는 MNIST와 같은 표준 벤치마크에서 최신 기술 대비 성능과 불확실성 정량화 측면에서 어떻게 비교되는가?

주요 결과

• 베이지안 신경미분방정식 프레임워크는 SGLD 샘플링을 통해 MNIST 데이터셋에서 테스트 앙상블 정확도 98.5%를 달성하였으며, 최신 기술과 경쟁 가능하다.
• SGHMC 기반의 베이지안 신경미분방정식은 컨볼루션 레이어와 결합되었을 때 테스트 앙상블 정확도 99.22%를 기록하여 MAP 추정치를 초월하고 편향을 감소시켰다.
• 베이지안 신경미분방정식에 변분 추론과 노멀라이징 플로우를 융합함으로써 물리계에서 예측 성능 향상과 정확한 불확실성 정량화가 가능해졌다.
• 베이지안 UDE는 포식자-피식자 모델과 SIR 유형의 전염병학적 모델 모두에서 누락된 항을 성공적으로 복원하였으며, 사후 샘플이 시간적·공간적 변동성과 관련된 불확실성을 잘 반영하였다.
• NUTS 샘플러는 다중모달 확률밀도에서 MAP 점이 낮은 확률 영역에 위치함에 따라 SGLD에 비해 더 높은 편향을 보였다.
• 이 프레임워크는 부분적으로 관측된 동역학계에서 모델 구조를 확률적으로 식별할 수 있게 하여, 불확실성 정량화 기능을 갖춘 새로운 과학적 머신러닝 도구를 제공한다.