Lab: Implementing Gradient Descent from Scratch

Implementing Gradient Descent in Python

Let's implement gradient descent step by step, starting with the simple function $J(x) = x^2$.

Step 1: Define the Function and Its Derivative

%%%python # Cost function def J(x): return x ** 2 # Gradient (derivative) def dJ(x): return 2 * x %%%

Step 2: Implement the Gradient Descent Loop

%%%python def gradient_descent(x0, lr, num_iters): """Run gradient descent. Args: x0: initial value lr: learning rate (alpha) num_iters: number of iterations Returns: history: list of (x, cost) at each step """ x = x0 history = [(x, J(x))] for i in range(num_iters): grad = dJ(x) # compute gradient x = x - lr * grad # update step cost = J(x) # compute new cost history.append((x, cost)) # Print progress every 5 steps if i % 5 == 0 or i == num_iters - 1: print(f"Step {i+1:3d}: x = {x:8.4f}, J(x) = {cost:10.6f}") return history %%%

Step 3: Run It

%%%python print("--- Gradient Descent on J(x) = x^2 ---") print(f"Starting at x0 = 4.0, learning rate = 0.1\n") history = gradient_descent(x0=4.0, lr=0.1, num_iters=30) final_x = history[-1][0] final_cost = history[-1][1] print(f"\nFinal: x = {final_x:.6f}, J(x) = {final_cost:.6f}") %%%

Expected output (selected steps):

Step 1: x = 3.2000, J(x) = 10.240000 Step 6: x = 1.0486, J(x) = 1.099511 Step 11: x = 0.3436, J(x) = 0.118060 Step 16: x = 0.1126, J(x) = 0.012677 Step 21: x = 0.0369, J(x) = 0.001361 Step 26: x = 0.0121, J(x) = 0.000146 Step 30: x = 0.0053, J(x) = 0.000028

Final: x = 0.005262, J(x) = 0.000028

Step 4: Visualize the Convergence

%%%python import matplotlib.pyplot as plt import numpy as np # Extract x and cost values from history x_vals = [h[0] for h in history] cost_vals = [h[1] for h in history] # Plot 1: Cost vs. iteration fig, axes = plt.subplots(1, 2, figsize=(12, 4)) axes[0].plot(cost_vals, "o-", markersize=3) axes[0].set_xlabel("Iteration") axes[0].set_ylabel("Cost J(x)") axes[0].set_title("Cost vs. Iteration") axes[0].grid(True, alpha=0.3) # Plot 2: Path on the function x_curve = np.linspace(-5, 5, 200) y_curve = x_curve ** 2 axes[1].plot(x_curve, y_curve, "b-", label="J(x) = x^2") axes[1].plot(x_vals, [x**2 for x in x_vals], "ro-", markersize=4, label="GD path") axes[1].set_xlabel("x") axes[1].set_ylabel("J(x)") axes[1].set_title("Gradient Descent Path") axes[1].legend() axes[1].grid(True, alpha=0.3) plt.tight_layout() plt.show() %%%

Comparing Learning Rates

Let's run gradient descent with three different learning rates and compare:

%%%python import matplotlib.pyplot as plt learning_rates = [0.01, 0.1, 0.5] colors = ["red", "blue", "green"] plt.figure(figsize=(8, 5)) for lr, color in zip(learning_rates, colors): history = gradient_descent(x0=4.0, lr=lr, num_iters=30) costs = [h[1] for h in history] plt.plot(costs, color=color, label=f"lr = {lr}", linewidth=2) plt.xlabel("Iteration") plt.ylabel("Cost J(x)") plt.title("Effect of Learning Rate on Convergence") plt.legend() plt.grid(True, alpha=0.3) plt.yscale("log") # log scale to see small values plt.tight_layout() plt.show() %%%

You will see that $\alpha = 0.5$ converges fastest, $\alpha = 0.1$ converges moderately, and $\alpha = 0.01$ is very slow.

Application: Linear Regression with Gradient Descent

Let's fit a line $y = wx + b$ to geoscience data (depth vs. temperature) using gradient descent:

%%%python import numpy as np import matplotlib.pyplot as plt # Geothermal data: depth (km) vs temperature (C) depth = np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]) temp = np.array([30, 45, 58, 72, 88, 97, 115, 128]) # Normalize depth for better convergence depth_norm = (depth - depth.mean()) / depth.std() # Initialize parameters w = 0.0 # slope b = 0.0 # intercept lr = 0.1 m = len(depth) # Gradient descent cost_history = [] for epoch in range(100): # Predictions y_pred = w * depth_norm + b # Cost (MSE) cost = (1 / m) * np.sum((y_pred - temp) ** 2) cost_history.append(cost) # Gradients dw = (2 / m) * np.sum((y_pred - temp) * depth_norm) db = (2 / m) * np.sum(y_pred - temp) # Update w = w - lr * dw b = b - lr * db print(f"Final: w = {w:.2f}, b = {b:.2f}") print(f"Final cost: {cost_history[-1]:.2f}") %%%

Visualize the Fit

%%%python fig, axes = plt.subplots(1, 2, figsize=(12, 4)) # Left: cost convergence axes[0].plot(cost_history) axes[0].set_xlabel("Epoch") axes[0].set_ylabel("MSE") axes[0].set_title("Training Loss") axes[0].grid(True, alpha=0.3) # Right: data + fitted line axes[1].scatter(depth, temp, c="steelblue", s=60, label="Data") x_line = np.linspace(0, 4.5, 100) x_norm = (x_line - depth.mean()) / depth.std() y_line = w * x_norm + b axes[1].plot(x_line, y_line, "r-", linewidth=2, label="GD fit") axes[1].set_xlabel("Depth (km)") axes[1].set_ylabel("Temperature (C)") axes[1].set_title("Geothermal Gradient Fit") axes[1].legend() axes[1].grid(True, alpha=0.3) plt.tight_layout() plt.show() %%%

References

• Harris, C.R., et al. (2020). Array programming with NumPy. Nature 585, 357–362.
• Pedregosa, F., Varoquaux, G., Gramfort, A., et al. (2011). Scikit-learn: Machine learning in Python. J. Mach. Learn. Res. 12, 2825–2830.
• Géron, A. (2022). Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow (3rd ed.), ch. 4 (training models, gradient descent). O’Reilly.