1.4 - L2 Regularization - Inside Deep Learning

In this notebook, we will not edit our model function $\hat{y}$ ; instead let’s add $\ell_2$ regularization (also denoted as $L_{2}$ ). Machine learning regularizers are techniques used to improve a model’s generalization and do not alter the way a model makes predictions.

This notebook takes the scratch model developed in 1.3 - Multioutput Linear Regression and only edits parameters updates, the rest remain unchanged. We are going to skip the explanations like multioutput task, weighted sum, gradients of weighted sum to focus on $\ell_2$ regularization.

Purpose of this Notebook:

Create a dataset for multioutput linear regression task
Add $\ell_2$ regularization into scratch model
Calculate the gradients from scratch
Implement gradient descent from scratch
Train our Perceptron
Compare our Perceptron to PyTorch’s built-in implementation

Setup¶

print('Start package installation...')

Start package installation...

%%capture
%pip install torch
%pip install scikit-learn

print('Packages installed successfully!')

Packages installed successfully!

import torch
from torch import nn

from platform import python_version
python_version(), torch.__version__

('3.14.4', '2.12.0+cu126')

# Set seeds for reproducibility
device = 'cpu'
torch.manual_seed(14)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(14)
    device = 'cuda'
    import numpy as np; np.random.seed(14)
import random; random.seed(14)
device

'cuda'

The add_to_class function is used to add new methods to a previously defined class; we do this to gradually enhance the class.

def add_to_class(Class):  
    """Register functions as methods in created class."""
    def wrapper(obj): setattr(Class, obj.__name__, obj)
    return wrapper

Dataset¶

from sklearn.datasets import make_regression
import random

N: int = 2_000  # number of samples
D: int = 5  # number of input features
C: int = 3  # number of output features

X, Y = make_regression(  # type: ignore
    n_samples=N,
    n_features=D,
    n_targets=C,
    n_informative=N - 1,
    bias=random.random(),
    noise=1
)

X = X.astype(np.float32)
Y = Y.astype(np.float32)

from sklearn.model_selection import train_test_split

X_train, X_valid, Y_train, Y_valid = train_test_split(
    X, Y,
    test_size=0.2,
    random_state=42,
    shuffle=True,
)

from torch.utils.data import TensorDataset

train_dataset = TensorDataset(
    torch.from_numpy(X_train),
    torch.from_numpy(Y_train)
)

valid_dataset = TensorDataset(
    torch.from_numpy(X_valid),
    torch.from_numpy(Y_valid)
)

from torch.utils.data import DataLoader

BATCH_SIZE: int = 32

train_loader = DataLoader(
    train_dataset,
    batch_size=BATCH_SIZE,
    shuffle=False,  # we want to ensure determinism
    pin_memory=True,
    drop_last=True,  # all batches have the same number of samples
)

valid_loader = DataLoader(
    valid_dataset,
    batch_size=BATCH_SIZE,
    shuffle=False,
    pin_memory=True,
    drop_last=False,  # set True to avoid stat bias
)

Scratch Model¶

class MultioutputRegression:
    def __init__(self, n_features: int, out_features: int):
        self.b = torch.randn(out_features).to(device)
        self.w = torch.randn(n_features, out_features).to(device)

    def copy_params(self, torch_layer: nn.modules.linear.Linear):
        """
        Copy the parameters from a module.linear to this model.

        Args:
            torch_layer: Pytorch module from which to copy the parameters.
        """
        self.b.copy_(torch_layer.bias.detach())
        self.w.copy_(torch_layer.weight.T.detach())

    def predict(self, x: torch.Tensor) -> torch.Tensor:
        """
        Predict the output for input x

        Args:
            x: Input tensor of shape (m_samples, d_features).

        Returns:
            y_pred: Predicted output tensor of shape (m_samples, c_features).
        """
        return torch.matmul(x, self.w) + self.b

    def mse_loss(self, y_true: torch.Tensor, y_pred: torch.Tensor):
        """
        MSE loss function between target y_true and y_pred.

        Args:
            y_true: Target tensor of shape (m_samples, c_features).
            y_pred: Predicted tensor of shape (m_samples, c_features).

        Returns:
            loss: MSE loss between predictions and true values.
        """
        return ((y_pred - y_true) ** 2).mean().item()

    def evaluate(self, x: torch.Tensor, y_true: torch.Tensor):
        """
        Evaluate the model on input x and target y_true using MSE.

        Args:
            x: Input tensor of shape (m_samples, d_features).
            y_true: Target tensor of shape (m_samples, c_features).

        Returns:
            loss: MSE loss between predictions and true values.
        """
        y_pred = self.predict(x)
        return self.mse_loss(y_true, y_pred)

    def update(self, x, y_true, y_pred, lr, lambda_):
        raise NotImplementedError

    def fit(self, train_loader: DataLoader,
        epochs: int, lr: float, lambda_: float,
        valid_loader: DataLoader):
        """
        Fit the model using gradient descent.

        Args:
            train_loader: Train dataloader.
            epochs: Number of epochs to fit.
            lr: Learning rate.
            lambda_: L2 regularization rate.
            valid_loader: Valid dataloader.
        """
        for epoch in range(epochs):
            # training epoch
            running_loss = 0.0
            for batch_x, batch_y in train_loader:
                # move only the current batch to VRAM
                batch_x = batch_x.to(device, non_blocking=True)
                batch_y = batch_y.to(device, non_blocking=True)

                # make predictions
                y_pred = self.predict(batch_x)

                running_loss += self.mse_loss(
                    batch_y, y_pred
                )

                self.update(
                    batch_x, batch_y,
                    y_pred, lr, lambda_
                )

            avg_loss = running_loss / len(train_loader)

            # validation epoch
            running_vloss = 0.0
            # disable gradient computation and reduce memory consumption.
            with torch.no_grad():
                for vbatch_x, vbatch_y in valid_loader:
                    vbatch_x = vbatch_x.to(device, non_blocking=True)
                    vbatch_y = vbatch_y.to(device, non_blocking=True)

                    vy_pred = self.predict(vbatch_x)

                    running_vloss += self.mse_loss(
                        vbatch_y, vy_pred
                    )
            avg_loss_v = running_vloss / len(valid_loader)

            print(f'epoch: {epoch} - MSE: {avg_loss:.4f} - vMSE: {avg_loss_v:.4f}')

Parameters update¶

Objective function¶

Instead of training the model with respect to the loss function $\mathcal{L}$ , we are going to use objective function $J$ . Usually, the objective function takes the form:

J(\theta) = \mathcal{L}(\theta) + R(\theta)

(1)

where $\theta$ are the parameters of the model, and $R$ is the regularization term.

Note: We do not use the objective function to evaluate the model; it is only used during the training step.

Remark: Some machine learning publications assign special meaning to the terms loss function, cost function and objective function. In these notebooks, we use the term objective function as the sum of a loss function plus a regularization method.

Regularizations methods are used to find a good balance between overfitting and underfitting. Without regularization, the model is trained solely on the training data, which means it may learn the noise and fail to approximate the underlying relationship or distribution of the data. With high regularization, the model will be unable to learn anything.

L2 regularization¶

$\ell_{2}$ (or just $L_{2}$ ) penalises the model more heavily if its weight values are larger; therefore, to minimise $J$ , the model must minimise both $\mathcal{L}$ and its weight values. $L_{2}$ regularizer is defined as:

\ell_{2}(\theta) = \frac{\lambda}{2} \| \theta \|^{2}_{2}

(2)

where $\lambda \geq 0$ is called regularization rate (or also regularization strength, regularization parameter) is a hyperparameter. Hyperparameters are parameters controlled by the developer (in this case you) not by the model. You can adjust the value of $\lambda$ . The ‘appropriate’ value depends on the specific scenario, such as the dataset, model, configuration, etc. Techniques such as k-fold cross-validation are used to find a good balance for $\lambda$ .

Remark: Generally, weights are used for $L_{2}$ , not bias.

For our case, the weight is of the form $\mathbf{W} \in \mathbb{R}^{d \times c}$ , then $L_{2}$ is:

\begin{align} \ell_{2}(\mathbf{W}) &= \frac{\lambda}{2} \left\| \mathbf{W} \right\|^{2}_{F} \\ &= \frac{\lambda}{2} \sum_{i=1}^{d} \sum_{j=1}^{c} w_{ij}^{2} \end{align}

(3)

where $d$ is the number of input features, and $c$ is the number of output features.

Note $\|A\|_{F}$ is called Frobenius norm.

Objective function derivative¶

\frac{\partial J}{\partial w_{rs}} = \frac{\partial \mathcal{L}}{\partial w_{rs}} + \frac{\partial \ell_{2}}{\partial w_{rs}}

(4)

Note: We shall omit the expansion of the first term on the right-hand side, as it has already been expanded in the previous notebook.

\begin{align} \frac{\partial \ell_{2}}{\partial w_{rs}} &= \frac{\lambda}{2} \sum_{i=1}^{d} \sum_{j=1}^{c} \frac{\partial}{\partial w_{rs}} \left( w_{ij}^{2} \right) \\ &= \lambda \sum_{i=1}^{d} \sum_{j=1}^{c} w_{ij} \frac{\partial w_{ij}}{\partial w_{rs}} \\ &= \lambda \sum_{i=1}^{d} \sum_{j=1}^{c} w_{ij} \delta_{ir} \delta_{js} \\ &= \lambda \sum_{i=1}^{d} w_{is} \delta_{ir} \\ &= \lambda w_{rs} \end{align}

(5)

for $r = 1, \ldots, d$ , and $s = 1, \ldots, c$ .

The vectorized form is:

\frac{\partial \ell_{2}}{\partial \mathbf{W}} = \lambda \mathbf{W}

(6)

Remark: $\nabla_{\mathbf{W}}\ell_2 \in \mathbb{R}^{d \times c}$ .

Final Gradients¶

\begin{align*} \frac{\partial J}{\partial \mathbf{W}} &= {\color{Orange} {\frac{\partial \mathcal{L}}{\partial \mathbf{W}}}} + {\color{Cyan} {\frac{\partial \ell_2}{\partial \mathbf{W}}}} \\ &= {\color{Orange} {\nabla_{\mathbf{W}} \mathcal{L}}} + {\color{Cyan} {\lambda \mathbf{W}}} \end{align*}

(7)

@add_to_class(MultioutputRegression)
def update(self, x: torch.Tensor, y_true: torch.Tensor,
           y_pred: torch.Tensor, lr: float, lambda_: float):
    """
    Update the model parameters.

    Args:
       x: Input tensor of shape (m_samples, d_features).
       y_true: Target tensor of shape (m_samples, c_features).
       y_pred: Predicted output tensor of shape (m_samples, c_features).
       lr: Learning rate. 
       lambda_: L2 regularization rate.
    """
    delta = 2 * (y_pred - y_true) / y_true.numel()

    # bias update (without modification)
    self.b -= lr * delta.sum(dim=0)

    # weight update
    l2_term = lambda_ * self.w
    loss_term = torch.matmul(x.T, delta)
    self.w -= lr * (loss_term + l2_term)  # new update

Scratch vs Torch.nn¶

PyTorch Module¶

class TorchLinearRegression(nn.Module):
    def __init__(self, d_features, c_out_features):
        super(TorchLinearRegression, self).__init__()
        self.layer = nn.Linear(
            d_features,
            c_out_features
        )
        self.loss = nn.MSELoss()

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

    @torch.inference_mode()
    def evaluate(self, x, y):
        self.eval()
        y_pred = self(x)
        return self.loss(y_pred, y).item()

    def fit(self, train_loader,
            epochs, lr, l2_lambda,
            valid_loader):
        # new edit here, only apply L2 to weight, not bias
        optimizer = torch.optim.SGD([
            {'params': self.layer.weight, 'weight_decay': l2_lambda},
            {'params': self.layer.bias}
        ], lr=lr)

        for epoch in range(epochs):
            self.train()  # use model.train() when training
            running_loss = 0.0  # train loss
            for batch_x, batch_y in train_loader:
                batch_x = batch_x.to(device, non_blocking=True)
                batch_y = batch_y.to(device, non_blocking=True)

                y_pred = self(batch_x)

                loss = self.loss(y_pred, batch_y)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()
            avg_loss = running_loss / len(train_loader)

            # validation epoch
            self.eval()  # set the model to evaluation mode
            with torch.no_grad():
                running_vloss = 0.0
                for vbatch_x, vbatch_y in valid_loader:
                    vbatch_x = vbatch_x.to(device, non_blocking=True)
                    vbatch_y = vbatch_y.to(device, non_blocking=True)

                    vy_pred = self(vbatch_x)
                    vloss = self.loss(vy_pred, vbatch_y)
                    running_vloss += vloss.item()
            avg_loss_v = running_vloss / len(valid_loader)

            print(f'epoch: {epoch} - MSE: {avg_loss:.4f} - vMSE: {avg_loss_v:.4f}')

torch_model = TorchLinearRegression(D, C).to(device)

model = MultioutputRegression(D, C)

Eval¶

We use a $L_{2}$ norm between PyTorch model and our Scratch model as parameter discrepancy.

def l2(pred, true_): 
    if isinstance(pred, (float, int)) and isinstance(true_, (float, int)):
        return abs(pred - true_)
    return torch.linalg.norm(pred - true_).item()

x_valid, y_valid = valid_dataset[:]
x_valid = x_valid.to(device, non_blocking=True)
y_valid = y_valid.to(device, non_blocking=True)

Copy Parameters¶

Copy the values of the PyTorch model parameters to our model.

model.copy_params(torch_model.layer)

Predictions post-copy¶

l2(
    model.predict(x_valid),
    torch_model(x_valid)
)

0.0

Training¶

LR: float = 0.01  # learning rate
EPOCHS: int = 16  # number of epochs
L2_LAMBDA: float = 0.01  # L2 regularization rate

model.fit(
    train_loader,
    EPOCHS, LR, L2_LAMBDA,
    valid_loader
)

epoch: 0 - MSE: 10516.6279 - vMSE: 6843.1859
epoch: 1 - MSE: 5384.1285 - vMSE: 3536.7894
epoch: 2 - MSE: 2777.5853 - vMSE: 1843.6148
epoch: 3 - MSE: 1446.6729 - vMSE: 971.4796
epoch: 4 - MSE: 762.7341 - vMSE: 519.0926
epoch: 5 - MSE: 408.5224 - vMSE: 282.4109
epoch: 6 - MSE: 223.3104 - vMSE: 157.2605
epoch: 7 - MSE: 125.3093 - vMSE: 90.2097
epoch: 8 - MSE: 72.6896 - vMSE: 53.7044
epoch: 9 - MSE: 43.9300 - vMSE: 33.4425
epoch: 10 - MSE: 27.8764 - vMSE: 21.9409
epoch: 11 - MSE: 18.6962 - vMSE: 15.2459
epoch: 12 - MSE: 13.3053 - vMSE: 11.2422
epoch: 13 - MSE: 10.0502 - vMSE: 8.7811
epoch: 14 - MSE: 8.0294 - vMSE: 7.2275
epoch: 15 - MSE: 6.7416 - vMSE: 6.2223

torch_model.fit(
    train_loader,
    EPOCHS, LR, L2_LAMBDA,
    valid_loader
)

epoch: 0 - MSE: 10516.6279 - vMSE: 6843.1859
epoch: 1 - MSE: 5384.1287 - vMSE: 3536.7896
epoch: 2 - MSE: 2777.5853 - vMSE: 1843.6148
epoch: 3 - MSE: 1446.6729 - vMSE: 971.4796
epoch: 4 - MSE: 762.7342 - vMSE: 519.0927
epoch: 5 - MSE: 408.5225 - vMSE: 282.4110
epoch: 6 - MSE: 223.3104 - vMSE: 157.2605
epoch: 7 - MSE: 125.3093 - vMSE: 90.2098
epoch: 8 - MSE: 72.6896 - vMSE: 53.7043
epoch: 9 - MSE: 43.9300 - vMSE: 33.4425
epoch: 10 - MSE: 27.8764 - vMSE: 21.9410
epoch: 11 - MSE: 18.6962 - vMSE: 15.2460
epoch: 12 - MSE: 13.3053 - vMSE: 11.2422
epoch: 13 - MSE: 10.0502 - vMSE: 8.7811
epoch: 14 - MSE: 8.0294 - vMSE: 7.2275
epoch: 15 - MSE: 6.7416 - vMSE: 6.2223

Predictions after Training¶

l2(
    model.predict(x_valid),
    torch_model(x_valid)
)

0.00010880782065214589

Bias Comparison¶

l2(
    model.b.clone(),
    torch_model.layer.bias.detach()
)

4.050985680237318e-08

Weight Comparison¶

l2(
    model.w.clone(),
    torch_model.layer.weight.detach().T
)

3.932099843950709e-06

In this notebook, we demonstrated the mathematical foundation and practical implementation of $L_2$ regularization. By incorporating the Frobenius norm of the weight matrix into the objective function, we introduced a penalty that scales with weight magnitude.

Remark: While $L_2$ regularization and weight decay are mathematically equivalent when using standard Stochastic Gradient Descent, this equivalence breaks down in adaptive gradient methods. In modern deep learning architectures, decoupled weight decay (such as the implementation found in the AdamW optimizer) is rigorously separated from the gradient of the loss function to achieve optimal generalization.