随机失活 - 从零开始大模型

上一章，我们给 GPT 模型加装了层归一化，保证了残差连接后的数值稳定和训练稳定。同时也提到了大规模训练的一个可能问题：过拟合（Overfitting）。

模型在训练集上表现很好，在测试集上却可能很差。模型记住了过多的训练数据的细节，而不是真正地学会了语言规律。就像一个学生只是死记硬背了所有的练习题，考试时遇到新题就傻眼了。

随机失活（Dropout）是防止过拟合的经典技术：训练时，随机丢弃一部分神经元的输出（设为 0），让模型不能过度依赖任何单个神经元：

训练时：随机将一定比例 $p$ （比如 0.1）的神经元的输出设为 0，剩下的输出除以 $1 - p$ 来保持期望值不变；
推理时：不丢弃任何神经元，使用全部输出。

加入随机失活的 GPT 基础模块示意图：随机失活位于嵌入层之后，以及注意力层和前馈网络层的残差连接之前

图例：加入随机失活的 GPT 基础模块

import numpy as np

from src.core import *

np.random.seed(42)

数据集¶

字符编码数据集¶

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):
        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(tokens[i + 1: i + self.context_size + 1])
        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 Tril(Layer):

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

    def forward(self, x: Tensor):
        keep = np.tril(np.ones(x.shape[-2:]))
        p = Tensor(np.where(keep, x.data, self.value))

        def gradient_fn():
            x.grad += p.grad * keep

        return p.attach(gradient_fn, {x})

归一化层¶

class LayerNorm(Layer):

    def __init__(self, normalized_size, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = Tensor(np.ones(normalized_size, dtype=DTYPE))
        self.bias = Tensor(np.zeros(normalized_size, dtype=DTYPE))

    def forward(self, x: Tensor):
        mean = np.mean(x.data, axis=-1, keepdims=True)
        var = np.var(x.data, axis=-1, keepdims=True, ddof=0)
        norm = (x.data - mean) / np.sqrt(var + self.eps)
        p = Tensor(self.weight.data * norm + self.bias.data)

        def gradient_fn():
            axis = tuple(range(p.grad.ndim - 1)) if p.grad.ndim > 1 else None
            self.weight.grad += np.sum(p.grad * norm, axis=axis)
            self.bias.grad += np.sum(p.grad, axis=axis)
            grad = p.grad * self.weight.data
            grad_mean = np.mean(grad, axis=-1, keepdims=True)
            norm_mean = np.mean(grad * norm, axis=-1, keepdims=True)
            x.grad += (grad - grad_mean - norm * norm_mean) / np.sqrt(var + self.eps)

        return p.attach(gradient_fn, {x})

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

随机失活层¶

随机失活层（Dropout 层）在训练时随机丢弃一部分神经元的输出：

前向传播：
- 训练时：生成一个和输入值格式相同的随机掩码，以概率 $1 - p$ 保留（缩放因子为 $1 / (1 - p)$ ），以概率 $p$ 丢弃（值为 0）；
- 推理时：直接返回输入值；
反向传播：梯度只通过被保留的神经元回传，被丢弃的神经元梯度为 0。

💡 提示：将保留的输入值除以 $1 - p$ ，是为了保持输出值的整体期望值不变，这样训练时和推理时的输出尺度一致。

class Dropout(Layer):

    def __init__(self, prob=0.1):
        super().__init__()
        self.prob = prob

    def forward(self, x: Tensor):
        # 推理时不失活
        if not self.training or self.prob == 0:
            return x

        # 随机掩码
        keep_prob = 1.0 - self.prob
        mask = (np.random.rand(*x.shape) < keep_prob).astype(DTYPE) / keep_prob
        # 创建输出值张量
        p = Tensor(x.data * mask)

        # 输出值张量的梯度（闭包）函数
        def gradient_fn():
            x.grad += p.grad * mask

        # 设置输出值张量梯度（闭包）函数和父节点列表
        return p.attach(gradient_fn, {x})

激活函数¶

高斯误差线性单元激活函数（GELU）¶

class GELU(Layer):

    def __init__(self):
        super().__init__()
        self.c = np.sqrt(2.0 / np.pi)

    def forward(self, x: Tensor):
        tanh = np.tanh(self.c * (x.data + 0.044715 * x.data ** 3))
        a = Tensor(0.5 * x.data * (1.0 + tanh))

        def gradient_fn():
            grad = 0.5 * (1.0 + tanh) + 0.5 * x.data * (1.0 - tanh ** 2) * self.c * (1.0 + 3.0 * 0.044715 * x.data ** 2)
            x.grad += a.grad * grad

        return a.attach(gradient_fn, parents=[x])

损失函数¶

GPT 交叉熵损失函数¶

class GPTLoss(Loss):

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

        flat_softmax = softmax.reshape(-1, softmax.shape[-1])
        flat_y = y.data.reshape(-1).astype(np.int64)
        n = len(flat_y)
        rows = np.arange(n)

        log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
        ce = Tensor(0 - np.sum(log) / n)

        def gradient_fn():
            flat_grad = flat_softmax.copy()
            flat_grad[rows, flat_y] -= 1
            p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n

        return ce.attach(gradient_fn, parents=[p])

模型¶

GPT 嵌入层¶

在词向量和位置向量合并以后，添加随机失活。

class GPTEmbedding(Composite):

    def __init__(self, vocab_size, context_size, embedding_size):
        self.embedding = Embedding(vocab_size, embedding_size)
        self.positional_embedding = Embedding(context_size, embedding_size)
        # 随机失活层
        self.dropout = Dropout()

        super().__init__([self.embedding,
                          self.positional_embedding,
                          self.dropout])

    def forward(self, x: Tensor):
        token = self.embedding(x)
        position = self.positional_embedding(Tensor(range(x.shape[1])))
        # 随机失活
        return self.dropout(token + position)

GPT 注意力层¶

在注意力层，随机失活添加在残差连接前面。

class GPTAttention(Composite):

    def __init__(self, embedding_size, heads=1):
        self.embedding_size = embedding_size
        self.heads = heads

        self.normalize = LayerNorm(embedding_size)
        self.query = Linear(embedding_size, embedding_size)
        self.key = Linear(embedding_size, embedding_size)
        self.value = Linear(embedding_size, embedding_size)
        self.mask = Tril()
        self.softmax = Softmax()
        self.output = Linear(embedding_size, embedding_size)
        # 随机失活层
        self.dropout = Dropout()

        super().__init__([self.normalize,
                          self.query,
                          self.key,
                          self.value,
                          self.mask,
                          self.softmax,
                          self.output,
                          self.dropout])

    def forward(self, x: Tensor):
        norm = self.normalize(x)
        head_embedding_size = self.embedding_size // self.heads
        multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
        shape = (-1, x.shape[1], self.embedding_size)

        query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
        value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
        scores = query @ key * scale
        weights = self.softmax(self.mask(scores))
        # 随机失活，然后残差连接
        return x + self.dropout(self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape)))

GPT 前馈网络层¶

同样的，在前馈网络层，随机失活也添加在残差连接前面。

class GPTFeedForward(Composite):

    def __init__(self, embedding_size):
        self.normalize = LayerNorm(embedding_size)
        self.input = Linear(embedding_size, embedding_size * 4)
        self.gelu = GELU()
        self.output = Linear(embedding_size * 4, embedding_size)
        # 随机失活层
        self.dropout = Dropout()

        super().__init__([self.normalize,
                          self.input,
                          self.gelu,
                          self.output,
                          self.dropout])

    def forward(self, x: Tensor):
        norm = self.normalize(x)
        h = self.gelu(self.input(norm))
        # 随机失活，然后残差连接
        return x + self.dropout(self.output(h))

GPT Transformer 模块¶

class GPTTransformer(Composite):

    def __init__(self, embedding_size, heads):
        self.attention = GPTAttention(embedding_size, heads)
        self.feed_forward = GPTFeedForward(embedding_size)

        super().__init__([self.attention,
                          self.feed_forward])

    def forward(self, x: Tensor):
        x = self.attention(x)
        return self.feed_forward(x)

GPT 输出层¶

class GPTOutput(Composite):

    def __init__(self, embedding_size, vocab_size):
        self.normalize = LayerNorm(embedding_size)
        self.output = Linear(embedding_size, vocab_size)

        super().__init__([self.normalize,
                          self.output])

    def forward(self, x: Tensor):
        norm = self.normalize(x)
        return self.output(norm)

GPT 模型¶

class GPT(Composite):

    def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
        self.embedding = GPTEmbedding(vocab_size, context_size, embedding_size)
        self.transformers = [GPTTransformer(embedding_size, heads) for _ in range(blocks)]
        self.output = GPTOutput(embedding_size, vocab_size)

        super().__init__([self.embedding] + self.transformers + [self.output])

    def forward(self, x: Tensor):
        x = self.embedding(x)
        for layer in self.transformers:
            x = layer(x)
        return self.output(x)

GPT 训练器¶

class GPTTrainer(Trainer):

    def train(self, dataset, epochs, scheduler=None):
        dataset.train()
        self.model.train()

        steps = 0
        for epoch in range(epochs):
            for i in range(len(dataset)):
                if scheduler is not None:
                    self.optimizer.lr = scheduler.step(steps)

                feature, label = dataset[i]

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

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

        predictions = []
        total_loss = 0.0
        with Tensor.no_grad():
            for i in range(len(dataset)):
                feature, label = dataset[i]
                prediction = self.model(feature)
                loss = self.loss_fn(prediction, label)
                predictions.append(prediction)
                total_loss += float(loss.data)

        loss = total_loss / len(dataset)
        dataset.train()
        self.model.train()
        return predictions, loss

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

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

                last_logits = logits.data[0, -1]
                exp = np.exp(last_logits - np.max(last_logits))
                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 = 2

上下文尺寸¶

CONTEXT_SIZE = 32

嵌入尺寸¶

EMBEDDING_SIZE = 64

头数¶

HEADS = 2

Transformer 模块数¶

BLOCKS = 2

建模¶

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = GPTTrainer(model, loss_fn, optimizer)

训练¶

scheduler = WarmupCosineScheduler(LEARNING_RATE, EPOCHS * len(dataset), 100, LEARNING_RATE / 10)
trainer.train(dataset, EPOCHS, scheduler)

评估¶

prediction, loss = trainer.test(dataset)
print(f'prediction:\t{len(prediction)} steps, each {prediction[0].data.shape}\nloss:\t{loss}')

prediction:	1742 steps, each (4, 32, 65)
loss:	1.8636507043608568

推理¶

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

ROMEO:
O:
O deseep would my lethis this they-anio, you hale should, gentis,
I'll I
Boy, this I have of del mogriess.

Secondself faic, not lifect were Clains; i' be So thour man I complo, and do hand
Mistrue such to which a city his old vice frieunts not man,
An

本章，我们给 GPT 模型添加了随机失活：在训练时随机丢弃一部分神经元的输出，防止过拟合。

至此，预训练的三大工程优化技术就配齐了：

残差连接：让梯度顺畅无阻，防止梯度消失；
层归一化：让数值稳定，训练稳定；
随机失活：防止过拟合。

有了这些工程优化技术，模型可以堆得更高、训练得更稳、泛化更好。

课后练习¶

试着把 Dropout 的概率 $p$ 从 0.1 改成 0.3 甚至 0.5，观察训练效果如何？概率是不是越大越好？