Gradient Descent Optimization

TensorFlow - Gradient Descent Optimization

Gradient Descent is a fundamental optimization algorithm used to find the minimum of a function. In machine learning, it’s widely used to minimize a model’s loss function by iteratively adjusting the model’s parameters (weights and biases).

The core idea is to take steps in the direction opposite to the gradient (or approximate gradient) of the function at the current point, as this is the direction of steepest descent. The size of these steps is determined by the learning rate.

Let’s implement a simple gradient descent optimization in TensorFlow 2.x to minimize the function f(x) = (log(x))^2. The minimum of this function occurs at x = 1, where f(1) = (log(1))^2 = 0.

Step 1: Setup Variables and Optimizer

First, we import TensorFlow and define the variable x that we want to optimize. We also define our loss function and the optimizer.

import tensorflow as tf

# Define the variable to optimize. Initialize it to a value, e.g., 2.0.
# We use tf.Variable because its value needs to change during optimization.
x = tf.Variable(2.0, name='x', dtype=tf.float32)

# Define the loss function: f(x) = (log(x))^2
# Note: log(x) is typically natural logarithm (ln(x)) in TensorFlow.
# Ensure x remains positive for log(x) to be defined.
def loss_function():
    # tf.math.log computes the natural logarithm.
    # Adding a small epsilon can improve stability if x might approach zero,
    # but for x starting at 2.0 and moving towards 1.0, it's generally not an issue.
    return tf.square(tf.math.log(x))

# Define the optimizer. We'll use Stochastic Gradient Descent (SGD).
# The learning rate controls the step size.
learning_rate = 0.5
optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate)

Step 2: Perform Optimization Steps

Next, we write a loop to perform the optimization. In each step, we calculate the loss, compute gradients of the loss with respect to x using tf.GradientTape, and then apply these gradients using the optimizer to update x.

def optimize(num_steps=10):
    print(f"Starting at: x = {x.numpy():.4f}, loss = (log(x))^2 = {loss_function().numpy():.4f}")

    for step in range(num_steps):
        # Open a GradientTape context to record operations for automatic differentiation.
        with tf.GradientTape() as tape:
            # Calculate the loss for the current value of x.
            current_loss = loss_function()

        # Calculate the gradients of the loss with respect to our variable [x].
        # tape.gradient will return a list of gradients, one for each trainable variable.
        gradients = tape.gradient(current_loss, [x])

        # Apply the gradients to the variable(s) using the optimizer.
        # optimizer.apply_gradients expects a list of (gradient, variable) pairs.
        optimizer.apply_gradients(zip(gradients, [x]))

        # Print the progress
        print(f"Step {step+1:2d}: x = {x.numpy():.4f}, loss = {loss_function().numpy():.4f}")

# Run the optimization
optimize(num_steps=10)

# For further learning:
# - tf.GradientTape: https://www.tensorflow.org/guide/autodiff
# - Optimizers in Keras: https://www.tensorflow.org/api_docs/python/tf/keras/optimizers

The output will show the value of x and the loss (log(x))^2 at each step. You should observe that x iteratively moves closer to 1.0, and consequently, the loss value decreases towards 0. This demonstrates the gradient descent algorithm minimizing the function. The exact values and speed of convergence depend on the initial value of x, the learning rate, and the function itself.