Linear regression from scratch: Batch, stochastic & mini-batch

I wanted to understand how Linear Regression works so I built it from scratch using Numpy and Pandas.
This is a breakdown of my process. I'll start with explaining the algorithm at a high level, then we'll do some calculations, and finally write some code.

The line of best fit

Linear regression works by predicting the line of best fit through a set of data points.
The formula for that line is:

y = wx + b

where:

• w is the weight. It controls how much the input (x), affects the output. If you increase x by one unit, the output changes by w.
• x is a data point. In housing dataset, x could be the number of bedrooms.
• b is the bias. If you didn't have it the line would be forced to start from (0,0) which is not realistic. No house cost nothing unfortunately. The bias gives the line a starting point that doesn't depend on x.

Training a linear regression model means finding the values of w and b than make predictions that are as close as possible to the true values.

The loss function

So how do you know your prediction is close enough? You need a way to measure it. That's what a loss function does.

The most common choice is mean squared error, or MSE. It's the average of the squared errors across your dataset.

Start with the error for a single point. This is the difference between the true and the predicted value.

error = yᵢ - ŷᵢ
where:

• yᵢ is the true value
• ŷᵢ is the predicted value

The loss(MSE) is the average of the squared errors across all n points.
L = (1/n) * Σ(yᵢ - ŷᵢ )²

The reason we're squaring the error is for two reasons

First, it stops the positive and negative errors canceling each other when you sum them. Second, it penalises larger errors more heavily than smaller ones. You could use absolute values instead (that's mean absolute error, or MAE), but MAE treats a big miss and a small miss more proportionally. Squaring makes the model care more about the outliers.

Gradient descent

To reduce the loss, the model adjusts w and b step by step. It does this by calculating the partial derivative of the loss with respect to each one, then subtracting a fraction of that value from the current weight and bias. This is gradient descent.

A partial derivative tells you how much a function's output will be affected if you change a variable by a tiny amount, holding everything else constant. This means:

• ∂L/∂w t tells you how much the loss will change if you change the weight
• ∂L/∂b tells you how much the loss will change if you change the bias.

How to calculate the partial derivative of loss w.r.t weight

Recall that ŷᵢ = wxᵢ + b so you can write the loss function as,
L = (1/n) * Σ(yᵢ - (wxᵢ + b) )²

This is a function inside another function so you'll use the chain rule to calculate the derivative.
The chain rule say take the derivative of the outer function and multiply it by the inner function as is, then multiple that result by the derivative of the inner function.
*- The outer function is ()². Its derivative with respect to w is 2().

• The inner function is yᵢ - (wxᵢ + b). Its derivative with respect to w is -xᵢ (yᵢ and b don't depend on w, so they drop to 0). Putting that together: 2(yᵢ - (wxᵢ + b)) * -xᵢ = -2xᵢ(yᵢ - (wxᵢ + b)) = -2xᵢ(yᵢ - ŷᵢ)

Averaging over all n points:
∂L/∂w = -(2/n) * Σ(yᵢ - ŷᵢ) * xᵢ

How to calculate the partial derivative of loss w.r.t bias

Same approach but this time with respect to bias:

• The outer function's derivative with respect to b is 2().
• The inner function's derivative with respect to b is -1 (yᵢ and wxᵢ don't depend on b, so they drop to 0).

2(yᵢ - (wxᵢ + b)) * -1
= -2(yᵢ - (wxᵢ + b))
= -2(yᵢ - ŷᵢ)

Averaging over all n points:
∂L/∂b = (-2/n) * Σ(yᵢ - ŷᵢ)

How to update the weight and bias

Now that you have both partial derivatives, the next step is to update the weights and bias in a direction that minimizes the loss. That means moving opposite the gradient, which is why you subtract the derivative rather than add it.
You wouldn't subtract the derivative directly, though. It can be large enough that you miss the optimum entirely. So you scale it down first, using a value called the learning rate.

This is a value you tune by experimenting. Too large you end up with very large weights and overshoot the minimum. Too small it takes forever to arrive at a small loss value.

new_weight = w - (∂L/∂w) * learning_rate
new_bias = b - (∂L/∂b) * learning_rate

Now, plug these new values into the ŷ = wx + b to get updated predictions, then recalculate the loss with those new predictions to see if it reduced.

All together now

The full algorithm repeats this cycle until the loss is small enough or stops improving meaningfully between iterations

• Initialize weights and bias to a random value or even 0
• Calculate predictions
• Calculate loss
• Calculate the partial derivatives of loss function w.r.t weight and bias
• Update bias and weight
• Recalculate predictions
• Recalculate loss
• Repeat ### Batch vs. Stochastic vs. Mini-batch

The above algorithm uses the entire dataset to calculate the gradients (the derivatives) on every iteration. This is called batch gradient descent. It thorough but slow if you have a large number of data points since every iteration requires summing over the whole dataset.

Stochastic gradient descent(SGD) takes the opposite approach. It updates the weight and bias using a single random data point per iteration. There's no summation so each step is faster. But if that one point happens to be unusual (an outlier) then the update the model calculates from it will steer w and b in a slightly wrong direction compared to what would happen if it had looked at the whole dataset. Next step, it might get a totally different (also slightly wrong) point, so the path zigzags rather than moving smoothly downhill.

Mini-batch gradient descent is the middle ground for batch and SGD. Instead of using the entire dataset or a single random data point, you use a small batch (for example 32 points) per iteration. It's faster than batch (only summing over 32 points) and steadier than SGD since averaging over 32 points smooths out the effect of any one outlier. This balance is why its the default choice for most ML training loops.

How to build a linear regression model from scratch

Let's now turn the math into code.

Start by importing Numpy and pandas.

import numpy as np
import pandas as pd

I'll create one LinearRegression class with a separate method for batch, stochastic, and mini-batch gradient descent.

Batch gradient descent from scratch

class LinearRegression:
    """
    A linear regression model trained using:
        - Batch gradient descent
    """

    def __init__(self):
        self.weights = None
        self.bias = None
        self.mse = []
        self.log = []

    def _initialize_params(self, n_features):
        """

        Initialize weights and bias to zeros based on the number of features.

        Parameters:

            n_features (int): The number of features in the dataset.

        """

        self.weights = np.zeros(n_features)
        self.bias = 0.0

    def fit_batch_gradient_descent(self, X,y, lr=0.05, epochs=100):
        """

        Train the model using batch gradient descent.
        Parameters:

            X (ndarray): Inputs
            y(ndarray): Output
            lr(float): Learning rate
            epochs(int): Number of iterations
        """

        self._initialize_params(X.shape[1])  # iniitalize weights and bias.

        N = len(X)  # number of samples
        y_pred = np.dot(X, self.weights) + self.bias # Initial predictions

        # Calculate loss and the partial derivatives of the loss function with respect to weights and bias

        for i in range(epochs):
             # Partial derivatives of loss w.r.t w and b
            dw = -2/N * np.dot(X.T, (y - y_pred)) 
            db = -2/N * (y - y_pred).sum() 

            # Update weights and bias
            self.weights -= lr * dw
            self.bias -= lr * db
            
            # Recalculate predictions with the new weights and bias
            y_pred = np.dot(X, self.weights) + self.bias
            
            # Calculate loss
            loss = ((y - y_pred) ** 2).mean()  # Mean squared error

            # Log this epoch
            self.log.append((self.weights.copy(), self.bias))
            self.mse.append(loss)

The fit_batch_gradient_descent takes four arguments:

• X, the independent variables used to make predictions
• y, the target variable (what the model will predict)
• lr, the learning rate
• epochs, the number of times the algorithm repeats It starts by setting the weights and bias to zero then loops for the number of epochs you've specified. Inside the loop, it calculates the partial derivatives using the formulas we derived earlier:

∂L/∂w = -(2/n) * Σ(yᵢ - ŷᵢ) * xᵢ -> -2/N * np.dot(X.T, (y - y_pred))
∂L/∂b = (-2/n) * Σ(yᵢ - ŷᵢ) -> -2/N * (y - y_pred).sum()

You have to transpose X (X.T) because np.dot does matrix multiplication and that requires the number of columns in the first matrix to match the number of rows in the second. X has shape ( ) while the error term (y - y_pred) has shape (samples, ). Transposing X to (features, samples) lines the dimensions up.

Once you've calculated the derivatives, the method:

• Updates the weights and bias
• Recalculates predictions using the new weights and bias
• Calculates the loss from the updated predictions
• Logs that epoch's loss plus the weights and bias. You can use this data to plot the training loss later.

Stochastic gradient descent

SGD updates the weights and bias using just one random data point at a time.

def fit_sgd(self, X,y,lr=0.05, epochs=100):
    self._initialize_params(X.shape[1])

    N = len(X)

    for _ in range(epochs):
        # Get random indices
        indices = np.random.permutation(N)
        for i in indices:
            xi = X[i:i+1]
            yi = y[i:i+1]

            y_pred = np.dot(xi, self.weights) + self.bias

            dw = -2 * xi.flatten() * (yi - y_pred)
            db = -2 * (yi - y_pred).sum()

            self.weights -= lr * dw
            self.bias -= lr * db
            self.log.append((self.weights.copy(), self.bias))

        y_pred_epoch = np.dot(X, self.weights) + self.bias
        loss = ((y-y_pred_epoch) **2).mean()

        self.mse.append(loss)

Each epoch starts with np.random.permutation which shuffles the row indices so the model doesn't see the data in the data in the same order every time. Then it loops through the shuffled indices one at a time. For each point, it:

• Calculate a prediction using just that one row
• Calculate the gradients of the loss w.r.t to weight and bias
• Update the weights and bias before moving to the next data point. Notice there's no /n averaging in the gradient formulas unlike batch gradient descent (You calculate each update form a single point so n=1) ### Mini-batch gradient descent Mini-batch gradient descent splits the data into multiple batches and updates the weights and bias oncer per batch.

def fit_mini_batch_gd(self, X,y,batch_size,lr=0.05, epochs=100):
    """Train the model using mini-batch gradient descent.
    Parameters:
        X (ndarray): Inputs
        y (ndarray): Output
        batch_size (int): Number of samples per batch
        epochs (int): Number of passes over the dataset
        lr (float): Learning rate
    """
    self._initialize_params(X.shape[1])

    N = len(X)

    for _ in range(epochs):
        indices = np.random.permutation(N)
        X_shuffled = X[indices]
        y_shuffled = y[indices]

        for start in range(0, N, batch_size):
            end = start + batch_size

            # Get a batch
            X_batch = X_shuffled[start:end]
            y_batch = y_shuffled[start:end]
            n_batch = len(X_batch)
            y_pred = np.dot(X_batch, self.weights) + self.bias

            # Calculate derivatives w.r.t. weights and bias
            dw = -2 / n_batch * np.dot(X_batch.T, (y_batch - y_pred))
            db = -2 / n_batch * (y_batch - y_pred).sum()

            # Update weight and bias
            self.weights -= lr * dw
            self.bias -= lr * db

        # Calculate predictions after each epoch
        y_pred_epoch = np.dot(X, self.weights) + self.bias

        # Calculate loss
        loss = ((y - y_pred_epoch) ** 2).mean()

        # Log loss, weights, and bias
        self.mse.append(loss)
        self.log.append((self.weights.copy(), self.bias))

Same shuffling trick to start, using np.random.permutation so the algorithm works through the data in a different order each epoch. Then it splits the shuffled dataset into consecutive chunks of batch_size, using range(0, N, batch_size).

For each batch, it:

1. Calculates predictions using just that batch's rows
2. Calculates the gradients of the loss with respect to weight and bias, averaged over the batch (note the /n_batch, since we're back to averaging over more than one point)
3. Updates the weights and bias immediately

Once it goes through every batch in the dataset, the method calculates predictions and loss across the entire dataset using the epoch's final weights and bias, then logs it.

If you set the batch_size to 1, this method will behave like SGD. If you use a batch_size set to the full length of the dataset, it behaves like batch gradient descent.

Make predictions

The predict method takes the final weights and bias from training and applies the same formula, ŷ = wx + b, to whatever new data you pass in.

def predict(self, X):
    """
    Predict target values for new input data using the learned weights and biases
    Parameters:
        X(ndarray): New features
    Return:
        Predictions (ndarray)
    """

    predictions = np.dot(X, self.weights) + self.bias

    return predictions

Example: Predict student performance

To see how my model works, I used it to predict student performance based on a dataset from Kaggle.

The dataset is clean, but it needs two main changes:

• The Extracturicular Activities column was Yes/No so I encoded it to 0,1.
• The columns were on different scales so I standardized them I also split the data into a training and test set.

df = pd.read_csv('Student_Performance.csv')

# Encode data
df['Extracurricular Activities'] = df['Extracurricular Activities'].map({'Yes': 1, 'No': 0})

# Define X and y
X = df.drop('Performance Index', axis=1)  # drop the column we're predicting
y = df['Performance Index']

# Split the data - test size is 20% of the total sample
X_train = X.sample(frac=0.8, random_state=42)
X_test = X.drop(X_train.index)

y_train = y.sample(frac=0.8, random_state=42)
y_test = y.drop(y_train.index)

# Scale the features - mean to 0, std to 1
X_train_mean = X_train.mean(axis=0)
X_train_std = X_train.std(axis=0)

X_train_scaled = (X_train - X_train_mean) / X_train_std
X_test_scaled = (X_test - X_train_mean) / X_train_std

# The model expects numpy arrays
X_train_scaled = X_train_scaled.to_numpy()
y_train = y_train.to_numpy()

X_test_scaled = X_test_scaled.to_numpy()
y_test = y_test.to_numpy()

Train with both three methods:

`python

Batch gradient descent

Make predictions

batch_model = LinearRegression()

batch_model.fit_gradient_descent(X_train_scaled, y_train, lr=0.05, epochs=100)

train_predictions = batch_model.predict(X_train_scaled)
test_predictions = batch_model.predict(X_test_scaled)

batch_train_mse = ((y_train - train_predictions) ** 2).mean()
batch_test_mse = ((y_test - test_predictions) ** 2).mean()

print("MSE for batch model")
print(f'Training mse: {batch_train_mse}\n Test Mse: {batch_test_mse}')

batch_model.visualize_loss()

`

I passed the scaled data, learning rate and epoch count into the fit method, then generated predictions on both the training and test sets. Comparing those two MSE values tells you whether the model is overfitting, if training MSE is much lower than test MSE, the model has memorised the training data rather than learned anything generalisable.

At first, I used lr=0.01 and epochs=50 and the errors were very large . The training MSE was 458 and test MSE of 460. MSE for batch model Training mse: 458.02 Test Mse: 459.74.

So I increased the learning rate to 0.05 and epochs to 100 to give the model enough time to minimize the loss. That brought the training MSE to 4.14 and a test MSE to 4.18. These errors are close together so I don't have an overfitting problem (basically, the model is generalizing well on unseen data)

To see this more clearly, I added a visualize_loss method to the class:
`python
def visualize_loss(self):
plt.plot(self.mse)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training loss vs iterations')
`

Below is the loss curve for batch GD at lr=0.05, epochs=100

Compare the same curve at the original lr=0.01, epochs=50, it never plateaus:

When I used the same learning rate and epoch count to train the SGD model, I got a larger MSE error. The training MSE was 5.43 and test MSE was 5.37. And as expected, the training took longer.
`python

Stochastic gradient descent

sgd_model = LinearRegression()
sgd_model.fit_sgd(X_train_scaled, y_train)

sgd_train_predictions = sgd_model.predict(X_train_scaled)
sgd_test_predictions = sgd_model.predict(X_test_scaled)

sgd_train_mse = ((y_train - sgd_train_predictions) ** 2).mean()
sgd_test_mse = ((y_test - sgd_test_predictions) ** 2).mean()
`
The loss curve was also very zig-zaggy.

Increasing the learning rate, made the error worse but dropping it to 0.01 brought the MSE much closer to batch's 4.14.

This tracks with what I said earlier. SGD updates on the raw error from a single point rather than an average. If that point is an outlier, its error will be bigger than the average across the entire dataset. If you now multiply that error with the same learning rate you used in batch GD, you get a big number that can swing the weights to far in one direction. That's why its better to start with a smaller learning rate.

I tried pushing further, down to lr=0.001 and epochs=20, but it never beat batch's numbers.

Next, I trained the mini-batch model with the same settings 0.05 learning rate, 100 epochs, and a batch size of 32.
`python

Mini-batch gradient descent

mini_batch_model = LinearRegression()
mini_batch_model.fit_mini_batch_gd(X_train_scaled, y_train, batch_size=32, lr=0.05, epochs=100)

mb_train_predictions = mini_batch_model.predict(X_train_scaled)
mb_test_predictions = mini_batch_model.predict(X_test_scaled)

mini_batch_train_mse = ((y_train - mb_train_predictions) ** 2).mean()
mini_batch_test_mse = ((y_test - mb_test_predictions) ** 2).mean()
`

Out of the gate, the MSE was much closer to batch. The training MSE was 4.31 and the test was 4.45. But the loss curve was noisy which told me the learning rate was too large.

Reducing the learning rate to 0.01, gave a much smoother curve.

The model was clearly reaching its minimum well before 100 epochs, so I cut that down to 20:

I also wanted to see what changing the batch size itself would do. Smaller batches produced noisier curves, closer to SGD, since there are fewer points to average the gradient over. Larger batches did the opposite.
Here's batch_size=16:

And batch_size=64, no jagged lines at all:

SGD and mini-batch eventually converged to similar train and test MSE values, but even after tuning learning rate, epochs and batch size, neither quite matched batch's result. Mini-batch was the fastest of the three, though, and got close enough to batch's accuracy that the speed trade-off is worth it.