Generate toy data

Linear regression is a fundamental supervised machine learning algorithm for modeling the relationship between input features and a continuous target variable. It assumes a linear relationship between the inputs and output and finds the optimal parameters (weights and bias) that minimize the mean squared error between the predicted and actual target values on a training dataset.

Despite its simplicity, linear regression is a powerful modeling technique that is widely used across many domains, from sales forecasting to analyzing scientific experimental results. It‘s an important foundational concept to master before moving on to more complex algorithms.

In this post, we‘ll take a deep dive into linear regression and the gradient descent optimization algorithm that‘s commonly used to train it, focusing on implementing them from scratch using PyTorch. Let‘s get started!

The Mathematical Formulation of Linear Regression

In linear regression, we model the target variable y as a linear function of the input features X:

y = Xw + b

where:

• y is the predicted target (a vector)
• X is the matrix of input feature values
• w is the vector of weights (model parameters)
• b is the bias or intercept term (a scalar)

The goal is to find the optimal values for the weights w and bias b that minimize the mean squared error (MSE) between the predicted targets y_pred and the actual targets y_true:

MSE = (1/n) * Σ(y_pred – y_true)^2

By convention, we add an extra feature column of all 1‘s to the input matrix X, so the bias term gets absorbed into the weights vector w. The linear regression equation then simplifies to:

y = Xw

This allows us to solve for the weights using matrix operations.

Coding Up Linear Regression in PyTorch

Let‘s see how to implement linear regression in PyTorch, taking advantage of its automatic differentiation capabilities. We‘ll start by importing the libraries we need:

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

Next we‘ll generate some toy linear data by adding random noise to a line:

def make_data(w, b, n_samples):
X = torch.normal(0, 1, (n_samples, len(w)))
y = torch.matmul(X, w) + b
y += torch.normal(0, 0.01, y.shape)
return X, y

true_w = torch.tensor([2, -3.4])
true_b = 4.2
n_samples = 100

X, y = make_data(true_w, true_b, n_samples)

Now we can define our model, which is just a linear layer with 2 input features and 1 output:

class LinearRegression(nn.Module):
def init(self, n_features):
super().init()
self.linear = nn.Linear(n_features, 1)

def forward(self, x):
    return self.linear(x)

model = LinearRegression(2)

The constructor creates a linear layer with 2 inputs and 1 output, and the forward method simply passes the input through this layer.

We‘ll also need a loss function, for which we‘ll use mean squared error:

criterion = nn.MSELoss()

Training with Gradient Descent

With our data and model ready, we can now train the model using the gradient descent optimization algorithm.

The idea behind gradient descent is simple yet powerful: to find the minimum of a function, take steps in the direction of the negative gradient, as this points in the direction of steepest descent.

For linear regression, the function we‘re minimizing is the MSE loss with respect to the weights and bias. We start with randomly initialized parameters. Then, for a certain number of iterations, we:

1. Make predictions using the current parameters
2. Calculate the MSE loss between predictions and true targets
3. Backpropagate to calculate the gradients of the loss with respect to the parameters
4. Take a step in the negative gradient direction to update the parameters
5. Reset the gradients to zero for the next iteration

Here‘s how this looks in code:

epochs = 100
learning_rate = 1e-1

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

for epoch in range(epochs):

y_pred = model(X)

# Compute loss
loss = criterion(y_pred, y)

# Backpropagation to get gradients 
loss.backward()

# Update parameters using gradient descent
optimizer.step()

# Reset gradients to zero
optimizer.zero_grad()

if (epoch+1) % 10 == 0:
    print(f‘Epoch: {epoch+1}, Loss: {loss.item():.4f}‘)

Some key things to note:

• We‘re using PyTorch‘s SGD optimizer, passing it the model‘s learnable parameters and the learning rate
• The learning rate controls the size of the steps we take in the negative gradient direction. It‘s a key hyperparameter that needs to be tuned for each problem.
• We run gradient descent for a fixed number of epochs (full passes through the training data). We could also use a different criterion, like stopping when the loss falls below a certain threshold.
• The loss decreases over the epochs as the model parameters are updated to better fit the data

Visualizing the Results

To visualize how well our trained model fits the data, let‘s plot the learned regression line against the original data points:

with torch.no_grad():
y_pred = model(X)

plt.figure(figsize=(8, 6))
plt.scatter(X[:, 0], y, label=‘data‘)
plt.plot(X[:, 0], y_pred, ‘r-‘, label=‘model‘)
plt.legend()
plt.show()

We can see that the model has learned to fit a line that nicely captures the linear trend in the noisy data.

It‘s also informative to plot the loss curve to check that it‘s decreasing over time:

plt.figure(figsize=(8, 6))
plt.plot(range(epochs), losses)
plt.xlabel(‘Epoch‘)
plt.ylabel(‘MSE Loss‘)
plt.show()

If the loss curve is unstable or not decreasing, that‘s a sign that the learning rate may be too high and is causing the optimizer to overshoot the minimum.

Limitations of Linear Regression and Gradient Descent

While linear regression is a useful algorithm, it does have some key limitations to be aware of:

• It assumes a linear relationship between the input features and output variable, which is often an oversimplification. Many real-world datasets have nonlinear or more complex patterns.

• It‘s sensitive to outliers, which can unduly influence the learned parameters and decrease model performance.

• It can‘t handle categorical features natively and requires transforming them (e.g. one-hot encoding) which can increase dimensionality.

• It tends to underfit on more complex datasets and may have high bias.

Gradient descent also has some failure modes. If the loss function has local optima, it can get stuck there instead of finding the global optimum. Additionally, it can be sensitive to the choice of learning rate – too low and it will take forever to converge, too high and it will diverge or "explode".

There are various techniques to improve gradient descent, like using momentum, adaptive learning rates (Adagrad, Adam), or second order methods like Newton‘s method that use second derivative information to help point in a better direction.

For more complex datasets, we can explore polynomial regression, regularization techniques (ridge, lasso), or switch to nonlinear algorithms like decision trees, random forests, or neural networks.

Comparing to Scikit-Learn

We can check our results against the optimized linear regression implementation in scikit-learn:

from sklearn.linear_model import LinearRegression as SklearnLR

sk_model = SklearnLR()
sk_model.fit(X, y)

sk_preds = sk_model.predict(X)

plt.figure(figsize=(8, 6))
plt.scatter(X[:, 0], y)
plt.plot(X[:, 0], y_pred.detach(), label=‘PyTorch‘)
plt.plot(X[:, 0], sk_preds, label=‘scikit-learn‘)
plt.legend()
plt.show()

print(f‘PyTorch MSE: {criterion(y_pred, y).item():.4f}‘)
print(f‘scikit-learn MSE: {criterion(torch.tensor(sk_preds), y).item():.4f}‘)

The results are essentially identical, which is reassuring! The small differences are likely due to randomness in initialization and the stochastic nature of mini-batch gradient descent.

Real-World Applications

Linear regression is used extensively in various fields to model and understand linear relationships:

• In finance, it‘s used to predict stock prices, analyze trends, and build trading strategies
• In social sciences, it‘s used to study relationships between variables like income and education level
• In medicine, it‘s used to relate patient characteristics to disease risk or treatment efficacy
• In business, it‘s used to forecast sales, understand factors driving churn, and set pricing

Although it‘s a simple algorithm, its versatility and interpretability make it a valuable tool in any data scientist‘s repertoire. I encourage you to try applying linear regression to a dataset of your own to solidify your understanding!

Conclusion

In this post, we took a deep dive into linear regression and gradient descent, focusing on building them from scratch in PyTorch. We discussed the mathematical formulation, coded up a working example, and visualized the results.

We saw how gradient descent iteratively updates the model parameters to minimize the loss, and how the learning rate affects convergence. We also highlighted some limitations of linear regression and failure modes of gradient descent to be aware of.

There are many directions you can go from here – trying out different optimizers, switching to mini-batch gradient descent, or moving on to more advanced algorithms. I hope this post gave you a solid foundation to build upon.

Happy learning, and I look forward to seeing what you create with PyTorch and gradient descent!