Cross-entropy loss function - From Neuron to GPT

In the previous two chapters, encoding and embedding, we solved the problem of making a network model understand text, and built a simple text generation model.

A text generation model is essentially a multi-class classification model, so we put a Softmax activation function after the output layer of the network, which successfully turned the model’s numerical output into a probability distribution. But in actual training we did not see a clear effect. The root of the problem is not the Softmax activation function itself, but the way we measure error: the mean squared error loss function.

The mean squared error loss function measures how far the model’s prediction is from the target by the squared error between the prediction and the label, and gives different gradients accordingly to adjust the model parameters. The larger the loss, the larger the gradient.

But the Softmax activation function turns the raw predictions into a relative probability distribution.

When a prediction is far from the label, the result of Softmax is very close to 0 or 1 and its curve becomes very flat, so the probability distribution changes very little. This is gradient saturation. It means that for large prediction errors, Softmax has a clear weakening effect. If we keep using the mean squared error loss function in this situation, the gradient is weakened and training stalls.

To resolve the logical contradiction of “the worse the prediction, the smaller the correction”, we introduce a new evaluation criterion for classification tasks: the cross-entropy loss function. It works perfectly with the Softmax activation function. When the model’s prediction is far off, it provides a strong corrective push that makes up for the weakening effect of Softmax, and it becomes gentle when the prediction is close to correct.

from abc import ABC, abstractmethod

import numpy as np

np.random.seed(42)

Tensor¶

class Tensor:

    def __init__(self, data):
        self.data = np.array(data)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = set()

    def backward(self):
        if self.gradient_fn is not None:
            self.gradient_fn()

        for p in self.parents:
            p.backward()

    def __str__(self):
        return f'Tensor({self.data})'

Dataset¶

Abstract class¶

class Dataset(ABC):

    def __init__(self, batch_size=1):
        self.batch_size = batch_size
        self.load()
        self.train()

    @abstractmethod
    def load(self):
        pass

    def train(self):
        self.data = self.train_data

    def eval(self):
        self.data = self.test_data

    def all(self):
        x, y = self.data
        return Tensor(x), Tensor(y)

    def __len__(self):
        x, *_ = self.data
        return len(x) // self.batch_size

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y = self.data
        return Tensor(x[s]), Tensor(y[s])

Character-encoding dataset¶

class CharDataset(Dataset):

    def __init__(self, filename, batch_size=1, context_size=64, stride=None, split=0.9):
        self.filename = filename
        self.context_size = context_size
        self.stride = stride if stride is not None else context_size // 2
        self.split = split
        super().__init__(batch_size)

    def load(self):
        with open(self.filename, encoding="utf-8") as f:
            text = f.read()

        self.vocab = sorted(set(text))
        self.vocab_size = len(self.vocab)
        self.stoi = {ch: i for i, ch in enumerate(self.vocab)}
        self.itos = {i: ch for i, ch in enumerate(self.vocab)}
        self.tokens = self.encode(text)

        split = int(len(self.tokens) * self.split)
        self.train_data = self._pack(self.tokens[:split])
        self.test_data = self._pack(self.tokens[split:])

    def _pack(self, tokens):
        onehot = np.eye(self.vocab_size)
        x, y = [], []
        for i in range(0, len(tokens) - self.context_size - 1, self.stride):
            x.append(tokens[i: i + self.context_size])
            y.append(onehot[tokens[i + self.context_size]])
        return x, y

    def encode(self, symbols):
        return [self.stoi[s] for s in symbols]

    def decode(self, tokens):
        return "".join(self.itos[t] for t in tokens)

Model¶

Abstract layer¶

class Layer(ABC):

    def __call__(self, x: Tensor):
        return self.forward(x)

    @abstractmethod
    def forward(self, x: Tensor):
        pass

    @property
    def parameters(self):
        return []

Linear layer¶

class Linear(Layer):

    def __init__(self, in_size, out_size):
        self.weight = Tensor(np.random.randn(out_size, in_size) * np.sqrt(2 / in_size))
        self.bias = Tensor(np.zeros(out_size))

    def forward(self, x: Tensor):
        p = Tensor(x.data @ self.weight.data.T + self.bias.data)

        def gradient_fn():
            self.weight.grad += p.grad.T @ x.data
            self.bias.grad += np.sum(p.grad, axis=0)
            x.grad += p.grad @ self.weight.data

        p.gradient_fn = gradient_fn
        p.parents = [x]
        return p

    @property
    def parameters(self):
        return [self.weight, self.bias]

Sequential layer¶

class Sequential(Layer):

    def __init__(self, layers):
        self.layers = layers

    def forward(self, x: Tensor):
        for l in self.layers:
            x = l(x)
        return x

    @property
    def parameters(self):
        return [p for l in self.layers for p in l.parameters]

Embedding layer¶

class Embedding(Layer):

    def __init__(self, vocab_size, embedding_size, std=0.02):
        super().__init__()
        self.weight = Tensor(np.random.randn(vocab_size, embedding_size) * std)

    def forward(self, x: Tensor):
        p = Tensor(self.weight.data[x.data])

        def gradient_fn():
            np.add.at(self.weight.grad, x.data, p.grad)

        p.gradient_fn = gradient_fn
        p.parents = [self.weight]
        return p

    @property
    def parameters(self):
        return [self.weight]

Mean pooling layer¶

class MeanPool(Layer):

    def forward(self, x: Tensor):
        p = Tensor(np.mean(x.data, axis=1))

        def gradient_fn():
            x.grad += p.grad[:, None, :] / x.data.shape[1]

        p.gradient_fn = gradient_fn
        p.parents = [x]
        return p

Activation functions¶

Rectified linear unit (ReLU)¶

class ReLU(Layer):

    def forward(self, x: Tensor):
        a = Tensor(np.maximum(0, x.data))

        def gradient_fn():
            x.grad += a.grad * (a.data > 0)

        a.gradient_fn = gradient_fn
        a.parents = [x]
        return a

Softmax¶

class Softmax(Layer):

    def __init__(self, axis=-1):
        super().__init__()
        self.axis = axis

    def forward(self, x: Tensor):
        exp = np.exp(x.data - np.max(x.data, axis=self.axis, keepdims=True))
        a = Tensor(exp / np.sum(exp, axis=self.axis, keepdims=True))

        def gradient_fn():
            grad = np.sum(a.data * a.grad, axis=self.axis, keepdims=True)
            x.grad += a.data * (a.grad - grad)

        a.gradient_fn = gradient_fn
        a.parents = [x]
        return a

Loss functions¶

We first extract an abstract class for loss functions, which makes it easy to add new ones.

Abstract class¶

The abstract class defines only one abstract method:

loss (loss method): computes the loss.

class Loss(ABC):

    def __call__(self, p: Tensor, y: Tensor):
        return self.loss(p, y)

    @abstractmethod
    def loss(self, p: Tensor, y: Tensor):
        pass

Mean squared error loss function¶

Then we change the mean squared error loss function to inherit from the abstract class.

class MSELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        mse = Tensor(np.mean(np.square(y.data - p.data)))

        def gradient_fn():
            p.grad += -2 * (y.data - p.data) / y.data.size

        mse.gradient_fn = gradient_fn
        mse.parents = [p]
        return mse

Cross-entropy loss function¶

The core idea of cross-entropy is:

the closer the predicted probability is to the label, the smaller the loss;
the further the predicted probability is from the true label, the loss grows sharply (logarithmically).

Forward propagation: the cross-entropy loss is defined as:

L = -\sum_{i=1}^{n} y_i \log(a_i)

(1)

Where:

$y_i$ : the probability of the label (one-hot encoding);
$a_i$ : the class probability of the activation (the output of the Softmax activation function).

The log function changes slowly when the probability is close to 1 and sharply when it is close to 0. It is the natural enemy of the exponential in Softmax. When the two are differentiated together, the gradient that vanished in the saturated region of Softmax is amplified again by the log function.

Gradient computation: implementations of the cross-entropy loss function usually include the Softmax computation. This greatly simplifies the gradient computation, and the formula is:

\delta = a_i - y_i

(2)

Where:

$y_i$ : the probability of the label (one-hot encoding);
$a_i$ : the class probability of the prediction after Softmax.

In the code, the loss is the average over a batch of samples, so the gradient is also divided by the batch size.

Parent list (parents): the prediction (without the Softmax activation function).

Parameter list (parameters): none.

💡 Tip: The cross-entropy loss function already includes the Softmax computation. So there is no need for an output layer activation function when you use it.

class CELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        # Built-in Softmax
        exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        # Create the loss tensor
        log = np.log(np.clip(softmax, 1e-10, 1.0))
        ce = Tensor(-np.mean(np.sum(y.data * log, axis=-1)))

        # Gradient (closure) function of the loss tensor
        def gradient_fn():
            p.grad += (softmax - y.data) / y.data.shape[0]

        ce.gradient_fn = gradient_fn
        # Parent list of the loss tensor
        ce.parents = [p]
        return ce

Binary cross-entropy loss function¶

In a binary classification problem, we usually care about only one predicted probability. We can no longer use a vector activation function such as Softmax; instead we use the Sigmoid activation function to map the output into the interval $(0, 1)$ .

Forward propagation: the binary cross-entropy loss is defined as:

L = -(y \log(a) + (1 - y) \log(1 - a))

(3)

Where:

$y$ : the probability of the label (0 or 1);
$a$ : the probability of the activation (the output of the Sigmoid activation function).

This formula measures the difference between the predicted probability and the label. The closer $a$ is to $y$ , the smaller the loss.

Gradient computation: the derivative with respect to the activation $a$ :

\delta = \frac{a - y}{a (1 - a)}

(4)

Where:

$y$ : the probability of the label (0 or 1);
$a$ : the probability of the activation (the output of the Sigmoid activation function).

The binary cross-entropy loss function is always used together with the Sigmoid activation function. This formula exactly cancels the gradient formula of Sigmoid, which solves the gradient saturation of Sigmoid for large predicted values.

Parent list (parents): the activation.

Parameter list (parameters): none.

💡 Tip: The binary cross-entropy loss function does not include the Sigmoid computation. So you still need an output layer activation function when you use it.

class BCELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        # Clip the values to prevent overflow
        clipped = np.clip(p.data, 1e-7, 1 - 1e-7)
        # Create the loss tensor
        bce = Tensor(-np.mean(y.data * np.log(clipped) + (1 - y.data) * np.log(1 - clipped)))

        # Gradient (closure) function of the loss tensor
        def gradient_fn():
            p.grad += (clipped - y.data) / (clipped * (1 - clipped)) / y.data.size

        bce.gradient_fn = gradient_fn
        # Parent list of the loss tensor
        bce.parents = [p]
        return bce

Optimizer (stochastic gradient descent)¶

class SGDOptimizer:

    def __init__(self, parameters, lr):
        self.parameters = parameters
        self.lr = lr

    def zero_grad(self):
        for p in self.parameters:
            p.grad = np.zeros_like(p.data)

    def step(self):
        for p in self.parameters:
            p.data -= p.grad * self.lr

Trainer¶

class Trainer:

    def __init__(self, model, loss_fn, optimizer):
        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    def train(self, dataset, epochs):
        dataset.train()

        for epoch in range(epochs):
            for i in range(len(dataset)):
                feature, label = dataset[i]

                self.optimizer.zero_grad()
                prediction = self.model(feature)
                loss = self.loss_fn(prediction, label)
                loss.backward()
                self.optimizer.step()

    def test(self, dataset):
        dataset.eval()

        feature, label = dataset.all()
        prediction = self.model(feature)
        loss = self.loss_fn(prediction, label)
        return prediction, loss

    def generate(self, dataset, prompt, steps=256):
        tokens = dataset.encode(prompt)

        for _ in range(steps):
            window = tokens[-dataset.context_size:]
            feature = Tensor([window])
            prediction = self.model(feature)

            exp = np.exp(prediction.data[0] - np.max(prediction.data[0]))
            probs = exp / np.sum(exp)
            token = np.random.choice(len(probs), p=probs)
            tokens.append(token)

        return dataset.decode(tokens)

Hyperparameters¶

Learning rate¶

The cross-entropy loss function is sensitive to large changes in the gradient, which can easily cause oscillation, divergence, or failure to converge. We can choose a smaller learning rate.

LEARNING_RATE = 0.001

Batch size¶

BATCH_SIZE = 4

Number of epochs¶

EPOCHS = 10

Context size¶

CONTEXT_SIZE = 32

Embedding size¶

EMBEDDING_SIZE = 16

Modeling¶

Let’s try the cross-entropy loss function and see how training goes.

Remember to remove the Softmax activation function after the output layer, because our cross-entropy loss function already includes the Softmax computation.

DATA_FILE = "../../../../data/tinyshakespeare.txt"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = Sequential([
    # Embedding layer
    Embedding(dataset.vocab_size, EMBEDDING_SIZE),
    # Mean pooling layer
    MeanPool(),
    # Hidden layer
    Linear(EMBEDDING_SIZE, EMBEDDING_SIZE * 2),
    # Hidden layer activation function
    ReLU(),
    # Output layer
    Linear(EMBEDDING_SIZE * 2, dataset.vocab_size)
])
# Cross-entropy loss function
loss_fn = CELoss()
optimizer = SGDOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = Trainer(model, loss_fn, optimizer)

Training¶

trainer.train(dataset, EPOCHS)

Evaluation¶

prediction, loss = trainer.test(dataset)
print(f'prediction:\t{prediction.data.shape}\nloss:\t{loss}')

prediction:	(6970, 65)
loss:	Tensor(3.3545101560265147)

We cannot compare the loss directly with the previous chapter, because a different loss function is used.

Inference¶

print(trainer.generate(dataset, prompt="ROMEO:"))

ROMEO:taots uaa e AVrsQsso nunaelf Urs uamoruUs nomAs  oesc sh ogiesso snnbto et vnlLEuSpd 
hwaaagetff yA middp rldeAofUiyeta 'd,
tso   ote fablo 
y:   draengss
peyeyoAhys na
d:rar
R:-  hJdasteY u OOW p l  t&rh uEs s
gottrt,PRcdh n afrol h bd Iuua  du o wdpal' u

With the new cross-entropy loss function, the generated text does not look very different. This is because an important problem is still unsolved: we build the features with a bag of words or mean pooling, which ignores the word order of the text.

Exercises¶

The cross-entropy loss function already includes the Softmax of the output layer. Try adding a Softmax activation function after the output layer anyway, and see how the loss and the final inference result behave.