Optimization — Descent Algorithms

Many algorithms used in Machine Learning are based on basic mathematical optimization methods. Discovering these algorithms directly in the context of machine learning can be confusing because of all the prerequisites. I think it's a good idea to see these algorithms free of any context in order to get a better understanding of these techniques.

Descent Algorithms

Descent algorithms are meant to minimize a given function. These algorithms proceed iteratively, which means they successively improve their current solution based on the previous one.

What if I want to maximize a function ?

Maximizing [begin-latex-inline]f(x)[end-latex-inline] is the same as minimizing [begin-latex-inline]-f(x)[end-latex-inline].

Thus, our problem is to find the [begin-latex-inline]x[end-latex-inline] that minimizes a given [begin-latex-inline]f(x)[end-latex-inline]:

[begin-latex]\left\{ \begin{array}{ll} f: \R^n \to \R \\ x^* = \underset{x}{\argmin}f(x) \end{array} \right.[end-latex]

Note that the input to the function may be a vector of any dimension.

Descent algorithms consist of building a sequence [begin-latex-inline]\{x_k\}[end-latex-inline] that will converge towards [begin-latex-inline]x^*=\argmin f(x)[end-latex-inline]. The sequence is built by successively correcting the current guess this way:

[begin-latex]x_{k+1} = x_k + d[end-latex]

Where [begin-latex-inline]k[end-latex-inline] is the iteration, and [begin-latex-inline]d[end-latex-inline] is a vector similar in size to [begin-latex-inline]x[end-latex-inline] called the descent vector. The algorithm keeps applying this update until the norm of the gradient (derivative) is small enough, since it reaches zero on any minimum or maximum.

x = x_init
while norm(gradient(f(x))) > epsilon:
    x += d

We will see 3 different descent vectors:

Newton's direction
Gradient's direction
Gradient + Optimal step size direction

But before, we need to define a function that we will try to minimize during our experiments.

Rosenbrock Function

I chose the Rosenbrock function, but you may find many others.

[begin-latex]f(x,y) = (a - x)^2 + b(y - x^2)^2[end-latex]

This function has a global minimum at [begin-latex-inline](x, y)=(a, a^2)[end-latex-inline] where [begin-latex-inline]f(a, a^2) = 0[end-latex-inline]. We will use [begin-latex-inline]a=2[end-latex-inline], [begin-latex-inline]b=100[end-latex-inline].

We will also need two other pieces of information, the gradient of that function as well as the hessian matrix.

[begin-latex]\nabla_{x,y} f = \begin{bmatrix} \dfrac{\partial f}{\partial x} \\[3ex] \dfrac{\partial f}{\partial y} \end{bmatrix} = \begin{bmatrix} 2(x - a) - 4bx(y - x^2) \\[3ex] 2b(y - x^2) \end{bmatrix}[end-latex]

[begin-latex]H_f = \begin{bmatrix} \dfrac{\partial^2 f}{\partial x^2} & \dfrac{\partial^2 f}{\partial x \partial y} \\[3ex] \dfrac{\partial^2 f}{\partial y \partial x} & \dfrac{\partial^2 f}{\partial y^2} \end{bmatrix} = \begin{bmatrix} 2 - 4b(y - 3x^2) & -4bx \\[3ex] -4bx & 2b \end{bmatrix}[end-latex]

import numpy as np


def rosenbrock(X, a=2, b=100):
    x, y = X
    return (a - x) ** 2 + b * (y - x**2) ** 2


def rosenbrock_grad(X, a=2, b=100):
    x, y = X
    return np.array([2 * (x - a) - 4 * b * x * (y - x**2), 2 * b * (y - x**2)])


def rosenbrock_hess(X, a=2, b=100):
    x, y = X
    return np.matrix([[2 - 4 * b * (y - 3 * x**2), -4 * b * x], [-4 * b * x, 2 * b]])

From now on, I will refer to the function input vector as [begin-latex-inline]x[end-latex-inline], akin to the problem definition.

Newton's Direction

Newton's direction is the following:

[begin-latex]d = - H_f^{-1}(x) \cdot \nabla_{x}f(x)[end-latex]

So the update is:

[begin-latex]x_{k+1} = x_k - H_f^{-1}(x_k) \cdot \nabla_{x}f(x_k)[end-latex]

This would be in Python:

def newton(x_init, grad, hess, epsilon=1e-10, max_iterations=1000):
    x = x_init
    for i in range(max_iterations):
        x = x - np.linalg.solve(hess(x), grad(x))
        if np.linalg.norm(grad(x)) < epsilon:
            return x, i + 1
    return x, max_iterations


# The Rosenbrock function takes 2 inputs (x, y)
x_init = np.zeros(2)
x_min, it = newton(x_init, rosenbrock_grad, rosenbrock_hess)
print("x* =", x_min)
print("Rosenbrock(x*) =", rosenbrock(x_min))
print("Grad Rosenbrock(x*) =", rosenbrock_grad(x_min))
print("Iterations =", it)

You will notice 2 differences with what I proposed:

I've introduced a max_iterations instead of just doing a while loop, so that we exit the program if we don't converge
The descent vector is not the dot product we expected: np.linalg.solve(hess(x), grad(x)). Actually it is: the inverse of a matrix can be expensive to compute (probably not in this case since we only have 2 variables), so in practice we solve the whole dot product directly:

[begin-latex]\begin{align*} & d = H_f^{-1}(x) \cdot \nabla_{x}f(x) \\[3ex] \iff & H_f(x) \cdot d = \nabla_{x}f(x) \end{align*}[end-latex]

Numpy's np.linalg.solve method will find [begin-latex-inline]d[end-latex-inline] given [begin-latex-inline]H_f(x)[end-latex-inline] and [begin-latex-inline]\nabla_{x}f(x)[end-latex-inline]. Let's see how this performs!

x* = [2. 4.]
Rosenbrock(x*) = 0.0
Grad Rosenbrock(x*) = [0. 0.]
Iterations = 2

The algorithm converged in only 2 iterations! That's really fast. You might think:

The initial [begin-latex-inline]x[end-latex-inline] is very close to the target [begin-latex-inline]x^*[end-latex-inline], which makes the problem easy!

Let's try with some other values, x_init = np.array([50, -30]): the algorithm terminates in 5 iterations.

This algorithm is called the Newton's Method and all descent algorithms used in Machine Learning are modifications of this method! It's kind of the mother formula. The reason it's so fast is that it uses second order information (the hessian matrix) which gives the algorithm information not only about the slope of the function at a given point, but also its curvature around that point. Think of it as being able to see a little further away on the function around a given point.

I don't like giving formulas out of the blue with no explanations, so I'll try to give some kind of proof for how to get this result.

Newton's Method: Finding Roots

There's another method that bares the name Newton's method. This one finds the roots of a function [begin-latex-inline]f(x)[end-latex-inline]; that is, the [begin-latex-inline]x[end-latex-inline] for which [begin-latex-inline]f(x)=0[end-latex-inline]. It does so by iteratively approximating the root by the intersection of the [begin-latex-inline]x[end-latex-inline]-axis with the tangent of the function at the current guess.

We have that:

[begin-latex]\begin{align*} & f'(x_n) = \dfrac{\Delta y}{\Delta x} = \dfrac{0 - f(x_n)}{x_{n+1} - x_n} \\[3ex] \iff & x_{n+1} = x_n - \dfrac{f(x_n)}{f'(x_n)} \end{align*}[end-latex]

How does that relate to minimizing a function?

Although this method does not find the minimum of a function directly, we know that a function is at an extremum (minimum or maximum) if the derivative is equal to zero! So we can apply this method, not over [begin-latex-inline]f(x)[end-latex-inline], but [begin-latex-inline]f'(x)[end-latex-inline]! In which case it becomes:

[begin-latex]\begin{align*} & x_{n+1} = x_n - \dfrac{f'(x_n)}{f''(x_n)} \\[3ex] \iff & x_{n+1} = x_n - \dfrac{1}{f''(x_n)} * f'(x_n) \end{align*}[end-latex]

Does that remind you anything?

[begin-latex]x_{n+1} = x_n - H_f^{-1}(x_n) \cdot \nabla_{x}f(x_n)[end-latex]

We got back to the 1d case of the more general formula with the hessian matrix and the gradient!

Also, this is why Newton's method can converge towards minimums OR maximums.

This method is clearly super fast to execute, but unfortunately not scalable. Since it uses the hessian matrix, we need to compute [begin-latex-inline]n^2[end-latex-inline] derivatives if the function has [begin-latex-inline]n[end-latex-inline] parameters (in modern neural networks, we have about hundreds of billions of parameters!). This is without even mentioning the cost of inverting the hessian matrix!

This is why other methods exist. All of them, in some way, try to approximate the second order information the hessian provides.

Gradient's Direction

If you did some machine learning, this formula will be familiar. The gradient direction:

[begin-latex]d = - \alpha \nabla_x f(x_k), \quad \alpha \in \R[end-latex]

This method approximates the hessian matrix by a simple scalar, called the learning rate. This has the advantage of being a lot faster to compute since it only uses the gradient.

def gradient_descent(x_init, grad, alpha=0.001, epsilon=1e-10, max_iterations=1000):
    x = x_init
    for i in range(max_iterations):
        x = x - alpha * grad(x)
        if np.linalg.norm(grad(x)) < epsilon:
            return x, i + 1
    return x, max_iterations


x_init = np.zeros(2)
x_min, it = gradient_descent(x_init, rosenbrock_grad)
print("x* =", x_min)
print("Rosenbrock(x*) =", rosenbrock(x_min))
print("Grad Rosenbrock(x*) =", rosenbrock_grad(x_min))
print("Iterations =", it)

And the result:

x* = [1.040607  1.0791156]
Rosenbrock(x*) = 0.9218391753965667
Grad Rosenbrock(x*) = [-0.35898539 -0.74946671]
Iterations = 1000

The algorithm did not converge to the exact solution after 1000 iterations! In fact, even after 100k iterations it didn't. That is the problem with gradient descent: it's cheap but slow, and may even get stuck in some local minima. Although, getting stuck in local minima tends to be less of a problem the more parameters the function has, due to an effect called the blessing of dimensionality.

Why Does It Work?

The gradient is by definition a vector pointing to the direction where the function increases the most around a given point. By subtracting the gradient from the parameters we follow the opposite direction: where the function decreases the most around a given point.

Here's an example with a small learning rate using the function [begin-latex-inline]f(x,y)=x^2+y^2[end-latex-inline].

Using a higher learning rate makes the algorithm bounce since the update is a lot larger:

Gradient's Direction + Optimal Step Size

One improvement to the classical gradient descent is to use a variable learning rate at each iteration. In fact, we can find the best possible learning rate! We want to find [begin-latex-inline]\alpha^*[end-latex-inline] that minimizes the [begin-latex-inline]f(x)[end-latex-inline] after the gradient descent step as much as possible.

[begin-latex]\alpha^* = \underset{\alpha}{\argmin} f(x_k - \alpha \nabla_x f(x_k))[end-latex]

Notice that [begin-latex-inline]x_k[end-latex-inline] and [begin-latex-inline]\nabla_x f(x_k)[end-latex-inline] are constants with respect to [begin-latex-inline]\alpha[end-latex-inline], therefore we just have a function of [begin-latex-inline]\alpha[end-latex-inline] to minimize:

[begin-latex]q(\alpha) = f(x_k - \alpha \nabla_x f(x_k))[end-latex]

How would we find the [begin-latex-inline]\alpha[end-latex-inline] that minimizes [begin-latex-inline]q(\alpha)[end-latex-inline] ?

Gradient descent? We could, but while we're at it, let's learn a new method: Golden Section Search.

GSS aims at finding the extremum (minimum or maximum) of a unimodal function inside a specified interval. A unimodal function is one that has a single minimum or maximum (a single change of direction from decreasing to increasing or vice-versa).

We're usually looking for an [begin-latex-inline]\alpha[end-latex-inline] in the range [begin-latex-inline][0,1][end-latex-inline] (and under certain assumptions we can show that [begin-latex-inline]q[end-latex-inline] is unimodal), which makes this algorithm suitable for the task.

def gss(f, a, b, tol=1e-7):
    phi = (np.sqrt(5) + 1) / 2
    d = b - (b - a) / phi
    c = a + (b - a) / phi

    while abs(d - c) > tol:
        if f(d) < f(c):
            b = c
        else:
            a = d

        d = b - (b - a) / phi
        c = a + (b - a) / phi

    return (a + b) / 2

This algorithm is like a binary search but it maintains the golden ratio [begin-latex-inline]\phi[end-latex-inline] between the intervals it's checking. To run the algorithm for a function [begin-latex-inline]f[end-latex-inline] in the interval [begin-latex-inline][a,b][end-latex-inline], the basic idea is the following:

Evaluate two other points [begin-latex-inline]c[end-latex-inline] and [begin-latex-inline]d[end-latex-inline] in [begin-latex-inline][a,b][end-latex-inline]
If [begin-latex-inline]f(c) < f(d)[end-latex-inline], then the new search interval is [begin-latex-inline][a,d][end-latex-inline]
If [begin-latex-inline]f(c) > f(d)[end-latex-inline], then the new search interval is [begin-latex-inline][c,b][end-latex-inline]

Now that we are able to find the best [begin-latex-inline]\alpha[end-latex-inline], we can implement gradient descent with optimal step size!

def gradient_descent_optimal(x_init, f, grad, epsilon=1e-10, max_iterations=1000):
    x = x_init
    for i in range(max_iterations):
        q = lambda alpha: f(x - alpha * grad(x))
        alpha = gss(q, 0, 1)
        x = x - alpha * grad(x)
        if np.linalg.norm(grad(x)) < epsilon:
            return x, i + 1
    return x, max_iterations


x_init = np.zeros(2)
x_min, it = gradient_descent_optimal(x_init, rosenbrock, rosenbrock_grad, max_iterations=30000)
print("x* =", x_min)
print("Rosenbrock(x*) =", rosenbrock(x_min))
print("Grad Rosenbrock(x*) =", rosenbrock_grad(x_min))
print("Iterations =", it)

We get the following result:

x* = [1.98408891 3.93658492]
Rosenbrock(x*) = 0.00025321978916729786
Grad Rosenbrock(x*) = [-0.0128689  -0.00477632]
Iterations = 30000

This converges towards the solution after 30k iterations, but at least it does converge unlike pure gradient descent. Although you do get a descent result even after a smaller number of iterations!

More Functions

Here are some other interesting functions you can try those methods on.

Himmelblau

[begin-latex]f(x,y) = (x^2 + y - 11)^2 + (x + y^2 - 7)^2[end-latex]

The Himmelblau's function has one local maximum [begin-latex-inline]f(-0.270845, -0.923039)=181.617[end-latex-inline] and four identical local minimums:

[begin-latex-inline]f(3.0,2.0)=0[end-latex-inline]
[begin-latex-inline]f(-2.805118,3.131312)=0[end-latex-inline]
[begin-latex-inline]f(-3.779310,-3.283186)=0 [end-latex-inline]
[begin-latex-inline]f(3.584428,-1.848126)=0[end-latex-inline]

Easom

[begin-latex]f(x,y) = -\cos(x)\cos(y)e^{-(x-\pi)^2 - (y-\pi)^2}[end-latex]

The Easom function has a global minimum where [begin-latex-inline]f(\pi, \pi) = -1[end-latex-inline].

Follow Up

Now that you've seen the principles behind those famous optimization methods, I hope neural networks are making more sense! Check out my other article where I build a neural network library from scratch!

Neural Network From Scratch Build your own machine learning library in Python omaraflak.com