Adam and AdamW step

ImplementMedium25 min11 checks

optimizersadamweight-decay

Lesson: OptimizersCheat sheetTrack: Deep learning coreTrack: Optimizers

Problem

Adam keeps two running averages for each parameter: the mean of the gradient and the mean of the squared gradient. It divides the first by the square root of the second, so every parameter gets its own step size. AdamW adds weight decay that does not pass through these averages. You write one Adam step as a pure function: it takes the state and returns a new state.

Task

Write two functions.

  1. init_adam_state(params) returns the start state: a dict with "m" and "v" as float zero arrays of the shape of params, and "t" equal to the int 0.
  2. adam_step(params, grads, state, lr, beta1, beta2, epsilon, weight_decay) does one update and returns (new_params, new_state).

One step computes, in this order:

t \leftarrow t + 1
m \leftarrow \beta_1 m + (1 - \beta_1) g, \qquad v \leftarrow \beta_2 v + (1 - \beta_2) g^2
\hat m = \frac{m}{1 - \beta_1^t}, \qquad \hat v = \frac{v}{1 - \beta_2^t}
\theta \leftarrow \theta (1 - \eta \lambda) - \eta \frac{\hat m}{\sqrt{\hat v} + \epsilon}

The counter grows before the bias correction, so the first step uses t = 1. All operations are element by element. The decay factor multiplies the old parameters, as in the AdamW class of the lesson code. With weight_decay = 0 the step is plain Adam.

Input and output

NameTypeShapeMeaning
paramsfloat arrayany shapethe parameters \theta before the step
gradsfloat arraysame as paramsthe gradient g
statedict"m" and "v": float arrays of the shape of params. "t": int, the number of steps done
lrfloatscalarthe learning rate \eta, default 0.001
beta1floatscalardefault 0.9
beta2floatscalardefault 0.999
epsilonfloatscalardefault 1e-8, added after the square root
weight_decayfloatscalarthe decoupled decay \lambda, default 0.0
  • adam_step returns a tuple. The first item is a float array with the shape of params. The second item is a new dict with the keys "m", "v" and "t".

Constraints

  • Arrays have at most 100 elements.
  • Do not change params, grads or the arrays in state. Return new arrays.
  • Use numpy only.

Examples

pythonparams = np.array([1.0, -2.0, 0.5])
grads = np.array([0.1, -0.3, 2.0])
state = init_adam_state(params)
params, state = adam_step(params, grads, state, lr=0.1)
params
state
textarray([ 0.90000001, -1.9       ,  0.4       ])
{'m': array([ 0.01, -0.03,  0.2 ]), 'v': array([1.e-05, 9.e-05, 4.e-03]), 't': 1}
pythonparams, state = adam_step(params, np.array([0.2, 0.1, -1.0]), state, lr=0.1)
params
state
textarray([ 0.80348181, -1.85997815,  0.3733663 ])
{'m': array([ 0.029, -0.017,  0.08 ]), 'v': array([4.999e-05, 9.991e-05, 4.996e-03]), 't': 2}
pythonparams = np.array([1.0, -2.0, 0.5])
adam_step(params, np.array([0.1, -0.3, 2.0]), init_adam_state(params), lr=0.1, weight_decay=0.5)[0]
textarray([ 0.85000001, -1.8       ,  0.375     ])

Hints

4 hints

    Solution and notes

    Locked

    Pass every check with Submit to open the reference solution and the notes.

    Run it on your computer

    Clone the repository, install numpy and run the same checks in a terminal:

    python3 scripts/practice/run.py adam-step my_answer.py
python3 scripts/practice/run.py adam-step my_answer.py --examples

    Start my_answer.py from starter.py.

    Your code

    Python loads on the first run

    Your code is saved in this browser as you type. On a phone you can read and edit the code, but the checks are easier to use on a desktop.

    Keyboard shortcuts
    Ctrl Enter
    Run the examples
    Ctrl Shift Enter
    Submit
    Tab / Shift Tab
    Indent or outdent the lines
    Esc then Tab
    Move the focus out of the editor
    Alt H
    Show the next hint

    On a Mac, use Cmd in place of Ctrl.

    Select Run to check the examples, or Submit to run every check.