무한 깊이 베이지안 신경망을 위한 네스테로프 가속 기법

본 논문은 연속 깊이 신경망인 SDE‑BNN에 네스테로프 가속 기법(NAG)을 도입하고, 함수 평가 횟수(NFE)에 기반한 잔차 스킵 연결을 설계함으로써 학습·추론 속도를 크게 높이고, 정확도와 불확실성 추정 성능을 동시에 개선한 새로운 모델인 Nesterov‑SDEBNN을 제안한다.

저자: Chenxu Yu, Wenqi Fang

무한 깊이 베이지안 신경망을 위한 네스테로프 가속 기법
본 논문은 연속‑시간 딥러닝 모델인 Neural ODE와 그 확장인 Neural SDE에 베이지안 신경망(BNN)의 불확실성 추정 메커니즘을 결합한 SDE‑BNN이 높은 함수 평가 횟수(NFE)와 수치적 불안정성으로 인해 실용성에 제약이 있다는 문제점을 인식하고, 이를 해결하기 위한 새로운 프레임워크인 Nesterov‑SDEBNN을 제안한다. 먼저 기존 SDE‑BNN의 수학적 배경을 정리한다. Neural SDE는 상태 hₜ 에 대해 drift f_w(·)와 diffusion g_θ(·)를 포함하는 확률 미분 방정식으로 표현되며, BNN은 가중치 wₜ 를 확률 과정으로 모델링한다. 변분 추론(VI) 하에서 ELBO를 최적화함으로써 posterior over wₜ 를 학습한다. 그러나 이 과정은 매 시간 스텝마다 drift와 diffusion을 평가해야 하므로 NFEs가 급증하고, 적응형 솔버 사용 시에도 수렴이 불안정해지는 단점이 있다. 이를 극복하기 위해 두 가지 핵심 기법을 도입한다. 첫 번째는 Nesterov 가속(NAG)이다. 기존 1차 SDE‑BNN을 2차 미분 방정식 형태인 d²hₜ/dt² + 3/t·dhₜ/dt + f_θ(hₜ,t)=0 로 변형하고, 이를 1차 시스템으로 재구성해 보조 변수 xₜ 와 momentum mₜ 를 도입한다. 이때 mₜ 는 관성항을 포함해 업데이트되며, Nesterov‑형 가속을 통해 최적화 경로가 더 빠르게 수렴한다. 두 번째는 NFE‑의존 잔차 스킵 연결이다. 기존 Nesterov‑ODE에서는 매 drift 평가 시마다 잔차 ξ·hₜ 를 바로 더했지만, SDE 솔버는 시간 스텝마다 독립적으로 동작하므로 이러한 직접 연결이 정보 흐름을 충분히 활용하지 못한다. 논문은 NFE가 홀수일 때 현재 hidden state hₜ 를 캐시하고, 짝수 NFE가 될 때 이 캐시를 drift mₜ 에 주입하는 “스킵” 메커니즘을 설계한다. 수식적으로는 dmₜ/dt = −mₜ − σ(f_w(hₜ,t) + ε·ξ·h_temp) 로 표현되며, ε와 h_temp은 NFE의 홀짝에 따라 0/1 혹은 이전 hₜ 값을 선택한다. 이 설계는 두 단계에 걸쳐 hₜ 와 hₜ₊Δt 를 모두 활용해 mₜ₊2Δt 를 계산함으로써 gradient 소실을 완화하고, NFEs를 크게 늘리지 않으면서도 정보 재사용을 가능하게 한다. 베이지안 측면에서는 가중치 wₜ 에 Ornstein‑Uhlenbeck(OU) 과정을 사전으로 두어 장기 안정성을 보장하고, posterior drift를 작은 신경망 NN_ϕ 으로 파라미터화한다. 이렇게 정의된 posterior SDE는 drift와 diffusion이 모두 학습 가능하며, ELBO는 로그우도와 Girsanov 변환 기반의 정규화 항(∥u∥²)으로 구성된다. Monte‑Carlo 샘플링을 통해 하나의 SDE 솔버 호출만으로 가중치 샘플링·전방 전파·손실 계산을 동시에 수행한다. 실험에서는 네 가지 베이스라인을 사용한다. 1‑D 회귀에서는 비단조적인 노이즈 데이터를 대상으로 Nesterov‑SDEBNN이 기존 SDE‑BNN과 동일한 예측 정확도와 불확실성 구간을 유지함을 확인한다. 이미지 분류에서는 MNIST와 CIFAR‑10을 대상으로 고정‑스텝과 적응‑스텝 두 가지 통합 솔버 설정에서 성능을 비교한다. 결과는 Nesterov‑SDEBNN이 정확도(예: CIFAR‑10 88.36 % vs 87.60 %), AUROC(85.61 % vs 83.05 %), 그리고 NLL(5.97 × 10⁻¹ vs 9.89 × 10⁻¹) 모두에서 우수함을 보이며, NFEs는 평균 30 % 이상 감소한다. 학습 곡선 역시 초기 단계부터 빠르게 수렴하는 모습을 보여, 가속 효과가 실질적으로 작동함을 입증한다. Walker2D와 같은 연속 제어 시뮬레이션에서도 비슷한 경향이 관찰되었다. 논문의 기여는 다음과 같다. (1) Nesterov 가속을 SDE‑BNN에 직접 통합해 2차 동역학을 구현하고, (2) NFE‑의존 잔차 스킵 연결을 설계해 정보 재사용과 수치 안정성을 동시에 달성했으며, (3) 변분 추론 프레임워크 내에서 하나의 SDE 솔버 호출만으로 효율적인 학습 파이프라인을 제공했다. 한계점으로는 Nesterov‑SDEBNN이 아직 다양한 도메인(예: 고차원 물리 시뮬레이션, 의료 데이터)에서 검증되지 않았으며, NFE‑의존 스킵 연결이 솔버 종류와 하이퍼파라미터에 따라 민감할 수 있다는 점을 들 수 있다. 또한 2차 동역학에 대한 이론적 수렴 보장이 완전히 증명되지 않아, 향후 수학적 분석이 필요하다. 그럼에도 불구하고, 복잡한 베이지안 연속‑시간 모델에 간단하고 효과적인 가속 기법을 도입함으로써 실용성을 크게 향상시킨 점은 학계·산업 모두에 의미 있는 진전이다.

원본 논문

고화질 논문을 불러오는 중입니다...

댓글 및 학술 토론

Loading comments...

의견 남기기