""" --- title: Adam Optimizer summary: A simple PyTorch implementation/tutorial of Adam optimizer --- # Adam Optimizer This is a [PyTorch](https://pytorch.org) implementation of popular optimizer *Adam* from paper [Adam: A Method for Stochastic Optimization](https://papers.labml.ai/paper/1412.6980). *Adam* update is, \begin{align} m_t &\leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t \\ v_t &\leftarrow \beta_2 v_{t-1} + (1 - \beta_2) \cdot g_t^2 \\ \hat{m}_t &\leftarrow \frac{m_t}{1-\beta_1^t} \\ \hat{v}_t &\leftarrow \frac{v_t}{1-\beta_2^t} \\ \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \end{align} where $\alpha$, $\beta_1$, $\beta_2$ and $\epsilon$ are scalar hyper parameters. $m_t$ and $v_t$ are first and second order moments. $\hat{m}_t$ and $\hat{v}_t$ are biased corrected moments. $\epsilon$ is used as a fix for division by zero error, but also acts as a form of a hyper-parameter that acts against variance in gradients. Effective step taken assuming $\epsilon = 0$ is, $$\Delta t = \alpha \cdot \frac{\hat{m}_t}{\hat{v}_t}$$ This is bounded by, $$\vert \Delta t \vert \le \alpha \cdot \frac{1 - \beta_1}{\sqrt{1-\beta_2}}$$ when $1-\beta_1 \gt \sqrt{1-\beta_2}$ and $$\vert \Delta t\vert \le \alpha$$ otherwise. And in most common scenarios, $$\vert \Delta t \vert \approx \alpha$$ """ import math from typing import Dict, Any, Tuple, Optional import torch from labml import tracker from torch import nn from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay class Adam(GenericAdaptiveOptimizer): """ ## Adam Optimizer We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html) to implement the Adam optimizer. """ def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-16, weight_decay: WeightDecay = WeightDecay(), optimized_update: bool = True, defaults: Optional[Dict[str, Any]] = None): """ ### Initialize the optimizer * `params` is the list of parameters * `lr` is the learning rate $\alpha$ * `betas` is a tuple of ($\beta_1$, $\beta_2$) * `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update` * `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html) * `optimized_update` is a flag whether to optimize the bias correction of the second moment by doing it after adding $\epsilon$ * `defaults` is a dictionary of default for group values. This is useful when you want to extend the class `Adam`. """ defaults = {} if defaults is None else defaults defaults.update(weight_decay.defaults()) super().__init__(params, defaults, lr, betas, eps) self.weight_decay = weight_decay self.optimized_update = optimized_update def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter): """ ### Initialize a parameter state * `state` is the optimizer state of the parameter (tensor) * `group` stores optimizer attributes of the parameter group * `param` is the parameter tensor $\theta_{t-1}$ """ # This is the number of optimizer steps taken on the parameter, $t$ state['step'] = 0 # Exponential moving average of gradients, $m_t$ state['exp_avg'] = torch.zeros_like(param, memory_format=torch.preserve_format) # Exponential moving average of squared gradient values, $v_t$ state['exp_avg_sq'] = torch.zeros_like(param, memory_format=torch.preserve_format) def get_mv(self, state: Dict[str, Any], group: Dict[str, Any], grad: torch.Tensor): """ ### Calculate $m_t$ and and $v_t$ * `state` is the optimizer state of the parameter (tensor) * `group` stores optimizer attributes of the parameter group * `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$ """ # Get $\beta_1$ and $\beta_2$ beta1, beta2 = group['betas'] # Get $m_{t-1}$ and $v_{t-1}$ m, v = state['exp_avg'], state['exp_avg_sq'] # In-place calculation of $m_t$ # $$m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t$$ m.mul_(beta1).add_(grad, alpha=1 - beta1) # In-place calculation of $v_t$ # $$v_t \leftarrow \beta_2 v_{t-1} + (1 - \beta_2) \cdot g_t^2$$ v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) return m, v def get_lr(self, state: Dict[str, any], group: Dict[str, any]): """ ### Get learning-rate This returns the modified learning rate based on the state. For *Adam* this is just the specified learning rate for the parameter group, $\alpha$. """ return group['lr'] def adam_update(self, state: Dict[str, any], group: Dict[str, any], param: torch.nn.Parameter, m: torch.Tensor, v: torch.Tensor): """ ### Do the *Adam* parameter update * `state` is the optimizer state of the parameter (tensor) * `group` stores optimizer attributes of the parameter group * `param` is the parameter tensor $\theta_{t-1}$ * `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$. This computes the following \begin{align} \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \end{align} Since $\alpha$, $\beta_1$, $\beta_2$ and $\epsilon$ are scalars and others are tensors we modify this calculation to optimize the computation. \begin{align} \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \\ \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{m_t / (1-\beta_1^t)}{\sqrt{v_t/(1-\beta_2^t)} + \epsilon} \\ \theta_t &\leftarrow \theta_{t-1} - \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}} \\ \end{align} where $$\hat{\epsilon} = (1-\beta_2^t) \epsilon$$ is what we should specify as the hyper-parameter. """ # Get $\beta_1$ and $\beta_2$ beta1, beta2 = group['betas'] # Bias correction term for $\hat{m}_t$, $1 - \beta_1^t$ bias_correction1 = 1 - beta1 ** state['step'] # Bias correction term for $\hat{v}_t$, $1 - \beta_2^t$ bias_correction2 = 1 - beta2 ** state['step'] # Get learning rate lr = self.get_lr(state, group) # Whether to optimize the computation if self.optimized_update: # $\sqrt{v_t} + \hat{\epsilon}$ denominator = v.sqrt().add_(group['eps']) # $\alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}$ step_size = lr * math.sqrt(bias_correction2) / bias_correction1 # $\theta_t \leftarrow \theta_{t-1} - \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot # \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}}$ param.data.addcdiv_(m, denominator, value=-step_size) # Computation without optimization else: # $\frac{\sqrt{v_t}}{\sqrt{1-\beta_2^t}} + \epsilon$ denominator = (v.sqrt() / math.sqrt(bias_correction2)).add_(group['eps']) # $\frac{\alpha}{1-\beta_1^t}$ step_size = lr / bias_correction1 # $\theta_t \leftarrow \theta_{t-1} - \alpha \cdot # \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$ param.data.addcdiv_(m, denominator, value=-step_size) def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter): """ ### Take an update step for a given parameter tensor * `state` is the optimizer state of the parameter (tensor) * `group` stores optimizer attributes of the parameter group * `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$ * `param` is the parameter tensor $\theta_{t-1}$ """ # Calculate weight decay grad = self.weight_decay(param, grad, group) # Get $m_t$ and $v_t$ m, v = self.get_mv(state, group, grad) # Increment $t$ the number of optimizer steps state['step'] += 1 # Perform *Adam* update self.adam_update(state, group, param, m, v)