기술 요약: 대규모 베이지안 심층 신경망을 위한 점수 기반 변분 사후 추론

문제 제기

베이지안 심층 신경망 (BNN) 은 점수 추정 심층 학습에 비해 강력한 불확실성 정량화, 노이즈 저항성, 과적합 완화 등 상당한 장점을 제공합니다. 그러나 BNN 에서 정확한 사후 추론은 일반적으로 계산적으로 불가능합니다. 변분 추론 (VI) 은 타겟 사후분포 $\pi(\theta)$를 계산 가능한 분포 $q_\lambda(\theta)$로 근사함으로써 추론을 최적화 문제로 재구성하는 지배적인 근사 방법입니다.

현재 지배적인 접근법인 ELBO 기반 VI(역 KL 발산을 최소화) 는 타겟 밀도의 모든 모드를 포착하지 못하는 잘 알려진 '모드 붕괴 (mode collapsing)' 문제에 시달립니다. 반면, 타겟 점수 $\nabla_\theta \log \pi(\theta)$와 모델 점수 $\nabla_\theta \log q_\lambda(\theta)$ 간의 차이를 최소화하는 점수 기반 VI 방법은 이론적으로 모드 붕괴를 해결할 수 있으나, 기존 점수 기반 방법들은 대규모 BNN 에는 확장성이 부족합니다. 구체적으로 다음과 같은 한계가 있습니다:

1. 계산적 비실용성: 가우시안 점수 매칭 (GSM) 및 배치 - 매칭 (BaM) 과 같은 방법들은 상태 차원에 대해 3 차 또는 2 차적으로 스케일링되는 행렬 역연산이나 곱셈에 의존하므로, 고차원 신경망 파라미터에는 적용이 불가능합니다.
2. 확률적 경사와의 비호환성: 기존 점수 기반 방법들은 일반적으로 정확한 점수를 요구하거나, 대규모 데이터셋 학습에 필수적인 확률적 경사에서 유도된 노이즈가 있는 편향되지 않은 미니배치 점수를 처리할 수 없습니다.
3. 재파라미터화 제약: 많은 방법들은 재파라미터화된 샘플링에 의존하는데, 이는 고차원에서 계산 비용이 많이 들거나 불안정할 수 있습니다.

방법론

저자들은 비전 트랜스포머 (ViT) 와 ResNet 을 포함한 대규모 BNN 으로 확장 가능한 새로운 근사 확률적-경사 점수 기반 VI 알고리즘을 제안합니다.

핵심 최적화 목적 함수

이 방법은 각 반복 $t$에서 다음 최적화 문제를 풀어 변분 분포 $q_t(\theta)$를 $q_{t+1}(\theta)$로 반복적으로 업데이트합니다:

$$q_{t+1} = \arg \min_{q \in \mathcal{Q}} \mathbb{E}_{\theta \sim q_t(\theta)} \left[ \alpha_t \|\nabla_\theta \log q(\theta) - \nabla_\theta \log q_t(\theta)\|_2^2 + \|\nabla_\theta \log q(\theta) - s(\theta)\|_2^2 \right]$$

여기서:

• $s(\theta) = \nabla_\theta \log \pi(\theta)$는 타겟 점수입니다.
• $\alpha_t$는 0 에서 상수 $\le 1$까지 선형적으로 스케줄링되는 적응형 근사 영향 파라미터입니다.
• 첫 번째 항은 새로운 반복이 점수 측면에서 현재 반복과 가깝게 유지되도록 강제하는 근사 페널티입니다.
• 두 번째 항은 점수 매칭 손실입니다.

주요 기술적 혁신

1. 노이즈가 있는 미니배치 점수 처리: 이전 점수 기반 방법들이 정확한 점수를 요구한 것과 달리, 이 접근법은 확률적 경사를 통해 얻은 노이즈가 있는 편향되지 않은 미니배치 점수 $\hat{s}(\theta)$를 활용합니다. 손실 함수 내 점수 항의 선형성으로 인해, 노이즈가 있는 점수라도 변분 파라미터 $\lambda$에 대한 손실의 기울기는 편향되지 않은 추정치로 유지됩니다.
2. 재파라미터화 회피: 근사 항을 통해 기대값을 변수 $q$가 아닌 고정된 현재 분포 $q_t$에 대해 취할 수 있게 합니다. 이는 (ADVI 에서 사용되는) 재파라미터화 트릭의 필요성을 우회하고, 전통적인 피셔/점수 방법이 요구하는 $\theta$에 대한 2 차 미분을 방지하여 계산 오버헤드와 메모리 풋프린트를 크게 줄입니다.
3. 확장성: 행렬 역연산이 필요한 공분산 유도 노름 대신 유클리드 노름을 사용하고 확률적 경사를 지원함으로써, 이 방법은 수억 개의 파라미터를 가진 네트워크 (예: ~305M 파라미터를 가진 ViT-L-32) 로 확장 가능합니다.

이론적 수렴

저자들은 변분 가족이 매끄러운 벡터장을 표현하기에 충분히 풍부하다는 가정 하에, 반복이 증가함에 따라 변분 분포의 점수가 타겟 점수로 수렴한다는 수렴 정리 (정리 3.1) 를 제시합니다. 표현 풍부성 가정이 이상적이기는 하나, 저자들은 참된 사후분포와 변분 가족이 모두 가우시안인 경우 이 가정이 정확히 성립한다고 언급합니다.

주요 기여

1. 새로운 알고리즘: 점수 매칭 손실과 근사 페널티를 결합한 근사 확률적-경사 점수 매칭 VI 알고리즘을 도입했습니다.
2. 대규모 BNN 으로의 확장성: 이 방법은 대규모 심층 네트워크 (ResNet, ViT) 와 대규모 데이터셋에서 효과적으로 작동함이 입증된 최초의 점수 기반 VI 접근법으로, GSM 과 BaM 의 계산 병목 현상을 극복했습니다.
3. 노이즈에 대한 강인성: 이 알고리즘은 노이즈가 있는 편향되지 않은 미니배치 점수를 효과적으로 처리하여 전체 배치 점수 계산 없이 확률적 경사를 사용할 수 있게 합니다.
4. 재파라미터화 제거: 이 방법은 재파라미터화 트릭을 피하여 최적화 지형을 단순화하고 ELBO 기반 ADVI 대비 효율성을 향상시킵니다.

실험 결과

저자들은 합성 데이터, PosteriorDB 벤치마크, 그리고 실제 대규모 작업에서 이 방법을 평가했습니다:

• 합성 데이터: 가우시안 및 가우시안 혼합 타겟에서 제안된 방법은 ADVI, GSM, BaM 에 비해 더 빠른 수렴과 초기 파라미터 선택 (특히 작은 공분산 고유값과 관련하여) 에 대한 더 나은 강인성을 보여주었습니다. 또한 GSM/BaM 이 수치적 문제로 인해 실패한 모델 불일치 및 고차원 상태 공간에서도 성공적으로 처리했습니다.
• MNIST (소규모 BNN): 저데이터 영역에서 이 방법은 ADVI 와 유사한 테스트 오차율을 달성했으나, 더 낮은 기대 보정 오차 (ECE) 와 음의 로그 가능도 (NLL) 를 통해 우월한 불확실성 정량화를 입증했습니다.
• 대규모 시각 인식 (ResNet-101 및 ViT-L-32): Pets, Flowers, Aircraft, DTD 와 같은 데이터셋에서 이 방법은 ADVI 와 동등한 테스트 예측 정확도를 유지하면서 불확실성 지표 (ECE 및 NLL) 에서 일관되게 ADVI 를 능가했습니다. 경쟁 방법들 (GSM, BaM, 전통적 점수 매칭) 은 메모리 제약이나 수치적 폭발로 인해 실패했습니다.
• 시계열 예측: Koopa 네트워크 아키텍처를 사용하여 이 방법은 여러 데이터셋 (ECL, ETTh2, Ex. Rate, ILI, Traffic, Weather) 에서 MSE, MAE, NLL 측면에서 ELBO 기반 VI 및 비베이지안 베이스라인을 능가했습니다.
• 정규화 흐름: 이 방법을 정규화 흐름 변분 밀도로 확장하여 시계열 작업에서 ADVI 대비 일관된 개선을 보여주었습니다.
• 효율성: 이전 및 현재 반복을 모두 유지함에도 불구하고, 실제 소요 시간과 GPU 메모리 풋프린트는 ADVI 와 비교 가능하여 대규모 학습에 대한 실용적 타당성을 입증했습니다.

의의 및 주장

이 논문은 대규모 베이지안 심층 학습을 위한 지배적인 ELBO 기반 VI 에 대한 확장 가능하고 유연한 대안을 제공한다고 주장합니다. 역 KL 목적 함수에 내재된 모드 붕괴 문제를 해결하고 이전 점수 기반 방법들의 계산적 한계를 극복함으로써, 제안된 알고리즘은 비전 트랜스포머와 같은 현대적이고 고용량의 아키텍처에서 엄격한 베이지안 추론을 가능하게 합니다.

저자들은 그들의 접근 방식이 다양한 벤치마크에서 경험적으로 효과적이며 예측 성능을 희생하지 않고 더 나은 불확실성 정량화를 제공한다고 강조합니다. 이론적 수렴은 변분 가족의 표현력에 대한 강력한 가정에 의존하지만, 경험적 결과는 실제에서의 방법 유용성을 뒷받침한다고 겸손하게 언급합니다. 이 작업은 점수 기반 VI 를 Gaussian-only 또는 전체 배치 요구사항의 이전 기술의 제약을 넘어선 베이지안 심층 학습을 위한 실현 가능하고 확장 가능한 경로로 위치시킵니다.