Giving Different Weights Different Step Scales

Some parameters repeatedly receive larger gradients than others. Adam keeps separate running estimates of gradient direction and squared magnitude for each parameter.

13 sections · 4 questions · 3 coding challenges

1 / 13

The arrow keys move between sections too.

Remembering Both Direction and Scale

BOTIE

Adaptive subgradient methods used gradient history to change parameter-specific scales. Kingma and Ba’s Adam combined running first and second moment estimates with a correction for their zero initialisation.

Adam is an update rule, not a substitute for choosing a loss or checking generalisation. Our examples calculate the state explicitly; they do not claim that Adam always outperforms other optimisers.

Sources for this history

Should Every Weight Use the Same Effective Step?

Two model weights can receive gradients with very different sizes. A single learning rate may make one change abruptly while the other barely moves.

Adam keeps a record of the gradients it has seen for each parameter. One running estimate smooths the gradients themselves. Another tracks their squared magnitudes and sets the scale used in the update. We will calculate these separately so we can see what each contributes.

The interactive calculation starts with already bias-corrected moments.

Step 1 of 3

Observe gradients with different scales

Each parameter keeps its own entries.

An invented example for this lesson.
Parameter	Gradients
Parameter 1	2
Parameter 2	-4

Parameter 1
Gradients
2
Parameter 2
Gradients
-4

Step 2 of 3

Keep two running arrays

One tracks direction; the other tracks squared magnitude.

An invented example for this lesson.
Parameter	First moment	Second moment
Parameter 1	0.2	0.004
Parameter 2	-0.4	0.016

Parameter 1
First moment
0.2
Second moment
0.004
Parameter 2
First moment
-0.4
Second moment
0.016

Step 3 of 3

Correct the initial bias

At the first step the corrected moments recover the gradient and its square.

An invented example for this lesson.
Parameter	Corrected direction	Corrected square
Parameter 1	2	4
Parameter 2	-4	16

Parameter 1
Corrected direction
2
Corrected square
4
Parameter 2
Corrected direction
-4
Corrected square
16

Adaptive updates retain a separate scale for each parameter

Magnitude scales each direction separately.

Do equally large directions produce equally large steps?

The stored second moments measure different recent gradient magnitudes. Change one coordinate's scale and inspect its update while holding the direction fixed.

BOTIE

Botie already calculated bias-corrected moments and now needs the parameter update.

Choose an example

Parameters

Corrected first moments

Corrected second moments

Learning rate

Positive epsilon

Updated parameter values

00.9
11.1

Follow the operations

1. Parameters
1. 01
2. 11
2. Learning rate
```
0.1
```
3. Corrected first moments
1. 02
2. 1-4
4. Corrected second moments
1. 04
2. 116
5. Square root
1. 02
2. 14
6. Positive epsilon
```
0
```
7. Add
1. 02
2. 14
8. Divide
1. 01
2. 1-1
9. Multiply
1. 00.1
2. 1-0.1
10. Subtract
1. 00.9
2. 11.1
11. Result
1. 00.9
2. 11.1

This final stage assumes valid corrected moments. The full optimiser also maintains state and counts updates.

1. Give Each Stored Array One Meaning

The first moment m is an exponential average of gradients. The second moment v is an exponential average of squared gradients, so negative gradients still contribute positive magnitude information.

Both arrays start at zero and match the parameter shape. Beta1 controls direction memory; beta2 controls squared-magnitude memory. These are different jobs and normally use different factors.

The second moment is not the gradient variance because it has not subtracted the squared mean.

2. Correct for Starting at Zero

Early moving averages are pulled toward their initial zero values. At update number t, divide each moment by one minus its retention factor raised to t. Count t from one for the first update.

m = beta1*m_old + (1-beta1)*g
v = beta2*v_old + (1-beta2)*g²
m_hat = m / (1-beta1^t)
v_hat = v / (1-beta2^t)

The correction depends on the update count, which must be saved when resuming.

3. Scale the Direction by Its Magnitude

Divide corrected direction by the square root of corrected second moment, plus a small epsilon outside the root. Multiply by the learning rate and subtract from the parameter.

At the first step with a nonzero gradient and tiny epsilon, the update is close to a sign-based step. Later steps depend on both histories, so interpreting every Adam update as just a sign is wrong.

parameter_new = parameter - rate × m_hat / (sqrt(v_hat) + epsilon)

Epsilon prevents division by zero and also changes scaling for very small gradients.

Use the retained scales for two coordinates

Start after moment estimation so the scale adjustment is easy to see. The directions have equal magnitudes, but the second coordinate has a larger second moment.

Python

import numpy as np
weight = np.array([1., 1.])
mean_gradient = np.array([2., -2.])
second_moment = np.array([4., 16.])
step = 0.1 * mean_gradient / (np.sqrt(second_moment) + 1e-8)
print(np.round(step, 3).tolist())
print(np.round(weight - step, 3).tolist())

Output

[0.1, -0.05]
[0.9, 1.05]

The second coordinate moves less despite an equally large direction. Practice also constructs the moments and applies the correction for their zero initialisation.

4. Check State, Convention, and Validation

Adam implementations can differ in epsilon placement and weight-decay handling. State the rule used when comparing a hand calculation with a library. AdamW applies a decoupled decay rule; it is not automatically identical to adding an L2 gradient.

A stable-looking training curve does not guarantee useful predictions. Compare validation performance, numerical finiteness, and repeatability, and retain the two moments along with weights when continuing a run.

The workshop starts from corrected moments so learners can inspect the final scaling operation; the coding example calculates the full first update.

5. Watch a Gradient Change Direction

The first update cannot show why Adam keeps a history. Give one weight a positive gradient, then a negative gradient. The first moment retains some of the old direction instead of changing instantly. The second moment retains squared magnitudes, so a sign change does not cancel it.

Write the two histories on separate rows. With beta1 equal to 0.5, a first gradient of two produces a first moment of one. A next gradient of minus two brings that moment below zero, but only partway towards the new gradient.

m₀ = 0
m₁ = 0.5 × 0 + 0.5 × 2 = 1
m₂ = 0.5 × 1 + 0.5 × (-2) = -0.5

The second moment uses 2² and (-2)², both 4.

A retained direction can disagree with the newest gradient. Inspect that disagreement before assuming the optimizer has made an arithmetic error.

6. Compare the Whole Training Path

Choose a small quadratic loss with coordinates on different scales. Start each optimizer from the same weights and feed it the gradient at its own current position. Those gradients will diverge as the methods take different routes.

Record loss, parameter values and step size for every update. An adaptive method may make steadier progress on a badly scaled problem, but a badly chosen rate can still make it unstable. Compare several rates on training data and use separate validation behaviour when choosing a procedure for a prediction task.

Example loss: L(w) = 0.5 × (w₀² + 10 × w₁²)
Gradient: [w₀, 10 × w₁]

This is a controlled optimization example with a known minimum. It lets us check state handling before blaming a dataset or architecture.

7. Resume the Same Optimizer, Not Just the Same Weights

Saving weights preserves the model’s current prediction rule. Continuing Adam also requires its first moment, second moment and update count. If any are reset, the next update follows a different history.

Store arrays as arrays of the original shape, and restore the counter as an integer. A checkpoint should also state the learning rate, retention factors and epsilon convention. Compare an uninterrupted run with a run stopped and resumed midway. Their next weights should agree within numerical tolerance.

A restart is a useful behavioural test because the final printed output alone can hide an accidental reset of the optimizer state.