交叉熵损失函数 - 从零开始大模型

通过前面两个章节，编码和嵌入，我们解决了让网络模型看懂文字的难题，构建了一个简单的文本生成模型。

文本生成模型本质上是一个多元分类模型，因此我们在网络模型的输出层后使用了 Softmax 激活函数，成功将模型的数值输出转化为概率分布。但在实际训练中，我们并没有看到明显效果。问题的根源不在于 Softmax 激活函数本身，而在于我们衡量误差的方式：均方误差损失函数。

均方误差损失函数是根据预测值和标签的平方误差来衡量模型预测的差距，并据此给出不同的梯度来调整模型参数。损失值越大、梯度也越大。

但是 Softmax 激活函数将实际预测值转换成了相对的概率分布。

当预测值远离标签的时候，Softmax 激活函数的结果非常接近于 0 或者 1，曲线变得非常平缓，因此给出的概率分布变化很小，出现梯度饱和现象。这相当于对于较大的预测误差，Softmax 激活函数有明显的弱化作用。这种情况下，继续使用均方误差损失函数，会造成梯度减弱，训练停滞不前。

为了解决这种“预测越错，修正越小”的逻辑矛盾，我们为分类任务引入一套新的评价标准：交叉熵损失函数（Cross-Entropy Loss）。它能与 Softmax 激活函数完美配合，在模型预测偏差较大时提供巨大的修正动力，来弥补 Softmax 激活函数的弱化效果，而在预测接近正确时则变得温和。

from abc import ABC, abstractmethod

import numpy as np

np.random.seed(42)

张量¶

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})'

数据集¶

抽象类¶

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])

字符编码数据集¶

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)

模型¶

抽象层¶

class Layer(ABC):

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

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

    @property
    def parameters(self):
        return []

线性层¶

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]

顺序层¶

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]

嵌入层¶

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]

平均池化层¶

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

激活函数¶

线性整流函数（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（损失函数）：计算损失值。

class Loss(ABC):

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

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

均方误差损失函数¶

然后将均方误差损失函数修改为从抽象类继承。

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

交叉熵损失函数¶

交叉熵的核心思想是：

预测概率越接近标签，损失值越小；
预测概率越偏离真实标签，损失值急剧增长（对数级增长）。

前向传播：交叉熵损失的定义是：

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

(1)

这里：

$y_i$ ：标签的概率（独热编码）；
$a_i$ ：激活值的类别概率（Softmax 激活函数的输出值）。

这里的 log 运算当概率接近 1 时变化缓慢，而在概率接近 0 时变化剧烈。可以说是 Softmax 激活函数的指数运算的“天敌”。当两者结合进行求导时，原本在 Softmax 饱和区消失的梯度，会被 log 函数重新放大。

梯度计算：交叉熵损失函数的实现，通常将 Softmax 激活函数的计算包含在内。这样做的目的，就是可以大大简化梯度计算，公式为：

\delta = a_i - y_i

(2)

这里：

$y_i$ ：标签的概率（独热编码）；
$a_i$ ：预测值经 Softmax 计算后的类别概率。

代码中的损失值是一批样本的平均值，所以梯度还要再除以批大小。

父节点列表（parents）：（无 Softmax 激活函数的）预测值。

参数列表（parameters）：无。

💡 提示：交叉熵损失函数已经包括 Softmax 激活函数的计算在内。因此，使用的时候就不再需要输出层激活函数。

class CELoss(Loss):

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

        # 创建损失值张量
        log = np.log(np.clip(softmax, 1e-10, 1.0))
        ce = Tensor(-np.mean(np.sum(y.data * log, axis=-1)))

        # 损失值张量的梯度（闭包）函数
        def gradient_fn():
            p.grad += (softmax - y.data) / y.data.shape[0]

        ce.gradient_fn = gradient_fn
        # 损失值张量的父节点列表
        ce.parents = [p]
        return ce

二元交叉熵损失函数¶

在处理二元分类问题时，我们通常只关心一个预测概率。这时我们就不能再使用 Softmax 这样的向量激活函数了，而是使用 Sigmoid 激活函数将输出映射到 $(0, 1)$ 的区间里。

前向传播：二元交叉熵损失的定义是：

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

(3)

这里：

$y$ ：标签的概率（0 或者 1）；
$a$ ：激活值的概率（Sigmoid 激活函数的输出值）。

这个公式衡量预测概率与标签之间的差异，当 $a$ 越接近 $y$ 时损失值越小。

梯度计算：对激活值 $a$ 求导：

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

(4)

这里：

$y$ ：标签的概率（0 或者 1）；
$a$ ：激活值的概率（Sigmoid 激活函数的输出值）。

二元交叉熵损失函数总是和 Sigmoid 激活函数配合使用。这个公式和 Sigmoid 激活函数的梯度公式正好相抵消，从而解决 Sigmoid 激活函数在遇到大的预测值时梯度饱和的问题。

父节点列表（parents）：激活值。

参数列表（parameters）：无。

💡 提示：二元交叉熵损失函数不包括 Sigmoid 激活函数的计算在内。因此，使用的时候仍然需要输出层激活函数。

class BCELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        # 裁剪数值，防止溢出
        clipped = np.clip(p.data, 1e-7, 1 - 1e-7)
        # 创建损失值张量
        bce = Tensor(-np.mean(y.data * np.log(clipped) + (1 - y.data) * np.log(1 - clipped)))

        # 损失值张量的梯度（闭包）函数
        def gradient_fn():
            p.grad += (clipped - y.data) / (clipped * (1 - clipped)) / y.data.size

        bce.gradient_fn = gradient_fn
        # 损失值张量的父节点列表
        bce.parents = [p]
        return bce

优化器（随机梯度下降）¶

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

训练器¶

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)

超参数¶

学习率¶

交叉熵损失函数对大的梯度变化比较敏感，容易造成震荡、发散或者不收敛的情况。我们可以选择使用小一些的学习率。

LEARNING_RATE = 0.001

批大小¶

BATCH_SIZE = 4

轮数¶

EPOCHS = 10

上下文尺寸¶

CONTEXT_SIZE = 32

嵌入尺寸¶

EMBEDDING_SIZE = 16

建模¶

让我们尝试使用一下交叉熵损失函数，看看训练效果会如何。

这里要记得去掉输出层的 Softmax 激活函数，因为我们的交叉熵损失函数已经包含 Softmax 激活函数的计算了。

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = Sequential([
    # 嵌入层
    Embedding(dataset.vocab_size, EMBEDDING_SIZE),
    # 平均池化层
    MeanPool(),
    # 隐藏层
    Linear(EMBEDDING_SIZE, EMBEDDING_SIZE * 2),
    # 隐藏层激活函数
    ReLU(),
    # 输出层
    Linear(EMBEDDING_SIZE * 2, dataset.vocab_size)
])
# 交叉熵损失函数
loss_fn = CELoss()
optimizer = SGDOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = Trainer(model, loss_fn, optimizer)

训练¶

trainer.train(dataset, EPOCHS)

评估¶

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

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

我们并不能直接和上一章对比损失值，因为使用的是不同的损失函数。

推理¶

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

换了新的交叉熵损失函数后，从模型生成的文本并不能看出太大的区别。这是因为我们还有一个重要的问题没有解决：就是我们使用词袋或者平均池化来构建特征，忽略了文本的词序。

课后练习¶

交叉熵损失函数已经包括输出层 Softmax 激活函数，试一试如果在模型的输出层后仍然加上了 Softmax 激活函数，损失值和最后的推理结果会有怎样的表现？