从零训练GPT

学完 Andrej Karpathy 的《Let's build GPT》和 3B1B 的《Visualizing Attention》，茅塞顿开。

但是我感觉自己并没有理解透彻，因此在本文中，我对两个教程的内容进行综合，同时加上我自己的理解。

在本文中，我以《从零训练GPT》作为主题，进行串讲。

语料准备

首先，第一步是训练语料准备。《Let's build GPT》采用 Tiny Shakespeare 语料集。这个语料集有点无聊，我打算使用中文语料进行训练。

加载语料代码如下：

# 加载语料数据集
with open('path/to/my/corpus.txt', 'r', encoding='utf-8') as f:
    text = f.read()
print('语料集的长度为：', len(text))
print('打印前100个字符：', text[:100])

Tokenizer

有了语料之后，接下来需要对其进行编码（Tokenize），将自然语言转为机器可以理解的形式。

在《Let's build GPT》中，由于内容都是英文字母加标点，因此采用统计方式，统计出 Tiny Shakespeare 中所有出现的字符种类，建了一个字典。然后，以字符为单位，以字符在字典中的 index 作为 Token。

这种做法，对于教学演示来说，是直观的。但是在实际中不会采用这种方式。

在实际中，自然语言不限于英文字母，比如中文。同时，以字符为单位效率也太多。

业界的 Tokenizer 会采用拆分子词的方式，进行 Token 转化，实现更高的效率、更好的效果。

那 Tokenizer 又是怎么来的呢？Tokenizer 也需要经过训练，但并非深度学习训练，而是基于频率计数的方式，对高频词素进行识别、聚类。经典的分词算法包括 byte-pair encoding。

在这里，我使用 OpenAI 的 tiktoken。这是 OpenAI 开源的 BPE 分词器，GPT-4 同款哦！

使用 tiktoken 将书记的文本转为 Token 列表：

import tiktoken
enc = tiktoken.get_encoding('cl100k_base')
data = torch.tensor(enc.encode(text), dtype=torch.long).to('cuda')
print('编码后的语料集长度为：', len(data))
print('打印前100个字符：', data[:100])

# Let's now split up the data into train and validation sets
n = int(0.9*len(text)) # first 90% will be train, rest val
train_data = data[:n]
val_data = data[n:]

这段代码将文本转换为 token 序列，并进一步转为 PyTorch 的长整型张量，存储在 GPU 上，这样能充分利用 GPU 加速训练。

同时，在上面代码中，还完成了训练集、验证集的拆分：

训练集：90% 的数据用于 GPT 训练
验证集：10% 数据不参与训练，用于衡量 GPT 的训练效果

这样，在训练时，每隔一定的训练周期，可以用验证集对模型误差进行验证，以更加客观得了解训练效果。

超参数

代码中还涉及到一系列超参数，在本节中统一给出。这些参数决定了我们的模型规模、以及训练参数。

# 超参数

# tiktoken 词表大小
vocab_size = enc.n_vocab # 词表大小
print('词表大小为：', vocab_size)

# 词嵌入空间的维度
n_embd = 128

# dropout概率
# 在训练时随机屏蔽部分神经网络
# 能够提升训练效果
dropout = 0.2

# 一次训练迭代中用于更新模型权重的样本数量。
batch_size = 32
# 模型预测时可以考虑的最大上下文长度。
block_size = 128

# 多头注意力头数
n_head = 16
# Transformer 模型中 Block 数量
n_layer = 16

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 每次训练多少个迭代
max_iters = 500
# 每隔多少个迭代，用验证集验证一次
eval_interval = 200
eval_iters = 50

# 学习率
learning_rate = 1e-3

词嵌入层

输入的 token 序列首先会经过词嵌入层，映射为一个固定维度的稠密向量，这一步称为 Embedding。

词向量与 token 有什么区别？词向量是嵌入空间中的一个向量，这个向量捕捉了 token 的语义信息。

嵌入空间，或者称为语义空间，让我非常着迷，它是意思的空间。

我将嵌入空间的维度设置为 128 维（n_embd）。

在这个空间内，每个向量，都将表示一个含义，通俗来说，就是一个意思。

在 Transformer Module 中，首先实现 Embedding 部分，这里直接使用 PyTorch 的 torch.nn.Embedding。

class Transformer(nn.Module):
    def __init__(self):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx and targets are both (B,T) tensor of integers
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
        x = tok_emb + pos_emb # (B,T,C)

理解 Embedding

torch.nn.Embedding 其实本质上就是一个查找表(lookup table)，或者说它是一个大矩阵。

矩阵的每一行对应一个离散 token（如词或字符）的嵌入向量。

具体来说，假设词表大小为 V，嵌入向量维度为 D，那么 Embedding 层就维护一个大小为V x D的矩阵(权重矩阵)。

输入给 Embedding 层的是一个大小为 N 的索引数组，其中每个元素是一个取值范围为 0 到 V-1 的整数，表示词表中的某个 token。

Embedding层读取权重矩阵的对应行,返回一个N x D 的嵌入向量数组。

用数学表达式可以写作: $e_{i} = W^{T} o_{i}$

其中 $o_{i}$ 是第i个输入 token 的 one-hot 向量, $W$ 是 Embedding 矩阵, $e_{i}$ 是第i个token的嵌入向量。

所以 torch.nn.Embedding 实现的词嵌入本质上是对一个 V x D 矩阵的查表操作。其好处为：

灵活性：允许每个token单独学习自己的表示，不同词的相似度可以直接体现在嵌入向量的距离上。
参数效率：虽然维护一个大矩阵，但是每次前向运算只需要简单查表，没有复杂的矩阵乘法运算。
可解释性：学习到的 Embedding 矩阵可以单独抽出来分析每个词的密集表示，做聚类、可视化等后续处理。

对 Embedding 进行训练

Embedding 层的目的是为每个离散的 token 学习一个连续的、密集的向量表示。这个表示应该能够体现 token 的语义信息。

换句话说，我们希望语义相近的词（如"国王"和"王后"）有相近的嵌入向量，而语义无关的词（如"国王"和"足球"）的嵌入向量应该相差较大。

那么，Embedding 矩阵是如何被训练的呢？关键在于[000.wiki/反向传播|反向传播]。

具体来说，Embedding层通常会被放在神经网络的底层。上层网络基于 Embedding 层的输出做一些操作（如 Transformer 中的自注意力和前馈计算），然后基于任务（如语言建模）计算损失函数。接下来通过反向传播计算损失函数对每个参数的梯度，并用优化器（optimizer）更新所有参数，包括 Embedding 矩阵在内。

举例来说，在 Transformer 的语言模型任务中，我们希望根据前面的词预测下一个词。如果模型预测错了，就会产生较大的损失。这个损失不仅会传导到 Transformer Block 中的参数,也会传导到底层的 Embedding 矩阵。

如果某两个词（如"国王"和"王后"）经常在相似的上下文中出现，那么语言模型就应该给它们相近的预测概率。这就要求它们的嵌入向量足够接近，使得Transformer的上层计算（注意力、前馈等）能给出相近的输出。

因此，通过不断地预测和纠错，Embedding矩阵可以被调整，最终将语义相关的 token 映射到接近的向量，而将无关的token映射到相差较大的向量。

位置编码

Transformer模型本身是没有捕捉序列顺序的能力的，所有的 token 都是同时输入到模型中的，因此需要引入位置编码。

位置编码（Positional Encoding）的目的是在 Transformer 模型中引入序列中每个 token 的位置信息。经过位置编码之后，Transformer 就能分得出来输入序列谁先谁后了。

接下来，将介绍两种位置编码方案。

正余弦位置编码

在 Transformer 的原始论文 "Attention is All You Need" 中，作者使用了基于正弦和余弦函数的位置编码方法。具体公式可参见位置编码。

正余弦位置编码方法有以下特点：

对于任意固定的偏移 k，PE(pos+k) 可以表示为 PE(pos) 的线性函数。这意味着模型可以容易地学习到相对位置，因为位置 pos+k 的编码可以通过位置 pos 的编码的线性变换得到。
正弦和余弦函数的周期性使得模型可以推广到比训练序列更长的序列。
每个位置都有一个唯一的编码，且编码的维度与嵌入向量的维度相同。

在模型中，位置编码会与词嵌入直接相加，然后作为输入送到 Transformer 的编码器和解码器中。这样，Transformer 的自注意力机制就可以利用位置信息了。

具体的代码实现如下：

让我们在你现有的 Transformer 类的基础上,修改 forward 函数,实现正余弦位置编码:

class Transformer(nn.Module):
    def __init__(self):
        super().__init__()
        # 词嵌入 Embedding
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # ...


    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx and targets are both (B,T) tensor of integers
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        
        # 计算正余弦位置编码
        position_encodings = get_sinusoid_encoding_table(T, n_embd) # (T,C)
        position_encodings = position_encodings.to(device) # 移动到与 tok_emb 相同的设备上
        
        x = tok_emb + position_encodings # (B,T,C)
        # ...

# 辅助函数:生成正余弦位置编码
def get_sinusoid_encoding_table(sequence_length, hidden_size):
    def cal_angle(position, hid_idx):
        return position / np.power(10000, 2 * (hid_idx // 2) / hidden_size)
    
    def get_posi_angle_vec(position):
        return [cal_angle(position, hid_j) for hid_j in range(hidden_size)]

    sinusoid_table = np.array([get_posi_angle_vec(pos_i) for pos_i in range(sequence_length)])
    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1
    
    return torch.FloatTensor(sinusoid_table)

在 forward 函数中,我们调用 get_sinusoid_encoding_table 函数计算正余弦位置编码。这个函数根据序列长度 T 和嵌入维度 n_embd 生成一个形状为 (T,C) 的位置编码矩阵。

我们将生成的位置编码移动到与词嵌入 tok_emb 相同的设备上(CPU 或 GPU)，以便后续计算。

将词嵌入 tok_emb 与位置编码 position_encodings 相加,得到最终的输入表示 x，形状为 (B,T,C)。

在辅助函数 get_sinusoid_encoding_table 中，我们根据正余弦位置编码的公式,计算出每个位置的编码向量。奇数维度使用正弦函数,偶数维度使用余弦函数。

Embedding 位置编码

在《Let's build GPT》中，使用了另一种位置编码方法——采用 Embedding 来实现位置编码：

在该方法下，位置编码是通过 nn.Embedding 实现的，这其实就是一种可学习的位置编码。与词嵌入类似，它也维护一个矩阵，每一行对应一个位置的编码。在前向传播时，你输入位置索引，它就返回对应的位置编码向量。

基于 Embedding 的位置编码，具体实现如下：

class Transformer(nn.Module):
    def __init__(self):
        super().__init__()
        # 词嵌入 Embedding
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # 位置编码 Embedding
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # ...

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx and targets are both (B,T) tensor of integers
        # 首先对 Token 进行词嵌入操作，得到词向量
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        # 之后根据位置编码的 Embedding，得到位置编码向量
        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
        # 对两者进行叠加，得到的向量，就是 Transformer 的输入
        x = tok_emb + pos_emb # (B,T,C)
	    # ...

Tranformer Blocks

Transformer Block 是 Transformer模型的基本组成单元。一个 Transformer 模型通常由多个（n_layer = 16）Transformer Block 堆叠而成，每个 Block 都包含了自注意力机制和前馈神经网络，以及一些其他的关键组件。

一个Transformer Block主要由以下几个部分组成:

多头自注意力层（Multi-Head Attention）：这是 Transformer Block 的核心。它允许模型在不同的表示子空间中关注输入序列的不同部分，有效地捕捉序列内和序列间的依赖关系。多头注意力将在后面的章节中详细介绍。
前馈神经网络（Feed Forward Network）: 这是一个简单的全连接神经网络，它对注意力层的输出进行进一步的变换,增强了模型的表达能力。
残差连接（Residual Connection）和层归一化（Layer Normalization）: 这些是现代神经网络设计中常用的技巧。残差连接使得信息可以更容易地在网络中流动，缓解了深度网络训练的难度。层归一化则有助于稳定和加速模型的训练。

在 Transformer Block 的前向传播过程中，输入序列 $X$ 首先经过多头自注意力层,得到注意力输出 $Z$ :

Z = M u l t i H e a d A t t e n t i o n (X)

然后，注意力输出 $Z$ 经过一个残差连接和层归一化,再传入前馈神经网络，得到Block的最终输出 $Y$ :

Y = L a y e r N o r m (F e e d F o r w a r d (Z) + Z)

这个过程可以被重复多次，形成一个深度的 Transformer 模型。

Transformer Block 的设计充分利用了自注意力机制的优势，使得模型能够高效地处理变长序列，并捕捉序列内和序列间的复杂依赖关系。

此外，Transformer Block 中的计算主要是矩阵乘法，非常适合在现代的 GPU 硬件上并行加速。这些特点使得 Transformer 模型能够在机器翻译、语言建模、问答系统等多个自然语言处理任务上取得优异的表现。

在代码实现上，Transformer Block 封装为名为 Block 的模块，在 Transformer 中，以 nn.Sequential 的方式，将多个 Block 连接到一起：

class Transformer(nn.Module):
    def __init__(self):
        super().__init__()
        # 词嵌入 Embedding
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # 位置编码 Embedding
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # 将多个 Block 连接到一起
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        # ...

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx and targets are both (B,T) tensor of integers
        # 首先对 Token 进行词嵌入操作，得到词向量
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        # 之后根据位置编码的 Embedding，得到位置编码向量
        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
        # 对两者进行叠加，得到的向量，就是 Transformer 的输入
        x = tok_emb + pos_emb # (B,T,C)
        # blocks 对输入进行多层级的运算
        x = self.blocks(x) # (B,T,C)
        # ...

Transformer Block

在上一节方式，将 Block 视作黑盒，介绍了 Transformer 中将 Block 进行多层连接的过程。在本节中，将讲解 Block 的内部。

Block 内部组成部分已经在上一节中介绍过了，包括：多头自注意力层、前馈神经网络、残差连接以及层归一化。

这里直接给出 Block 实现：

class Block(nn.Module):
    """ Transformer block: communication followed by computation """

    def __init__(self, n_embd, n_head):
        # n_embd: embedding dimension, n_head: the number of heads we'd like
        super().__init__()
        head_size = n_embd // n_head
        # 多头自注意力
        self.sa = MultiHeadAttention(n_head, head_size)
        # 前馈神经网络
        self.ffwd = FeedFoward(n_embd)
        # 输入的层归一化
        self.ln1 = nn.LayerNorm(n_embd)
        # 输出的层归一化
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
	    # 首先对输入进行层归一化处理
	    # 之后，进行多头自注意力运算
	    # 之后，有一个 `x +`，这是进行了一次残差连接
        x = x + self.sa(self.ln1(x))
        # 对残差连接的输出，再一次进行层归一化处理
        # 之后，进入前馈神经网络
        # 之后，又进行了一次残差连接
        x = x + self.ffwd(self.ln2(x))
        # 最终得到 Block 输出
        return x

对 Block 的直观理解

Block 的输入是 token 嵌入向量的序列，Block 的输出也是嵌入向量的序列，但两者在语义表示上有着本质的区别。

通过 Embedding 层得到的初始嵌入向量序列，主要捕捉到了每个 token 相对独立的语义信息。但自然语言中，token 的意义往往需要在特定的依赖关系下才能准确解读。这正是 Block 发挥作用的地方。

经过一次 Block 运算：

首先，嵌入向量序列会通过 Block 中的多头自注意力机制进行变换。在语言模型任务中，多头自注意力通常采用 masked 的方式，使得每个 token 的嵌入向量只能与它之前的 token 进行交互。通过这种交互,每个 token 的表示都融合了它所能“看到”的其他 token 的信息。多头自注意力中的不同头，可以学习到 token 之间的不同依赖关系，比如语义关联、句法关系等。
接下来，自注意力的输出会经过一个前馈神经网络。这个前馈网络可以被视作是一种特征提取器，它对自注意力输出的 token 表示进行进一步的非线性变换，提取出更加抽象和高层次的特征。这增强了模型的表达能力。
在这两个主要组件之外，Block 中还使用了Layer Normalization 和残差连接。Layer Normalization 有助于稳定深度网络的训练，而残差连接使得梯度能够更容易地流通到网络的浅层，缓解了深度模型的优化难题。

经过 Block 的处理，输出的嵌入向量序列相比原始的词嵌入，能够更好地反映 token 在特定上下文中的语义。

值得一提的是，Transformer 模型通常由多个 Block 堆叠而成。每一层 Block 都在前一层的基础上，学习到更深层次的 token 表示。随着 Block 的层数增加，模型可以建模越来越复杂的依赖关系，token 的表示也变得更加丰富和抽象。

需要注意的是，虽然 Block 的堆叠确实能学习到更深层次的表示，但模型的性能并非仅仅取决于 Block 的层数。模型的表现还受其他因素影响，比如模型规模、训练数据质量、超参数选择等。

3B1B 的《Visualizing Attention》中举过一个形象的例子：

对于“塔”，Embedding 会嵌入到空间中，表示“塔”的意思的一个向量
对于“埃菲尔铁塔”，经过 Block 运算，会对初始嵌入向量进行调整，调整到那个表示埃菲尔铁塔意思的位置。毕竟，埃菲尔铁塔的有它专属的含义。
对于“埃菲尔铁塔钥匙扣”，经过 Block 运算，发现它并非完全是埃菲尔铁塔，则会调整到另外一个位置，在那个位置，既包含了埃菲尔铁塔的意思，也包含了钥匙扣的意思。

这就是 Block 的作用，非常奇妙，这这种“意思”空间的投射让我着迷。

多头注意力

在《Transformer Block》一节中，看到 Block 内部包括 MultiHeadAttention 和 FeedFoward。

前馈神经网络是最基础的神经网络形式，可点击进入其笔记了解。本节中，将介绍 MultiHeadAttention 多头注意力机制，这是 Transformer 模型的核心。

所谓多头注意力，就是将多个单头注意力组合起来。组合多少个单头注意力，我们称之为注意力头数，由参数（n_head = 16）控制。

那什么又是单头注意力呢？在下一节中将进行介绍，在本节中，我们将其视为一个黑盒的神经网络模块。我们重点看如何组装单头注意力，形成多头注意力的：

class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

其中，多头注意力的实现主要分为三个步骤:

初始化多个单头注意力（Head）
并行计算多个单头注意力
将多个单头注意力的输出拼接并线性变换

具体来说：

【步骤1】初始化多个单头注意力：

在 __init__ 方法中，我们首先根据指定的头数 num_heads 和每个头的维度 head_size 来初始化多个单头注意力:

self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])

这里使用了 PyTorch 的 nn.ModuleList 类来封装多个 Head 实例。每个 Head 实例代表一个单头注意力，它的输入输出维度为 head_size。这样，我们就得到了一个包含 num_heads 个单头注意力的列表 self.heads。

除此之外，__init__ 方法中还初始化了一个线性变换 self.proj 和一个 dropout 层 self.dropout，它们将在第三步中用到。

【步骤2】并行计算多个单头注意力：

在 forward 方法中，我们对输入 x 并行地应用多个单头注意力，并将结果拼接起来:

out = torch.cat([h(x) for h in self.heads], dim=-1)

这行代码使用了 Python 的列表推导式，对 self.heads 中的每一个单头注意力 h，我们都将其应用于输入 x，得到一个形状为 (B, T, head_size) 的张量。然后，我们使用 torch.cat 函数将这些张量在最后一个维度(dim=-1) 上拼接起来,得到一个形状为 (B, T, num_heads * head_size) 的张量 out。

这一步的关键在于，我们让每个单头注意力独立地计算，它们各自关注输入序列的不同方面。这种并行计算的方式提高了模型的表达能力。

【步骤3】线性变换：

最后,我们使用初始化的线性变换 self.proj 和 dropout 层 self.dropout 对拼接后的张量 out 进行变换:

out = self.dropout(self.proj(out))

这里的线性变换 self.proj 可以将 out 的维度从 num_heads * head_size 变换回 n_embd，与多头注意力的输入维度保持一致。接着，我们使用 dropout 层对线性变换的输出进行随机遗弃，这是一种常见的正则化手段，可以减少过拟合。

至此，多头注意力的计算就完成了。让我们回顾一下整个过程：首先初始化多个单头注意力，然后并行地计算它们，接着将结果拼接起来并经过线性变换，最终得到了多头注意力的输出。

这种多头注意力的设计使得模型可以在不同的子空间里学习到输入序列的不同表示,提高了模型的表达能力和泛化能力。它是 Transformer 模型的一个关键组成部分。

单头注意力

现在，我们了解了多头注意力机制，但是对于单头注意力还是一头雾水，要知道，这是 Transformer 模型最精华的部分。

从 Transformer 模型提出的论文《Attention is All You Need》的标题即可看出，Attention 注意力机制的关键性。

注意力机制的灵感来自于人类的视觉注意力机制。当我们观察一幅图像或一个场景时，我们的视觉系统并不是平等地处理每一个细节，而是会有选择性地关注某些部分。例如，当我们看一张人脸的照片时，我们通常会更多地关注眼睛、鼻子、嘴巴等面部特征，而相对忽略其他次要的细节。这种选择性关注的机制使我们能够高效地提取关键信息，而不被大量无关的细节所淹没。

在自然语言处理任务中，注意力机制的作用与之类似。当我们处理一个句子或一段文本时，其中某些词或短语通常比其他部分更重要，包含了更多的信息。注意力机制允许模型去学习如何区分重要的信息和次要的信息，并根据这些重要的信息来做出判断或预测。

具体来说，注意力机制的核心思想可以用"查询-键-值（Query-Key-Value）"的框架来描述:

查询（Query）：我们想要关注的内容。在自然语言处理任务中，查询通常是我们当前正在处理的词或句子。
键（Key）：我们用来判断其他信息是否重要的参考。键可以是句子中的其他词，或者是来自其他来源的信息。
值（Value）:我们想要提取的信息。值通常与键是一一对应的。

注意力机制的过程可以概括为：对于每个查询，我们用它去和所有的键进行比较，计算出每个键与查询的相关性或重要性。然后，我们用这些重要性作为权重，对相应的值进行加权求和，得到最终的注意力结果。这个结果就是模型认为对于当前的查询最重要的信息。

通过这种方式，注意力机制使模型能够动态地调整对不同部分信息的关注，从而更好地理解和处理复杂的语言数据。

在 Transformer 模型中，注意力机制被用于捕捉句子内部和句子之间的依赖关系。例如,当模型处理一个代词时（如"it"）,注意力机制允许模型去关注前面提到的相关名词,从而正确地理解代词的指代对象。

单头注意力的数学推导

注意力机制的核心思想是，当查询（query）某个值（value）时，通过与一系列键（key）进行相似度计算来得到每个值的重要性权重，然后用这些权重对相应的值进行加权求和，得到最终的注意力结果。这个过程可以用数学公式表示为:

A t t e n t i o n (Q, K, V) = s o f t m a x (\frac{Q K^{T}}{\sqrt{d_{k}}}) V

其中， $Q$ 表示查询（query）， $K$ 表示键（key）， $V$ 表示值（value）， $d_{k}$ 是键向量的维度。

在单头注意力中，Q、K、V 实际上都是矩阵，而不是向量。它们的形状都是（B, T, head_size），其中 B 是 batch size,T 是序列长度，head_size 是每个注意力头的维度。

这里的 head_size 与嵌入维度 n_embd 有关。具体来说，在多头注意力中，我们将嵌入维度 n_embd 均分给每个注意力头。所以，每个头的维度 head_size = n_embd // num_heads。

Q、K、V 矩阵是从输入序列的词嵌入矩阵 X 计算得到的。设输入序列的词嵌入矩阵为 X,其形状为 (B, T, n_embd),即对于每个 batch 中的每个序列，我们都有 T 个 n_embd 维的词嵌入向量。

在自注意力机制中，查询、键、值都来自同一个输入序列。具体来说，对于输入序列的每个位置，们使用三个学习到的矩阵 $W_{Q}$ 、 $W_{K}$ 、 $W_{V}$ 将其映射为查询向量、键向量和值向量:

Q = X W_{Q}, K = X W_{K}, V = X W_{V}

其中， $X$ 表示输入序列的嵌入表示。

换算成更加熟悉的代码表示：

Q = X @ W_Q
K = X @ W_K  
V = X @ W_V

其中，W_Q、W_K、W_V 的形状都是 (n_embd, head_size)。这个映射过程可以理解为，我们为每个注意力头学习一组特定的查询、键、值表示。

所以，Q、K、V 矩阵可以看作是从原始词嵌入空间映射到注意力头特定的查询、键、值空间。这种映射允许每个注意力头关注输入序列的不同方面。

在实际实现中，我们通常会一次性计算整个序列的注意力。这时， $Q$ 、 $K$ 、 $V$ 都是矩阵，注意力计算可以用矩阵乘法高效地实现:

A t t e n t i o n (Q, K, V) = s o f t m a x (\frac{Q K^{T}}{\sqrt{d_{k}}}) V

其中， $Q$ 、 $K$ 、 $V$ 的形状分别为 $(B, T, d_{k})$ 、 $(B, T, d_{k})$ 、 $(B, T, d_{v})$ ，注意力输出的形状为 $(B, T, d_{v})$ 。

对于每个查询向量（Q 矩阵中的每一行），我们计算它与所有键向量（K 矩阵中的所有行）的相似度,得到一个注意力分数向量。这个过程可以用矩阵乘法高效地实现，换成代码表示:

scores = Q @ K.transpose(-1, -2)

这里的 scores 矩阵形状为 (B, T, T)，表示对于每个 batch 中的每个查询位置，我们都计算了它与所有键位置的相似度。

然后，我们对这些分数应用 softmax，得到注意力权重矩阵:

weights = softmax(scores, dim=-1)

最后，我们用这些权重对值矩阵 V 进行加权求和:

output = weights @ V

这里的 output 形状为 (B, T, head_size)，表示对于每个 batch 中的每个位置，我们都计算了一个 head_size 维的注意力输出向量。

计算注意力分数的过程可以解释为，对于每个查询向量，我们计算它与所有键向量的点积相似度，然后除以 $\sqrt{d_{k}}$ 进行缩放(这是为了防止点积结果过大)。接着，我们对这些相似度进行 softmax 归一化，得到注意力权重。最后，我们用这些权重对值向量进行加权求和，得到注意力输出。

单头 output 与多头组装与 output

单头注意力的输出可以看作是对输入序列的一种新的表示。具体来说，对于序列中的每个位置，注意力机制根据这个位置与其他位置的相关性，有选择性地聚合了序列中的信息。所以，注意力输出向量可以看作是蕴含了序列上下文信息的表示。

从直觉上理解，注意力机制让模型能够在编码每个位置时"参考"序列中的其他位置，就像人在理解一个词的意思时,会联系上下文一样。那么，注意力输出向量可以看作是一种"上下文感知"的表示，相比原始的词嵌入，它融合了更多的上下文信息。

多头注意力机制通过引入多个注意力“头”，允许模型在不同的子空间中以不同的方式关注序列的不同部分。每个头可以学习到不同的关注模式。例如，有的头可能更关注局部信息，有的头可能更关注全局。多头注意力的输出是将这些不同的表示拼接起来。所以，它可以看作是一种“多视角”的序列表示，综合了多个注意力头的关注结果。

在多头注意力之后进行的线性变换（proj），可以看作是对多头注意力输出的一种“整合”。这个变换将不同头的输出映射到同一个空间，并调整其维度与输入的词嵌入维度一致。这使得注意力层的输出可以与输入相加（残差连接）,也可以传递给下一个注意力层。

Transformer Block 的输出是对输入序列的一种新的表示。

但是，Block的输出并不是直接用于预测下一个词的概率分布。==这个任务是在Transformer的最后一层，通过一个线性层（lm_head）将最后一个Block的输出映射到词表大小的空间来实现的。

所以，Transformer Block 的作用是逐层地对输入序列进行变换，使其蕴含越来越丰富的上下文信息，但它并不直接做出预测。可以说，Block是在为最后的预测任务提供更好的序列表示。每一个Block的输出，都可以看作是一个更高层次的序列表示,捕捉了更复杂的上下文依赖。

自注意力

对于注意力机制，分为两种形式：自注意力、交叉注意力。对两者的介绍如下：

自注意力（Self-Attention） 是指序列内部的注意力计算。在自注意力中，Query、Key、Value 矩阵都来自同一个输入序列。也就是说，序列中的每个位置都要和序列中的每个位置（包括自己）计算注意力。这使得序列中的每个位置都能够“关注”到序列中的任意一个位置。通过自注意力，模型能够学习序列内部的依赖关系，捕捉序列的内部结构。

交叉注意力（Cross-Attention）则是在两个不同序列之间进行注意力计算。在交叉注意力中，Query 矩阵来自一个序列（通常称为“目标序列”）,而 Key 和 Value 矩阵来自另一个序列（通常称为“源序列”）。这使得目标序列中的每个位置都能够"关注"到源序列中的任意一个位置。通过交叉注意力，模型能够学习两个序列之间的对应关系，实现信息的传递和融合。

在 GPT（Generative Pre-trained Transformer）模型中,主要使用的是自注意力机制。这是因为 GPT 是一个单向语言模型,它的任务是根据给定的前文，预测下一个词。在这个任务中，我们只需要考虑序列内部的依赖关系，即每个词只需要关注它前面的词。使用自注意力机制，GPT 可以高效地建模序列内部的长距离依赖，学习词与词之间的关系。

具体来说,在 GPT 中，每个 Transformer Block 都包含一个自注意力层。在这个层中，输入序列首先被映射为 Query、Key、Value 矩阵，然后通过自注意力计算得到一个新的序列表示。这个新的表示融合了序列中每个位置与其他位置的关系信息。

通过多个 Transformer Block 的堆叠，GPT 可以逐层地提取序列的高级特征，建立起复杂的词与词之间的关系。最后,通过一个线性层将最后一个 Block 的输出映射到词表大小的空间，得到下一个词的概率分布。

需要注意的是，虽然 GPT 主要使用自注意力，但这并不意味着交叉注意力在自然语言处理中不重要。事实上，在一些需要建立两个序列之间关系的任务中，如机器翻译、问答系统等，交叉注意力发挥着关键作用。比如在 Transformer 的原始论文中，编码器-解码器结构中就大量使用了交叉注意力。

Masked 掩码注意力

在自回归语言模型如 GPT 中，我们使用自注意力机制来建模序列内部的依赖关系。然而，单纯的自注意力机制存在一个问题：它允许每个位置去 attend to 序列中的所有位置，包括它后面的位置。这在语言建模任务中是不合理的，因为在预测下一个词时，我们只能使用当前词及其之前的上下文，而不能"窥视"未来的词。

为了解决这个问题，我们引入了Masked 掩码注意力机制。具体来说，我们定义了一个attention mask矩阵，用于指示每个位置在计算注意力时，哪些位置是被允许 attend to 的。这个mask矩阵通常是一个上三角矩阵，对角线及以下的元素为1（表示允许 attend to），而上三角部分的元素为0（表示不允许attend to）。

举一个具体的例子，假设我们有一个序列 "I love machine learning"，我们要预测序列中每个位置的下一个词。当我们预测 "love" 后面的词时,我们只能利用 "I" 和 "love" 的信息，而不能利用后面的"machine learning"。所以，理想的attention mask应该是这样的:

其中，第 i 行表示位置 i 在计算注意力时，哪些位置是被允许 attend to 的。可以看到，对于每个位置，它只允许 attend to 它自己和它前面的位置。

在实际实现中，我们通常不直接使用 0 和 1，而是将 mask 中的 0 替换为一个非常大的负值（如-1e10）。这是因为我们最终会将 mask 与注意力分数相加，mask 中的 0 会让对应位置的注意力分数变得极小，经过 softmax 后，这些位置的注意力权重就会变为0。而使用一个大的负值，则可以更加稳定地实现这一效果。

通过引入 attention mask，我们巧妙地引入了位置的先后顺序信息，使得模型在预测下一个词时，只能利用当前词及其之前的上下文，这与语言模型的要求完全吻合。同时，Masked Self-Attention 的计算过程与普通的 Self-Attention 非常相似,只是多了一个加上 mask 的步骤，因此可以很容易地集成到 Transformer 模型中。

Masked Self-Attention 是自回归语言模型如 GPT 的核心，它允许模型以一种自回归的方式，逐词地预测整个序列。通过堆叠多个 Masked Self-Attention 层，模型可以建立起复杂的长距离依赖关系，从而生成连贯且语法正确的文本。可以说，Masked Self-Attention 是 GPT 等生成式预训练模型能够取得巨大成功的关键所在。

如何理解 GPT 中的自回归

在上一节中，我们提到 GPT 是一个自回归语言模型。那么，"自回归"在这里具体指的是什么呢?

“自回归”（Auto-regressive）是一个来自统计学的概念。在统计学中，如果一个变量的当前值可以被它自己的历史值很好地预测，我们就说这个变量是“自回归”的。

在语言模型的上下文中，“自回归”指的是一种特定的概率建模方式。具体来说，就是模型通过已经生成的词，来预测下一个最可能的词。这个过程可以用如下的概率公式表示:

P (x_{1}, x_{2}, . . ., x_{n}) = P (x_{1}) P (x_{2} | x_{1}) P (x_{3} | x_{1}, x_{2}) . . . P (x_{n} | x_{1}, . . ., x_{n - 1})

其中， $x_{1}, x_{2}, . . ., x_{n}$ 表示一个长度为 n 的文本序列，每个 $x_{i}$ 表示序列中的一个词。

这个公式告诉我们，序列的联合概率 $P (x_{1}, x_{2}, . . ., x_{n})$ 可以被分解为一系列条件概率的乘积。每个条件概率 $P (x_{i} | x_{1}, . . ., x_{i - 1})$ 表示在给定前 i-1 个词的条件下，第 i 个词是 $x_{i}$ 的概率。

GPT 模型就是以这种自回归的方式工作的。具体来说:

在训练阶段，GPT 的目标是最大化序列的联合概率 $P (x_{1}, x_{2}, . . ., x_{n})$ 。它通过最小化每个位置的交叉熵损失来实现这一点。
在生成（推断）阶段，GPT 通过已经生成的词，逐词地预测下一个词。在生成第 i 个词时，模型基于前 i-1 个词的上下文，计算词表中每个词作为第 i 个词的条件概率 $P (x_{i} | x_{1}, . . ., x_{i - 1})$ ，然后从这个概率分布中采样得到第 i 个词。

这种自回归的生成过程使得 GPT 可以生成连贯且语法正确的文本。因为在生成每个词时，模型都会考虑之前生成的所有词，这使得生成的文本在局部和全局层面上都有很好的一致性。

从模型结构的角度看，GPT 中的 Masked Self-Attention 机制起到了关键作用。通过引入 attention mask，GPT 在计算第 i 个词的表示时，只允许它 attend to 前 i-1 个词，而不能访问后面的词。这正是自回归的要求。

单头注意力的代码实现

经过上面的理论梳理，在本节中给出单头注意力（Head）的代码实现：

class Head(nn.Module):
    """ 单头自注意力机制 """

    def __init__(self, head_size):
        super().__init__()
        # 定义 key, query, value 的线性变换
        # 这里的线性变换相当于将输入 x 映射到 key, query, value 空间
        # 映射的维度由 head_size 定义,通常 head_size = n_embd // n_head
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        
		# 定义注意力掩码矩阵的上三角矩阵
        # 这个矩阵用于在计算注意力时,屏蔽掉后面的位置,实现因果注意力
        # 这里使用 register_buffer 是为了将这个矩阵注册为模型的一部分,但不作为参数进行优化
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

		# dropout 层,用于在训练时随机丢弃一部分注意力权重,提高模型的泛化能力
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
	    # x 的维度: (batch_size, seq_length, n_embd)
        B,T,C = x.shape

		# 计算 key, query, value
        # k, q, v 的维度: (batch_size, seq_length, head_size)
        k = self.key(x)   # (B,T,C)
        q = self.query(x) # (B,T,C)
        
        # 计算注意力分数 (attention scores)
        # 这里先计算 q 和 k 的点积,然后除以 head_size 的平方根进行缩放
        # 这个缩放操作是为了让注意力分数的方差在不同的 head_size 下保持稳定
        # wei 的维度: (batch_size, seq_length, seq_length)
        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)

		# 应用注意力掩码
        # 这里使用 masked_fill 函数,将 tril 矩阵中为 0 的位置 (代表要被屏蔽的位置) 
        # 在 wei 中对应位置的值设置为负无穷大
        # 这样在计算 softmax 时,这些位置的注意力分数就会变成 0
        # 这实现了因果注意力,即每个位置只能 attend to 它前面的位置
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)

		# 对注意力分数应用 softmax,得到注意力权重
        wei = F.softmax(wei, dim=-1) # (B, T, T)

		# 应用 dropout
        wei = self.dropout(wei)
        
        # 根据注意力权重聚合值 (value)
        # v 的维度: (batch_size, seq_length, head_size)
        v = self.value(x) # (B,T,C)

		# 注意力加权求和
        # out 的维度: (batch_size, seq_length, head_size)
        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
        return out

这段代码实现了单头自注意力机制,主要分为以下几个步骤:

初始化阶段：
- 定义了三个线性变换：key， query，value，用于将输入 x 映射到 key，query，value 空间。
- 定义了一个上三角矩阵 tril，用于在计算注意力时实现因果注意力 (即屏蔽掉后面的位置)。
- 定义了一个 dropout 层，用于在训练时随机丢弃一部分注意力权重，提高模型的泛化能力。
前向传播阶段：
- 首先，通过线性变换计算出 key，query，value。
- 然后，通过计算 query 和 key 的点积并缩放，得到注意力分数。
- 接着，应用注意力掩码 (即上三角矩阵 tril)，屏蔽掉后面的位置，实现因果注意力。
- 对注意力分数应用 softmax，得到注意力权重。
- 应用 dropout 到注意力权重上。
- 最后，根据注意力权重对 value 进行加权求和，得到输出。

单头注意力的输出可以看作是对输入序列的一种新的表示。具体来说，对于序列中的每个位置，注意力机制根据这个位置与其他位置的相关性，有选择性地聚合了序列中的信息。

所以，注意力输出向量可以看作是蕴含了序列上下文信息的表示。从直觉上理解，注意力机制让模型能够在编码每个位置时"参考"序列中的其他位置，就像人在理解一个词的意思时，会联系上下文一样。

举个例子，假设我们有这样一个句子:"The animal didn't cross the street because it was too tired"，它的意思是"这只动物没有过马路，因为它太累了"。

当我们在编码"it"这个词时，人类很自然地理解"it"指代的是"the animal"，因为我们能够根据上下文去推断。

注意力机制可以让模型也有类似的能力。当模型在处理"it"这个词时，注意力机制可以让"it"的表示向量中融合"the animal"的信息，从而更好地理解"it"在这个句子中的实际意思。

所以，注意力的输出可以看作是一种"上下文感知"的表示，它相比原始的词嵌入，融合了更多的上下文信息，从而能够更好地表示词在特定语境下的意义。

Block 前馈网络层

在 Transformer 模型中，每个 Block 除了包含自注意力层（Self-Attention）外，还包含一个由两层全连接层组成的前馈神经网络（Feed Forward Network）。尽管自注意力机制是 Transformer 的核心创新，但从参数量的角度来看，这个前馈神经网络实际上占据了 Transformer 参数的大部分。

这个前馈神经网络的结构非常简单：第一层将输入维度从 n_embd 映射到一个更高的维度（通常是 4 * n_embd），然后通过一个非线性激活函数（通常是 ReLU）；第二层再将维度映射回 n_embd。

class FeedFoward(nn.Module):
    """ a simple linear layer followed by a non-linearity """

    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

那么，这个看似简单的前馈神经网络在 Transformer 中究竟扮演了什么角色呢？

从直觉上理解，自注意力层允许模型在不同的位置之间建立直接的依赖关系，使得每个位置的表示能够融合其他位置的信息。但是，这种交互是在相同的表示空间内进行的，因为 Q、K、V 矩阵都来自同一个输入 X。前馈神经网络的作用，就是在自注意力交互之后，为模型提供一种 "空间转换" 的能力，让表示在另一个空间中进行变换和学习。

具体来说，通过第一个线性层，模型可以将每个位置的表示投影到一个更高维的空间。在这个空间中，模型可以学习到一些在原始空间中可能难以学习的特征和关系。然后，通过 ReLU 激活函数引入非线性，增强模型的表达能力。最后，通过第二个线性层，模型再将高维表示映射回原始的维度。

从功能上看，前馈神经网络可以被视为对自注意力层输出的一种 "增强" 或 "细化"。它给了模型一个机会，去学习如何在另一个维度上理解和处理自注意力的结果。

此外，前馈神经网络还起到了增加模型容量和可学习参数的作用。更多的参数意味着模型可以学习到更多的模式和关系。

当然，过多的参数也可能带来过拟合的风险。为了缓解这一问题，Transformer 在前馈神经网络之后引入了 Dropout 层。Dropout 随机地屏蔽一部分神经元，迫使网络学习到更加稳健和泛化的特征，减少过拟合。

综上所述，尽管前馈神经网络的结构非常简单，但它在 Transformer 中扮演了不可或缺的角色。它与自注意力机制互补，提供了另一种形式的表示学习和处理，增强了模型的容量和表达能力。同时，它也是 Transformer 参数规模的主要贡献者。理解前馈神经网络的作用和机制，对于深入理解 Transformer 模型至关重要。

预测下一个 Token

在理解了 Transformer 对 Blocks 的堆叠组装后，我们已经知道，Transformer Block 的作用是逐层地对输入序列进行变换，使其蕴含越来越丰富的上下文信息。但是，这只是故事的一半。为了完成 GPT 的使命，即预测下一个 Token，我们还需要最后一步：将这个高级的序列表示转化为下一个 Token 的概率分布。

这一步是由 Transformer 类中的 lm_head 完成的。lm_head 是一个简单的线性层（全连接层），它接受 Transformer 最后一个 Block 的输出（经过层归一化），并将其映射到一个大小为 vocab_size 的向量，这个向量表示词表中每个 Token 作为下一个 Token 的得分（logit）。

self.lm_head = nn.Linear(n_embd, vocab_size)

在前向传播过程中，lm_head 是这样被使用的：

# 堆叠 Blocks 运算
x = self.blocks(x) # (B,T,C)
# 层归一化
x = self.ln_f(x) # (B,T,C)
# 输入本节中的全连接层，转为整个词典中，对下一个 Token 的得分
logits = self.lm_head(x) # (B,T,vocab_size)

这里的关键是理解 logits 的含义。对于批次中的每个序列（共 B 个），对于序列中的每个位置（共 T 个），我们都得到了一个大小为 vocab_size 的向量。这个向量可以被解释为：在当前位置，基于之前的所有 Token，词表中的每个 Token 作为下一个 Token 出现的可能性大小。

举个例子，假设我们的词表大小为 10000，那么在某个位置，lm_head 的输出可能是这样的：

[-2.1, 5.6, -0.7, ..., 3.2] # 共 10000 个数

这个向量告诉我们，在当前位置，第 0 号 Token 作为下一个 Token 出现的得分是-2.1，第 1 号 Token 的得分是 5.6,以此类推。

为了将这些得分转化为实际的概率，我们需要应用 softmax 函数：

probs = F.softmax(logits, dim=-1) # (B, T, vocab_size)

softmax 函数将这些得分转化为正数，并且保证它们的和为 1，因此可以被解释为概率。

有了这个概率分布，我们就可以进行实际的预测了。最简单的方法是选择概率最大的 Token：

next_token = torch.argmax(probs, dim=-1) # (B, T)

这给了我们下一个最可能的 Token。但是，这种贪心的方法可能会导致生成的文本缺乏多样性。一种更好的方法是从概率分布中随机采样：

next_token = torch.multinomial(probs, num_samples=1) # (B, T, 1)

这种方法允许模型生成多样化的文本，虽然可能不总是选择概率最大的 Token。

所以，通过 lm_head 和 softmax，我们实现了从 Transformer 的高级表示到下一个实际 Token 的转化。这一过程，连同 Transformer Block 的自注意力计算，构成了 GPT 生成文本的完整流程。

用最后一个位置计算下一个 Token

上一节说到，Transformer Blocks 之后的全连接层，得到对于序列中的每个位置（共 T 个）的词典 Token 可性能向量。

对于推理场景（给定一句话，预测下一个 Token）来说，我们实际上只需要序列的最后一个位置的概率，对其采样即可。

因为我们要根据上下文生成下一个词，而这个"下一个词"，就对应着上下文序列的下一个位置，即最后一个位置的下一个位置。

所以，我们只需要取出最后一个位置的 logits 向量，将其转化为概率分布，然后从这个分布中采样，就可以得到下一个生成的 Token 了。

在代码实现中，这一步通常是这样的：

# logits 的形状：(B, T, vocab_size)
# 取出最后一个时间步的 logits，形状变为 (B, vocab_size)
logits = logits[:, -1, :] 

# 应用 softmax 得到概率分布，形状为 (B, vocab_size)
probs = F.softmax(logits, dim=-1) 

# 从概率分布中采样，得到下一个 token，形状为 (B, 1)
next_token = torch.multinomial(probs, num_samples=1)

这段代码首先取出 logits 张量的最后一个时间步（logits[:, -1, :]），得到一个形状为 (B, vocab_size) 的张量，表示在上下文的最后一个位置，每个词作为下一个词的得分。

然后，它对这些得分应用 softmax 函数，将其转化为概率分布。最后，它使用 torch.multinomial 函数从这个概率分布中采样，得到下一个生成的 Token。

这个采样过程可以被看作是模型在给定上下文下，对下一个词的“猜测”。通过多次重复这个过程（每次都将新生成的 Token 附加到上下文的末尾），我们就可以让模型生成一段连贯的文本。

交叉熵误差的计算

在上文中，我们提到 Transformer 的输出 logits 是一个形状为 (B, T, vocab_size) 的张量，表示对于每个批次中的每个序列的每个位置，模型预测下一个 token 是词表中每个 token 的得分。

如何将其与真是语料作对比，形成交叉熵误差，得出损失函数中的损失，供优化器进行反向传播学习呢？

要计算交叉熵误差，我们需要将这些 logits 与真实的下一个 token 进行比较。这里的关键是理解，我们的目标是最大化真实的下一个 token 的概率。

假设我们的训练数据是这样的:

# 输入序列
inputs = [
    [2, 5, 1, 9, ...],
    [3, 2, 7, 6, ...],
    ...
]

# 目标序列（即每个输入 token 的下一个 token）
targets = [
    [5, 1, 9, 3, ...],
    [2, 7, 6, 4, ...],
    ...
]

注意，targets 序列实际上就是 inputs 序列向右移动一位的结果（最后一个 token 可以忽略，因为我们不需要预测序列结束后的下一个 token）。

在模型的前向传播过程中,我们会得到 logits:

logits, loss = model(inputs, targets)

这里的 logits 是模型的预测，形状为 (B, T, vocab_size)，而 targets 是真实的下一个 token,形状为 (B, T)。

为了计算交叉熵误差，我们首先需要将 logits 和 targets 的形状进行调整:

B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)

这里我们将 logits 重塑为 (B*T, C) 的形状,将 targets 重塑为 (B*T,) 的形状。这样做的目的是将每个位置的预测和目标都看作是独立的样本。

接下来,我们可以直接使用 PyTorch 的 cross_entropy 函数来计算交叉熵误差:

loss = F.cross_entropy(logits, targets)

这个函数会在内部进行以下操作:

对 logits 应用 softmax 函数，得到每个词的预测概率。
对每个位置，取出目标词的预测概率。
对这些概率取负对数，得到交叉熵损失。
对所有位置的交叉熵损失取平均，得到最终的损失值。

直观地理解，交叉熵损失衡量的是模型的预测与真实目标之间的差异。如果模型对正确的下一个词给出了高概率，那么损失就会很低;反之，如果模型给出了错误的预测，损失就会很高。

在训练过程中，我们的目标就是最小化这个交叉熵损失。PyTorch 的优化器（如 Adam）会自动计算损失对模型参数的梯度，并根据梯度更新模型参数，使损失逐步降低。

这个过程会反复进行多个 epoch，直到模型在验证集上的性能不再提升。这意味着模型已经学会了根据上下文预测下一个词。

所以，交叉熵损失是连接模型预测和真实目标的桥梁。它提供了一种衡量模型性能的方法，并指引模型通过梯度下降来学习和改进。

直观理解 Token 预测

前面从理论上介绍了，生成场景下，Transformer Blocks 之后的全连接层，将经过多层 Block 运算后，饱含含义的序列向量，转为对于序列中的每个位置（共 T 个）的词典 Token 可性能向量。我们该如何对这一过程建立直觉理解呢？

首先，我们需要理解，经过多层 Transformer Block 处理后的序列向量，已经不仅仅代表了单个词的意思，而是融合了上下文的信息。对于序列中的每个位置，其向量表示已经考虑了这个位置之前的所有词，以及这些词之间的关系。这种表示是高度抽象和富有语义的。

接下来，全连接层的作用可以被看作是一种"解码"或"翻译"的过程。它将这些抽象的、上下文相关的表示，映射到了一个更加具体的空间——词表空间。在这个空间中，每个维度对应着词表中的一个 Token。

具体来说，全连接层包含了一个权重矩阵，其维度为 (n_embd, vocab_size)。对于序列中的每个位置,它的向量表示(维度为 n_embd)都会与这个权重矩阵相乘。这个乘积的结果，就是一个维度为 vocab_size 的向量，其中每个元素表示了在当前位置,对应的 Token 作为下一个词出现的得分或者说是"倾向性"。

从直觉上理解，这个权重矩阵扮演了一个"解释者"或"翻译官"的角色。它了解每个 Token 在词表空间中的"位置",以及这个位置与 Transformer 产生的上下文表示之间的对应关系。通过与上下文向量的乘积，它实际上是在评估，在当前的上下文下,每个 Token 作为下一个词出现的合理性。

举一个具体的例子,假设我们有这样一个句子："The cat sat on the ..."。在处理到最后一个词 "the" 时，Transformer 产生的上下文向量已经编码了前面所有词的信息，以及 "the" 在这个上下文中的角色。它可能表示一种"期待"或"引入"的语义，暗示着后面应该出现一个名词。

全连接层接收到这个信息后，会去查询它的权重矩阵。它会发现，像 "mat"、"chair"、"couch" 这样的名词，其对应的权重可能较大，因为它们经常出现在 "the" 的后面，是合理的延续。而像 "happy"、"run"、"blue" 这样的形容词或动词，其权重则可能较小。

最后，这个权重向量会被转化为一个概率分布（通过 softmax 函数）。这个分布告诉我们,在给定的上下文下，每个 Token 作为下一个词的可能性有多大。我们可以根据这个分布，选择最可能的 Token（如 "mat"），或者从中随机采样，作为模型生成的下一个词。

所以，从抽象的上下文表示，到具体的下一词预测，全连接层起到了一个至关重要的桥梁作用。它将 Transformer 捕捉到的语义信息，解码为了我们可以直接理解和使用的词表征。这个过程，就像是将一种抽象的语言(上下文语义)翻译成了另一种具体的语言(词)。

通过这种方式，GPT 模型可以根据上下文，生成连贯、合理、富有语义的文本。每一次的预测，都建立在对之前所有词的理解之上;而每一个新生成的词，又会反过来影响后续的预测。这就是 GPT 能够生成如此神奇和逼真文本的秘密所在。

单层全连接够用吗？

我看到，最后的这一层全连接，仅是一个单层感知（nn.Linear），我想知道这样够用吗？会不会过于单薄？如果我换成多层感知会如何。或者说，最后一层的神经网络形态，对于 GPT 会有什么样的影响？

首先，我们要理解这个全连接层在整个 GPT 模型中的作用。如前所述，它的任务是将 Transformer 产生的高维、抽象的上下文表示，映射到词表维度的空间，从而生成下一个词的概率分布。从这个角度看，单层全连接似乎已经足够完成这个任务了。因为本质上，它只需要学习一个线性变换，将上下文向量空间线性地映射到词表空间。

但是，如果我们增加全连接层的层数，使其成为一个[000.wiki/多层感知机|多层感知机],会有什么影响呢？主要有以下几点：

增加了模型的复杂度和表达能力：多层感知机引入了非线性变换(如 ReLU)，使得模型可以学习到更加复杂的上下文与词之间的关系。也就是说，它可以捕捉到一些单层线性变换无法表达的模式。这在某些情况下可能是有益的，尤其是当上下文和词之间的关系非常复杂，不能被简单的线性变换所捕捉时。
增加了计算量和训练难度：每增加一层全连接，都会引入额外的矩阵乘法运算，增加计算开销。同时，由于引入了更多的参数（权重和偏置）和非线性变换，也可能使得模型的优化和收敛更加困难。这意味着我们可能需要更大的训练数据,更长的训练时间，以及更细致的超参数调优。
可能引入过拟合的风险：更加复杂的模型，虽然有更强的表达能力，但也更容易过拟合训练数据。过拟合意味着模型可以很好地记忆训练样本，但是在新的、看不见的数据上，性能反而会下降。这在训练数据不够大或者多样的情况下尤其容易发生。
需要更多的调参和实验：当我们改变模型结构时，最优的超参数设置也可能会改变。这意味着我们可能需要重新进行超参数搜索，如学习率、批次大小、dropout 率等。这是一个试错和优化的过程，需要大量的实验和计算资源。

综上所述，增加最后一层全连接的层数,是一把双刃剑。它可以提升模型的表达能力，捕捉更复杂的模式;但同时也可能带来计算开销、优化难度、过拟合风险等问题。

无论如何，GPT 的成功，主要还是归功于其 Transformer 部分的革新（如 Self-Attention、Layer Normalization 等）。最后一层全连接，虽然重要，但相对而言只是锦上添花。我们在优化模型时,更多的精力应该放在 Transformer 部分的改进上。

完整的 Transformer 结构

经过前面小节，我们对 Transformer 整体架构的各个部分进行了详细介绍。在本节中，给出 Transformer 的模块的完整代码，作为本文的结束：

class Transformer(nn.Module):
    def __init__(self):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd) # final layer norm
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx and targets are both (B,T) tensor of integers
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
        x = tok_emb + pos_emb # (B,T,C)
        x = self.blocks(x) # (B,T,C)
        x = self.ln_f(x) # (B,T,C)
        logits = self.lm_head(x) # (B,T,vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):
            # crop idx to the last block_size tokens
            idx_cond = idx[:, -block_size:]
            # get the predictions
            logits, loss = self(idx_cond)
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx

喜欢我文章的朋友请随缘打赏，鼓励我创作更多更好的作品！