AdamW / Learning Rate / Clipping

Three functions in optim.py turn a batch's mean gradient into changed parameters. This page documents them as they are coded. train.main runs them in this order:

train.pymainGitHub

62            norm = clip_grad_norm(grads, config.grad_clip)63            step += 164            lr = lr_schedule(step, config.warmup_steps, total_steps,65                             config.lr_max, config.lr_min)66            opt.step(params, grads, lr)

So clipping comes first, then the learning rate, then the optimizer step. How the Model Learns shows how they fit into one whole step. optim.py has no other optimizer, schedule or clipping mode.

The sections below are in the order of the code in optim.py.

AdamW

What is it?

AdamW updateImplemented

Optimizer

Adam.step changes every parameter once per batch. For each key it updates two moving averages, m and v, corrects them for their zero start, and divides one by the square root of the other. Keys with two or more dimensions also get weight decay. The class is named Adam, and train.main gives it a nonzero weight_decay, so what runs is AdamW.

Why does the model need it?

The gradient says which way each number should move, but its size varies a lot from one number to another and from one batch to the next. Dividing by the typical size in v gives each number a move of about lr, and m smooths the direction over recent batches. Weight decay shrinks the matrices slightly on each step, separately from those averages.

m and v

The class keeps two arrays per parameter, both starting at zero (optim.py:8). In general, these are called the first moment and the second moment. The code calls them m and v:

m is momentum. It is a running average of the gradient, so it holds the direction of recent batches. beta1 is how much of the old average it keeps on each step.
v is the second moment. It is a running average of the squared gradient, so it holds the typical size of the gradient for each number. beta2 is how much of the old average it keeps.

The update divides mhat by the square root of vhat. A number whose gradient is always large gets a large vhat, and the division brings its move back to about lr. So does a number whose gradient is always small. This is what the comment on that line says: every number moves about lr per step.

Mathematics

Powers, squares, square roots and divisions here apply to each number of the array separately.

Moving average of the gradient #

m_t = \beta_1 m_{t-1} + (1 - \beta_1)\, g_t

Moving average of the squared gradient #

v_t = \beta_2 v_{t-1} + (1 - \beta_2)\, g_t^2

m and v start at zero, so on early steps they are too small. Bias correction divides them by the amount they fall short. self.t goes up by one at the start of each call, so t is 1 on the first step and 1 - beta1 ** t is never zero.

Bias correction of m #

\hat m_t = \frac{m_t}{1 - \beta_1^{\,t}}

Bias correction of v #

\hat v_t = \frac{v_t}{1 - \beta_2^{\,t}}

The update before decay #

u_t = \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}

Decoupled weight decay #

u_t \leftarrow u_t + \lambda\, w_{t-1} \qquad \text{only where } w \text{ has two or more dimensions}

The parameter update, in place #

w_t = w_{t-1} - \eta_t\, u_t

Decay is decoupled: it is added to update after m and v are computed, so it never enters those averages, and clipping never sees it. Since the gradient plays no part in it, a parameter whose gradient is zero still shrinks. Embeddings shows this on the unused rows of pos_emb.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$w$	One parameter array	`params[k]`	same as the parameter
$g_t$	Its gradient at this step, after averaging over the batch and after clipping	`g = grads[k]`	same as the parameter
$m_t$	Running average of the gradient, which holds the direction	`self.m[k]`	same as the parameter
$v_t$	Running average of the squared gradient, which holds the typical size	`self.v[k]`	same as the parameter
$\hat m_t, \hat v_t$	The same averages after the correction for their zero start	`mhat, vhat`	same as the parameter
$u_t$	The update	`update`	same as the parameter
$t$	Step counter. It is 1 on the first call	`self.t`	scalar
$\beta_1, \beta_2$	How much of the old average each step keeps	`self.b1 from config.beta1, self.b2 from config.beta2`	scalar
$\epsilon$	Keeps the division defined when v is 0	`self.eps from config.adam_eps`	scalar
$\lambda$	The weight-decay rate	`self.wd from config.weight_decay`	scalar
$\eta_t$	The learning rate for this step	`lr`	scalar

Which parameters get weight decay

The rule is one line, params[k].ndim >= 2 (optim.py:27). The rule depends only on the number of dimensions, so the two embedding tables get decay too, although config.py calls this "weight matrices". The table lists every name init_params creates. "Weight decay" comes from running the real Adam.step on all ones with zero gradients, where only decay can change a value.

Key	Weight decay	`ndim`	Shape	Shape in config.py values
`tok_emb`	Yes	2	`(V, d)`	`(V, 32)`
`pos_emb`	Yes	2	`(T_max, d)`	`(32, 32)`
`b{l}.ln1_g`	No	1	`(d,)`	`(32,)`
`b{l}.ln1_b`	No	1	`(d,)`	`(32,)`
`b{l}.wq`	Yes	2	`(d, d)`	`(32, 32)`
`b{l}.bq`	No	1	`(d,)`	`(32,)`
`b{l}.wk`	Yes	2	`(d, d)`	`(32, 32)`
`b{l}.bk`	No	1	`(d,)`	`(32,)`
`b{l}.wv`	Yes	2	`(d, d)`	`(32, 32)`
`b{l}.bv`	No	1	`(d,)`	`(32,)`
`b{l}.wo`	Yes	2	`(d, d)`	`(32, 32)`
`b{l}.bo`	No	1	`(d,)`	`(32,)`
`b{l}.ln2_g`	No	1	`(d,)`	`(32,)`
`b{l}.ln2_b`	No	1	`(d,)`	`(32,)`
`b{l}.ff1_w`	Yes	2	`(d, F)`	`(32, 128)`
`b{l}.ff1_b`	No	1	`(F,)`	`(128,)`
`b{l}.ff2_w`	Yes	2	`(F, d)`	`(128, 32)`
`b{l}.ff2_b`	No	1	`(d,)`	`(32,)`
`lnf_g`	No	1	`(d,)`	`(32,)`
`lnf_b`	No	1	`(d,)`	`(32,)`
`head_w`	Yes	2	`(d, V)`	`(32, V)`
`head_b`	No	1	`(V,)`	`(V,)`

With num_layers = 2, the b{l}. keys appear once per block, so 15 tensors get decay and 23 do not.

Tiny numerical example

Two steps, with and without decayIllustrative ExampleThe numbers of test_adam_two_step_hand_trace: start at 1, gradients 0.5 then 0.25, and lr 0.1, which is far above lr_max so that the moves are easy to see. Other settings are beta1 = 0.9, beta2 = 0.999, adam_eps = 1e-8 and weight_decay = 0.01, the values in config.py. w has two dimensions and b has one.

Value	Code	Step 1, `w`	Step 1, `b`	Step 2, `w`	Step 2, `b`
Before	`params[k]`	`1`	`1`	`0.899`	`0.9`
Gradient	`g = grads[k]`	`0.5`	`0.5`	`0.25`	`0.25`
Learning rate	`lr`	`0.1`	`0.1`	`0.1`	`0.1`
Counter	`self.t`	`1`	`1`	`2`	`2`
Average	`self.m[k]`	`0.05`	`0.05`	`0.07`	`0.07`
Average of squares	`self.v[k]`	`0.00025`	`0.00025`	`0.00031225`	`0.00031225`
Corrected	`mhat`	`0.5`	`0.5`	`0.36842105`	`0.36842105`
Corrected	`vhat`	`0.25`	`0.25`	`0.1562031`	`0.1562031`
Before decay	`mhat / (np.sqrt(vhat) + self.eps)`	`0.99999998`	`0.99999998`	`0.93217962`	`0.93217962`
Update	`update`	`1.00999998`	`0.99999998`	`0.94116962`	`0.93217962`
After	`params[k] -= lr * update`	`0.899`	`0.9`	`0.80488304`	`0.80678204`

Read the first column of each pair against the second. w and b start equal and get the same gradients, so self.m[k] and self.v[k] stay equal. They differ only in update, because w also has self.wd * params[k]. The values for b round to the 0.9000 and 0.8068 that test_adam_two_step_hand_trace asserts.

On the first step, the correction makes mhat equal to g and vhat equal to g squared. So update before decay is close to the sign of g, whatever its size. A test checks this (tests/test_optim.py:21).

One step for one number

One step for one numberIllustrative ExampleThe fields start at the two-step example's first step, and you can edit them. To see that example's second step, copy self.m[k], self.v[k], self.t and params[k] from the result back into the fields, and set grads[k] to its second gradient. The betas, eps and weight_decay stay at their config.py values.

params[k]

grads[k]

self.m[k]

self.v[k]

self.t

Dimensions

Code, in order	Result
`self.t += 1`	`1`
`self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * g`	`0.05`
`self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * g * g`	`0.00025`
`mhat = self.m[k] / (1 - self.b1 ** self.t)`	`0.5`
`vhat = self.v[k] / (1 - self.b2 ** self.t)`	`0.25`
`update = mhat / (np.sqrt(vhat) + self.eps)`	`0.99999998`
`update = update + self.wd * params[k]`	`1.00999998`
`params[k] -= lr * update`	`0.899`

Tensor shapes

There is one set of arrays per key, and each has the shape of its parameter. V depends on the corpus, so it stays a symbol.

One key with two dimensions, head_w, whose update includes decayoptim.py:28

grads[k](d, V)(32, V)
self.m[k](d, V)(32, V)
self.v[k](d, V)(32, V)
update(d, V)(32, V)
params[k](d, V)(32, V)

d = emb_dim = 32. V depends on the corpus. It is not a value in config.py.

One key with one dimension, head_b, whose update has no decayoptim.py:25

grads[k](V,)
self.m[k](V,)
self.v[k](V,)
update(V,)
params[k](V,)

V depends on the corpus. It is not a value in config.py.

Actual implementation

optim.py__init__GitHub

7    def __init__(self, params, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.0):8        self.m = {k: np.zeros_like(v) for k, v in params.items()} # running average of the gradient (which direction to go)9        self.v = {k: np.zeros_like(v) for k, v in params.items()} # running average of gradient squared (how big the gradients usually are)10        self.t = 0 # step counter, needed for bias correction11        self.b1, self.b2, self.eps, self.wd = beta1, beta2, eps, weight_decay

optim.pystepGitHub

13    def step(self, params, grads, lr):14        self.t += 1 # first step is t=1, so 1 - b1**t is never zero1516        for k in params:17            g = grads[k]1819            self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * g # mostly old direction, a little new gradient20            self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * g * g # same idea, but for the squared gradient2122            mhat = self.m[k] / (1 - self.b1 ** self.t) # m starts at 0 so it is too small early on, this scales it back up23            vhat = self.v[k] / (1 - self.b2 ** self.t)2425            update = mhat / (np.sqrt(vhat) + self.eps) # direction divided by typical size, so every number moves about lr per step2627            if params[k].ndim >= 2: # weight decay only for matrices, not biases or layer norm gamma/beta28                update = update + self.wd * params[k] # shrink the weight a little, outside the m/v part (this is the "W" in AdamW)2930            params[k] -= lr * update # changes the array in place, nothing to return

params[k] -= lr * update changes the array in place, and step returns nothing. train.main creates the object once (train.py:51) and passes it the same params dictionary on every step.

The default for weight_decay in the class is 0.0, but train.main passes 0.01. So the pipeline always applies decay.

Code and math connection

Equation	Role	Lines	Code
Moving average of the gradient	Computes it	optim.py:19	`self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * g # mostly old direction, a little new gradient`
Moving average of the squared gradient	Computes it	optim.py:20	`self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * g * g # same idea, but for the squared gradient`
Bias correction of m	Computes it	optim.py:22	`mhat = self.m[k] / (1 - self.b1 ** self.t) # m starts at 0 so it is too small early on, this scales it back up`
Bias correction of v	Computes it	optim.py:23	`vhat = self.v[k] / (1 - self.b2 ** self.t)`
The update before decay	Computes it	optim.py:25	`update = mhat / (np.sqrt(vhat) + self.eps) # direction divided by typical size, so every number moves about lr per step`
Decoupled weight decay	Computes it	optim.py:27	`if params[k].ndim >= 2: # weight decay only for matrices, not biases or layer norm gamma/beta`
Decoupled weight decay	Computes it	optim.py:28	`update = update + self.wd * params[k] # shrink the weight a little, outside the m/v part (this is the "W" in AdamW)`
Decoupled weight decay	Tests it	tests/test_optim.py:29	`assert np.allclose(params["b"], 1.0) # untouched`
The parameter update, in place	Computes it	optim.py:30	`params[k] -= lr * update # changes the array in place, nothing to return`
The parameter update, in place	Calls it	train.py:66	`opt.step(params, grads, lr)`
The parameter update, in place	Tests it	tests/test_optim.py:13	`assert np.isclose(params["w"][0], 0.8068, atol=1e-4)`

Where is it used?

Source

optim.py, lines 13-30, step

PDF reference

Section 4.6, Eq. 55, 56, 57, p. 13

Used by

Called by: train.main

Tested by: tests.test_optim.test_adam_first_step_is_lr_times_sign_of_gradient, tests.test_optim.test_adam_two_step_hand_trace, tests.test_optim.test_weight_decay_only_touches_matrices, tests.test_transformer.test_overfit_one_window

Tests it: tests/test_optim.py, line 29 (test_weight_decay_only_touches_matrices), for Decoupled weight decay
Calls it: train.py, line 66 (main), for The parameter update, in place
Tests it: tests/test_optim.py, line 13 (test_adam_two_step_hand_trace), for The parameter update, in place

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Moving average of the gradient	Computes it	`optim.py`	`step`	19	GitHub: optim.py, line 19
Moving average of the squared gradient	Computes it	`optim.py`	`step`	20	GitHub: optim.py, line 20
Bias correction of m	Computes it	`optim.py`	`step`	22	GitHub: optim.py, line 22
Bias correction of v	Computes it	`optim.py`	`step`	23	GitHub: optim.py, line 23
The update before decay	Computes it	`optim.py`	`step`	25	GitHub: optim.py, line 25
Decoupled weight decay	Computes it	`optim.py`	`step`	27	GitHub: optim.py, line 27
Decoupled weight decay	Computes it	`optim.py`	`step`	28	GitHub: optim.py, line 28
The parameter update, in place	Computes it	`optim.py`	`step`	30	GitHub: optim.py, line 30

Learning-rate schedule

What is it?

Learning-rate scheduleImplemented

Optimizer

lr_schedule gives the learning rate for one step. Up to warmup_steps it rises in a straight line, reaching lr_max exactly at step warmup_steps. After that it follows a cosine curve down to lr_min, reaching it exactly at step total_steps.

Why does the model need it?

In general, large moves right after initialisation can do harm, and smaller moves near the end make the parameters settle. The schedule gives small steps first, the largest step at the end of warmup, and smaller steps toward the end.

Mathematics

Linear warmup #

\eta(t) = \eta_{\max} \cdot \frac{t}{t_{\text{warm}}} \qquad t \le t_{\text{warm}}

How far through the decay #

f = \frac{t - t_{\text{warm}}}{t_{\text{total}} - t_{\text{warm}}} \qquad t > t_{\text{warm}}

Cosine decay #

\eta(t) = \eta_{\min} + \tfrac{1}{2} (\eta_{\max} - \eta_{\min}) \big(1 + \cos(\pi f)\big)

What does each symbol mean?

Symbol	Meaning	In code	Shape
$t$	The step. train.main adds one before the call, so the first value is 1	`step`	scalar
$t_{\text{warm}}$	The last step of the warmup	`warmup, from config.warmup_steps`	scalar
$t_{\text{total}}$	The last step of training	`total, from total_steps in train.main`	scalar
$\eta_{\max}$	The largest learning rate	`lr_max, from config.lr_max`	scalar
$\eta_{\min}$	The learning rate at the last step	`lr_min, from config.lr_min`	scalar
$f$	How far through the decay this step is, from just above 0 to 1	`frac`	scalar

warmup_steps is 50, lr_max is 1e-3 and lr_min is 1e-4 in config.py. total_steps is computed by train.main from the number of windows, so it depends on the corpus. Training Loop shows that calculation.

Tiny numerical example

Edit the four values to see the curve change. The table lists the values at the key steps.

total_steps

warmup_steps

lr_max

lr_min

lr against stepIllustrative Examplewarmup_steps, lr_max and lr_min start at their config.py values. total_steps depends on the corpus, so its starting value is hand-picked. Both axes are linear, and the y axis starts at 0.

0step200

Data table

step	lr
0	0
1	0.00002
2	0.00004
3	0.00006
4	0.00008
5	0.0001
6	0.00012
7	0.00014
8	0.00016
9	0.00018
10	0.0002
11	0.00022
12	0.00024
13	0.00026
14	0.00028
15	0.0003
16	0.00032
17	0.00034
18	0.00036
19	0.00038
20	0.0004
21	0.00042
22	0.00044
23	0.00046
24	0.00048
25	0.0005
26	0.00052
27	0.00054
28	0.00056
29	0.00058
30	0.0006
31	0.00062
32	0.00064
33	0.00066
34	0.00068
35	0.0007
36	0.00072
37	0.00074
38	0.00076
39	0.00078
40	0.0008
41	0.00082
42	0.00084
43	0.00086
44	0.00088
45	0.0009
46	0.00092
47	0.00094
48	0.00096
49	0.00098
50	0.001
51	0.0009999
52	0.0009996
53	0.0009991
54	0.0009984
55	0.0009975
56	0.0009965
57	0.0009952
58	0.0009937
59	0.000992
60	0.0009902
61	0.0009881
62	0.0009859
63	0.0009834
64	0.0009808
65	0.000978
66	0.000975
67	0.0009718
68	0.0009684
69	0.0009648
70	0.0009611
71	0.0009572
72	0.0009531
73	0.0009488
74	0.0009443
75	0.0009397
76	0.0009349
77	0.0009299
78	0.0009248
79	0.0009195
80	0.0009141
81	0.0009084
82	0.0009027
83	0.0008967
84	0.0008906
85	0.0008844
86	0.000878
87	0.0008715
88	0.0008648
89	0.000858
90	0.0008511
91	0.000844
92	0.0008368
93	0.0008295
94	0.0008221
95	0.0008145
96	0.0008068
97	0.000799
98	0.0007911
99	0.0007831
100	0.000775
101	0.0007668
102	0.0007585
103	0.0007501
104	0.0007416
105	0.000733
106	0.0007244
107	0.0007157
108	0.0007069
109	0.000698
110	0.0006891
111	0.0006801
112	0.000671
113	0.0006619
114	0.0006528
115	0.0006436
116	0.0006343
117	0.000625
118	0.0006157
119	0.0006064
120	0.000597
121	0.0005877
122	0.0005783
123	0.0005688
124	0.0005594
125	0.00055
126	0.0005406
127	0.0005312
128	0.0005217
129	0.0005123
130	0.000503
131	0.0004936
132	0.0004843
133	0.000475
134	0.0004657
135	0.0004564
136	0.0004472
137	0.0004381
138	0.000429
139	0.0004199
140	0.0004109
141	0.000402
142	0.0003931
143	0.0003843
144	0.0003756
145	0.000367
146	0.0003584
147	0.0003499
148	0.0003415
149	0.0003332
150	0.000325
151	0.0003169
152	0.0003089
153	0.000301
154	0.0002932
155	0.0002855
156	0.0002779
157	0.0002705
158	0.0002632
159	0.000256
160	0.0002489
161	0.000242
162	0.0002352
163	0.0002285
164	0.000222
165	0.0002156
166	0.0002094
167	0.0002033
168	0.0001973
169	0.0001916
170	0.0001859
171	0.0001805
172	0.0001752
173	0.0001701
174	0.0001651
175	0.0001603
176	0.0001557
177	0.0001512
178	0.0001469
179	0.0001428
180	0.0001389
181	0.0001352
182	0.0001316
183	0.0001282
184	0.000125
185	0.000122
186	0.0001192
187	0.0001166
188	0.0001141
189	0.0001119
190	0.0001098
191	0.000108
192	0.0001063
193	0.0001048
194	0.0001035
195	0.0001025
196	0.0001016
197	0.0001009
198	0.0001004
199	0.0001001
200	0.0001

Where	step	lr
Step 0, which train.main never uses	`0`	`0`
Step 1, the first one used	`1`	`0.00002`
The last warmup step	`50`	`0.001`
The first step on the cosine	`51`	`0.000999901`
The last step	`200`	`0.0001`

train.main adds one to step before calling the schedule (train.py:63), so step 0 is what the function would return, not a value training uses.

lr_max is reached at the last warmup step. At step warmup_steps, the test step <= warmup is still true (optim.py:33), so the function returns lr_max * step / warmup (optim.py:34), which is exactly lr_max. The next step is already on the cosine and slightly lower, as the table shows. The comment on lr_max in config.py says "reached right after warmup" (config.py:24). This page follows the code.

Tensor shapes

The schedule takes numbers and returns one number, so it has no tensors. Adam.step multiplies every update by that same number.

Actual implementation

optim.pylr_scheduleGitHub

32def lr_schedule(step, warmup, total, lr_max, lr_min):33    if step <= warmup:34        return lr_max * step / warmup # warm-up: ramp from near 0 up to lr_max, a fresh model is fragile3536    frac = (step - warmup) / (total - warmup) # how far through the decay we are, goes from 0 to 137    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * frac)) # cosine slide from lr_max down to lr_min

Code and math connection

Equation	Role	Lines	Code
Linear warmup	Computes it	optim.py:33	`if step <= warmup:`
Linear warmup	Computes it	optim.py:34	`return lr_max * step / warmup # warm-up: ramp from near 0 up to lr_max, a fresh model is fragile`
How far through the decay	Computes it	optim.py:36	`frac = (step - warmup) / (total - warmup) # how far through the decay we are, goes from 0 to 1`
Cosine decay	Computes it	optim.py:37	`return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * frac)) # cosine slide from lr_max down to lr_min`
Cosine decay	Calls it	train.py:64	`lr = lr_schedule(step, config.warmup_steps, total_steps,`
Cosine decay	Tests it	tests/test_optim.py:35	`assert np.allclose(got, [1e-4, 5e-4, 1e-3, 5.5e-4, 1e-4])`

Where is it used?

Source

optim.py, lines 32-37, lr_schedule

PDF reference

Section 4.6, Eq. 58, p. 14

Used by

Called by: train.main

Tested by: tests.test_optim.test_lr_schedule_values

Calls it: train.py, line 64 (main), for Cosine decay
Tests it: tests/test_optim.py, line 35 (test_lr_schedule_values), for Cosine decay

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Linear warmup	Computes it	`optim.py`	`lr_schedule`	33	GitHub: optim.py, line 33
Linear warmup	Computes it	`optim.py`	`lr_schedule`	34	GitHub: optim.py, line 34
How far through the decay	Computes it	`optim.py`	`lr_schedule`	36	GitHub: optim.py, line 36
Cosine decay	Computes it	`optim.py`	`lr_schedule`	37	GitHub: optim.py, line 37

Gradient clipping

What is it?

Gradient clippingImplemented

Optimizer

clip_grad_norm treats every gradient of every parameter as one long vector and measures its length. If that length is above the limit, it scales every tensor by the same factor, in place, so the length becomes the limit and the direction stays. It returns the length from before clipping.

Why does the model need it?

The comment on grad_clip in config.py gives the reason, that one bad batch must not blow up the model. A batch with an unusually large gradient still moves the parameters in its direction, but with a limited length.

Mathematics

Global norm, over all tensors #

\lVert g \rVert = \sqrt{\sum_{k} \sum_{i} g_{k,i}^2}

Scaling by one factor #

g_k \leftarrow g_k \cdot \frac{c}{\lVert g \rVert} \qquad \text{for every key } k, \text{ only when } \lVert g \rVert > c

What does each symbol mean?

Symbol	Meaning	In code	Shape
$g_k$	The gradient tensor of one parameter	`grads[k]`	same as the parameter
$g_{k,i}$	One number of that tensor	not named in the code	scalar
$\lVert g \rVert$	The length of all gradients together. This is the value train.py prints as \|g\|	`norm`	scalar
$c$	The limit	`max_norm, from config.grad_clip`	scalar

Tiny numerical example

Change the values or the limit. The two tensors start with one length of 3 and one of 4, which gives a norm of 5. With a limit of 1.0, each is scaled by the same factor.

Tensor a

Tensor b

max_norm

norm is above max_norm, so every tensor is multiplied by max_norm / norm = 0.2.

Length of the whole gradientIllustrative ExampleHand-picked values that you can edit. max_norm starts at grad_clip from config.py.

norm before5.0000

max_norm1.0000

norm after1.0000

Data table

Label	Length
norm before	5.0000
max_norm	1.0000
norm after	1.0000

Tensor	Before	After
`a`	`3, 0`	`0.6, 0`
`b`	`0, 4`	`0, 0.8`

norm = 5, the value clip_grad_norm returns
max_norm / norm = 0.2, the factor for every tensor

Tensor shapes

Clipping changes values, not shapes. head_w as one exampleoptim.py:48

grads[k] before(d, V)(32, V)
grads[k] after(d, V)(32, V)

d = emb_dim = 32. V depends on the corpus. It is not a value in config.py.

Actual implementation

optim.pyclip_grad_normGitHub

40def clip_grad_norm(grads, max_norm):41    total = 0.042    for k in grads:43        total += (grads[k] ** 2).sum() # add up the squares of every gradient number in every tensor44    norm = np.sqrt(total) # length of the whole gradient as one big vector4546    if norm > max_norm:47        for k in grads:48            grads[k] *= max_norm / norm # shrink every tensor by the same factor, direction stays, length becomes max_norm4950    return float(norm) # the length before clipping, train.py prints it as |g|

The function has two results. It changes grads in place, and it returns the norm from before clipping as a Python float. train.main stores that as norm and prints it as |g| (train.py:67-68). So |g| can be above grad_clip, although the gradients that reach opt.step have a length of at most grad_clip.

It runs right after batch_loss_and_grads and before the learning rate and opt.step, so m and v only ever see clipped gradients.

Code and math connection

Equation	Role	Lines	Code
Global norm, over all tensors	Computes it	optim.py:43	`total += (grads[k] ** 2).sum() # add up the squares of every gradient number in every tensor`
Global norm, over all tensors	Computes it	optim.py:44	`norm = np.sqrt(total) # length of the whole gradient as one big vector`
Global norm, over all tensors	Tests it	tests/test_optim.py:41	`assert np.isclose(norm, 5.0) # norm BEFORE clipping`
Scaling by one factor	Computes it	optim.py:46	`if norm > max_norm:`
Scaling by one factor	Computes it	optim.py:48	`grads[k] *= max_norm / norm # shrink every tensor by the same factor, direction stays, length becomes max_norm`
Scaling by one factor	Calls it	train.py:62	`norm = clip_grad_norm(grads, config.grad_clip)`
Scaling by one factor	Tests it	tests/test_optim.py:49	`assert np.allclose(grads["a"], [0.3, 0.4])`

Where is it used?

Source

optim.py, lines 40-50, clip_grad_norm

PDF reference

Section 4.6, Eq. 59, p. 14

Used by

Called by: train.main

Tested by: tests.test_optim.test_clip_leaves_small_gradients_alone, tests.test_optim.test_clip_rescales_all_tensors_together

Tests it: tests/test_optim.py, line 41 (test_clip_rescales_all_tensors_together), for Global norm, over all tensors
Calls it: train.py, line 62 (main), for Scaling by one factor
Tests it: tests/test_optim.py, line 49 (test_clip_leaves_small_gradients_alone), for Scaling by one factor

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Global norm, over all tensors	Computes it	`optim.py`	`clip_grad_norm`	43	GitHub: optim.py, line 43
Global norm, over all tensors	Computes it	`optim.py`	`clip_grad_norm`	44	GitHub: optim.py, line 44
Scaling by one factor	Computes it	`optim.py`	`clip_grad_norm`	46	GitHub: optim.py, line 46
Scaling by one factor	Computes it	`optim.py`	`clip_grad_norm`	48	GitHub: optim.py, line 48