专栏名称: AI开发者
AI研习社,雷锋网旗下关注AI开发技巧及技术教程订阅号。
目录
相关文章推荐
宝玉xp  ·  对 MCP 协议最形象生动的比喻:MCP ... ·  2 天前  
爱可可-爱生活  ·  HieroLM: ... ·  2 天前  
爱可可-爱生活  ·  [CL]《HieroLM: ... ·  2 天前  
李继刚  ·  Claude Prompt:言外之意 ·  2 天前  
李继刚  ·  Claude Prompt:言外之意 ·  2 天前  
51好读  ›  专栏  ›  AI开发者

GPT-2没什么神奇的,PyTorch 就可以复现代码

AI开发者  · 公众号  · AI  · 2020-02-25 18:32

正文

点击上方“蓝字”关注“AI开发者”


欢迎来到「带注释的 GPT-2」。

我读过的最精彩、解释最清楚的文章之一是「The Annotated Transformer」 https://nlp.seas.harvard.edu/2018/04/03/attention.html 。它引起了前所未有的关注,一个简单的想法就是用一个文件注释你所需要的代码。

我在机器学习方面的经验让我意识到,当你将事情写成代码的时候,其实现和秘密会变得更清晰,而不再是魔法了。

魔法没有什么神奇的。魔术师只是理解一些简单的东西,而这些东西对未经训练的观众来说似乎并不简单。一旦你学会了操作魔术卡片,你也可以变魔术。

——Jeffrey Friedl 在《Mastering Regular Expressions》一书中写道

GPT-2 一开始看起来像是魔术,它的看起来太美丽了,但希望我能为你解释魔术,在你读完这篇文章时揭示所有的技巧,这是我的目标。使那些热衷于理解 GPT-2 模型是如何工作的人更好理解。

注:几乎所有代码都是从Hugging Face( https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_gpt2.py )的 GPT-2 实现中复制、启发和引用的,只保留了简单的基本要素。如果你想在并行 GPU 上训练 GPT-2 模型,在微调时保存检查点,在多个 CPU 上运行推理任务等等,我建议你使用 Hugging Face API。最近,Hugging Face 发布了一个简单的教程,它教你如何做到这一点: https://huggingface.co/blog/how-to-train

在这篇文章中,我并不是在试图重新发明轮子,而是仅仅把一系列已经存在的优秀资源放在一起,让读者更容易掌握 GPT-2。我让读者在他们所选择的任何领域进一步建立这些基础。

你不能在弱小的基础上建造一座伟大的建筑。如果你要有一个强大的上层建筑,你必须有坚实的基础。

——Gordon B. Hinckley


学习本教程的先决条件

本文假设读者对注意力机制和 tansformers 有着扎实的理解。GPT-2 采用 12 层的,仅有解码器的 transformer 架构。如果你想复习一下或了解注意力机制和 transformers,这里有一个很好的资源列表:

  • Jay Alammar 的 The illustrated Transformer: http://jalammar.github.io/illustrated-transformer/

  • 哈佛大学 The Annotated Transformer: https://nlp.seas.harvard.edu/2018/04/03/attention.html

  • Rachel Thomas 和 Jeremy Howard的 Introduction to the Transformer: https://www.youtube.com/watch?v=AFkGPmU16QA&list=PLtmWHNX-gukKocXQOkQjuVxglSDYWsSh9&index=18&t=0s

如果你刚刚开始你的 NLP 之旅,或者你是一个专家,我绝对会推荐 Rachel Thomas 和 Jeremy Howard 教授的 fast.ai NLP 课程( https://www.fast.ai/2019/07/08/fastai-nlp/ )。本课程从基础开始,包括使用朴素贝叶斯和 Logistic 回归进行语义分类,接着是 RNN,后面还讨论了迁移学习、ULMFiT、Seq2Seq 翻译和 transformers 等。它是 fast.ai 团队免费提供的优秀资源。

另一个关于 GPT-2 本身的优秀资源,是 Jay Alammar 的 The Illustrated GPT-2( http://jalammar.github.io/illustrated-gpt2/ )。本文从语言模型的基本介绍开始,以一种非常容易理解的方式逐步解释 GPT-2 模型。我强烈建议读者阅读这篇文章。

哈佛大学 The Annotated Transformer 使用 PyTorch 实现了完整的 transformer 架构,是深入理解 transformer 的好方法。

然后,让我们在这些优秀的现有资源的基础上,用代码实现 GPT-2 吧~


摘要

自然语言处理任务,如问答、机器翻译、阅读理解等,通常是在特定任务的数据集上进行有监督的学习。我们证明,当语言模型在一个名为 WebText 的数百万网页的新数据集上训练时,它开始学习这些任务,而不需要任何明确的监督。我们最大的模型,GPT-2,是一个 1.5B 参数的 transformer,它可以获得最先进的语言建模成果,但仍然不适合 WebText。模型中的示例反映了这些改进,并包含连贯的文本段落。这些发现为构建语言处理系统提供了一条有希望的途径,该系统可以从自然发生的演示中学习执行任务。

Zero-shot 设置是不微调语言模型并直接在目标数据集上运行推理的设置。例如,在 WebText 上预览一个 LM,并直接尝试预测 Amazon 影评数据集的下一个单词。


模型架构(GPT-2)

我们的 LM 使用基于 transformer 的架构。该模型主要遵循 OpenAI GPT 模型的细节,并进行了一些修改。层规范化被移动到每个子块的输入,类似于预激活剩余网络,并且在最终的自关注块之后添加了额外的层规范化。我们在初始化时将剩余层的权重按 1/√N 的因子进行缩放,其中 N 是剩余层的数量。词汇量扩大到 50257 个单词。我们还将上下文大小从 512 增加到 1024 个,并使用更大的批大小——512。


模型规格(GPT)

我们的模型基本上遵循了最初 transformer 的工作原理。我们训练了一个 12 层的只解码的 transformer,它有隐藏的自注意力头(768 维状态和 12 个注意力头)。对于位置前馈网络,我们使用了 3072 维的内部状态。我们使用 Adam 优化方案,最大学习速率为 2.5e-4。学习速率在前 2000 次更新中从零线性增加,并使用余弦调度将其退火为 0。我们在 64 个随机抽样的小批量、512 个令牌的连续序列上训练了 100 个阶段。由于 layernorm 在整个模型中广泛使用,简单的 N(0,0.02)权重初始化就足够了。我们使用了一个 bytepair 编码(BPE)词汇表。我们还采用了在中提出的 L2 正则化的改进版本,在所有非偏倚或增益权重上的 w=0.01。对于激活函数,我们使用高斯误差线性单位(GELU)。


导入

import torchimport copyimport torch.nn as nnimport torch.nn.functional as Ffrom torch.nn.modules import ModuleListfrom torch.nn.modules.normalization import LayerNormimport numpy as npimport osfrom tqdm import tqdm_notebook, trangeimport logginglogging.basicConfig(level = logging.INFO)logger = logging.getLogger()


GPT-2 内部的 transformer 解码器

要重用用于描述 transformer 的术语,注意是一个查询(Q)和一组键(K)和值(V)对的函数。为了处理更长的序列,我们修改了 transformer 的多头自注意力机制,通过限制 Q 和 K 之间的点积来减少内存使用:

注意力是查询、键和值的组合

class Conv1D(nn.Module):    def __init__(self, nx, nf):        super().__init__()        self.nf = nf        w = torch.empty(nx, nf)        nn.init.normal_(w, std=0.02)        self.weight = nn.Parameter(w)        self.bias = nn.Parameter(torch.zeros(nf))
def forward(self, x): size_out = x.size()[:-1] + (self.nf,) x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight) x = x.view(*size_out) return x


CONV1D 层解释

CONV1D 层本身可以看作是一个线性层。本质上,它是投射一个初始张量 x(最终尺寸为 x.size(-1))并传递给它,最终尺寸为 self.nf。

下面是相同的输出示例:

d_model = 768conv1d  = Conv1D(d_model, d_model*3)x       = torch.rand(1,4,d_model) #represents a sequence of batch_size=1, seq_len=4 and embedding_sz=768, something like "Hello how are you"x       = conv1d(x)x.shape
>> torch.Size([1, 4, 2304])

如上例所示,CONV1D 返回的张量的最终维数是初始大小的 3 倍。我们这样做是为了能够将输入转换为查询、键和值矩阵。

然后可以检索查询、键和值矩阵,如下所示:

query, key, value = x.split(d_model, dim=-1)
query.shape, key.shape, value.shape>> (torch.Size([1, 4, 768]), torch.Size([1, 4, 768]), torch.Size([1, 4, 768]))

将输入转换为 Q、K 和 V 矩阵的另一种方法是必须有单独的 Wq、Wk 和 Wv 矩阵。我已经在这篇文章底部的附加部分解释了这一点。我发现这种方法更直观、更具相关性,但在本文中我们使用了 CONV1D 层,因为我们重用了 Hugging Face 的 CONV1D 预训练权重。


前向层解释

class FeedForward(nn.Module):    def __init__(self, dropout, d_model=768, nx=768*4):        super().__init__()        self.c_fc    = Conv1D(d_model, nx)        self.c_proj  = Conv1D(nx, d_model)        self.act     = F.gelu        self.dropout = nn.Dropout(dropout)
def forward(self, x): return self.dropout(self.c_proj(self.act(self.c_fc(x))))

在 Jay Alammar 的文章中有一个很好的解释,也就是上面提到的,输入是如何先经过注意力层,然后再进入前向层的。前馈网络是一个正常的网络,它接受来自注意力层(768)的输出,将其投射到 nx(768×4)维,添加一个激活函数 self.act(GELU),将其投射回 d_model (768) 并添加 dropout(0.1)。


注意力层解释

下面的摘录是从论文上摘取的:https://arxiv.org/abs/1706.03762

  • 标度点产品注意力

我们称我们的注意力为「标度点产品注意力」。输入包括维度 dk 的查询和键以及维度 dv 的值。我们使用所有键计算查询的点积,用√dk除以每个键,然后应用 softmax 函数获得值的权重。

在实际应用中,我们同时计算一组查询的注意力函数,将它们组合成一个矩阵 Q,并将键和值组合成矩阵 K 和 V。我们将输出矩阵计算为:

输出矩阵为 Q、K 和 V 的组合

最常用的两个注意力函数是加性注意力函数和点积(乘法)力函数注意。除了比例因子 1/√dk 外,点积注意力与我们的算法相同。附加注意力使用具有单个隐藏层的前馈网络计算兼容性函数。虽然二者在理论复杂度上相似,但在实际应用中,点积注意力速度更快,空间效率更高,因为它可以使用高度优化的矩阵乘法码来实现。当 dk 值较小时,两种机制的表现相似,但在 dk 值较大时,加性注意力优于点积注意力。我们怀疑,对于 dk 的较大值,点积在数量上增长较大,将 softmax 函数推入具有极小梯度的区域。为了抵消这一影响,我们将网点产品缩放至 1/√dk。

为了在代码中实现注意力层,我们首先利用 CONV1D 层,得到前面解释的 Q、K 和 V 矩阵。

一旦我们有了 Q、K 和 V 矩阵,我们就可以使用函数 _attn 来执行注意力。此函数复制了上述注意力点积公式。

class Attention(nn.Module):    def __init__(self, d_model=768, n_head=12, n_ctx=1024, d_head=64, bias=True, scale=False):        super().__init__()        self.n_head  = n_head        self.d_model = d_model        self.c_attn  = Conv1D(d_model, d_model*3)        self.scale   = scale        self.softmax = nn.Softmax(dim=-1)        self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))        self.dropout = nn.Dropout(0.1)        self.c_proj  = Conv1D(d_model, d_model)
def split_heads(self, x): "return shape [`batch`, `head`, `sequence`, `features`]" new_shape = x.size()[:-1] + (self.n_head, x.size(-1)//self.n_head) x = x.view(*new_shape) return x.permute(0, 2, 1, 3)
def _attn(self, q, k, v, attn_mask=None): scores = torch.matmul(q, k.transpose(-2, -1)) if self.scale: scores = scores/math.sqrt(v.size(-1)) nd, ns = scores.size(-2), scores.size(-1) if attn_mask is not None: scores = scores + attn_mask scores = self.softmax(scores) scores = self.dropout(scores) outputs = torch.matmul(scores, v) return outputs
def merge_heads(self, x): x = x.permute(0, 2, 1, 3).contiguous() new_shape = x.size()[:-2] + (x.size(-2)*x.size(-1),) return x.view(*new_shape)
def forward(self, x): x = self.c_attn(x) #new `x` shape - `[1,3,2304]` q, k, v = x.split(self.d_model, dim=2) q, k, v = self.split_heads(q), self.split_heads(k), self.split_heads(v) out = self._attn(q, k, v) out = self.merge_heads(out) out = self.c_proj(out) return out

另一种实现注意力的方法在本博客底部的附加部分进行了说明。我发现它更直观,更容易与研究论文进行比较。它利用线性层而不是 CONV1D 将输入转换为 Q、K 和 V 矩阵。我们之所以没有使用它,是因为我们使用了预训练的权重,从 Hugging Face 转换为一维层。


多头注意力

下面一段是从论文「Attention is all you need」上摘取的。

我们发现,使用不同的、学习到的线性映射将查询、键和值分别线性映射到 dk、dk 和 dv 维度更好。然后,在这些查询、键和值的隐射版本中,我们并行地执行注意力函数,生成 dv 维输出值。这些值被连接起来,然后再次进行映射,得到最终值,如下图所示:

多头注意力机制允许模型在不同的位置共同关注来自不同表示子空间的信息。

多头注意力等式

在这项工作中,我们使用了 h=8 个平行的注意力层,或者说头。其中,我们使用的都是 dk=dv=dmodel/h=64。由于每个头的维数减少,总的计算成本与全维度的单头部注意的计算成本相似。

不要被这个弄糊涂了,本质上,我们所做的就是给 Q,K 和 V 矩阵增加一个维数。也就是说,如果这些矩阵之前的大小是 [1, 4, 768],表示 [bs, seq_len, d_model],则这些矩阵被投影到[bs, n_head, seq_len, d_model//n_head],大小为 [1, 12, 4, 64]。GPT-2 使用 12 个平行头。我们将 Q,K,V 矩阵分解到 split_heads 函数中。最后,当我们通过应用并行注意力得到一个输出时,我们将它连接到合并头中,返回到维度矩阵 [bs,seq_len,d_model]。


代码中的 GPT-2 模型体系结构







请到「今天看啥」查看全文