Adamax

The problem, from first principles

Adam adapts learning rates per parameter using first and second moment estimates. Adamax is a variant that replaces the L2-norm-based second moment with an L-infinity norm operation.

Instead of dividing by sqrt(v), Adamax divides by the maximum of all past absolute gradients seen.

The updates are:

m_{\text{new}} = \beta_1 \cdot m + (1-\beta_1) \cdot \text{grad}

u_{\text{new}} = \max(\beta_2 \cdot u, |\text{grad}|)

\text{param}_{\text{new}} = \text{param} - \text{lr} \cdot \frac{m_{\text{new}}}{u_{\text{new}} + \text{eps}}

Where:

m is the first moment (exponential moving average of gradients).
u is the infinity norm (max absolute gradient observed, decayed).
beta1 controls first moment decay (typically 0.9).
beta2 controls infinity norm decay (typically 0.999).
eps provides numerical stability.

From theory to code

Implement:

adamax_step(
    param,
    grad,
    m,
    u,
    lr,
    beta1,
    beta2,
    eps,
)

The function returns:

(new_param, new_m, new_u)

The signature is provided in the editor.

First update the first moment estimate. Then update the infinity norm. Finally, compute the scaled update using both.

Constraints

param, grad, m, and u must have the same shape. Mismatch raises ValueError.
lr must be greater than 0. Otherwise raise ValueError.
beta1 and beta2 must both be in [0, 1). Otherwise raise ValueError.
eps must be greater than 0. Otherwise raise ValueError.
Inputs must not be modified.
Return new arrays.
m and u start at zeros.

Hints

Hint 1

Compute m_new first as an exponential moving average of the gradient.

Hint 2

Compute u_new as the element-wise maximum of beta2 * u and the absolute gradient.

Hint 3

After computing m_new and u_new, scale the parameter update by m_new divided by (u_new + eps).

The problem, from first principles

From theory to code

Constraints

Hints

Keywords