2.1 - Multiclass classification - Inside Deep Learning

We are going to add a crucial element: the activation function. This function will allow us to modify the output to suit our problem, in this case the classification of multiple classes.

The softmax function will allow us to convert an input into the probability of remaining in the classes.

We can interpret the perceptron with softmax as a dense layer and an activation layer, this interpretation will be useful later in chapter 3.

Purpose of this Notebook:

The purposes of this notebook are:

Create a dataset for classification regression task
Create our own Perceptron class from scratch
Add Softmax function as activation function from scratch
Calculate the gradient descent from scratch
Train our Perceptron
Compare our Perceptron to the one prebuilt by PyTorch
[Extra] Calculate the gradient descent by another way

import torch
from torch import nn

from platform import python_version
python_version(), torch.__version__

('3.14.7', '2.12.0+cpu')

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
device

'cpu'

torch.set_default_dtype(torch.float64)

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

Dataset¶

create dataset¶

\begin{align*} \mathbf{X} &\in \mathbb{R}^{m \times n} \\ \mathbf{Y} &\in \mathbb{R}^{m \times n_{1}} \end{align*}

(1)

where $n_{1}$ is the number of classes.

from sklearn.datasets import make_classification

M: int = 10_100 # number of samples
N: int = 5 # number of input features
CLASSES: int = 3 # number of classes

X, Y = make_classification(
    n_samples=M, 
    n_features=N, 
    n_classes=CLASSES, 
    n_informative=N - 1, 
    n_redundant=0
)

print(X.shape)
print(Y.shape)

(10100, 5)
(10100,)

one hot encoding¶

Y_hat = nn.functional.one_hot(
    torch.tensor(Y, device=device).long(), 
    CLASSES
).type(torch.float32)
Y_hat.shape

torch.Size([10100, 3])

split dataset into train and valid¶

X_train = torch.tensor(X[:100], device=device)
X_valid = torch.tensor(X[100:], device=device)
X_train.shape, X_valid.shape

(torch.Size([100, 5]), torch.Size([10000, 5]))

Y_train, Y_valid = Y_hat[:100], Y_hat[100:]
Y_train.shape, Y_valid.shape

(torch.Size([100, 3]), torch.Size([10000, 3]))

delete raw dataset¶

del X
del Y
del Y_hat

Model¶

weights and bias¶

\begin{align*} \mathbf{W} &\in \mathbb{R}^{n \times n_{1}} \\ \mathbf{b} &\in \mathbb{R}^{n_{1}} \end{align*}

(2)

class SoftmaxClassifier:
    def __init__(self, n_features: int, n_classes: int):
        self.w = torch.randn(n_features, n_classes, device=device)
        self.b = torch.randn(n_classes, device=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().clone())
        self.w.copy_(torch_layer.weight.T.detach().clone())

weighted sum and softmax function¶

weighted sum

\mathbf{Z}(\mathbf{X}) = \mathbf{X} \mathbf{W} + \mathbf{b} \\ \mathbf{Z} : \mathbb{R}^{m \times n} \rightarrow \mathbb{R}^{m \times n_{1}}

(3)

softmax function

\sigma(\mathbf{z}_{i,:})_{j} = \frac{\exp(z_{ij})} {\sum_{k=1}^{n_{1}}(\exp(z_{ik}))} \in \mathbb{R}^{+}

(4)

then

\sigma(\mathbf{z}_{i,:}) = \begin{bmatrix} \sigma(\mathbf{z}_{i,:})_{1} & \sigma(\mathbf{z}_{i,:})_{2} & \cdots & \sigma(\mathbf{z}_{i,:})_{n_{1}} \end{bmatrix}

(5)

therefore

\mathbf{\Sigma(Z)} = \begin{bmatrix} \sigma(\mathbf{z}_{1,:}) \\ \sigma(\mathbf{z}_{2,:}) \\ \vdots \\ \sigma(\mathbf{z}_{m,:}) \end{bmatrix} \\ \mathbf{\Sigma} : \mathbb{R}^{m \times n_{1}} \rightarrow \mathbb{R}^{m \times n_{1}}

(6)

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

    Args:
        x: Input tensor of shape (n_samples, n_features).

    Returns:
        y_pred: Predicted output tensor of shape (n_samples, n_classes).
    """
    # weighted sum
    z = torch.matmul(x, self.w) + self.b
    # avoid underflow and overflow
    z_norm = z - torch.max(z, dim=1, keepdims=True)[0]
    # softmax function
    z_exp = torch.exp(z_norm)
    return z_exp / z_exp.sum(1, keepdims=True) # y_pred

Cross-entropy loss¶

Cross-entropy loss

L(\mathbf{\hat{Y}}) = - \frac{1}{m} \sum_{i=1}^{m} \sum_{j=1}^{n_{1}}( y_{ij} \log_{e}(\hat{y}_{ij}) ) \\ L : \mathbb{R}^{m \times n_{1}} \rightarrow \mathbb{R}

(7)

Remark: for this case $\mathbf{\hat{Y}}$ is $\mathbf{\Sigma(Z)}$ .
It is not mandatory to use softmax for cross-entropy loss, but some modules like PyTorch require softmax to use cross-entropy loss.

Vectorized form

L(\mathbf{\hat{Y}}) = - \frac{1}{m} \sum_{i=1}^{m} \left( \mathbf{y}_{i,:}^\top \log_{e}(\mathbf{\hat{y}}_{i,:}) \right)

(8)

or

L(\mathbf{\hat{Y}}) = - \frac{1}{m} \text{sum} \left( \mathbf{Y} \odot \log_{e}(\mathbf{\hat{Y}}) \right)

(9)

@add_to_class(SoftmaxClassifier)
def cross_entropy_loss(self, y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
    """
    CE loss function between target y_true and y_pred.

    Args:
        y_true: Target tensor of shape (n_samples, n_classes).
        y_pred: Predicted tensor of shape (n_samples, n_classes).

    Returns:
        loss: CE loss between predictions and true values.
    """
    loss = y_true * torch.log(y_pred)
    return - loss.sum().item() / len(y_true)

@add_to_class(SoftmaxClassifier)
def evaluate(self, x: torch.Tensor, y_true: torch.Tensor) -> float:
    """
    Evaluate the model on input x and target y_true using CE.

    Args:
        x: Input tensor of shape (n_samples, n_features).
        y_true: Target tensor of shape (n_samples, n_classes).

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

Gradient¶

Cross-entropy derivative¶

\begin{align*} \frac{\partial L}{\partial \hat{y}_{pq}} =& -\frac{1}{m} \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial}{\partial \hat{y}_{pq}} \left( y_{ij} \log_{e}(\hat{y}_{ij}) \right) \\ &= -\frac{1}{m} \left(\frac{y_{pq}}{\hat{y}_{pq}} \right) \end{align*}

(10)

for all $p = 1, \ldots, m$ and $q = 1, \ldots, n_{1}$ .

Remark: $\hat{y}_{pq}$ must be different of 0, $\hat{y}_{pq} \neq 0$ . Softmax returns positive real values, $\sigma(z) \in \mathbb{R}^{+}$ .

In general

\frac{\partial L}{\partial \hat{\mathbf{Y}}} = -\frac{1}{m} \left( \mathbf{Y} \oslash \hat{\mathbf{Y}} \right)

(11)

Note: $\oslash$ is element-wise divide.

softmax derivative¶

\begin{align*} \frac{\partial L}{\partial z_{pq}} =& -\frac{1}{m} \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial}{\partial z_{pq}} \left( y_{ij} \log_{e}(\hat{y}_{ij}) \right) \\ =& \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial L}{\partial \sigma_{ij}} \frac{\partial \sigma_{ij}}{\partial z_{pq}} \end{align*}

(12)

for all $p = 1, \ldots, m$ and $q = 1, \ldots, n_{1}$ .

where

\frac{\partial \sigma_{ij}}{\partial z_{pq}} = \begin{cases} \sigma(z_{pq})(1 - \sigma(z_{pq})) & \text{if } i=p, j=q \\ -\sigma(z_{pq}) \sigma(z_{ij}) & \text{if } i=p, j \neq q \\ 0 & \text{otherwise} \end{cases}

(13)

therefore

\begin{align*} \frac{\partial L}{\partial z_{pq}} =& \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial L}{\partial \sigma_{ij}} \frac{\partial \sigma_{ij}}{\partial z_{pq}} \\ =& \sum_{j=1}^{n_{1}} \frac{\partial L}{\partial \sigma_{pj}} \begin{cases} \sigma(z_{pq})(1 - \sigma(z_{pq})) & \text{if } j=q \\ -\sigma(z_{pq}) \sigma(z_{pj}) & \text{if } j \neq q \end{cases} \end{align*}

(14)

Check softmax function and its derivative for more information about the softmax derivative.

In general

\frac{\partial L}{\partial \mathbf{Z}} = \mathbf{\Sigma} \odot \left( \frac{\partial L}{\partial \mathbf{\Sigma}} - \left( \frac{\partial L}{\partial \mathbf{\Sigma}} \odot \mathbf{\Sigma} \right) \mathbf{1} \right)

(15)

where $\mathbf{1} \in \mathbb{R}^{n_{1} \times n_{1}}$ .

weighted sum derivative¶

respect to bias¶

\begin{align*} \frac{\partial L}{\partial b_{q}} =& -\frac{1}{m} \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial}{\partial b_{q}} \left( y_{ij} \log_{e}(\hat{y}_{ij}) \right) \\ &= \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial L}{\partial z_{ij}} \frac{\partial z_{ij}}{\partial b_{q}} \\ &= \sum_{i=1}^{m} \frac{\partial L}{\partial z_{iq}} \end{align*}

(16)

for all $q = 1, \ldots, n_{1}$ .

In general

\frac{\partial L}{\partial \mathbf{b}} = \mathbf{1} \frac{\partial L}{\partial \mathbf{Z}}

(17)

where $\mathbf{1} \in \mathbb{R}^{m}$ .

respect to weight¶

\begin{align*} \frac{\partial L}{\partial w_{pq}} =& -\frac{1}{m} \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial}{\partial w_{pq}} \left( y_{ij} \log_{e}(\hat{y}_{ij}) \right) \\ &= \sum_{i=1}^{m} \sum_{j=1}^{n_{1}} \frac{\partial L}{\partial z_{ij}} \frac{\partial z_{ij}}{\partial w_{pq}} \\ &= \sum_{i=1}^{m} x_{ip} \frac{\partial L}{\partial z_{iq}} \end{align*}

(18)

for all $p = 1, \ldots, m$ and $q = 1, \ldots, n_{1}$ .

In general

\frac{\partial L}{\partial \mathbf{W}} = \mathbf{X}^\top \frac{\partial L}{\partial \mathbf{Z}}

(19)

@add_to_class(SoftmaxClassifier)
def update(self, x: torch.Tensor, y_true: torch.Tensor, 
           y_pred: torch.Tensor, lr: float) -> None:
    """
    Update the model parameters.

    Args:
       x: Input tensor of shape (n_samples, n_features).
       y_true: Target tensor of shape (n_samples, n_classes).
       y_pred: Predicted output tensor of shape (n_samples, n_classes).
       lr: Learning rate. 
    """
    # cross entropy der
    delta = -(y_true / y_pred) / len(y_true)
    # softmax der
    delta = y_pred * (delta - (delta * y_pred).sum(axis=1, keepdims=True))
    # weighted sum der
    self.b -= lr * delta.sum(axis=0)
    self.w -= lr * (x.T @ delta)

metric: accuracy¶

@add_to_class(SoftmaxClassifier)
def accuracy(self, y_true, y_pred) -> float:
    preds = y_pred.argmax(axis=-1)
    compare = (y_true.argmax(axis=-1) == preds).type(torch.float32)
    return compare.mean().item()

fit (train)¶

@add_to_class(SoftmaxClassifier)
def fit(self, x_train: torch.Tensor, y_train: torch.Tensor, 
        epochs: int, lr: float, batch_size: int, 
        x_valid: torch.Tensor, y_valid: torch.Tensor) -> None:
    """
    Fit the model using gradient descent.

    Args:
        x_train: Input tensor of shape (n_samples, num_features).
        y_train: Target tensor one hot of shape (n_samples, n_classes).
        epochs: Number of epochs to train.
        lr: learning rate).
        batch_size: Int number of batch.
        x_valid: Input tensor of shape (n_valid_samples, num_features).
        y_valid: Input tensor one hot of shape (n_valid_samples, n_valid_classes).
    """
    for epoch in range(epochs):
        loss = []
        for batch in range(0, len(y_train), batch_size):
            batch_end = batch + batch_size

            y_pred = self.predict(x_train[batch:batch_end])
            loss.append(self.evaluate(
                x_train[batch:batch_end], 
                y_train[batch:batch_end]
            ))

            self.update(
                x_train[batch:batch_end], 
                y_train[batch:batch_end], 
                y_pred, lr
            )

        loss = round(sum(loss) / len(loss), 4)
        loss_v = round(self.evaluate(x_valid, y_valid), 4)
        acc = round(self.accuracy(y_valid, self.predict(x_valid)), 4)
        print(f'epoch: {epoch} - CE: {loss} - CE_v: {loss_v} - acc_v: {acc}')

Scratch vs nn¶

nn model¶

Important: nn.CrossEntropyLoss applies Softmax to input

class TorchSoftmax(nn.Module):
    def __init__(self, n_features, n_out_features):
        super(TorchSoftmax, self).__init__()
        self.layer = nn.Linear(n_features, n_out_features, device=device)
        self.soft = nn.Softmax(dim=1)
        self.loss = nn.CrossEntropyLoss()

    def forward(self, x):
        z = self.layer(x)
        return self.soft(z)
    
    def evaluate(self, x, y):
        self.eval()
        with torch.no_grad():
            y_pred = self.layer(x)
            # do not use self.soft because nn.CrossEntropyLoss already uses softmax
            return self.loss(y_pred, y).item()
    
    def fit(self, x, y, epochs, lr, batch_size, x_valid, y_valid):
        optimizer = torch.optim.SGD(self.parameters(), lr=lr)
        for epoch in range(epochs):
            loss_t = []
            for batch in range(0, len(y), batch_size):
                batch_end = batch + batch_size

                y_pred = self.layer(x[batch:batch_end])
                loss = self.loss(y_pred, y[batch:batch_end])
                loss_t.append(loss.item())

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

            loss_t = round(sum(loss_t) / len(loss_t), 4)
            loss_v = round(self.evaluate(x_valid, y_valid), 4)
            print(f'epoch: {epoch} - CE: {loss_t} - CE_v: {loss_v}')

torch_model = TorchSoftmax(N, CLASSES)

scratch model¶

model = SoftmaxClassifier(N, CLASSES)

evals¶

import MAPE modified¶

# This cell imports torch_mape 
# if you are running this notebook locally 
# or from Google Colab.

import os
import sys

module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)

try:
    from tools.torch_metrics import torch_mape as mape
    print('mape imported locally.')
except ModuleNotFoundError:
    import subprocess

    repo_url = 'https://raw.githubusercontent.com/PilotLeoYan/inside-deep-learning/main/content/tools/torch_metrics.py'
    local_file = 'torch_metrics.py'
    
    subprocess.run(['wget', repo_url, '-O', local_file], check=True)
    try:
        from torch_metrics import torch_mape as mape # type: ignore
        print('mape imported from GitHub.')
    except Exception as e:
        print(e)

mape imported locally.

predict¶

mape(
    model.predict(X_valid),
    torch_model(X_valid)
)

2.283273950528358

copy parameters¶

model.copy_params(torch_model.layer)
parameters = (model.b.clone(), model.w.clone())

predict after copy parameters¶

mape(
    model.predict(X_valid),
    torch_model(X_valid)
)

3.302948811450026e-17

CE¶

mape(
    model.evaluate(X_valid, Y_valid),
    torch_model.evaluate(X_valid, Y_valid)
)

0.0

train¶

LR = 0.01
EPOCHS = 16
BATCH = len(X_train) // 3

torch_model.fit(
    X_train, Y_train, 
    EPOCHS, LR, BATCH, 
    X_valid, Y_valid
)

epoch: 0 - CE: 1.572 - CE_v: 1.8893
epoch: 1 - CE: 1.4875 - CE_v: 1.8393
epoch: 2 - CE: 1.419 - CE_v: 1.7933
epoch: 3 - CE: 1.3625 - CE_v: 1.7506
epoch: 4 - CE: 1.3148 - CE_v: 1.7106
epoch: 5 - CE: 1.2736 - CE_v: 1.673
epoch: 6 - CE: 1.2374 - CE_v: 1.6375
epoch: 7 - CE: 1.2051 - CE_v: 1.6037
epoch: 8 - CE: 1.176 - CE_v: 1.5717
epoch: 9 - CE: 1.1494 - CE_v: 1.5411
epoch: 10 - CE: 1.1249 - CE_v: 1.512
epoch: 11 - CE: 1.1024 - CE_v: 1.4842
epoch: 12 - CE: 1.0814 - CE_v: 1.4577
epoch: 13 - CE: 1.062 - CE_v: 1.4324
epoch: 14 - CE: 1.0438 - CE_v: 1.4083
epoch: 15 - CE: 1.0268 - CE_v: 1.3852

model.fit(
    X_train, Y_train, 
    EPOCHS, LR, BATCH, 
    X_valid, Y_valid
)

epoch: 0 - CE: 1.572 - CE_v: 1.8893 - acc_v: 0.1408
epoch: 1 - CE: 1.4875 - CE_v: 1.8393 - acc_v: 0.1506
epoch: 2 - CE: 1.419 - CE_v: 1.7933 - acc_v: 0.161
epoch: 3 - CE: 1.3625 - CE_v: 1.7506 - acc_v: 0.1701
epoch: 4 - CE: 1.3148 - CE_v: 1.7106 - acc_v: 0.179
epoch: 5 - CE: 1.2736 - CE_v: 1.673 - acc_v: 0.1862
epoch: 6 - CE: 1.2374 - CE_v: 1.6375 - acc_v: 0.193
epoch: 7 - CE: 1.2051 - CE_v: 1.6037 - acc_v: 0.1992
epoch: 8 - CE: 1.176 - CE_v: 1.5717 - acc_v: 0.2063
epoch: 9 - CE: 1.1494 - CE_v: 1.5411 - acc_v: 0.2142
epoch: 10 - CE: 1.1249 - CE_v: 1.512 - acc_v: 0.2205
epoch: 11 - CE: 1.1024 - CE_v: 1.4842 - acc_v: 0.2255
epoch: 12 - CE: 1.0814 - CE_v: 1.4577 - acc_v: 0.2317
epoch: 13 - CE: 1.062 - CE_v: 1.4324 - acc_v: 0.2384
epoch: 14 - CE: 1.0438 - CE_v: 1.4083 - acc_v: 0.2462
epoch: 15 - CE: 1.0268 - CE_v: 1.3852 - acc_v: 0.2556

predict after train¶

mape(
    model.predict(X_valid),
    torch_model.forward(X_valid)
)

1.35383281674615e-16

weight¶

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

1.5600348830805582e-16

bias¶

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

7.66750107073557e-17

Compute gradient with einsum¶

Gradient descent is

\frac{\partial L}{\partial \mathbf{W}} = \frac{\partial L}{\partial \mathbf{\Sigma}} \frac{\partial \mathbf{\Sigma}}{\partial \mathbf{Z}} \frac{\partial \mathbf{Z}}{\partial \mathbf{W}}

(20)

and

\frac{\partial L}{\partial \mathbf{b}} = \frac{\partial L}{\partial \mathbf{\Sigma}} \frac{\partial \mathbf{\Sigma}}{\partial \mathbf{Z}} \frac{\partial \mathbf{Z}}{\partial \mathbf{b}}

(21)

where their shapes are

\begin{align*} \frac{\partial L} {\partial \mathbf{W}} &\in \mathbb{R}^{n \times n_{1}} \\ \frac{\partial L} {\partial \mathbf{b}} &\in \mathbb{R}^{n_{1}} \\ \frac{\partial L} {\partial \mathbf{\Sigma}} &\in \mathbb{R}^{m \times n_{1}} \\ \frac{\partial \mathbf{\Sigma}} {\partial \mathbf{Z}} &\in \mathbb{R}^{(m \times n_{1}) \times (m \times n_{1})} \\ \frac{\partial \mathbf{Z}} {\partial \mathbf{W}} &\in \mathbb{R}^{(m \times n_{1}) \times (n \times n_{1})} \\ \frac{\partial \mathbf{Z}} {\partial \mathbf{b}} &\in \mathbb{R}^{(m \times n_{1}) \times n_{1}} \end{align*}

(22)

Then we have 2 cases

\frac{\partial \mathbf{\Sigma}(\mathbf{Z})_{i,:}} {\partial \mathbf{Z}_{p=i,:}}

(23)

and

\frac{\partial \mathbf{\Sigma}(\mathbf{Z})_{i,:}} {\partial \mathbf{Z}_{p\neq i,:}}

(24)

First case

\frac{\partial \mathbf{\Sigma}(\mathbf{Z})_{i,:}} {\partial \mathbf{Z}_{p=i,:}} = \text{diag}(\sigma(\mathbf{Z}_{i,:})) - \sigma(\mathbf{Z}_{i,:}) \sigma(\mathbf{Z}_{i,:})^\top

(25)

Second case

\frac{\partial \mathbf{\Sigma}(\mathbf{Z})_{i,:}} {\partial \mathbf{Z}_{p \neq i,:}} = \mathbf{0}

(26)

Weighted sum derivative

\frac{\partial \mathbf{Z}}{\partial \mathbf{W}} = \mathbb{I} \otimes \mathbf{X}

(27)

\frac{\partial z_{ij}}{\partial b_{p}} = \begin{cases} 1 & \text{if } j=p \\ 0 & \text{if } j\neq p \end{cases}

(28)

for all $i = 1, \ldots, m$ and $j, p = 1, \ldots, n_{1}$

therefore using Einstein summation

\begin{align*} {\color{Orange} {\frac{\partial L}{\partial \mathbf{Z}}}} &= {\color{Lime} {\frac{\partial L}{\partial \mathbf{\Sigma}}}} {\color{Cyan} {\frac{\partial \mathbf{\Sigma}}{\partial \mathbf{Z}}}} \\ &\in \mathbb{R}^{ {\color{Lime} {(m \times n_{1})}} \times {\color{Cyan} {(m \times n_{1} \times m \times n_{1})}}} \\ &\in \mathbb{R}^{\color{Orange} {{m \times n_{1}}}} \end{align*}

(29)

\begin{align*} {\color{Magenta} {\frac{\partial L}{\partial \mathbf{b}}}} &= {\color{Orange} {\frac{\partial L}{\partial \mathbf{Z}}}} {\color{Cyan} {\frac{\partial \mathbf{Z}}{\partial \mathbf{b}}}} \\ &\in \mathbb{R}^{ {\color{Orange} {(m \times n_{1})}} \times {\color{Cyan} {(m \times n_{1} \times n_{1})}}} \\ &\in \mathbb{R}^{\color{Magenta} {n_{1}}} \end{align*}

(30)

and

\begin{align*} {\color{Magenta} {\frac{\partial L}{\partial \mathbf{W}}}} &= {\color{Orange} {\frac{\partial L}{\partial \mathbf{Z}}}} {\color{Cyan} {\frac{\partial \mathbf{Z}}{\partial \mathbf{W}}}} \\ &\in \mathbb{R}^{ {\color{Orange} {(m \times n_{1})}} \times {\color{Cyan} {(m \times n_{1} \times n \times n_{1})}}} \\ &\in \mathbb{R}^{\color{Magenta} {n \times n_{1}}} \end{align*}

(31)

Model¶

class EinsumSoftmaxClassifier(SoftmaxClassifier):
    def update(self, x: torch.Tensor, y_true: torch.Tensor,
           y_pred: torch.Tensor, lr: float) -> None:
        """
        Update the model parameters.

        Args:
            x: Input tensor of shape (n_samples, n_features).
            y_true: Target tensor of shape (n_samples, n_classes).
            y_pred: Predicted output tensor of shape (n_samples, n_classes).
            lr: Learning rate. 
        """
        m, n_classes = y_true.shape
        # cross entropy der
        delta = -(y_true / y_pred) / m
        # softmax der
        diag_a = torch.diag_embed(y_pred)
        outer_a = torch.einsum('ij,ik->ijk', y_pred, y_pred) 
        soft_der = torch.zeros(
            (m, n_classes, m, n_classes), 
            dtype=y_pred.dtype, 
            device=device
        )
        idx = torch.arange(m, device=device)
        soft_der[idx, :, idx, :] = diag_a - outer_a
        delta = torch.einsum('pq,pqij->ij', delta, soft_der)
        # weighted sum der
        self.b -= lr * delta.sum(axis=0)
        
        identity = torch.eye(n_classes, device=device)
        w_der = torch.kron(
            x.unsqueeze(1).unsqueeze(3), 
            identity.unsqueeze(0).unsqueeze(2)
        )
        w_der = torch.einsum('pq,pqij->ij', delta, w_der)
        self.w -= lr * w_der

einsum_model = EinsumSoftmaxClassifier(N, CLASSES)
einsum_model.b.copy_(parameters[0])
einsum_model.w.copy_(parameters[1])

tensor([[-0.2256, -0.1411,  0.3085],
        [ 0.0310, -0.1917,  0.2306],
        [-0.2142,  0.4443,  0.3190],
        [ 0.3056, -0.2796, -0.1243],
        [ 0.2665,  0.3218, -0.4009]])

einsum_model.fit(
    X_train, Y_train, 
    EPOCHS, LR, BATCH, 
    X_valid, Y_valid
)

epoch: 0 - CE: 1.572 - CE_v: 1.8893 - acc_v: 0.1408
epoch: 1 - CE: 1.4875 - CE_v: 1.8393 - acc_v: 0.1506
epoch: 2 - CE: 1.419 - CE_v: 1.7933 - acc_v: 0.161
epoch: 3 - CE: 1.3625 - CE_v: 1.7506 - acc_v: 0.1701
epoch: 4 - CE: 1.3148 - CE_v: 1.7106 - acc_v: 0.179
epoch: 5 - CE: 1.2736 - CE_v: 1.673 - acc_v: 0.1862
epoch: 6 - CE: 1.2374 - CE_v: 1.6375 - acc_v: 0.193
epoch: 7 - CE: 1.2051 - CE_v: 1.6037 - acc_v: 0.1992
epoch: 8 - CE: 1.176 - CE_v: 1.5717 - acc_v: 0.2063
epoch: 9 - CE: 1.1494 - CE_v: 1.5411 - acc_v: 0.2142
epoch: 10 - CE: 1.1249 - CE_v: 1.512 - acc_v: 0.2205
epoch: 11 - CE: 1.1024 - CE_v: 1.4842 - acc_v: 0.2255
epoch: 12 - CE: 1.0814 - CE_v: 1.4577 - acc_v: 0.2317
epoch: 13 - CE: 1.062 - CE_v: 1.4324 - acc_v: 0.2384
epoch: 14 - CE: 1.0438 - CE_v: 1.4083 - acc_v: 0.2462
epoch: 15 - CE: 1.0268 - CE_v: 1.3852 - acc_v: 0.2556

mape(
    einsum_model.w.clone(),
    torch_model.layer.weight.detach().T
)

2.1556410508718776e-16

mape(
    einsum_model.b.clone(),
    torch_model.layer.bias.detach()
)

7.56436911839403e-17