{"id":1287,"date":"2018-11-03T19:19:37","date_gmt":"2018-11-03T19:19:37","guid":{"rendered":"http:\/\/www.philippeadjiman.com\/blog\/?p=1287"},"modified":"2025-07-18T13:15:13","modified_gmt":"2025-07-18T13:15:13","slug":"visualising-sgd-with-momentum-adam-and-learning-rate-annealing","status":"publish","type":"post","link":"https:\/\/philippeadjiman.com\/blog\/2018\/11\/03\/visualising-sgd-with-momentum-adam-and-learning-rate-annealing\/","title":{"rendered":"Visualising SGD with Momentum, Adam and Learning Rate Annealing"},"content":{"rendered":"<\/span>\n
\n\n\n
\"\"<\/figure>\n\n\n[Full code on my github\u00a0<\/span>here<\/a>\u00a0. To see it from mobile, once you land on github, click on \u201cDesktop Version\u201d ]<\/strong>\n\nAt the very heart of the model training procedure of almost every modern machine learning or deep learning algorithm applied to big enough data, you\u2019ll find Stochastic Gradient Descent (SGD).\n\nThe best part of SGD is its simplicity. As\u00a0<\/span>Francois Chollet<\/a>\u00a0<\/span>would say, it is made of a small set of high school-level ideas put together. But it does not make it less powerful and beautiful.\n\nIn this post we\u2019ll implement from scratch SGD and some optimizations around it like Momentum, Adam and learning rate annealing, and we\u2019ll apply it on some very simple generated toy data in order to visually compare them together with some animated graph in python. In the post, we\u2019ll only show\u00a0 some snippets of a subset of the code, check\u00a0<\/span>here<\/a>\u00a0<\/span>for the full code.\n

Vanilla SGD<\/h2>\nFirst we generate some data. We\u2019ll take on purpose the simplest and almost smallest data set ever, by simply generating 20 random points from a linear function ax+b.\n
\n
\n\n\n\n
\n
\n
np.random.seed(<\/code>7<\/code>)<\/code><\/div>\n
a_real <\/code>=<\/code> 1.5<\/code><\/div>\n
b_real <\/code>=<\/code> -<\/code>28<\/code><\/div>\n
xlim <\/code>=<\/code> [<\/code>-<\/code>10<\/code>,<\/code>10<\/code>]<\/code><\/div>\n
x_gen <\/code>=<\/code> np.random.randint(low<\/code>=<\/code>xlim[<\/code>0<\/code>], high<\/code>=<\/code>xlim[<\/code>1<\/code>], size<\/code>=<\/code>20<\/code>)<\/code><\/div>\n
y_real <\/code>=<\/code> a_real<\/code>*<\/code>x_gen <\/code>+<\/code> b_real<\/code><\/div>\n
plt.plot(x_gen, y_real, <\/code>'bo'<\/code>)<\/code><\/div>\n<\/div><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<\/div>\n\"\"<\/a>\n\nSo we\u2019ll start with some initial a and b (say a=0 and b=0) and the goal of SGD is to find alone the real a and b (a=1.5 and b=-28 in our example) which were used to generate those few data points.\u00a0 To do so, we simply need to minimize a cost function on the data, in our case\u00a0\u00a0<\/span>\u2211<\/span>(<\/span>\ud835\udc65<\/span>,<\/span>\ud835\udc66<\/span>)<\/span>\u2208<\/span>\ud835\udc51<\/span>\ud835\udc4e<\/span>\ud835\udc61<\/span>\ud835\udc4e<\/span>\ud835\udc60<\/span>\ud835\udc52<\/span>\ud835\udc61<\/span><\/span><\/span><\/span>(<\/span>\ud835\udc4e<\/span>\ud835\udc65<\/span>+<\/span>\ud835\udc4f<\/span>\u2212<\/span>\ud835\udc66<\/span>)<\/span>2<\/span><\/span><\/span><\/span><\/span><\/span><\/nobr><\/span>\u00a0<\/span>. SGD achieves that by simply following the negative of the gradient (negative because the gradient is the direction of the steepest increase of the function and we\u2019re looking for the minimum of the cost function).\n\nSo basically, the vanilla SGD parameter update is simply:\n
\n
\n\n\n\n
\n
\n
param <\/code>+<\/code>=<\/code> -<\/code>lr<\/code>*<\/code>dx<\/code><\/div>\n<\/div><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<\/div>\nwith lr being the learning rate, and dx being the gradient of the cost function relative to the corresponding param you want to update (in our case, only a or b).\n\nHow to compute dx ? if our hypothesis function was a deep neural network you could simply apply the chain rule multiple times (a.k.a backpropagation) via e.g.\u00a0<\/span>pytorch\u2019s autograd<\/a>\u00a0, but in our case we can simply compute the analytical gradient of the cost function\u00a0 w.r.t. a and b, or just use wolfram alpha (like\u00a0<\/span>this<\/a>\u00a0<\/span>) if you\u2019re lazy.\n\nThis inevitably leads you to a full implementation of SGD (for our example) in less than 10 lines of code:\n
\n
\n\n\n\n
\n
1<\/div>\n
2<\/div>\n
3<\/div>\n
4<\/div>\n
5<\/div>\n
6<\/div>\n
7<\/div>\n
8<\/div><\/td>\n
\n
\n
def<\/code> sgd(X,Y, a, b, lr, epochs<\/code>=<\/code>1<\/code>):<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>for<\/code> e <\/code>in<\/code> range<\/code>(epochs):<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>for<\/code> x_ , y_ <\/code>in<\/code> zip<\/code>(X,Y):<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>a <\/code>=<\/code> a <\/code>-<\/code> lr<\/code>*<\/code>2<\/code>*<\/code>x_<\/code>*<\/code>(a<\/code>*<\/code>x_<\/code>+<\/code>b<\/code>-<\/code>y_)<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>b <\/code>=<\/code> b <\/code>-<\/code> lr<\/code>*<\/code>2<\/code>*<\/code>(a<\/code>*<\/code>x_<\/code>+<\/code>b<\/code>-<\/code>y_)\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0 <\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>return<\/code> a,b<\/code><\/div>\n
a,b <\/code>=<\/code> gradient_descent(x_gen,y_real,<\/code>0<\/code>,<\/code>0<\/code>,<\/code>0.001<\/code>,epochs<\/code>=<\/code>150<\/code>) <\/code><\/div>\n
print<\/code>(<\/code>\"a={:.3f} , b={:.3f}\"<\/code>.<\/code>format<\/code>(a,b))<\/code><\/div>\n<\/div><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<\/div>\nwhich prints:\n
\n
\n\n\n\n
\n
\n
a=1.501 , b=-27.913<\/code><\/div>\n<\/div><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<\/div>\nNot that far from our real a=1.5 and b= -28 .\n\nNote that we could make this code much more efficient by vectorizing it but we keep it dumb on purpose to observe easily how simple it is (and also add metrics to be able to monitor and visualize the gradients steps, c.f. later section). In the code\u00a0we treat one data point at a time (batch size = 1) , and going over all the data points multiple times (as many time as specified in the epoch parameter).\n\nWe can keep track of the a and b updates after each epoch and animate the evolution (see the section named \u201cSimple standalone Animation code\u201d in the\u00a0<\/span>notebook<\/a>\u00a0<\/span>to see how to generate such an animation):\n\n\"\"<\/a>\n\nNote how a and b are converging slowly but surely to the real values (of a = 1.5 and b = -28).\n

Momentum<\/h2>\nSo, the vanilla gradient descent is converging surely towards the optimal, but also rather slowly as we see above, taking the same small step (in the right direction) at every iteration. More than that, here we have a nice and easy convex cost function, but in case of ravines, SGD becomes even more slower by taking hesitant steps toward the optimum.\n\nTo improve that, the momentum update is taking advantage of the history of previous gradient steps directions in order to make more aggressive steps when the gradient direction seems stable, and slows it down when it is starting to go in multiple directions, inspired from the velocity principle in physics.\n\nSo, instead of of doing the vanilla update rule of SGD (param += -lr*dx) , the momentum update is actually replacing the dx part by a decaying average of previous gradients. The average is controlled by a parameter beta , and the gradient is replaced by a linear interpolation of previous gradient\u2019s update and current gradient. It gives the simple following code:\n
\n
\n\n\n\n
\n
1<\/div>\n
2<\/div>\n
3<\/div>\n
4<\/div>\n
5<\/div>\n
6<\/div>\n
7<\/div>\n
8<\/div>\n
9<\/div>\n
10<\/div>\n
11<\/div>\n
12<\/div><\/td>\n
\n
\n
def<\/code> sgd_momentum(X,Y, a, b, lr,\u00a0 beta, epochs<\/code>=<\/code>1<\/code>):<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>avg_ga <\/code>=<\/code> 0<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>avg_gb <\/code>=<\/code> 0<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>for<\/code> e <\/code>in<\/code> range<\/code>(epochs):<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>for<\/code> x_ , y_ <\/code>in<\/code> zip<\/code>(X,Y):<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>de_da <\/code>=<\/code> 2<\/code>*<\/code>x_<\/code>*<\/code>(a<\/code>*<\/code>x_<\/code>+<\/code>b<\/code>-<\/code>y_)<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>de_db <\/code>=<\/code> 2<\/code>*<\/code>(a<\/code>*<\/code>x_<\/code>+<\/code>b<\/code>-<\/code>y_) <\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>avg_ga <\/code>=<\/code> avg_ga<\/code>*<\/code>beta <\/code>+<\/code> (<\/code>1.0<\/code>-<\/code>beta)<\/code>*<\/code>de_da<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>avg_gb <\/code>=<\/code> avg_gb<\/code>*<\/code>beta <\/code>+<\/code> (<\/code>1.0<\/code>-<\/code>beta)<\/code>*<\/code>de_db<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>a <\/code>=<\/code> a <\/code>-<\/code> lr<\/code>*<\/code>avg_ga<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/code>b <\/code>=<\/code> b <\/code>-<\/code> lr<\/code>*<\/code>avg_gb\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0 <\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>return<\/code> a,b<\/code><\/div>\n<\/div><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<\/div>\nWhat is interesting is to compare the evolution of the gradient per method, to see if we do see the expected smoother evolution of the gradient update each iterations. This is how compares the vanilla SGD v.s momentum gradient updates on the first learned parameter (the a):\n\n\"\"<\/a>\n\nAnd on the second one (the b):\n\n\"\"<\/a>\n\nYeah, i know it looks a bit like a\u00a0\"Image, but the point is that the momentum update is doing its job: it is not going crazy in all direction like the raw gradient, but it is smoothing it, based on previous iterations. For such a simple convex error function like ours, it does not really matters (and it won\u2019t make a dramatic difference in terms of how fast we converge as we\u2019ll see below), but we can easily understand how in a very bumpy loss function surface, this could be a great advantage to surf around those rather than entering eyes closed inside each small ravine.\n

Adam, Learning Rate Annealing and other SGD optimisations<\/h2>\nIn the same spirit as the momentum update, many different methods are exposing multiple variations of how to modify the SGD update rule, in order to converge faster and better to the optimal parameters that the model is trying to learn.\n\nWe won\u2019t get into details of all those great optimizations because there are already excellent posts\/video around that topic, e.g. from Sebastian Rudder\u00a0<\/span>here<\/a>\u00a0or Andrej Karpathy\u00a0\u00a0<\/span>here<\/a>\u00a0or the fantastic video (and corresponding\u00a0<\/span>excel file)<\/a>\u00a0<\/span>by Jeremy Howard . But we\u2019ll do mention a couple of important concepts that those methods are using:\n

Per coordinate adaptive learning rate<\/h4>\nThe idea is that the size of the steps taken in gradient descent should be adapted for each learned parameter separately. Intuitively, the idea would be to make the learning rate smaller as a function of how much data was observed for the specific corresponding parameter. First popular proposed method is\u00a0AdaGrad<\/a>\u00a0and then Adadelta and RMSProp are some evolution of it, then\u00a0 Adam (Adaptive moment estimation) is combining that idea with momentum, then other methods are improving on top of Adam etc.. Again, c.f. the links above and the excel file, they are the best reference for all those. You can also find an illustration of how to apply and implement (in few lines of code) this concept in the context of logistic regression in my blog post\u00a0<\/span>here<\/a>.\n

Learning rate annealing (with restarts)<\/h4>\nThis simple concept is also one the most effective tricks you can find in the deep learning world.\u00a0The idea is that when you start your search of the optimal parameter, you can afford doing some big jump, but the more you progress towards your minimum, the more you want to make smaller steps to nail it down (and not miss it by too big steps), and thus you progressively reduce your learning rate.\u00a0<\/span>You can also combine that idea by reinitialising your learning rate to its highest value from time to time (this is called SGD with restarts) to find more general optimums. More on that in the fantastic\u00a0fast.ai course<\/a>\u00a0and also in\u00a0that great blog post<\/a>.\u00a0<\/span>\n

Implementation and Visual comparisons on our simple example<\/h2>\nIn\u00a0<\/span>that notebook<\/a>, I\u2019ve implemented a few functions allowing to simulate, visualize, debug, investigate, experiment with few variations of SGD: vanilla, momentum, adam and adam with learning rate annealing .\n\nThe implementation can easily be extended with any function (not only a linear function as in our example), only the derivative needs to be provided as a function (which itself could be automated using e.g.\u00a0<\/span>pytorch\u2019s autograd<\/a>), although the visualisations are adapted only for functions in the domain\u00a0<\/span><\/span>\u211d<\/span><\/span><\/span>\ud835\udc5b<\/span><\/span>\u2192<\/span>\u211d<\/span><\/span><\/span><\/span><\/span><\/span><\/span><\/nobr><\/span>\u00a0<\/span>.\n\nBelow, we\u2019re showing some of the output given when calling this code with:\n
\n
\n\n\n\n
\n
1<\/div>\n
2<\/div>\n
3<\/div>\n
4<\/div>\n
5<\/div>\n
6<\/div>\n
7<\/div>\n
8<\/div>\n
9<\/div>\n
10<\/div>\n
11<\/div>\n
12<\/div>\n
13<\/div><\/td>\n
\n
\n
params,methods,x,y, \\<\/code><\/div>\n
loss_evolution_list,lr_evolution_list <\/code>=<\/code> \\<\/code><\/div>\n
compare_methods_and_plot([<\/code>1.5<\/code>,<\/code>-<\/code>28<\/code>],[<\/code>1<\/code>,<\/code>1<\/code>],lin,<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>linear_gradients,[<\/code>-<\/code>10<\/code>,<\/code>10<\/code>],<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>size_gen_data<\/code>=<\/code>30<\/code>,<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>epochs <\/code>=<\/code> 50<\/code>,<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>methods <\/code>=<\/code> [<\/code>\"SGD\"<\/code>,<\/code>\"Momentum\"<\/code>,<\/code>\"Adam\"<\/code>,<\/code>\"AdamAnn\"<\/code>],<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>methods_optim_params <\/code>=<\/code> [[<\/code>0.001<\/code>] , [<\/code>0.001<\/code>,<\/code>0.95<\/code>] , [<\/code>1<\/code>,<\/code>0.7<\/code>,<\/code>0.9<\/code>] , [<\/code>1<\/code>,<\/code>0.7<\/code>,<\/code>0.9<\/code>] ],<\/code><\/div>\n
\u00a0\u00a0\u00a0\u00a0<\/code>anim_interval_ms<\/code>=<\/code>100<\/code>,ylim_anim<\/code>=<\/code>[<\/code>-<\/code>50<\/code>,<\/code>15<\/code>])<\/code><\/div>\n
<\/div>\n
fig, ax <\/code>=<\/code> plt.subplots(figsize<\/code>=<\/code>(<\/code>12<\/code>, <\/code>6<\/code>))\u00a0\u00a0\u00a0 <\/code><\/div>\n
anim <\/code>=<\/code> draw_animation(params,methods, lin,x,y,<\/code>75<\/code>,ylim<\/code>=<\/code>[<\/code>-<\/code>50<\/code>,<\/code>15<\/code>])<\/code><\/div>\n
HTML(anim.to_jshtml())<\/code><\/div>\n<\/div><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<\/div>\nFirst let\u2019s observe the animation:\n\n\"\"<\/a>\n\nWe can note few things:\n