循环神经网络 - 从零开始大模型

目前为止，我们的简单文本生成模型在通往真正的语言理解的道路上，还存在两个瓶颈：

无法理解语序：“老鼠爱大米”和“大米爱老鼠”生成的上下文特征是一样的。词语出现的先后顺序这一重要信息被丢失了。
缺乏长程记忆：模型的视野被限制在一个狭小的上下文尺寸之内。超出了这个范围的文字，模型就完全无法建立起词义间的关联。比如这句歌词：“我爱你就像老鼠爱大米”。在上下文尺寸为 5 的情况下，就无法学习到“我爱你”和“老鼠爱大米”之间的关联。

怎样让机器拥有“长期记忆”，并不是一个新问题。

1990 年，语言学家埃尔曼（Elman）提出了简单循环神经网络（Simple RNN）。埃尔曼提出将隐状态（Hidden State）反馈到下一个时刻，这成为了现代循环神经网络（Recurrent Neural Network，RNN）的标准范式：模型不仅看当前输入，也参考过去的自己。

所谓隐状态，就是在网络模型中，输入数据经过层层变换，准备进入最后的输出层时的状态。隐状态保存了网络模型到此刻为止，学习到的全部知识的汇总，就是所谓的时间记忆（Temporal Memory）。

回头看我们的词向量网络模型，有三层：嵌入层、隐藏层和输出层。其中隐藏层（经过激活函数后）的输出就是隐状态。

简单地讲，循环神经网络把读取一个字符的过程称为一个时间步（Time Step），并把上一个时间步进入输出层前的隐状态，连接到当前时间步嵌入层构建的词向量中。这样，一个时间步将包括四个步骤：

嵌入层对输入数据进行初步学习，生成词向量；
连接词向量和上一个时间步的隐状态；
对叠加了时间记忆的词向量进一步学习，生成当前时间步的隐状态；
输出层生成预测值。

最终的输出数据包括预测值和隐状态两个部分。预测值用于损失函数开始梯度计算和反向传播；隐状态则参与到下一个时间步。

循环神经网络示意图：输入特征 x 和上一个时间步的隐状态 h 一起生成新的隐状态 h′，再由隐状态得到预测值 p′

图例：循环神经网络

这里：

$x$ ：（输入）特征；
$h$ ：（上一个时间步）隐状态；
$h^{\prime}$ ：（输出）隐状态；
$p^{\prime}$ ：（输出）预测值。

from abc import ABC, abstractmethod

import numpy as np

np.random.seed(42)

张量¶

反向函数¶

张量的反向函数是反向传播的最主要载体：我们在反向函数中遍历计算图中所有父节点，计算父节点梯度。

我们一直使用非常简洁的递归遍历算法。这在计算图中所有节点排成一条直线或者树状时有效，每个节点只会被访问一次。而循环神经网络模型恰恰打破了这个前提。RNN 里，隐状态 $h$ 既要送进输出层生成预测，又要传给下一个时间步继续更新，于是同一个节点会被两条路径各访问一次；而同一套计算逻辑又被反复调用几十次。结果是计算图里的同一个节点被多次访问，梯度被重复计算。

解决办法是换用拓扑排序算法（Topological Sort）：先把计算图里的所有节点按照依赖顺序排列，保证每个节点都排在它的所有父节点之后。这样我们从这个节点序列的末尾（损失值张量）开始，倒序依次访问每个节点，计算梯度。每个节点被且只被访问一次，梯度计算既不会重复也不会遗漏。

加法¶

RNN 的一个重要步骤就是将上一个时间步的隐状态 $h$ 加到新的特征嵌入向量上。这不是简单地把两个向量相加，而是需要像嵌入层、平均池化层那样处理前向传播中构建计算图，和反向传播中计算梯度的任务。

我们的解决方案是扩展张量的加法运算：在进行加法运算的同时，实现计算结果张量的梯度函数 gradient_fn 和父节点列表 parents。

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

        # 自身初始梯度为 1
        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])

字符编码数据集¶

前面几章里，我们一直采用下一个字符的独热编码作为样本的标签。

而 RNN 是对一个样本逐个字符处理的。例如我们的上下文尺寸是 32，那么每个样本是 32 个字符。RNN 从第一个字符开始，依次进行 32 次预测，每次要带上前一步的隐状态。因此一共预测出了 32 个字符，正好对应后移一位的 32 个字符。

因此，在 RNN 模型里，我们将采用：

$x$ （特征）：tokens[i: i + self.context_size]；
$y$ （标签）：tokens[i + 1: i + self.context_size + 1]。

就是说：如果特征是文本位置 1 到 32 的字符，那么标签就是位置 2 到 33 的字符，分别对应 RNN 32 次预测的结果。

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])
            # 标签：索引编码数组（后移 1 位）
            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]

组合抽象层¶

我们从顺序层抽象出组合抽象层（Composite Layer），只保留子层容器的功能。组合抽象层将成为之后复杂网络层的基础层。

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]

激活函数¶

线性整流函数（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

双曲正切函数（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

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

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

循环神经网络¶

循环单元¶

RNN 的核心是循环单元（Cell）。它只做一件事：把当前读到的字符 $x_t$ 和之前的记忆 $h_{t-1}$ 融合成新的记忆 $h_t$ 。

它内部有三个子层：

input（输入线性层）：处理当前读到的字符 $x_t$ ；
hidden（隐状态线性层）：处理之前的记忆 $h_{t-1}$ ；
tanh（激活函数）：将上面两者相加的结果压缩到 $(-1, 1)$ ，得到新的记忆 $h_t$ 。

用公式表示就是：

h_t = \tanh(W_{in} \, x_t + W_{hidden} \, h_{t-1} + b)

(1)

这里：

$h_{t-1}$ ：上一个时间步的记忆；
$h_t$ ：更新后（当前时间步）的记忆。

代码中两个线性层各有一个偏置，公式里的 $b$ 是两者之和。

记忆就这样一个字符、一个字符地往后传，循环（Recurrent）这个名字也由此而来。

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)

循环神经网络模型¶

把循环单元应用到一个完整的神经网络，就得到了 RNN。它处理一个完整的特征（一段字符序列）的方式是：

embedding（嵌入）：每个字符经过嵌入层转换成嵌入向量；
cell（循环）：和之前记忆（上一步的隐状态）融合，得到新的记忆（这一步的隐状态）；
output（输出）：新的记忆经过输出层，得到下一个字符的预测概率。

推理函数 forward 里的 for 循环，就是把特征的字符序列逐个字符送入同一个循环单元 cell，并把新的记忆传给下一个时间步。这样，每个时间步都在训练同一套循环单元参数。

💡 提示：所以 RNN 的训练过程是 3 层循环：轮次、迭代、时间步。也因此效率不高。

随时间反向传播¶

这种共享同一个单元、随时间步展开的算法结构，在反向传播时也会造成梯度沿着时间步一路回传到最开始，称为随时间反向传播（Backpropagation Through Time，BPTT）。

梯度计算要穿过几十甚至上百个时间步，特别容易造成梯度消失或者梯度爆炸。为此，RNN 训练中还可以采用一种称为 BPTT 截断（Truncated BPTT）的技术：不再让梯度穿过整个序列，而是将序列划分为较短的时间片段，只在有限的时间步内进行反向传播。这样可以降低计算和内存开销，但也意味着模型无法通过反向传播直接学习更远距离的依赖关系。

梯度消失和梯度爆炸是 RNN 难以学习长期依赖关系的重要原因之一。即使使用截断 BPTT，也无法从根本上解决这个问题。这也是后来 LSTM、GRU 等改进型 RNN，以及 Transformer 等架构相继发展的原因之一。

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

循环神经网络训练器¶

训练函数 train：RNN 的随时间步训练和之前的多层神经网络略有不同：模型一次返回的是一串预测（每个时间步一个），而不是一个预测。所以我们要把一次推理中每个时间步的损失累加起来，然后做一次反向传播。

这里我们就用上了张量里扩展的加法 __add__：

loss += self.loss_fn(...)

这样我们在累加每个时间步损失的同时，也将每个时间步的计算图串联起来。通过一次反向传播，梯度就会沿着计算图，穿过每个时间步，把梯度传回每一个参数。

生成函数 generate：生成函数也变成了逐字符处理：

先让模型把提示词 prompt 逐个处理，更新记忆；
根据当前记忆预测下一个字符，再次更新记忆；
重复第二步，并将每一步的预测串联成最终的输出文字。

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]))
                # 反向传播（梯度沿时间步传播，BPTT）
                loss.backward()
                # 更新参数
                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

建模¶

我们已经把 RNN 的所有子层在循环单元 Cell 和网络模型 RNN 中实现，因此只需要创建一个 RNN 的实例即可。

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = RNN(dataset.vocab_size, EMBEDDING_SIZE)
loss_fn = CELoss()
optimizer = SGDOptimizer(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(68.8332931423169)

推理¶

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

ROMEO:
Hat the will you at resters hex, peale Integ stair,
Shate whis. Bute to wain ray batter
Pon's that wigr,
Shist loble it a er? Andake stave ilis muthponl?

TUIEG:
Thald;
And her it fees sRish co'r, forg be cruls fort sad
zould fire I hork have me, me a Pro

这一次，我们终于让网络模型开始读懂词序：通过记忆（隐状态）机制，RNN 一个字符、一个字符地处理输入，每处理一个字符就更新一次记忆。

从上面的推理结果就可以看出，输出结果已经有单词的影子了。我们已经可以辨认出一些单词，这和之前完全是一堆乱字符的情况相比，是很大的进步。也说明 RNN 的确有效。

课后练习¶

把 Cell 里的激活函数从 Tanh 换成 ReLU，再训练一次，观察生成结果有什么不同？结合 Tanh 的输出结果范围 $(-1, 1)$ 考虑一下原因。