최적화 - RAdam (Rectified Adam)

Rectified Adam Optimizer란

• RAdm Optimizer는 기존의 Adam optimizer의 문제점을 개선한 최적화 방식이다. Adam Optimizer는 학습률이 기하급수적으로 늘어날 수도 있고, 학습을 너무 빨리 하여 local minima에 도달할 수 있다는 단점이 있다.
• RAdam은 학습이 너무 빠르다는 문제를 완화한 optimizer로 초기에 불안정하게 학습하는 Adam의 특성을 분석하여 고안되었다.
• 따라서 RAdam은 Adaptive Learning Rate의 신뢰도를 보정하는 보정 계수를 도입하였다. 이는 Variance rectification term을 통해 학습 초기에 Adam이 흔들리지 않도록 안정화 시킨다.
• 또한 논문에서는 두 가지 분산 감소 매커니즘을 평가한다
  • Adam-2k: 학습 초기에 2k 스텝 동안만 적응형 학습률 ($V_{t})을 계산하고 파라미트 업데이트를 하지만 모멘텀($m_{t})은 계산하지 않는다
  • Adam-eps : 큰 $\epsilon$을 사용하는 Adam 변형

수식 정리

Adam은 다음의 수식을 통해 1차, 2차 모멘텀을 추정한다.

Adam Optimizer Equation

 

(Adam Optimizer에 대한 설명은 아래 포스팅에 정리해 두었다.)

2025.03.05 - [Artificial intelligence/Deep Learning] - 최적화 - Adam Optimizer (Adaptive Moment Estimation)

최적화 - Adam Optimizer (Adaptive Moment Estimation)

 

RAdam에서는 $\hat{v_{t}}$의 분산이 클 경우 학습률을 줄이기 위해 다음을 도입:

RAdam Equation

 

여기서 $\alpha_{t}$는 Rectification Term을 고려해 조정된다. 신뢰 가능한 분산 추정치를 판별하기 위해 $p_{t}$를 도입했다.

Rectification Term

• $p_{\infty}$ : $2/1-\beta_{2} - 1$ : 이론상 최대값
• $t$ : 현재 step
• $\beta_{2}$: 모멘텀 파라미터

이때 적용 조건은 다음과 같다

• $p_{t} > 4$  : 분산이 신뢰할 수 있는 수준 , recitified term 적용
• $p_{t} < 4$ : 초기단계로서 Adam대신 SGD 사용

RAadm은 warm-up 없이도 학습률을 안정적으로 적용 가능하며 Adam의 빠른 수렴 속도와 SGD의 안정적 수렴 성질을 모두 지니고 있다. 또한 Adam보다 하이퍼 파라미터의 민감도가 낮아 튜닝이 더 편리하며 다양한 작업에서 성능향상이 확인되었다고 한다.

코드 구현

아래 구현은 pytorch로 구현한 RAdam optimizer이다.

import torch
from torch.optim.optimizer import Optimizer

class RAdam(Optimizer):
    def __init__(self, params,
                 lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
        defaults=dict(lr=lr,betas=betas,eps=eps,weight_decay=weight_decay)
        super().__init__(params, defaults)
    '''
    defaults : optimizer class dict
    params : 파라미터 리스트
    lr : 학습률
    betas : Tuple(beta1,beta2)
    eps: epsilon값
    '''

    def step(self):
        # self.param_groups는 optimizer에 선언되어있음
        for group in self.param_groups:
            for p in group['params']:
                # gradient가 None이면 skip
                if p.grad is None:
                    continue
                grad = p.grad.data
                state = self.state[p]
                
                # state 초기화
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2 = group['betas']
                state['step'] += 1
                t = state['step'] # 현재 스텝

                # 업데이트
                exp_avg.mul_(beta1).add_(grad,alpha=1-beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

                # Recitification term
                rho_inf = 2 / (1 - beta2) - 1
                rho_t = rho_inf - 2 * t * (beta2 ** t) / (1 - beta2 ** t)

                if rho_t > 4:
                    r_t = ((rho_t - 4) * (rho_t - 2) * rho_inf /
                           ((rho_inf - 4) * (rho_inf - 2) * rho_t)) ** 0.5
                    denom = exp_avg_sq.sqrt().add_(group['eps'])
                    step_size = group['lr'] * r_t / (1 - beta1 ** t)
                    p.data.addcdiv_(exp_avg, denom, value=-step_size)
                else:
                    step_size = group['lr'] / (1 - beta1 ** t)
                    p.data.add_(exp_avg, alpha=-step_size)