[Paper Review] The Bayesian Learning Rule
이 포스팅은 Khan 등 (2023)의 논문 “The Bayesian Learning Rule”를 읽고 정리한 글입니다.
Introduction
머신러닝의 목표는 결국 어떤 입력 $x$가 주어졌을 때 출력 $y$를 예측하는 unknown function $f(x)$를 추론하는 문제입니다.
\[y \approx f(x; \theta)\]여러 알고리즘, 예를 들어 경사 하강법(Gradient Descent)을 이용한 최적화부터 MCMC를 이용한 베이즈 추론까지, 모두 이 최적의 함수 $f$ 또는 $f$의 파라미터 $\theta$를 찾아가는 다양한 방법론으로 볼 수 있습니다.
| 지금까지 최적화(Optimization)와 베이즈 추론(Bayesian Inference)은 각기 다른 패러다임으로 발전해왔습니다. 최적화는 손실 함수를 최소화하는 단일 최적점 $\theta^*$를 찾는 데 집중하는 반면, 베이즈 추론은 데이터로부터 얻을 수 있는 파라미터 $\theta$의 전체 분포, 즉 사후 분포(Posterior Distribution) $p(\theta | D)$를 구하는 데 초점을 맞춥니다. |
하지만 두 방식을 포함한 많은 알고리즘들이 사실은 단 하나의 일반화된 학습 규칙의 특별한 경우로 표현될 수 있습니다.
이제 본격적으로 “베이지안 학습 규칙(Bayesian Learning Rule, BLR)”에 대한 이야기를 해보겠습니다.
Khan(2021)은 이 질문에 대한 답으로, 여러 머신러닝 알고리즘을 통합하는 단일 프레임워크인 Bayesian Learning Rule 을 제안했습니다.
핵심 아이디어는 다음과 같습니다.
$\rightarrow$ “Bayesian Inference의 Posterior Updating Problem을 Variational Inference와 Natural Gradient를 이용해 푸는 일반화된 Updating Rule을 만들고, 이 규칙의 특정 조건 하에서 다양한 알고리즘이 유도됨을 보인다.”
이 논문은 우리가 당연하게 사용하던 알고리즘들이 왜 그런 형태로 동작하는지에 대한 근본적인 통찰을 제공합니다.
The Bayesian Learning Rule Model
BLR 모델은 두 가지 핵심 개념 위에 세워져 있습니다: Bayesian Inference를 위한 Variational Approximation Objective와, 그 목표를 달성하기 위한 Update Mechanism입니다.
Bayesian Inference \& Variational Approximation
| BI의 최종 목표는 데이터 $D$가 주어졌을 때 파라미터 $\theta$의 사후 분포 $P(\theta | D)$를 찾는 것입니다. |
| 하지만 분모의 $P(D) = \int P(D | \theta)P(\theta)d\theta$ 항은 대부분의 현실 문제에서 계산이 불가능(intractable)합니다. |
| 이 문제를 해결하기 위해 Variational Inference 는 다루기 쉬운 근사 분포 $q(\theta)$를 도입하여, 이 분포가 실제 사후 분포 $P(\theta | D)$와 최대한 비슷해지도록 만듭니다. ‘비슷함’의 척도로는 KL-Divergenc) 이 사용되며, 우리의 목표는 이 KLD를 최소화하는 것입니다. |
이것이 바로 BLR이 달성하고자 하는 최적화 목표입니다.
Natural Gradient
KL Divergence를 최소화하기 위해 $q(\theta)$의 파라미터를 어떻게 업데이트해야 가장 효율적일까요?
- Standard Gradient: ‘파라미터 공간’에서 가장 가파른 방향으로 이동합니다. 이는 유클리드 기하학에 기반한 최단 경로입니다.
- Natural Gradient: ‘확률 분포 공간’에서 가장 가파른 방향으로 이동합니다. 이는 두 확률 분포 사이의 거리(ex. KLD)를 가장 빨리 줄이는 경로이며, Information Geometry에 기반합니다.
BLR은 이 Natural Gradient를 Update Mechanism으로 채택합니다. 이를 통해 변분 추론의 목표를 가장 효율적으로 달성하는 이론적 기반을 마련합니다.
Generalised Form of Bayesian Learning Rule
BLR 은 위 두 요소를 결합한 것입니다. 즉, 변분 추론의 KLD를 minimize하는 목표를 natural gradient를 이용해 푸는 generalised updating rule입니다.
이는 특정 데이터와 사전 믿음이 주어졌을 때, 우리의 믿음(근사 분포 $q$)을 어떻게 합리적으로 업데이트해야 하는지에 대한 일반 원리를 제공합니다. 이 규칙의 놀라운 점은, 근사 분포 $q(\theta)$를 어떤 형태로 가정하느냐에 따라 수많은 알고리즘과 같아진다는 것입니다.
from BLR to other algorithms
1. Optimization Algorithms
1.1. Gradient Descent & SGD
- 가정: 근사 분포 $q(\theta)$를 파라미터 $\mu$에 위치한 단일 포인트(Dirac delta function), 즉 $q(\theta) = \delta(\theta - \mu)$로 가정합니다.
-
유도: 이 가정 하에서 KL divergence minimization는 **로그 사후 확률 $\log P(\mu D)$를 최대화하는 MAP 문제와 동일해집니다. 이를 경사 상승법으로 풀면, 사전 분포를 무시할 경우 정확히 **경사 하강법(GD) 업데이트 규칙 $\mu_{t+1} \leftarrow \mu_t - \eta \nabla_\mu \text{Loss}(\mu_t)$가 됩니다. 전체 데이터가 아닌 미니배치를 사용하면 SGD가 됩니다.
1.2. Newton’s Method
- 가정: GD와 동일하게 $q(\theta) = \delta(\theta - \mu)$로 가정하여 MAP 문제를 풉니다.
- 유도: MAP의 목적 함수인 로그 사후 확률 $\mathcal{L}(\mu)$를 1차 미분(경사)이 아닌 2차 미분(헤시안 행렬 $H$)까지 고려하여 최적화하면 뉴턴 방법이 됩니다. 뉴턴 업데이트 규칙 $\mu_{t+1} \leftarrow \mu_t - \eta H^{-1} \nabla_\mu \mathcal{L}(\mu_t)$는 BLR 관점에서, 현재 지점 $\mu_t$ 주변에 국소적인 가우시안 근사를 적용하고 그에 대한 자연 경사 업데이트를 수행하는 것과 같습니다. 즉, 손실 함수의 곡률(curvature)을 이용해 더 빠르게 최적점에 도달하는 과정으로 해석됩니다.
1.3. Ridge Regression
- 가정: 모델이 선형($y = X\theta + \epsilon$)이고, 파라미터에 대한 사전 분포가 가우시안 분포 $p(\theta) = \mathcal{N}(0, \alpha^{-1}I)$ 라고 가정합니다.
-
유도: 이 또한 MAP 문제입니다. 로그 사후 확률은 $\log P(\theta D) \propto -|y - X\theta|^2 - \alpha|\theta|^2$ (상수 제외) 가 됩니다. 이 식을 $\theta$에 대해 미분하여 0으로 두면, 우리가 아는 릿지 회귀의 Closed-form solution $(X^T X + \lambda I)\theta = X^T y$ (단, $\lambda$는 $\alpha$와 관련)를 얻게 됩니다. 즉, 릿지 회귀는 BLR이 특정한 선형-가우시안 모델과 가우시안 사전 분포 하에서 내놓는 정확한 해입니다.
2. Deep Learning Techniques
2.1. Dropout
- 가정: 가중치 $W$에 대한 근사 분포 $q(W)$를, 각 가중치를 독립적인 베르누이 분포 $z_{ij} \sim \text{Bernoulli}(p)$와 곱한 형태라고 가정합니다. 즉, 학습 시마다 랜덤하게 일부 뉴런을 0으로 만드는 과정입니다.
- 유도: 이러한 형태의 근사 분포 $q(W)$를 가정하고 변분 추론의 목적 함수(ELBO)를 최대화하는 것은, 결과적으로 드롭아웃을 적용하여 신경망을 학습시키는 것과 수학적으로 동일한 효과를 낳습니다. 즉, 드롭아웃은 가중치에 대한 특정 베이지안 근사를 수행하는 방법으로 해석될 수 있습니다.
2.2. Binary Neural Networks
- 가정: 가중치 $\theta$가 ${-1, +1}$ 같은 이진 값만 가질 수 있다고 제약합니다. 근사 분포 $q(\theta)$ 역시 이 이진 상태에 대한 독립적인 분포 $q(\theta) = \prod_i q_i(\theta_i)$로 가정합니다.
- 유도: 이산적인(discrete) 파라미터에 대한 BLR을 적용하면 특정 형태의 업데이트 규칙이 나옵니다. 이 규칙은 이진 신경망 학습에 사용되는 Straight-Through Estimator와 같은 기법들과 깊은 연관성을 가집니다. 즉, BNN의 학습 과정은 이진 파라미터에 대한 변분 추론으로 이해될 수 있습니다.
3. Probabilistic Model Inference
3.1. EM Algorithm
- 가정: 모델에 관측되지 않는 잠재 변수(latent variable) $z$가 존재하고, 근사 분포가 $q(\theta, z) = q(\theta)q(z)$ 형태로 분리(factorize) 가능하다고 가정합니다.
- 유도: BLR의 목표인 ELBO를 이 분리된 $q$에 대해 좌표 상승법(coordinate ascent)으로 최적화하면 EM 알고리즘이 됩니다.
-
E-step: $q(\theta)$를 고정한 채 $q(z)$에 대해 최적화하면, $q(z)$는 실제 사후 분포 $p(z D, \theta)$와 같아집니다. 이는 잠재 변수의 기댓값을 계산하는 과정입니다. - M-step: 업데이트된 $q(z)$를 고정한 채 $q(\theta)$에 대해 최적화하면, 이는 완전한 데이터의 로그 가능도 기댓값을 최대화하는 과정입니다. 따라서 EM 알고리즘은 BLR의 한 종류인 평균장 변분 추론(Mean-field VI)의 특별한 경우입니다.
-
3.2. Stochastic Variational Inference(SVI) \& Adam
- 가정: EM과 달리 $q$의 분리 가능성을 가정하지 않지만, 미니배치(mini-batch) 데이터를 사용하여 목적 함수의 경사(gradient)를 확률적으로 추정합니다.
- 유도: BLR의 업데이트 규칙에서 사용되는 실제 경사를 미니배치 기반의 확률적 경사로 대체하면, 이것이 바로 Stochastic Variational Inference(SVI) 입니다. 위에서 다룬 Adam/RMSProp은 SVI를 수행하기 위한 효율적인 방법론 중 하나입니다. 즉, 가우시안 근사($q(\theta) = \mathcal{N}(\mu, \sigma^2)$)를 사용한 SVI에서, 자연 경사를 효율적으로 근사하는 과정이 Adam의 업데이트 규칙으로 귀결됩니다.
Significance of the BLR: A Unifying Perspective
BLR은 단지 여러 알고리즘을 유도하는 수학적 트릭이 아니라 머신러닝 알고리즘을 바라보는 근본적인 관점을 제시합니다.
- 통합적 이해: 개별적으로 존재하던 알고리즘들이 ‘Updating Bayesian Belif’라는 원리의 다른 표현임을 보여줌으로써, 깊고 일관된 이해를 가능케 합니다.
- 새로운 알고리즘 설계: $q(\theta)$에 대한 새로운 가정을 하거나, 계산 과정에서 다른 종류의 근사를 적용함으로써 기존에 없던 새로운 하이브리드 알고리즘을 체계적으로 설계할 수 있는 가능성을 제공합니다.
- 알고리즘의 ‘왜’에 대한 설명: “왜 Adam이 잘 작동하는가?”와 같은 질문에 “경험적으로 좋았다”를 넘어, “불확실성을 고려한 베이지안 학습 규칙의 효율적인 근사이기 때문”이라는 원리적 답변을 가능하게 합니다.
Conclusion
이번 포스팅에서는 [The Bayesian Learning Rule] 논문의 핵심 개념과 의의를 살펴보았습니다. BLR은 Variational Inference과 Natural Gradient를 결합하여, 수많은 머신러닝 알고리즘을 아우르는 Generalised Powerful Framework를 제시합니다.
긴 글 읽어주셔서 감사합니다.