[딥러닝] 딥러닝과 역전파의 핵심, Optimzier AdamW 위주로 정리

딥러닝 모델을 만들 때 이론적으로 가장 중요한 게 뭘까? 나는 단연코 역전파라고 생각한다. 물론 지금 우리는 pytorch나 tensorflow 같은 잘 만들어진 프레임워크에서 간단한 함수 몇 줄로 이 역전파 및 훈련이 가능하기 때문에, 실제 활용할 때는 그 중요성이 잘 느껴지지 않는다.

https://tiabet0929.tistory.com/79

[딥러닝] 역전파를 단 한 줄로 가능하게 해주는 backward() 함수 탐구

이전에 backward() 함수에 대해 탐구해보면서 언젠가 Optimizer에 대한 글도 적어야겠다고 마음먹었었는데, 최근 Paged AdamW, 8bit Optimizer 등 Optimzier 쪽도 발전한 내용이 굉장히 많아서 이 부분에 대한 글을 작성해보려고 한다. 즉 이번 글은 Optimizer과 AdamW의 기초가 될 내용만 다뤄보려고 한다.

 

Optimizer 언제 쓰는가

우선 Optimizer가 언제 쓰이는지와 그 역할부터 명확히 하고 넘어가고 싶다.

Optimizer는 앞서도 말했지만 '역전파를 할 때' 사용된다. Pytorch를 예로 들어서 말하면, backward()가 업데이트하고자 하는 값에 연결된 모든 변수들과 그 기울기를 다 추적해놓았으면, optimizer는 실제로 업데이트를 시켜주는 함수라고 이해하면 된다. 

 

import torch
import torch.optim as optim

model = SimpleLinearModel()

criterion = nn.MSELoss()  # 평균 제곱 오차 (Mean Squared Error)

optimizer = optim.SGD(model.parameters(), lr=0.01)  # SGD 사용, 학습률 0.01

x_train = torch.randn(100, 1)
y_train = 2 * x_train + 1 + 0.1 * torch.randn(100, 1)  # 약간의 노이즈 추가

num_epochs = 100
for epoch in range(num_epochs):
    # 순전파 (Forward pass)
    outputs = model(x_train)
    loss = criterion(outputs, y_train)

    # 역전파 및 최적화 (Backward and optimize)
    optimizer.zero_grad()  # 기존 기울기 초기화
    loss.backward()        # 기울기 계산
    optimizer.step()       # 파라미터 업데이트

 

코드를 통해 보면 이런거다. optimizer 안에는 model.parameters()가 들어가 있고, Loss함수(여기선 MSE)에 들어간 outputs를 통해 정답과 예측값 사이의 loss 함수가 정해지게 된다. 이후 loss.backward()로 loss와 연결된 모든 값들의 gradient를 계산, optimizer.step()으로 경사하강법에 의거 loss가 작아지는 방향으로 한 step 나아가게 되는 것이다. 중간에 optimizer.zero_grad() 함수는 이전 step에서 계산했던 기울기 정보를 초기화하는 단계이다. optimizer는 이전에 계산했던 값들을 저장해놓고 있기 때문에 매 epoch마다 optimizer의 기울기 정보를 초기화해줘야 한다.

역전파 (Backpropagation)

역전파도 간단하게만 언급하고 넘어가야 할 것 같다. 딥러닝 모델의 역전파는 경사하강법에 의해 이루어진다. 경사하강법이란 쉽게 말하면 loss가 작아지는 방향으로 가다 보면 언젠간 최고의 성능을 낼 수 있다라는 이론이다. 

경사하강법 설명

그러니까 이런 느낌이다. 따라서 우리는 loss가 작아지는 방향을 구하기 위해 미분을 하는 거고, 이 미분을 하는 대상인 활성화함수, Loss 함수 등의 선택이 딥러닝 모델의 성능을 가르는 경우가 많은 것이다. 경사하강법이 적용된 역전파를 수식으로 나타내자면 다음과 같다.

$$\theta \leftarrow \theta - \eta \frac{\partial \mathcal{L}}{\partial \theta}$$

$$\frac{\partial \mathcal{L}}{\partial \theta} = \frac{\partial \mathcal{L}}{\partial a} \times \frac{\partial a}{\partial z} \times \frac{\partial z}{\partial \theta}$$

즉 Loss에 대한 파라미터의 편미분 후 학습률 $\eta$를 곱해주어서 원래 파라미터에서 빼주면 역전파이다. 그리고 위 수식은 SGD, Stochastic Gradient Descent Optimizer라고도 불린다.

 

여기까지만 보면 "엥? 복잡한 Optimizer 왜 필요하지? 그냥 이렇게 역전파해주면 끝인데" 라는 의문을 품을 수 있는데, 위 식은 아주 기초적은 역전파 공식인 것이고 그 뒤로 파라미터 업데이트, 즉 학습을 잘 시키기 위해 발전된 부분이 많다. AdamW를 이해하기 위해 필요한 Momentum, RMSProp, Adam을 차근차근 정리해보겠다.

 

$$\begin{aligned} \theta_{t+1} &= \theta_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t \right) \end{aligned}$$

▲ AdamW에서 사용하는 역전파 수식

Momentum

$$\begin{aligned} v_t &= \beta v_{t-1} + (1 - \beta) \nabla_{\theta} \mathcal{L}(\theta_t) \\ \theta_{t+1} &= \theta_t - \eta v_t \end{aligned}$$

▲ Momentum 수식

 

Momentum 은 한국어로 하면 운동량, 정확히는 관성에 의해 발생하는 운동량이다. SGD는 단순히 그때그때의 Gradient만 가지고 다음 업데이트 방향을 정하기 때문에 목적지까지 빠르게 도달하지 못하며, Local Optimal에 갇히면 다시 탈출하지 못 할 가능성이 있다. 이를 보완하기 위해 적용한 것이 Momentum이다. 수식을 보면 알겠지만 Momentum은 이전의 Mometum과 Loss의 Gradient의 가중평균, 내지는 이동평균이다. 즉, Momentum은 업데이트될수록 이전의 값들을 기억하며, 여기서 값이란 결국 업데이트 방향이기 때문에 이전의 업데이트 방향들에 영향을 계속 누적하며 파라미터를 업데이트하는 방식이라 볼 수 있다. 이 방식으로 상기한 SGD의 문제점 2개, 즉 빠르게 Optimal에 도달하지 못하는 문제와 Local Optimal에 갇히면 탈출하지 못하는 문제를 해결할 수 있게 된다.(이론적으로) 이전에 빠르게 움직인 값들이 내재되어 있기 때문에 학습 순간의 Gradient가 작다고 해도 Momentum의 파라미터 업데이트는 계속해서 빠르게 움직일 수 있으며, 마찬가지로 값이 크게 바뀌고 있다면 Local Optimal을 탈출할 수 있기 때문.

 

https://paperswithcode.com/method/sgd-with-momentum

Papers with Code - SGD with Momentum Explained

이 글은 Momentum에 대해 자세히 설명해주는데, 글에 의하면 가중평균에서 $\beta$ 는 0.99, 0.9, 0.5 셋 중 하나의 값을 선택하는 것이 일반적이라고 한다.

 

RMSProp

$$\begin{aligned} s_t &= \beta s_{t-1} + (1 - \beta) \left( \nabla_{\theta} \mathcal{L}(\theta_t) \right)^2 \\ \theta_{t+1} &= \theta_t - \eta \frac{\nabla_{\theta} \mathcal{L}(\theta_t)}{\sqrt{s_t} + \epsilon} \end{aligned}$$

▲ RMSProp 수식

 

RMSProp은 Root Mean Square Propagation의 줄임말이다. RMSProp은 처음에 AdaGrad의 문제를 해결하기 위해 등장한 개념이다. AdaGrad까지 다 설명하긴 너무 길어지고, AdaGrad의 어떤 문제점이었는지만 간략히 설명하면 Step이 누적될수록 Gradient Vanishing이 심해지는 문제다. 즉 RMSProp은 Gradient가 너무 작아지지 않도록 방지하기 위해서 등장했다. 수식을 자세히 살펴보면 Momentum처럼 이동평균을 쓰는 것을 알 수 있다. 특이한 것은 기울기의 제곱과의 이동평균을 구하는 건데, 제곱을 해주는 건 AdaGrad에서 나왔던 내용이다. 그리고 이건 파라미트 업데이트 과정에서 Gradient 보정을 해주는 느낌으로 진행되어, Gradient가 너무 크면 파라미터 업데이트 크기를 줄이고, Gradient가 너무 작으면 파라미터 업데이트 크기를 늘리는 역할을 해주는 개념이 RMSProp이라고 정리할 수 있다. 여기서의 $\beta$도 0.9, 0.99 정도로 선택하는 것이 일반적이다.

 

Adam

$$\begin{aligned} m_t &= \beta_1 m_{t-1} + (1 - \beta_1) \nabla_{\theta} \mathcal{L}(\theta_t) \\ v_t &= \beta_2 v_{t-1} + (1 - \beta_2) \left( \nabla_{\theta} \mathcal{L}(\theta_t) \right)^2 \\ \hat{m}_t &= \frac{m_t}{1 - \beta_1^t} \\ \hat{v}_t &= \frac{v_t}{1 - \beta_2^t} \\ \theta_{t+1} &= \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \end{aligned}$$

▲ Adam(Adaptive Momentum estimation) 수식

아주 오랫동안 사랑받은 Optimizer인 Adam은 RMSProp + Momentum이라고 정리 가능하다. 수식을 살펴보자. $m$과 $v$는 1차, 2차 모멘텀이라고 부르긴 하지만 2차 모멘텀은 사실상 RMSProp의 $s$와 같은 역할을 한다. 파라미터 업데이트 단계를 보면 Momentum에서처럼 Gradient를 직접 학습률에 곱하는 것이 아닌 1차 모멘텀을 곱해주며, 이를 RMSProp처럼 2차 모멘텀으로 보정해주는 것이다. 그리고 RMSProp의 내용인 Gradient 보정이 Adaptive Learning Rate 효과를 갖는다고 하여 이름이 Adam이 된 것이다. 하이퍼파라미터의 디폴트 값은 $\beta_1$는 0.9, $\beta_2$는 0.999, $\epsilon$은 $10^{-8}$ 라고 한다.

AdamW

https://arxiv.org/abs/1711.05101

Decoupled Weight Decay Regularization

▲ AdamW 논문

 

짧게 쓴다고 썼는데 내용이 너무 길어졌다. 이제 대망의 AdamW 차례다. 상기한 수식의 전문을 다시 가져와보겠다.

 

$$\begin{aligned} g_t &= \nabla_{\theta} \mathcal{L}(\theta_t) \\ m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \\ \hat{m}_t &= \frac{m_t}{1 - \beta_1^t} \\ \hat{v}_t &= \frac{v_t}{1 - \beta_2^t} \\ \theta_{t+1} &= \theta_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t \right) \end{aligned}$$

 

Adam과 대비해보면 달라진건 딱 하나, 학습률과 곱해지는 인자에 파라미터 자체가 더해진다는 것이다. $\lambda \theta_t$ 부분인데, 이는 Weight Decay 항이라고 설명 가능하다. Weight Decay에 대해서도 설명할 부분이 많은데, 역시나 이번 포스팅의 목적하고는 조금 다르므로 간략하게만 정리하겠다. 

https://paperswithcode.com/method/weight-decay

Papers with Code - Weight Decay Explained

▲ Weight Decay 설명

Weight Decay는 정규화, Reguralization 기법의 일종이다. 정규화를 하는 이유는 어디서나 마찬가지지만 과적합 방지에 있다. 위와 같이 정규화를 해주지 않으면 모델의 예측 모형이 지나치게 극단적이게 되어서, 이를 완화해주는 데에 사용되는 것이다. 이 Weight Decay 방법은 이전부터 많이 쓰이던거라, SGD와 Adam에 적용되었는데 Adam을 사용하면서 Gradient에 정규화를 진행하면

$$\nabla_{\theta} \mathcal{L}(\theta_t) + \lambda \theta_t$$

이런 식으로 되는데 이는 Adam의 수식에서

$$\begin{aligned} \theta_{t+1} &= \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \end{aligned}$$

의 $\hat{m}_t$ 에 들어가게 된다. 이렇게 되면 정규화 요소인 분모에도 이 정규화가 또 영향을 받게 되는 거라, 정규화의 수학적인 일관성이 깨져버리게 된다.Weight Decay는 L2 Regularization과 수학적으로 동일하다. L2 Regularization은 목표가 파라미터를 작게 유지시키는 것이다. 하지만 Adam의 정규화는 큰건 작게, 작은건 크게 만드는 것이다. 그래서 이 둘이 곱해지면 수학적으로 충돌이 일어난다. 그래서 아예 Adam에는 정규화를 다른 식으로 적용해버리자 해서 나온 것이 AdamW이다.

 

$$\begin{aligned} \theta_{t+1} &= \theta_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t \right) \end{aligned}$$

 

다시 AdamW의 수식으로 돌아오자. 이 수식에서 왜 Weight Decay과 완전히 분리되냐면 괄호를 풀고 다시 엮어보면

 

$$\theta_{t+1} = (1 - \eta \lambda)\theta_t - \eta \cdot \text{AdamUpdate}$$

 

이렇게 써지기 때문이다. AdamUpdate는 당연히 순수 Adam에서의 업데이트 값들을 말한다. 이렇게 Adam with Weight Decay, AdamW가 완성되고 이 일련의 과정을 Decoupled Weight Decay Regularization, 원 논문의 제목이 된다.

 

AdamW에서 $\lambda$의 값은 디폴트가 0.01이다.

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)

 

Pytorch에서 코드로 불러오려면 이렇게 불러올 수 있다. 

 

결론 : AdamW는 여러 수학적 이론(어렵지는 않은)이 합쳐져 나온 현재 최고의 Optimizer

 

다음 포스팅에선 8bit Optimizer 혹은 PagedAdamW에 대해 다뤄보겠다.