企业官网建设流程全解析

GPT-2 from scratch with torch：用 torch 从零实现 GPT-2

原文：Posit AI Blog, GPT-2 from scratch with torch
作者：Sigrid Keydana（Posit）
发布日期：2023-06-20
授权说明：原文页 “Reuse” 部分说明，Text and figures are licensed under Creative Commons Attribution CC BY 4.0。原文同时说明，来自其他来源的复用图片不在该许可范围内，并会在图注中以 “Figure from …” 标识。本文为中文翻译，保留原文结构、代码、引用关系和图注语义，并按 CC BY 4.0 署名。

从零实现一个语言模型，可以说是形成准确直觉、理解其内部引擎如何工作的最佳方式。本文使用torch编写 GPT-2，也就是原始 GPT 的直接继任者。最终，你会得到一个 R 原生模型，它可以直接使用 Hugging Face 上已经预训练好的 GPT-2 模型权重。

目录

• 来源与资源
• 一个最小化 GPT-2
• 端到端使用：加载预训练权重

无论你如何看待大语言模型（LLM）——它们有益吗？危险吗？会不会像加密货币一样只是短暂风潮？——它们都已经存在，而且就在当下。这意味着，了解它们如何工作是一件好事；至于需要了解到什么层次，则取决于每个人自己的目标。

同一天，我还发布了 What are Large Language Models? What are they not?，那篇文章面向更广泛的读者。本文则想面向深度学习实践者，逐步走读一个用torch实现的 GPT-2（Radford et al. 2019）。GPT-2 是 OpenAI 一系列模型中的第二个：模型越来越大，训练语料也越来越庞大。你会看到，一个完整的模型实现不到 250 行 R 代码。

来源与资源

我要展示的代码来自minhub仓库。这个仓库本身也值得单独介绍。正如 README 所强调的：

minhub是一组深度学习模型的最小实现，灵感来自 minGPT。所有模型都设计成自包含、单文件、没有外部依赖，因此很容易复制并集成到你自己的项目中。

显然，这让它们成为很好的学习材料；但它的价值还不止于此。模型还提供了从 Hugging Face model hub 加载预训练权重的选项。更方便的是，你也不必担心如何正确做分词：直接从 Hugging Face 下载匹配的 tokenizer 即可。我会在本文最后一节展示它如何工作。正如minhubREADME 所说，这些功能由hfhub和tok两个包提供。

在minhub中，gpt2.R 基本上是 Karpathy 的 MinGPT 的移植版本。实现时也参考了 Hugging Face 更复杂的 GPT-2 实现。如果你想看 Python 代码走读，可以参考 https://amaarora.github.io/posts/2020-02-18-annotatedGPT2.html。那篇文章还整理了一批关于深度学习语言建模的博客和学习材料，这些资料在短短几年内已经逐渐成为“经典”。

一个最小化 GPT-2

总体架构

原始 Transformer（Vaswani et al. 2017）由 encoder stack 和 decoder stack 两部分组成，一个典型用例是机器翻译。后续模型会根据主要用途舍弃其中一侧结构。第一代 GPT 与 GPT-2 的差异主要是一些相对细微之处；在架构上，它只保留了 decoder stack。

由于每个 decoder block 内部都接入了“自注意力”，并且最开始还有 embedding 步骤，因此这并不是问题：外部输入与后续的内部表示在技术上没有本质区别。

下面是原始 GPT 论文（Radford and Narasimhan 2018）中的总体架构截图；这个结构对 GPT-2 仍然适用。token embedding 和 position embedding 之后，是 12 次重复的 transformer block（结构相同，但不共享权重）；最终由一个依赖任务的线性层构成模型输出。

在 gpt2.R 中，这个全局结构及其行为由nn_gpt2_model()定义。（代码被进一步模块化了，所以不要因为代码和截图不完全一一对应而困惑。）

首先，在initialize()中，我们定义模块：

self$transformer<-nn_module_dict(list(wte=nn_embedding(vocab_size,n_embd),wpe=nn_embedding(max_pos,n_embd),drop=nn_dropout(pdrop),h=nn_sequential(!!!map(1:n_layer,\(x)nn_gpt2_transformer_block(n_embd,n_head,n_layer,max_pos,pdrop))),ln_f=nn_layer_norm(n_embd,eps=1e-5)))self$lm_head<-nn_linear(n_embd,vocab_size,bias=FALSE)

这个模型最顶层的两个组件是transformer和lm_head，后者是输出层。这种代码层面的区分有重要语义含义，尤其体现在两个方面。

第一，transformer的定义用很简洁的方式传达了“什么构成了一个 Transformer”。之后接上的内容——在这里是lm_head——可以变化。

第二，更重要的是，这个区分反映了深度学习自然语言处理中的基本思路，或者说一种基本操作方式。学习分成两个步骤：第一步，也是不可或缺的一步，是学习关于“语言”本身的知识（这正是 LLM 做的事）；第二步资源消耗要小得多，即适配到具体任务，例如问答或文本摘要。

为了看清事情发生的顺序，以及每个模块发生多少次，我们看看forward()：

tok_emb<-self$transformer$wte(x)pos<-torch_arange(1,x$size(2))$to(dtype="long")$unsqueeze(1)pos_emb<-self$transformer$wpe(pos)x<-self$transformer$drop(tok_emb+pos_emb)x<-self$transformer$h(x)x<-self$transformer$ln_f(x)x<-self$lm_head(x)x

transformer中的所有模块都会被调用并执行一次；这包括h。但h本身是一个顺序模块，由多个 transformer block 构成。

这些 block 是模型核心，接下来我们就看它们。

Transformer block

在nn_gpt2_transformer_block()中，12 个 block 中的每一个都是这样定义的：

self$ln_1<-nn_layer_norm(n_embd,eps=1e-5)self$attn<-nn_gpt2_attention(n_embd,n_head,n_layer,max_pos,pdrop)self$ln_2<-nn_layer_norm(n_embd,eps=1e-5)self$mlp<-nn_gpt2_mlp(n_embd,pdrop)

在这个分辨率层次上，我们看到自注意力会在每一层重新计算；另一个构成性成分是前馈神经网络。此外，这里还有两个执行layer normalization的模块。Layer normalization 是 transformer block 中使用的归一化类型。

不同归一化算法的区别，往往在于它们对哪些维度取平均。Layer normalization（Ba, Kiros, and Hinton 2016）可能会让一些读者意外：它是对每个 batch item 做统计。也就是说，对于一个模块中的每个单元，都会有一个均值和一个标准差。其他维度，例如图像中的空间维度和通道维度，都会作为该 item 级统计计算的输入。

继续放大细节，我们很快会分别看注意力模块和前馈网络。不过在那之前，需要先看这些层如何被调用。forward()中发生的事情只有这些：

x<-x+self$attn(self$ln_1(x))x+self$mlp(self$ln_2(x))

这两行值得仔细读。这里并不是简单地把前一层输出依次送入下一层，而是在两个主要阶段外分别插入了 skip connection，也称 residual connection。其效果是：每个子模块并不替换传入的表示，而是用自己的“视角”去更新它。

近看 Transformer block：自注意力

在 GPT-2 的所有模块中，自注意力看起来最吓人。但这里使用的基本算法，仍然和 2014 年经典“点积注意力论文”（Bahdanau, Cho, and Bengio 2014）提出的思想一致：注意力被概念化为相似度，而相似度通过点积来度量。

一个容易让人困惑的点，是 self-attention 中的 “self”。这个术语首次出现在 Transformer 论文（Vaswani et al. 2017）中。那篇论文同时有 encoder stack 和 decoder stack。在那里，“attention” 指 decoder block 如何决定关注 encoder 阶段传来的信息；而 “self-attention” 是把这种技术用于 stack 内部时创造的术语，也就是在一个 stack 内部的 block 之间使用注意力。到了 GPT-2，只剩下现在显得有些重复命名的 self-attention。

从上文继续看，自注意力之所以显得复杂，有两个原因。

第一，是 Transformer 通过 query-key-value 框架引入了 token 的“三重化”¹。

第二，是多头注意力带来的额外批处理：每一层不只运行一个注意力计算过程，而是同时运行多个并行、独立的注意力头。走读代码时，我会指出这两点分别在什么地方出现。

我们还是从模块初始化开始。nn_gpt2_attention()这样列出它的组件：

# key, query, value projections for all heads, but in a batchself$c_attn<-nn_linear(n_embd,3*n_embd)# output projectionself$c_proj<-nn_linear(n_embd,n_embd)# regularizationself$attn_dropout<-nn_dropout(pdrop)self$resid_dropout<-nn_dropout(pdrop)# causal mask to ensure that attention is only applied to the left in the input sequenceself$bias<-torch_ones(max_pos,max_pos)$bool()$tril()$view(c(1,1,max_pos,max_pos))|>nn_buffer()

除了两个 dropout 层，我们还看到：

• 一个线性模块，它执行上面提到的“三重化”。注意，这不同于简单地复制三份相同的 token 表示。即使所有表示在初始时基本相同，例如随机初始化后如此，一旦模型开始训练，它们也不会继续保持相同。
• 一个名为c_proj的模块，用于执行最终仿射变换。要理解它具体做什么，需要看它如何被使用。
• 一个buffer。buffer 是模块状态的一部分，但不参与训练。这里的 buffer 确保注意力不会应用到“未来”的 previous-block output。基本做法是借助下三角矩阵，把未来 token mask 掉。

对于forward()，我们把它拆成容易消化的小块来看。

进入这个方法时，参数x的形状正如语言模型中所预期的那样：batch dimension × sequence length × embedding dimension。

x$shape [1] 1 24 768

接下来会发生两个批处理操作：第一，把表示三重化为 query、key、value；第二，腾出空间，让指定数量的 attention head 可以一次性并行计算。先列出完整代码，再解释。

# batch size, sequence length, embedding dimensionality (n_embd)c(b,t,c)%<-%x$shape# calculate query, key, values for all heads in batch and move head forward to be the batch dimc(q,k,v)%<-%((self$c_attn(x)$split(self$n_embd,dim=-1))|>map(\(x)x$view(c(b,t,self$n_head,c/self$n_head)))|>map(\(x)x$transpose(2,3)))

首先，调用self$c_attn()会为每个嵌入后的输入 token 生成 query、key 和 value 向量。split()把得到的矩阵拆成一个列表。然后map()负责第二个批处理操作：三个矩阵都被 reshape，并添加第四个维度；这个第四维用来表示 attention head。

注意，这与前面把 embedding 三倍扩展的乘法过程不同。这里是把已有内容分配给不同的 head，让每个 head 处理一个子集；每个子集大小与使用的 head 数量成反比。最后，map(\(x) x$transpose(2, 3))会交换 head 维度和 sequence-position 维度。

接下来是注意力本身的计算。

# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)att<-q$matmul(k$transpose(-2,-1))*(1/sqrt(k$size(-1)))att<-att$masked_fill(self$bias[,,1:t,1:t]==0,-Inf)att<-att$softmax(dim=-1)att<-self$attn_dropout(att)

首先计算 query 和 key 之间的相似度；矩阵乘法在这里实际上是批量点积。如果你对第一行最后的除法项感到好奇，这个缩放操作是 GPT-2 与前代 GPT 有所不同的少数方面之一。感兴趣的话，可以查看论文中相关考虑。

接着，应用前面提到的 mask，对得到的分数做归一化，并使用 dropout 正则化来鼓励稀疏性。

最后，计算出来的attention²需要传给后续层。这时 value 向量终于登场了——它是 query-key-value 三元组中我们还没看到发挥作用的成员。

y<-att$matmul(v)# (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)y<-y$transpose(2,3)$contiguous()$view(c(b,t,c))# re-assemble all head outputs side by side# output projectiony<-self$resid_dropout(self$c_proj(y))y

具体来说，这里的矩阵乘法会用attention对 value 向量加权并求和。这一步会同时作用于所有 attention head，也真正代表整个注意力算法的输出。

剩余步骤会恢复原始输入大小。这包括把所有 head 的结果依次对齐，然后应用线性层c_proj，确保这些结果不会被等同、独立地对待，而是以有用的方式组合起来。因此，这里的 projection 实际由一个机械步骤（view()）和一个“智能”步骤（c_proj()变换）共同构成。

近看 Transformer block：前馈网络（MLP）

与注意力模块相比，transformer block 的第二个核心组件nn_gpt2_mlp()并没有太多神秘之处。它确实“只是”一个 MLP，没有什么特殊技巧。不过仍有两点值得指出。

第一，你可能听说过 transformer block 中的 MLP 是 “position-wise” 的，并想知道这是什么意思。考虑一个 block 中发生的事情：

x<-x+self$attn(self$ln_1(x))x+self$mlp(self$ln_2(x))

MLP 的输入几乎直接来自注意力模块。我们已经看到，注意力模块返回的是形状为 [batch size,sequence length, embedding dimension] 的张量。在 MLP 内部，也就是它的forward()中，维度数量始终不变：

x|>self$c_fc()|># nn_linear(n_embd, 4 * n_embd)self$act()|># nn_gelu(approximate = "tanh")self$c_proj()|># nn_linear(4 * n_embd, n_embd)self$dropout()# nn_dropout(pdrop)

因此，这些变换会独立地应用于序列中的每个元素。

第二，由于这里是激活函数唯一出现的地方，可以顺便说一下。GeLU 表示 “Gaussian Error Linear Units”，由 Hendrycks and Gimpel（2020）提出。它的想法是把类似 ReLU 的激活效果与正则化/随机性结合起来。理论上，每个中间计算会根据它在高斯累积分布函数中的位置被加权，也就是根据它比其他值大多少或小多少来加权。实践中，正如模块实例化所显示的，会使用一个近似形式。

至此，GPT-2 的主角——重复出现的 transformer block——就讲完了。还剩两件事：之前发生了什么，之后又发生了什么。

从词到编码：token embedding 和 position embedding

严格说来，如果你按要求对输入数据做分词（使用 Hugging Face 中匹配的 tokenizer，见下文），最终得到的并不真的是“词”。但一个成熟事实仍然成立：如果模型要成功提取语言知识，就必须发生某种表示变换。

和许多基于 Transformer 的模型一样，GPT 家族用两种方式编码 token。第一种是 word embedding。回到最开始走读的顶层模块nn_gpt2_model()，我们可以看到：

wte=nn_embedding(vocab_size,n_embd)

这已经很有用，但由此产生的表示空间并不包含那些可能随序列位置变化的语义关系，例如句法规则或短语语用。第二种编码就是用来弥补这一点的。它被称为 position embedding，在nn_gpt2_model()中写作：

wpe=nn_embedding(max_pos,n_embd)

又一个 embedding layer？是的，不过这次嵌入的不是 token，而是一组预先指定的有效位置（在 GPT 中范围是 1 到 1024）。换句话说，网络需要学习序列中的位置意味着什么。这也是不同模型差异很大的领域。原始 Transformer 使用的是一种正弦位置编码；更新一些的改进可以在 GPT-NeoX 等模型中看到，例如 RoPE（Su et al. 2021）。

一旦两种编码都可用，它们会被直接相加，见nn_gpt2_model()$forward()：

tok_emb<-self$transformer$wte(x)pos<-torch_arange(1,x$size(2))$to(dtype="long")$unsqueeze(1)pos_emb<-self$transformer$wpe(pos)x<-self$transformer$drop(tok_emb+pos_emb)

得到的张量随后会传入一串 transformer block。

输出

应用完 transformer block 后，最后一个映射由lm_head完成：

x<-self$lm_head(x)# nn_linear(n_embd, vocab_size, bias = FALSE)

这是一个线性变换，会把内部表示映射回离散词表索引，并为每个索引分配一个分数。到这里，模型本身的最终动作就结束了。接下来由样本生成过程决定如何使用这些分数。换句话说，生成过程可以在几种成熟技术之间自由选择。下一节我们会看到一种相当标准的方法。

至此，模型走读结束。我略过了一些细节，例如权重初始化；如果你感兴趣，可以查看 gpt2.R。

端到端使用：加载预训练权重

很少有用户真的想从零开始训练 GPT-2。因此，我们来看如何快速把它设置好，用于样本生成。

创建模型、加载权重、获得 tokenizer

Hugging Face model hub 允许你直接从 GPT-2 页面访问并下载所有必需文件，包括 weights 和 tokenizer。所有文件都有版本管理；这里使用最新版本。

identifier<-"gpt2"revision<-"e7da7f2"# instantiate model and load Hugging Face weightsmodel<-gpt2_from_pretrained(identifier,revision)# load matching tokenizertok<-tok::tokenizer$from_pretrained(identifier)model$eval()

分词

decoder-only Transformer 类型的模型并不需要 prompt。但通常应用都会想把输入传给生成过程。有了tok，对输入做分词就非常方便：

idx<-torch_tensor(tok$encode(paste("No duty is imposed on the rich, rights of the poor is a hollow phrase...)","Enough languishing in custody. Equality"))$ids)$view(c(1,-1))idx

输出：

torch_tensor Columns 1 to 11 2949 7077 318 10893 319 262 5527 11 2489 286 262 Columns 12 to 22 3595 318 257 20596 9546 2644 31779 2786 3929 287 10804 Columns 23 to 24 13 31428 [ CPULongType{1,24} ]

生成样本

样本生成是一个迭代过程：模型最后一次预测会被追加到不断增长的 prompt 后面。

prompt_length<-idx$size(-1)for(iin1:30){# decide on maximal length of output sequence# obtain next prediction (raw score)with_no_grad({logits<-model(idx+1L)})last_logits<-logits[,-1,]# pick highest scores (how many is up to you)c(prob,ind)%<-%last_logits$topk(50)last_logits<-torch_full_like(last_logits,-Inf)$scatter_(-1,ind,prob)# convert to probabilitiesprobs<-nnf_softmax(last_logits,dim=-1)# probabilistic samplingid_next<-torch_multinomial(probs,num_samples=1)-1L# stop if end of sequence predictedif(id_next$item()==0){break}# append prediction to promptidx<-torch_cat(list(idx,id_next),dim=2)}

要查看输出，只需使用tok$decode()：

tok$decode(as.integer(idx))

输出：

[1] "No duty is imposed on the rich, rights of the poor is a hollow phrase... Enough languishing in custody. Equality is over"

如果想尝试文本生成，只要复制这个自包含文件，然后试试不同的采样参数即可。当然，也可以换不同的 prompt。

一如既往，感谢阅读！

封面照片来自 Marjan Blan on Unsplash。

参考文献

• Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. “Layer Normalization.” https://arxiv.org/abs/1607.06450.
• Bahdanau, Dzmitry, Kyunghyun Cho, and Yoshua Bengio. 2014. “Neural Machine Translation by Jointly Learning to Align and Translate.” http://arxiv.org/abs/1409.0473.
• Hendrycks, Dan, and Kevin Gimpel. 2020. “Gaussian Error Linear Units (GELUs).” https://arxiv.org/abs/1606.08415.
• Radford, Alec, and Karthik Narasimhan. 2018. “Improving Language Understanding by Generative Pre-Training.”
• Radford, Alec, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. “Language Models Are Unsupervised Multitask Learners.”
• Su, Jianlin, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. 2021. “RoFormer: Enhanced Transformer with Rotary Position Embedding.”arXiv Preprint arXiv:2104.09864.
• Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. “Attention Is All You Need.” https://arxiv.org/abs/1706.03762.

复用与引用

原文复用说明：Text and figures are licensed under Creative Commons Attribution CC BY 4.0。来自其他来源的复用图片不在该许可范围内，并会在图注中以 “Figure from …” 标识。

原文建议引用：

Keydana (2023, June 20). Posit AI Blog: GPT-2 from scratch with torch. Retrieved from https://blogs.rstudio.com/tensorflow/posts/2023-06-20-gpt2-torch/

BibTeX：

@misc{keydanagpt2, author = {Keydana, Sigrid}, title = {Posit AI Blog: GPT-2 from scratch with torch}, url = {https://blogs.rstudio.com/tensorflow/posts/2023-06-20-gpt2-torch/}, year = {2023} }