自适应矩估计优化器 - 从零开始大模型

到目前为止，我们一直使用随机梯度下降（SGD）优化器更新参数。它的算法逻辑非常简单：给梯度乘上一个固定的学习率，就完事了。

但是它也有一个毛病：一刀切。所有参数共用一个学习率，而且从头到尾不变。可是实际上，不同的参数、不同的训练阶段，需要的步幅是不一样的：

有的参数梯度一直很大，有的则一直很小。用同一个学习率，大的容易冲过头，小的又挪不动步；
训练初期通常可以大胆迈步，接近收敛时则需要小心微调。

这一章，我们将引入一个更聪明的优化器：自适应矩估计（Adaptive Moment Estimation，Adam）优化器。它能给每个参数单独地调整有效步幅。简单地讲：就是当梯度长时间偏小的时候，Adam 会适当调大有效步幅；当梯度长时间偏大的时候，Adam 会适当调小有效步幅。Adam 是当前训练神经网络模型最常用的优化器。

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):
        topo = []
        visited = set()
        stack = [(self, False)]

        while stack:
            node, expanded = stack.pop()
            if node in visited:
                continue

            if expanded:
                visited.add(node)
                topo.append(node)
            else:
                stack.append((node, True))
                for p in node.parents:
                    if p not in visited:
                        stack.append((p, False))

        self.grad = np.ones_like(self.data)
        for t in reversed(topo):
            if t.gradient_fn is not None:
                t.gradient_fn()

        for t in topo:
            t.gradient_fn = lambda: None
            t.parents = set()

    def __add__(self, other):
        p = Tensor(self.data + other.data)

        def gradient_fn():
            self.grad += p.grad
            other.grad += p.grad

        p.gradient_fn = gradient_fn
        p.parents = [self, other]
        return p

    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):
        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 Layer(ABC):

    def __call__(self, *args):
        return self.forward(*args)

    @abstractmethod
    def forward(self, *args):
        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 Composite(Layer, ABC):

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

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

顺序层¶

class Sequential(Composite):

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

嵌入层¶

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]

激活函数¶

双曲正切函数（Tanh）¶

class Tanh(Layer):

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

        def gradient_fn():
            x.grad += a.grad * (1 - a.data ** 2)

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

损失函数¶

抽象类¶

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 CELoss(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)

        if y.data.ndim == softmax.ndim:
            labels = y.data
        else:
            labels = np.zeros_like(softmax)
            rows = np.arange(y.data.size)
            labels[rows, y.data] = 1

        log = np.log(np.clip(softmax, 1e-10, 1.0))
        ce = Tensor(-np.mean(np.sum(labels * log, axis=-1)))

        def gradient_fn():
            p.grad += (softmax - labels) / labels.shape[0]

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

优化器¶

在引入新的优化器之前，我们也先提炼出一个优化器的抽象类。

抽象类¶

抽象类定义了两个函数：

step（步进函数）：这是一个抽象函数，根据梯度更新所有参数；
zero_grad（清零函数）：每次迭代前把所有参数的梯度归零。

class Optimizer(ABC):

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

    @abstractmethod
    def step(self):
        pass

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

随机梯度下降优化器¶

同时将现有的随机梯度下降优化器改为从抽象类继承，只需要实现步进函数即可。

class SGDOptimizer(Optimizer):

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

自适应矩估计优化器¶

自适应矩估计优化器名字里的矩（Moment）指的是梯度的两种统计量。Adam 为每个参数单独维护两个矩：

$m$ （一阶矩，First Moment）：梯度的指数加权移动平均（Exponentially Weighted Moving Average，EWMA）。它像是给梯度加了惯性，让更新方向更平稳，不会被单次的随机波动带偏；
$v$ （二阶矩，Second Moment）：梯度平方的指数加权移动平均，反映梯度的大小尺度。对于梯度一直很大的参数， $v$ 就很大，实际步幅会被自动调小；反之，对于梯度一直很小的参数， $v$ 就很小，实际步幅会被自动放大。

指数加权移动平均¶

加权移动平均的目的是让梯度（或者梯度平方）的变化更加平缓，公式为：

m_t = \beta_1 m_{t-1} + (1 - \beta_1)g_t

(1)

v_t = \beta_2 v_{t-1} + (1 - \beta_2){g_t}^2

(2)

这里：

$m_t$ ， $v_t$ ：当前第 $t$ 步的指数加权平均值；
$m_{t-1}$ ， $v_{t-1}$ ：上一步（第 $t-1$ 步）的指数加权平均值；
$g_t$ ， $g_t^2$ ：当前的梯度、梯度平方；
$\beta_1$ ， $\beta_2$ ：梯度的衰减率（超参数，通常为 0.9 和 0.999），决定了过去信息所占的权重。

步幅缩放¶

更新时，用 $m$ 决定方向，再除以 $\sqrt{v}$ 来缩放步幅：

p \;{-}{=}\; lr \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}

(3)

这里：

$p$ ：参数；
$lr$ ：学习率；
$\hat{m}$ ：梯度的指数加权平均值；
$\hat{v}$ ：梯度平方的指数加权平均值；
$\epsilon$ ：极小的数值，防止除以 0。

这里的 $\hat{m}$ 和 $\hat{v}$ 都需要偏差修正（Bias Correction）。

偏差修正¶

在训练初期， $m_0$ 和 $v_0$ 通常被初始化为 0。这会导致在最初的几步中， $m_t$ 和 $v_t$ 的值严重偏小（接近 0）。为了消除这个误差，通常需要进行偏差修正，以 $m$ 为例，公式为：

\hat{m_t} = \frac{m_t}{1 - {\beta_1}^t}

(4)

这里：

$\hat{m_t}$ ：修正后的指数加权平均值；
$m_t$ ：指数加权平均值；
$\beta_1$ ：梯度的衰减率；
$t$ ：优化器已经更新的步数。

当 $t$ 越来越大时， ${\beta_1}^t$ 会趋近于 0，分母趋近于 1，修正公式就会自然失效。

⚠️ 注意：这里的步数 $t$ 是优化器更新参数的次数，和上一章 RNN 里读取一个字符的“时间步”不是一回事。

梯度裁剪¶

自适应矩估计优化器里还有一个梯度裁剪（Gradient Clipping）函数 clip_grad_norm。

RNN 的梯度计算要穿过很多个时间步，偶尔会突然变得极大（梯度爆炸），把参数一下冲飞。梯度裁剪的作用就是：如果所有参数的梯度合在一起的总范数超过了阈值，就整体按比例缩小，把步幅限制在安全范围内。

class AdamOptimizer(Optimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8):
        super().__init__(parameters, lr)
        self.b1, self.b2 = betas
        self.eps = eps
        # 一阶矩
        self.m = [None] * len(parameters)
        # 二阶矩
        self.v = [None] * len(parameters)
        # 更新步数
        self.t = 0

    def step(self):
        self.t += 1
        for i, p in enumerate(self.parameters):
            if p is not None:
                if self.m[i] is None:
                    self.m[i] = np.zeros_like(p.data)
                    self.v[i] = np.zeros_like(p.data)

                # 计算一阶矩
                self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * p.grad
                # 计算二阶矩
                self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (p.grad ** 2)
                # 偏差修正
                m = self.m[i] / (1 - self.b1 ** self.t)
                v = self.v[i] / (1 - self.b2 ** self.t)
                # 权重衰减（AdamW）
                self._apply_weight_decay(p)
                # 更新参数
                p.data -= self.lr * m / (np.sqrt(v) + self.eps)

    # 权重衰减（钩子函数）
    def _apply_weight_decay(self, p):
        pass

    # 梯度裁剪函数
    def clip_grad_norm(self, max_norm=1.0):
        sq = 0.0
        for p in self.parameters:
            sq += float(np.sum(p.grad ** 2))

        if np.sqrt(sq) > max_norm > 0:
            scale = max_norm / (np.sqrt(sq) + 1e-6)
            for p in self.parameters:
                p.grad *= scale

权重衰减自适应矩估计优化器¶

权重衰减自适应矩估计（Adaptive Moment Estimation with Weight Decay，AdamW）优化器是 Adam 的一个小改进，也是训练 GPT 这类大模型的标配。它在 Adam 的基础上加了权重衰减（Weight Decay）：每次更新时，把权重本身也往 0 的方向拉一点点。

这样做的目的是防止过拟合：不让任何权重长得太大，模型就不容易“死记硬背”训练数据的细节。

要注意只对矩阵（ndim >= 2）做衰减，偏置等一维参数不衰减。这也是实践中的经验总结。

class AdamWOptimizer(AdamOptimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
        super().__init__(parameters, lr, betas, eps)
        self.weight_decay = weight_decay

    # 权重衰减
    def _apply_weight_decay(self, p):
        if p.data.ndim >= 2:
            p.data -= p.data * self.weight_decay * self.lr

循环神经网络¶

循环单元¶

class Cell(Composite):

    def __init__(self, in_size, out_size):
        self.input = Linear(in_size, out_size)
        self.hidden = Linear(out_size, out_size)
        self.tanh = Tanh()

        super().__init__([self.input,
                          self.hidden,
                          self.tanh])

    def forward(self, x: Tensor, h: Tensor):
        p = self.input(x)
        h = self.hidden(h)
        return self.tanh(p + h)

循环神经网络模型¶

class RNN(Composite):

    def __init__(self, vocab_size, embedding_size):
        self.embedding_size = embedding_size

        self.embedding = Embedding(vocab_size, embedding_size)
        self.cell = Cell(embedding_size, embedding_size * 2)
        self.output = Linear(embedding_size * 2, vocab_size)

        super().__init__([self.embedding,
                          self.cell,
                          self.output])

    def step(self, token: Tensor, h: Tensor):
        p = self.embedding(token)
        h = self.cell(p, h)
        return self.output(h), h

    def forward(self, x: Tensor, h: Tensor = None):
        batch_size, sequence_len = x.data.shape
        if h is None:
            h = Tensor(np.zeros((batch_size, self.embedding_size * 2)))

        ps = []
        for s in range(sequence_len):
            p, h = self.step(Tensor(x.data[:, s]), h)
            ps.append(p)

        return ps, h

训练器¶

抽象类¶

class Trainer(ABC):

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

    @abstractmethod
    def train(self, dataset, epochs):
        pass

    @abstractmethod
    def test(self, dataset):
        pass

循环神经网络训练器¶

在训练器里，我们在步进函数 step 前先调用梯度裁剪函数 clip_grad_norm，用来防止梯度爆炸。

class RNNTrainer(Trainer):

    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 = Tensor(0.0)
                for t in range(len(prediction)):
                    loss += self.loss_fn(prediction[t], Tensor(label.data[:, t]))
                loss.backward()
                # 梯度裁剪
                self.optimizer.clip_grad_norm()
                self.optimizer.step()

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

        feature, label = dataset.all()
        prediction, _ = self.model(feature)
        loss = Tensor(0.0)
        for t in range(len(prediction)):
            loss += self.loss_fn(prediction[t], Tensor(label.data[:, t]))
        return prediction, loss

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

        h = Tensor(np.zeros((1, self.model.embedding_size * 2)))
        for token in tokens:
            p, h = self.model.step(Tensor([token]), h)

        for _ in range(steps):
            exp = np.exp(p.data[0] - np.max(p.data[0]))
            probs = exp / np.sum(exp)
            token = np.random.choice(len(probs), p=probs)
            tokens.append(token)
            p, h = self.model.step(Tensor([token]), h)

        return dataset.decode(tokens)

超参数¶

学习率¶

LEARNING_RATE = 0.001

批大小¶

BATCH_SIZE = 4

轮数¶

EPOCHS = 2

上下文尺寸¶

CONTEXT_SIZE = 32

嵌入尺寸¶

EMBEDDING_SIZE = 16

建模¶

我们把优化器改成权重衰减自适应矩估计（AdamW）优化器。

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = RNN(dataset.vocab_size, EMBEDDING_SIZE)
loss_fn = CELoss()
# 权重衰减自适应矩估计优化器
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = RNNTrainer(model, loss_fn, optimizer)

训练¶

trainer.train(dataset, EPOCHS)

评估¶

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

prediction:	32 steps, each (6970, 65)
total loss:	Tensor(65.93744561104684)

推理¶

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

ROMEO:
I not beting you at sheenes hew so ene I's have boo?

GRUMPIO:
Westeesilv bees
tha my sapin, it and ofte, couther me it,
Ane sceat, bot, mance?

Prevees live hem?

GOMYe leald het emen, Beran: I but graed cruith, quep a you geed, I
Anthing will sunf a Sor

这一章，我们把优化器从“一刀切”的 SGD 升级到“量身定制”的 AdamW：

Adam：给每个参数单独地调整有效步幅（方向看动量，步幅看尺度）；
梯度裁剪：防止梯度爆炸；
权重衰减：防止过拟合。

💡 提示：总体来说，训练初期 Adam、AdamW 的收敛速度通常比 SGD 更快，但是最终的泛化能力可能不如经过精心调参的 SGD。在 Transformer 类模型上，AdamW 仍然是最常用的选择。

课后练习¶

试试把梯度裁剪函数 clip_grad_norm 的阈值 max_norm 改大或者改小，观察训练过程有什么不同？