딥러닝을 하다 보면 최적화쪽에서 large batch size가 small learning rate와 동일하다는 논리를 보곤 합니다. 이것은 실험적인 결과일까요? 아니면 이론적인 뼈대가 있는 것일까요? 이 글에서는, 이론적으로 large batch size가 small learning rate와 같은 효과를 줄 수 있다는 것을 보이기로 합니다.
Stochastic Gradient Descent의 조작
SGD를 다음처럼 씁시다:
\[\theta_{k+1} = \theta_{k} - \eta\nabla_{\theta_k}L_k\]이 때 $\theta_k$는 iteration $k$의 parameter이고, $\eta$는 learning rate이며 $L_k$는 $k$번째 iteration의 batch loss입니다. 이제 데이터가 ${(x_1,y_1),\cdots,(x_N,y_N)}$만큼 $N$개가 있다고 하면 total loss $L$은
\[L = \frac{1}{N}\sum_{i=1}^Nl(x_i,y_i)\]가 됩니다. Full gradient는 물론
\[\nabla L = \frac{1}{N}\sum_{i=1}^N\nabla l(x_i,y_i)\]가 될 것이고 mini-batch $B={(x_{k_1},y_{k_1}),\cdots(x_{k_B},y_{k_B})}$를 생각하되, $B\ll N$이라고 합시다.
평균과 공분산
$L_k$의 평균과 공분산은
\[\mathbb{E}[L_k] = \frac{1}{B}\sum_{j=1}^B\mathbb{E}[\nabla l(x_{k_j},y_{k_j})]=\frac{1}{B}\cdot B\cdot \nabla L=\nabla L\]과
\[\Sigma = \text{Cov}[\nabla L_K] = \mathbb{E}[(\nabla L_k - \nabla L)(\nabla L_k- \nabla L)^T]\]로 정의하게 되는 것으로부터
\[\text{Cov}[\nabla L_k]=\frac{1}{B}\Sigma\]가 됨을 알 수 있습니다.
Noise로 분해 조작
이제 mini-batch gradient를
\[\nabla L_k = \nabla L + (\nabla L_k-\nabla L)\]로 조작하고,
\[\xi_B:=\sqrt{B}(\nabla L_k-\nabla L)\]로 정의합시다. 그러면
평균이 0
\[\mathbb{E}[\xi_B]=\sqrt{B}(\nabla L - \nabla L)=0\]공분산이 $B$에 무관
\[\text{Cov}[\xi_B] = B\cdot \text{Cov}[\nabla L_k]=B\cdot\frac{1}{B}\Sigma=\Sigma\]가 됩니다. 즉, $\xi_B$는 $B$에 영향을 받지 않습니다. 따라서,
\[\nabla L_k = \nabla L + \frac{1}{\sqrt{B}}\xi_B\]가 됩니다. 여기서, $\xi_B$는 평균이 $0$, 공분산 $\Sigma$를 갖는 random vector입니다.
$\xi$는 Gaussian
$\xi_B$를 잘 살펴보면, $Y_i=\nabla L_k$이고 $\mu = \mathbb{E}[Y] = \nabla L$이라고 할 때
\[\xi_B = \sqrt{B}\bigg(\frac{1}{B}\sum_{j=1}^BY_j-\mu\bigg)=\frac{1}{\sqrt{B}}\sum_{j=1}^B(Y_j-\mu)\]가 되고, 정확히 multivariable central limit theorem에 의해서
\[\xi\xrightarrow{d}\mathcal{N}(0,\Sigma)\quad\text{as }B\to\infty\]가 됩니다. 즉,
\[\theta_{k+1}=\theta_k-\eta\nabla L_k\]는 full batch와 noise term
\[\begin{equation}\theta_{k+1} = \theta_k - \eta\nabla L + \frac{\eta}{\sqrt{B}}\xi,\quad\xi\sim\mathcal{N}(0,\Sigma)\end{equation}\]가 되는 것입니다.
Notes
이 유도는 논리적이 보이지만 실제로 몇 가지 가정에 기반하고 있습니다.
- Sampling with replacement를 가정합니다. 이는 $B\ll N$일때는 전혀 문제가 없습니다. 하지만 $N$이 적당히 작고 $B$가 적당히 큰 경우에는 도출이 힘든 가정입니다.
- 분산의 유한성. Heavy-tailed process와 같은 경우에는 본 mini-batch noise가 Gaussian이 되지 않을 수 있습니다.
- $B$가 너무 작지도 너무 크지도 않아야 함. $B=1$처럼 극단적으로 작은 경우에는 multivariate central limit theorem을 사용할 수가 없습니다. $N$이 아주 크고 $B$가 1보다는 상당히 크지만 $N$보다는 충분히 작은 경우에만 이 이야기를 할 수 있습니다.
- Large batch size가 small learning rate와 같다는 논리가 여기서 나옵니다. $(1)$식에서, $B$를 고정하고 $\eta$를 키우는 것과, $\eta$를 고정하고 $B$를 줄이는 것이 동일함을 알 수 있습니다.