대규모 점수 기반 변분 베이지안 신경망 추론

초록

본 논문은 베이지안 딥러닝에서 변분 추론을 수행하기 위한 새로운 점수 매칭 기반 방법을 제안한다. 기존 ELBO 방식과 달리 점수 차이를 최소화하는 근접(proximal) 손실을 사용해 재파라미터화 샘플링 없이도 미니배치의 잡음이 섞인 점수를 효율적으로 활용한다. Vision Transformer와 ResNet 등 대규모 모델에 적용 가능함을 실험을 통해 입증한다.

상세 분석

이 연구는 베이지안 신경망(BNN)의 변분 추론을 점수 매칭(score matching)으로 재구성한다는 점에서 혁신적이다. 기존 ELBO 기반 변분 방법은 재파라미터화 샘플링을 통해 기대값을 근사하지만, 고차원 파라미터 공간에서 샘플링 비용이 크게 늘어난다. 반면 점수 기반 방법은 목표 분포와 변분 분포의 로그밀도 미분(스코어)을 직접 맞추어 KL 발산을 최소화한다. 그러나 기존 점수 매칭 기법인 Gaussian Score Matching(GSM)과 Batch‑and‑Match(BaM)는 정확한 스코어를 필요로 하고, 행렬 역연산 등 O(d³) 복잡도가 발생해 대규모 네트워크에 적용하기 어렵다.

논문은 이러한 한계를 극복하기 위해 두 개의 손실을 결합한 근접 점수 매칭 프레임워크를 제안한다. 첫 번째 항은 현재 변분 분포 q_t와 다음 단계 q 사이의 스코어 차이를 제곱해 α_t 가중치로 제어하고, 두 번째 항은 변분 스코어와 실제(또는 미니배치 추정) 스코어 사이의 차이를 최소화한다. α_t는 0에서 1로 선형 증가하도록 스케줄링되어 초기에는 안정적인 업데이트를, 후기에선 목표 스코어에 더 가까워지도록 한다. 중요한 점은 기대값을 q_t에 대해 고정해 두어 θ에 대한 이중 미분을 피하고, 따라서 재파라미터화 샘플링이 필요 없으며, 변분 파라미터 λ에만 직접적인 그래디언트를 계산한다는 것이다.

또한 논문은 미니배치에서 얻는 잡음이 섞인 스코어 ˆs(θ)를 사용해도 손실의 그래디언트가 무편향(unbiased)임을 증명한다. 이는 스코어가 선형으로 λ에만 등장하기 때문에 가능하며, 기존 점수 매칭 방법이 미니배치 스코어를 다루지 못하던 문제를 해결한다. 수렴에 대해서는 완전한 정리를 제공하지 않지만, 완벽한 최적화와 풍부한 변분 가족 가정 하에 α_t = t/T 로 진행하면 q_T가 목표 분포 π에 수렴한다는 비공식적 논증을 제시한다.

실험에서는 toy Gaussian, Gaussian mixture, PosteriorDB 등 작은 규모 데이터에서 정확도와 수렴 속도를 검증하고, 대규모 비전 인식(ViT, ResNet)과 시계열 예측(DeepAR 계열) 작업에서 ADVI, GSM, BaM과 비교해 비슷하거나 더 나은 로그우도와 예측 불확실성을 보여준다. 특히 근접 점수 매칭은 메모리 사용량이 낮고, 재파라미터화 샘플링이 필요 없으므로 GPU 메모리 제한이 있는 대형 모델에서도 효율적으로 동작한다.

이러한 장점에도 불구하고 변분 가족을 Gaussian에 제한하지는 않지만, 실제 구현에서는 Gaussian 혹은 다중 가우시안 형태를 사용한다. 또한 수렴 이론이 아직 완전하지 않아, 향후 확률적 근접점 방법론과 결합한 엄밀한 분석이 필요하다. 전반적으로 이 논문은 대규모 베이지안 딥러닝에 실용적인 점수 기반 변분 추론 프레임워크를 제공하며, 기존 ELBO 기반 방법과 비교해 계산 효율성과 확장성에서 의미 있는 진전을 보여준다.