Adam on a badly scaled problem

A quadratic whose two coordinates differ in curvature by a factor of 1000. Fixed-step gradient descent must use a step small enough for the stiff coordinate, so it crawls along the flat one. Adam rescales each coordinate by its root-mean-square gradient and moves both at a similar speed.

import numpy as np

from mathematicskit.optimization import Adam, GradientDescent
from mathematicskit.optimization.visualizers.plots import plot_contour_path

f(x, y) = x^2 / 1000 + y^2

f = lambda x: x[0] ** 2 / 1000.0 + x[1] ** 2
grad = lambda x: np.array([2.0 * x[0] / 1000.0, 2.0 * x[1]])
x0 = np.array([-1.5, 2.5])

gd = GradientDescent(alpha=0.4, max_iter=300).minimize(f, grad, x0)
adam = Adam(alpha=0.05, max_iter=300).minimize(f, grad, x0)
print(f"after 300 steps: gradient descent x = {gd.x.round(4)}, Adam x = {adam.x.round(4) + 0.0}")

ax = plot_contour_path(f, gd, x_range=(-2.0, 0.5), y_range=(-1.0, 3.0), label="gradient descent")
plot_contour_path(f, adam, ax=ax, x_range=(-2.0, 0.5), y_range=(-1.0, 3.0), label="Adam (Kingma & Ba, 2014)")

after 300 steps: gradient descent x = [-1.1798  0.    ], Adam x = [0. 0.]

<Axes: title={'center': 'Optimizer iterate path'}, xlabel='x', ylabel='y'>