Advanced Optimization: Understanding and Implementing ADAM

Introduction to ADAM

Hello! Today, we will explore the ADAM (Adaptive Moment Estimation) algorithm. This advanced optimization algorithm is a favorite among machine learning practitioners as it combines the advantages of two other extensions of Stochastic Gradient Descent (SGD): Root Mean Square Propagation (RMSprop) and Adaptive Gradient Algorithm (AdaGrad). Our primary focus today is understanding ADAM, and we will also build it from scratch in Python to optimize multivariable functions.

Understanding ADAM

Before we dive into ADAM, let us recall that classic gradient descent methods like SGD and even sophisticated versions like Momentum and RMSProp have some limitations. These limitations relate to sensitivity to learning rates, the issue of vanishing gradients, and the absence of individual adaptive learning rates for different parameters. ADAM, a promising choice for an optimization algorithm, combines the merits of RMSProp and AdaGrad . It maintains a per-parameter learning rate adapted based on the average of recent magnitudes of the gradients for the weights (similar to RMSProp) and the average of recent gradients (like Momentum). This mechanism enables the algorithm to traverse quickly over the low gradient regions and slow down near the optimal points.

ADAM Mathematically

m_t = \beta_1 * m_{t-1} + (1 - \beta_1) * grad

ADAM in Python Code

Let's now consolidate the ADAM concept into Python code. We will define an ADAM function, which takes the gradients, the decay rates beta1 and beta2, a numerical constant epsilon, the learning rate, and previous estimates of m and v (initialized to 0) as input and returns the updated parameters, along with the updated m and v . Python def ADAM(beta1, beta2, epsilon, grad, m_prev, v_prev, learning_rate): # Update biased first-moment estimate m = beta1 * m_prev + (1 - beta1) * grad # Update biased second raw moment estimate v = beta2 * v_prev + (1 - beta2) * np.power(grad, 2) # Calculate updates updates = learning_rate * m / (np.sqrt(v) + epsilon) return updates, m, v def ADAM(beta1, beta2, epsilon, grad, m_prev, v_prev, learning_rate): # Update biased first-moment estimate m = beta1 * m_prev + (1 - beta1) * grad # Update biased second raw moment estimate v = beta2 * v_prev + (1 - beta2) * np.power(grad, 2) # Calculate updates updates = learning_rate * m / (np.sqrt(v) + epsilon) return updates, m, v v and m are initialized with zeros and therefore they are biased towards zero at the start of the optimization, especially when the decay rates are small (beta1 and beta2 close to 1). To counteract these biases, Adam also usually includes the correction terms m_hat and v_hat . These terms adjust m and v by an amount that lessens as the number of time steps increases: Python m_hat = m / (1 - np.power(beta1, epoch+1)) # Correcting the bias for the first moment v_hat = v / (1 - np.power(beta2, epoch+1)) # Correcting the bias for the second moment updates = learning_rate * m_hat / (np.sqrt(v_hat) + epsilon) return updates, m, v m_hat = m / (1 - np.power(beta1, epoch+1)) # Correcting the bias for the first moment v_hat = v / (1 - np.power(beta2, epoch+1)) # Correcting the bias for the second moment updates = learning_rate * m_hat / (np.sqrt(v_hat) + epsilon) return updates, m, v Note that we still return plain m and v .

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Application of ADAM on Multivariable Function Optimization

Now, let's test ADAM slightly by finding the minimum of a multivariable function f(x, y) = x^2 + y^2. The corresponding gradients are df/dx = 2*x and df/dy = 2*y. With an initial starting point at (x, y) = (3, 4), selected reasonable values for beta1=0.9, beta2=0.9999, epsilon=1e-8, learning_rate=0.02 and an epoch size of 150, we can start minimizing our function.

Python

def f(x, y):
    return x ** 2 + y ** 2

def df(x, y):
    return np.array([2 * x, 2 * y])

coordinates = np.array([3.0, 4.0])
learning_rate = 0.02
beta1 = 0.9
beta2 = 0.9999
epsilon = 1e-8
max_epochs = 150

m_prev = np.array([0, 0])
v_prev = np.array([0, 0])

for epoch in range(max_epochs + 1):
    grad = df(coordinates[0], coordinates[1])
    updates, m_prev, v_prev = ADAM(beta1, beta2, epsilon, grad, m_prev, v_prev, learning_rate)
    coordinates -= updates
    if epoch % 30 == 0:
        print(f"Epoch {epoch}, current state: {coordinates}")

def f(x, y):
    return x ** 2 + y ** 2

def df(x, y):
    return np.array([2 * x, 2 * y])

coordinates = np.array([3.0, 4.0])
learning_rate = 0.02
beta1 = 0.9
beta2 = 0.9999
epsilon = 1e-8
max_epochs = 150

m_prev = np.array([0, 0])
v_prev = np.array([0, 0])

for epoch in range(max_epochs + 1):
    grad = df(coordinates[0], coordinates[1])
    updates, m_prev, v_prev = ADAM(beta1, beta2, epsilon, grad, m_prev, v_prev, learning_rate)
    coordinates -= updates
    if epoch % 30 == 0:
        print(f"Epoch {epoch}, current state: {coordinates}")

The output of this code is the following:

Epoch 0, current state: [2.80000003 3.80000002]
Epoch 30, current state: [ 0.27175946 -0.35494334]
Epoch 60, current state: [-0.07373187 -0.06706317]
Epoch 90, current state: [-0.02001478  0.0301726 ]
Epoch 120, current state: [ 0.00082782 -0.0039881 ]
Epoch 150, current state: [ 0.00094425 -0.00038352]

Epoch 0, current state: [2.80000003 3.80000002]
Epoch 30, current state: [ 0.27175946 -0.35494334]
Epoch 60, current state: [-0.07373187 -0.06706317]
Epoch 90, current state: [-0.02001478  0.0301726 ]
Epoch 120, current state: [ 0.00082782 -0.0039881 ]
Epoch 150, current state: [ 0.00094425 -0.00038352]