Minimizing the Loss Function in Linear Regression

Finding the Optimal Slope and Intercept

Introduction

Linear regression is one of the most fundamental algorithms in machine learning, and at its heart lies a simple but powerful idea: find the line that best fits a set of data points. But what does “best fits” actually mean mathematically? The answer lies in minimizing a loss function, and the most common choice for regression problems is the Mean Squared Error (MSE).

In this article, we’ll build a concrete example from scratch, implement gradient descent in Python, and visualize how the algorithm converges toward the optimal slope and intercept — including a 3D visualization of the loss surface itself.

The Mathematical Formulation

Given a dataset of $n$ points $(x_i, y_i)$, we want to fit a line:

$$
\hat{y}_i = wx_i + b
$$

where $w$ is the slope and $b$ is the intercept. The Mean Squared Error loss function is defined as:

To minimize this loss, we use gradient descent. The partial derivatives of $L$ with respect to $w$ and $b$ are:

$$
\frac{\partial L}{\partial w} = -\frac{2}{n}\sum_{i=1}^{n}x_i(y_i - (wx_i + b))
$$

$$
\frac{\partial L}{\partial b} = -\frac{2}{n}\sum_{i=1}^{n}(y_i - (wx_i + b))
$$

At each iteration, we update the parameters using a learning rate $\eta$:

$$
w \leftarrow w - \eta \frac{\partial L}{\partial w}, \qquad b \leftarrow b - \eta \frac{\partial L}{\partial b}
$$

We repeat this process until the loss converges to a minimum, at which point $w$ and $b$ represent the best-fit line.

The Concrete Example

For this example, we generate synthetic data based on the true relationship $y = 3.5x + 7$ with added Gaussian noise, then use gradient descent to recover the slope (3.5) and intercept (7) purely from the noisy data.

Python Implementation

The code below performs the following steps:

  1. Generates synthetic noisy linear data.
  2. Implements a vectorized (NumPy-based) gradient descent algorithm for speed — avoiding slow Python for loops over individual data points.
  3. Tracks the loss history for convergence analysis.
  4. Computes the loss surface across a grid of $(w, b)$ values for visualization.
  5. Produces four plots: the fitted regression line, the loss convergence curve, a 3D loss surface, and a 2D contour map with the gradient descent path overlaid.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ---------------------------------------------------------
# 1. Generate synthetic data
# ---------------------------------------------------------
np.random.seed(42)

true_w = 3.5
true_b = 7.0
n_samples = 200

X = np.random.uniform(-10, 10, n_samples)
noise = np.random.normal(0, 4, n_samples)
Y = true_w * X + true_b + noise

# ---------------------------------------------------------
# 2. Vectorized gradient descent
# ---------------------------------------------------------
def compute_loss(w, b, X, Y):
    predictions = w * X + b
    return np.mean((Y - predictions) ** 2)

def gradient_descent(X, Y, w_init=0.0, b_init=0.0,
                      learning_rate=0.01, n_iterations=1000):
    w, b = w_init, b_init
    n = len(X)
    loss_history = []
    w_history = []
    b_history = []

    for i in range(n_iterations):
        predictions = w * X + b
        errors = Y - predictions

        dw = -(2 / n) * np.dot(X, errors)
        db = -(2 / n) * np.sum(errors)

        w -= learning_rate * dw
        b -= learning_rate * db

        loss = compute_loss(w, b, X, Y)
        loss_history.append(loss)
        w_history.append(w)
        b_history.append(b)

    return w, b, loss_history, w_history, b_history

learning_rate = 0.01
n_iterations = 500

final_w, final_b, loss_history, w_history, b_history = gradient_descent(
    X, Y, w_init=0.0, b_init=0.0,
    learning_rate=learning_rate, n_iterations=n_iterations
)

print(f"True parameters:      w = {true_w}, b = {true_b}")
print(f"Estimated parameters: w = {final_w:.4f}, b = {final_b:.4f}")
print(f"Final MSE loss: {loss_history[-1]:.4f}")

# ---------------------------------------------------------
# 3. Compute loss surface for visualization
# ---------------------------------------------------------
w_range = np.linspace(final_w - 5, final_w + 5, 100)
b_range = np.linspace(final_b - 15, final_b + 15, 100)
W_grid, B_grid = np.meshgrid(w_range, b_range)

Loss_grid = np.zeros_like(W_grid)
for i in range(W_grid.shape[0]):
    for j in range(W_grid.shape[1]):
        Loss_grid[i, j] = compute_loss(W_grid[i, j], B_grid[i, j], X, Y)

# ---------------------------------------------------------
# 4. Visualization
# ---------------------------------------------------------
plt.style.use('dark_background')
fig = plt.figure(figsize=(16, 12))

# --- Plot 1: Fitted regression line ---
ax1 = fig.add_subplot(2, 2, 1)
ax1.scatter(X, Y, color='#00d4ff', alpha=0.6, s=25, label='Data points')
x_line = np.linspace(X.min(), X.max(), 100)
y_line = final_w * x_line + final_b
ax1.plot(x_line, y_line, color='#ff6b6b', linewidth=2.5,
          label=f'Fitted line: y = {final_w:.2f}x + {final_b:.2f}')
ax1.set_xlabel('X', fontsize=12)
ax1.set_ylabel('Y', fontsize=12)
ax1.set_title('Linear Regression Fit', fontsize=14, fontweight='bold')
ax1.legend(fontsize=10)
ax1.grid(alpha=0.2)

# --- Plot 2: Loss convergence curve ---
ax2 = fig.add_subplot(2, 2, 2)
ax2.plot(loss_history, color='#ffd93d', linewidth=2)
ax2.set_xlabel('Iteration', fontsize=12)
ax2.set_ylabel('MSE Loss', fontsize=12)
ax2.set_title('Loss Convergence over Iterations', fontsize=14, fontweight='bold')
ax2.grid(alpha=0.2)

# --- Plot 3: 3D loss surface ---
ax3 = fig.add_subplot(2, 2, 3, projection='3d')
surf = ax3.plot_surface(W_grid, B_grid, Loss_grid, cmap='plasma',
                          alpha=0.85, edgecolor='none')
ax3.plot(w_history, b_history, loss_history, color='#00ff88',
          linewidth=2.5, label='Gradient descent path')
ax3.scatter([final_w], [final_b], [loss_history[-1]],
             color='red', s=60, label='Final minimum')
ax3.set_xlabel('Slope (w)', fontsize=10)
ax3.set_ylabel('Intercept (b)', fontsize=10)
ax3.set_zlabel('MSE Loss', fontsize=10)
ax3.set_title('3D Loss Surface with Descent Path', fontsize=14, fontweight='bold')
ax3.legend(fontsize=9)

# --- Plot 4: Contour map with descent path ---
ax4 = fig.add_subplot(2, 2, 4)
contour = ax4.contourf(W_grid, B_grid, Loss_grid, levels=40, cmap='plasma')
ax4.plot(w_history, b_history, color='#00ff88', linewidth=2,
          marker='o', markersize=2, label='Gradient descent path')
ax4.scatter([final_w], [final_b], color='red', s=80,
             marker='*', label='Final minimum', zorder=5)
ax4.set_xlabel('Slope (w)', fontsize=12)
ax4.set_ylabel('Intercept (b)', fontsize=12)
ax4.set_title('Loss Contour Map (Top View)', fontsize=14, fontweight='bold')
ax4.legend(fontsize=10)
fig.colorbar(contour, ax=ax4, label='MSE Loss')

plt.tight_layout()
plt.show()
True parameters:      w = 3.5, b = 7.0
Estimated parameters: w = 3.4844, b = 7.2644
Final MSE loss: 14.9412

Code Walkthrough

Data generation: We create 200 points along the line $y = 3.5x + 7$, then inject Gaussian noise with a standard deviation of 4. This simulates real-world measurement noise, giving gradient descent a genuine estimation problem rather than a trivial exact fit.

Vectorized gradient computation: Instead of looping over each data point with a Python for loop (which would be extremely slow for large datasets), the gradients are computed using np.dot(X, errors) and np.sum(errors). This leverages NumPy’s underlying C implementation, making the computation orders of magnitude faster than a pure Python loop — critical when scaling to larger datasets or more iterations.

Gradient descent loop: At each of the 500 iterations, predictions are computed for the entire dataset at once, the error vector is calculated, and both gradients ($\partial L/\partial w$ and $\partial L/\partial b$) are computed simultaneously. The parameters are then nudged in the direction that reduces the loss, scaled by the learning rate of 0.01.

Loss surface computation: To visualize the shape of the loss function itself, we build a grid of candidate $(w, b)$ pairs surrounding the final solution and compute the MSE loss at every grid point. This produces a bowl-shaped surface — a hallmark of MSE loss for linear regression, since it’s a convex quadratic function with a single global minimum.

Interpreting the Results

The top-left plot shows the raw noisy data alongside the line found by gradient descent — despite the noise, the algorithm recovers a slope and intercept very close to the true values of $w=3.5$ and $b=7$.

The top-right plot shows the loss dropping sharply in the first several iterations before flattening out, which is typical of gradient descent: large early steps followed by fine-tuning as the algorithm approaches the minimum.

The 3D surface plot is the most illuminating: it reveals the loss function as a smooth, convex bowl in $(w, b)$ space. The green trajectory traces the exact path taken by gradient descent, starting from $w=0, b=0$ and spiraling down toward the bottom of the bowl — the point where MSE is minimized.

The contour map gives a bird’s-eye view of the same bowl, making it easy to see how the descent path curves toward the minimum, following the steepest downhill direction at every step, which is exactly what the negative gradient represents.

Why This Matters

This simple example illustrates the core mechanism behind training almost every regression-based machine learning model, from simple linear regression to the first layer of a neural network. Understanding how MSE creates a convex loss landscape — and how gradient descent navigates it — provides the foundation for understanding more complex loss surfaces in deep learning, where the landscape is no longer a simple bowl but is optimized using the very same underlying principle: follow the gradient downhill.