1 / 1

AdamW from Scratch

JAX

Medium

Apply one AdamW update to a parameter pytree and return the new parameters together with the new optimizer state. The state is a dict with an integer step counter under step and two pytrees under mu and nu, shaped like params, that hold the running first and second moments of the gradient. With $t$ equal to step + 1, every leaf is updated as

\begin{aligned} m &= \beta_1 \, m_{\text{old}} + (1 - \beta_1) \, g, \qquad v = \beta_2 \, v_{\text{old}} + (1 - \beta_2) \, g^{2}, \\ \hat m &= \frac{m}{1 - \beta_1^{\,t}}, \qquad \hat v = \frac{v}{1 - \beta_2^{\,t}}, \qquad p_{\text{new}} = p - \eta \left( \frac{\hat m}{\sqrt{\hat v} + \epsilon} + \lambda \, p \right). \end{aligned}

Here p, g, $m_{\text{old}}$ and $v_{\text{old}}$ are matching leaves of params, grads, the mu tree and the nu tree. The symbol $\beta_1$ is b1, $\beta_2$ is b2, $\eta$ is learning_rate, $\epsilon$ is eps and $\lambda$ is weight_decay. Every operation is elementwise, and the decay term uses the parameter from before the update.

Your function is called inside jax.jit, so step and every hyperparameter arrive as traced scalars.

Return a dict with keys params and state. The value under params is the updated pytree. The value under state is a dict with keys step, holding step + 1 as an int32 JAX array of shape (), and mu and nu, holding the new $m$ and $v$ pytrees. Every returned pytree has the structure of params, and its leaves are float32 JAX arrays with the shapes of the matching parameter leaves.

Examples

Input: params = {"w": [1.0, -2.0], "b": 0.5}, grads = {"w": [0.5, -0.1], "b": 0.0}, state = {"step": 0, "mu": {"w": [0.0, 0.0], "b": 0.0}, "nu": {"w": [0.0, 0.0], "b": 0.0}}, learning_rate = 0.1, b1 = 0.9, b2 = 0.999, eps = 1e-8, weight_decay = 0.0

Output: {"params": {"b": 0.5, "w": [0.9, -1.9]}, "state": {"step": 1, "mu": {"b": 0.0, "w": [0.05, -0.01]}, "nu": {"b": 0.0, "w": [0.00025, 1e-05]}}}

Explanation: On the first step the bias corrections divide the moments by 0.1 and 0.001, which gives back the gradient and its square. Their ratio is about the sign of the gradient, so each entry of w moves by the learning rate 0.1 against its gradient, while b has a zero gradient and stays at 0.5.

Input: params = {"w": [[1.0, 0.5], [-0.5, 2.0]]}, grads = {"w": [[0.2, -0.4], [0.0, 1.0]]}, state = {"step": 4, "mu": {"w": [[0.1, -0.1], [0.05, 0.3]]}, "nu": {"w": [[0.01, 0.02], [0.001, 0.2]]}}, learning_rate = 0.01, b1 = 0.9, b2 = 0.99, eps = 1e-8, weight_decay = 0.1

Output: {"params": {"w": [[0.993141, 0.504304], [-0.507232, 1.993614]]}, "state": {"step": 5, "mu": {"w": [[0.11, -0.13], [0.045, 0.37]]}, "nu": {"w": [[0.0103, 0.0214], [0.00099, 0.208]]}}}

Input: params = {"layers": [{"w": [0.3, -0.3], "b": 0.1}, {"w": [1.5], "b": -1.0}]}, grads = {"layers": [{"w": [0.01, -0.02], "b": 0.5}, {"w": [-0.3], "b": 0.0}]}, state = {"step": 1, "mu": {"layers": [{"w": [0.001, 0.0], "b": 0.05}, {"w": [-0.03], "b": 0.0}]}, "nu": {"layers": [{"w": [0.00001, 0.00004], "b": 0.00025}, {"w": [0.00009], "b": 0.0}]}}, learning_rate = 0.05, b1 = 0.8, b2 = 0.95, eps = 0.01, weight_decay = 0.05

Output: {"params": {"layers": [{"b": 0.047402, "w": [0.281729, -0.283095]}, {"b": -0.9975, "w": [1.547675]}]}, "state": {"step": 2, "mu": {"layers": [{"b": 0.14, "w": [0.0028, -0.004]}, {"b": 0.0, "w": [-0.084]}]}, "nu": {"layers": [{"b": 0.012738, "w": [1.4e-05, 5.8e-05]}, {"b": 0.0, "w": [0.004586]}]}}}

Hint 1

jax.tree.map(f, params, grads, mu, nu) passes one matching leaf from each tree.

Hint 2

1 - b1 ** t works with a traced int32 t.

Hint 3

Compute the two bias corrections once, outside the tree map.

Requirements

Decay every leaf, including biases and scalars, unlike optimizers that skip them.

Constraints

params, grads, mu and nu share one pytree structure made of dicts and lists, and every leaf is float32 with between zero and three dimensions.
step is an integer between 0 and 1000, and when it is 0 every entry of mu and nu is zero.
Every entry of nu is non-negative.
learning_rate and weight_decay are non-negative, b1 satisfies 0 ≤ b1 < 1, b2 satisfies 0 < b2 < 1, and eps is positive.
In the Run panel, state is an object with the keys step, mu and nu. A JSON list whose items are all objects is read as a list node of a tree, and any other list or number is read as one array leaf.

Try Similar Problems

Global Norm Gradient ClippingMedium SGD with Momentum over a PytreeEasy Gradient Accumulation with ScanHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium

AdamW from Scratch

JAX

Medium

\begin{aligned} m &= \beta_1 \, m_{\text{old}} + (1 - \beta_1) \, g, \qquad v = \beta_2 \, v_{\text{old}} + (1 - \beta_2) \, g^{2}, \\ \hat m &= \frac{m}{1 - \beta_1^{\,t}}, \qquad \hat v = \frac{v}{1 - \beta_2^{\,t}}, \qquad p_{\text{new}} = p - \eta \left( \frac{\hat m}{\sqrt{\hat v} + \epsilon} + \lambda \, p \right). \end{aligned}

Your function is called inside jax.jit, so step and every hyperparameter arrive as traced scalars.

Examples

Output: {"params": {"b": 0.5, "w": [0.9, -1.9]}, "state": {"step": 1, "mu": {"b": 0.0, "w": [0.05, -0.01]}, "nu": {"b": 0.0, "w": [0.00025, 1e-05]}}}

Output: {"params": {"w": [[0.993141, 0.504304], [-0.507232, 1.993614]]}, "state": {"step": 5, "mu": {"w": [[0.11, -0.13], [0.045, 0.37]]}, "nu": {"w": [[0.0103, 0.0214], [0.00099, 0.208]]}}}

Hint 1

jax.tree.map(f, params, grads, mu, nu) passes one matching leaf from each tree.

Hint 2

1 - b1 ** t works with a traced int32 t.

Hint 3

Compute the two bias corrections once, outside the tree map.

Requirements

Decay every leaf, including biases and scalars, unlike optimizers that skip them.

Constraints

params, grads, mu and nu share one pytree structure made of dicts and lists, and every leaf is float32 with between zero and three dimensions.
step is an integer between 0 and 1000, and when it is 0 every entry of mu and nu is zero.
Every entry of nu is non-negative.
learning_rate and weight_decay are non-negative, b1 satisfies 0 ≤ b1 < 1, b2 satisfies 0 < b2 < 1, and eps is positive.
In the Run panel, state is an object with the keys step, mu and nu. A JSON list whose items are all objects is read as a list node of a tree, and any other list or number is read as one array leaf.

Try Similar Problems

params =

Accepts: object, array, number

grads =

Accepts: object, array, number

state =

Accepts: object

learning_rate =

Accepts: number

b1 =

Accepts: number

b2 =

Accepts: number

eps =

Accepts: number

weight_decay =

Accepts: number