AdaMax

Adam can be understood as updating weights inversely proportional to the scaled L2 norm (squared) of past gradients. AdaMax extends this to the so-called infinite norm (max) of past gradients.

The calculation of the infinite norm exhibits a stable behavior. Mathematically, the infinite norm can be viewed as,

$\$\$u\_t=\backslash beta2^\backslash infty\; v\{t-1\}\; +\; (1-\backslash beta2^\backslash infty\; v\{t-1\})|g\_t|^\backslash infty=max(\backslash beta2\; \backslash cdot\; v\{t-1\},|g\_t|)\$\$$

We see the mathematical equivalence of the calculation of the infinite norm with the maximum of the parameters calculated up to t. Now the update would be similar to Adam,

$\$\$\backslash theta\_\{t+1\}=\backslash theta\_t-\; \backslash eta\; \backslash cdot\; \backslash frac\{m\_t\}\{u\_t\}\$\$$

Again, here $\$\backslash eta\$$ is the base learning rate and $\$m\_t\$$ is the momentum similar to as discussed in Adam.

Major Parameters

• Base Learning Rate
• Weight Decay
• Epsilon Coefficient
• Betas

Betas

The exponential moving average and the infinite norm are calculated in Adamax. Mathematically, given by the formula,

$\$\$V\_\{dw\}=\backslash beta1\; \backslash cdot\; V\{dw\}+(1-\backslash beta\_1)\backslash cdot\; \backslash partial\; w\backslash \; u\_t=\backslash beta2^\backslash infty\; v\{t-1\}\; +\; (1-\backslash beta2^\backslash infty\; v\{t-1\})|g\_t|^\backslash infty\$\$$

Here $\$\backslash beta\_1\$$ and $\$\backslash beta\_2\$$ are the betas. They are the exponential decay rates of the first moment and the exponentially weighted infinity norm.

Code Implementation

      # importing the library
import torch
import torch.nn as nn

x = torch.randn(10, 3)
y = torch.randn(10, 2)

# Build a fully connected layer.
linear = nn.Linear(3, 2)

# Build MSE loss function and optimizer.
criterion = nn.MSELoss()

# Optimization method using Adamax
optimizer = torch.optim.Adamax(linear.parameters(), lr=0.002, 
betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

# Forward pass.
pred = linear(x)

# Compute loss.
loss = criterion(pred, y)
print('loss:', loss.item())

optimizer.step()