Featured image of post Transformer模型结构

Transformer模型结构

内容介绍

先简要概述构建和训练一个基于 Transformer 的自然语言处理模型的步骤,对其的详细描述和基于pytorch的代码实现例子及解析见后文。

1,分词,得到和词语对应的词表,参数:词表长度,vocab_size;还需要包含BOS(begin of sentence),EOS(end of sentence),UNK (Unknown) 未知词标记。

2,构建transformer架构。这个架构依靠注意力层以捕捉到上下文的语义关联,利用注意力层和全连接层的堆叠可以构造编码器/解码器。

参数包含:d_model(模型给每个词分配的向量的长度),num_heads(多头注意力的头数),num_layers=6,(编码器/解码器的堆叠层数),context_length(最大序列长度,超出此长度会被截断),等。

transformer结构,加上位置信息的词嵌入层,注意力层+全连接层的反复堆叠,layer norm和残差连接。具体结构可见https://www.runoob.com/pytorch/transformer-model.html。

这个网站也会给出简单的实现transfomer具体结构的代码。

3,经过模型训练,训练每个词对应的向量和模型参数。

参数包括,batch size,一次性放入多少个数据对/多少个样本。一般所有数据张量的第一个维度都是batchsize,bsz.

另外还有填充padding的概念,统一用填充标记PAD (Padding),将不同长度的序列统一到相同长度seq_length(当前 Batch 中最长的那个序列的长度),形成规则的张量以方便计算。如果Batch 中的最大序列长度 超过 max_seq_length就会被截断。


注意力层的核心原理:

注意力层是transformer架构的核心,是模型理解上下文语意的关键。注意力机制以Q K的点积生成叠加系数,V作为叠加向量,生成新的包含了上下文信息的向量。

以下以decoder only的架构为例。

embedding词嵌入层把 token id 映射为(batch_size, seq_length, d_model)的向量表示,这里d_model是模型给每个词分配的向量的长度,batchsize是模型同时处理的序列数目,seq_length是每个序列被填充后的长度。再通过绝对位置嵌入加上位置信息(也可以采用相对位置编码,直接旋转每层的Q K),得到第一层 Transformer block 的输入。

具体来说,每个token乘以每一个注意力层的W_Q W_K W_V矩阵后都对应于Q K V三个向量:Q:query, 查询,“我想要什么”。K,key,被查询,“我能提供什么”, V,value,输出,“我的实际内容是什么”。 任意两个token对应的Q K^T之间都可以计算出一个注意力得分(除以每个 head 的 key/query 维度d_k的根号,防止后续softmax时注意力过早变得极端;开根号是因为,点积的方差正比于维度,标准差正比于维度开根号),组成一个注意力矩阵。

模型在生成下一个词的时候,会用之前文本序列的最后一个token作为Q,用其对应的注意力矩阵的这一行作为叠加系数(经过softmax时e指数的归一化),对应的V作为叠加向量,生成一个新的向量:softmax(QK^⊤)V。

这个向量在经过一个 FFN 子层(通常由两次线性变换 + 非线性/门控构成)之后,在RMSNorm层归一化后,流到下一个注意力层(记为B层)作为输入,再乘以B层的W_Q W_K W_V矩阵得到Q K V,给出B层的输出向量…这个过程涉及KV cache,B层计算时用到的K V向量是之前时间步计算过的,需要存储和复用。

一直到最后transformer层的输出,最后一个向量(Hidden State)在RMSNorm层归一化后进入线性层映射到词表大小,以概率分布的形式给出下一个 Token。

在模型训练过程中,输入的文本序列里的每个词都依次作为Q,可以一次性算出所有词两两之间的注意力得分,但用于预测下一个token的时候不可以直接偷看结果,需要在注意力矩阵里面通过负无穷大掩盖这个词之后的序列,即因果掩码概念。 生成/推理阶段,由于每步只计算最后一个位置的注意力,通常不需要显式构造因果掩码;但在 prompt 的并行 prefill 阶段以及任何并行计算多个位置注意力的场景下,必须使用因果掩码来保证自回归约束。

实际的掩码来源有两个,一个是把同一batch都填充到最长序列长度时带来的(只看原本有内容的部分,其余部分填充掩码;对batch内不同序列都是不同的),另外一个就是这里的no peek/因果机制导致,不能偷看这个词之后的序列。这两个掩码的融合就是进行 逻辑与 (Logical AND) 运算,实现上常将两者转为加性 mask(无效处为大负数)并相加。


注意力层具体的维度:

设Z是输入给注意力层的序列嵌入(batch_size, seq_length, d_model)。Wq Wk矩阵维度是(d_model,dim1),$Q=Z W_q , K=Z W_k$。由于Q K需要点积,其维度必须一致。$V=Z W_v$,Wv矩阵维度是(d_model,dim2),得到的V维度是(batch_size, seq_length,dim2)。 通常情况下,dim1 和 dim2 被选择为相等。

如果有多个头,需要拼接得到V_all=(batch_size, seq_length,dim2×N_heads)。

之后需要混合多个头的贡献,还需要乘以矩阵Wo, Output=V_all × W_o,dim(Wo)=(dim2×N_heads,d_model)。这里有 dim2 * N_heads = d_model,即$\mathbf{W}_o$ 作用后得到的最终输出向量的维度,与该注意力层的输入向量维度 必须一致,因为之后需要残差连接,然后层归一化。

以及,GQA (Grouped Query Attention),K V的头数比Q少,减少GPU 每次需要从内存搬运到计算核心的KV cache的数据量。


全连接层

注意力层后面会跟两个全连接层nn.Linear,构成transformer层。

首先投影到高维,例如4×输出维度(4×d_model),以获得高维的复杂的特征表示,然后激活函数例如RELU变得非线性。之后必须投影回原来维度,因为也会残差连接,以保证了信息和梯度能够顺畅地流过许多层而不衰减。同时,保持输入输出维度相同,可以让所有 Transformer 层共享一个统一的输入/输出接口维度,使得整个架构具有模块化和可扩展性。

transformer层可以多次叠加,组成encoder或decoder层。这两个层的结构是很像的。

可以改一下全连接层,把FFN 都换成 MOE(Mixture of Experts),让模型在不同的token上使用不同的专家,并且每次只激活 1-2 个专家,扩大参数量,提供更强的拟合能力。


训练阶段的全序列并行与推理阶段的自回归解码(KV Cache)

文本序列经过分词,可以在嵌入层用张量表示;此时也加上位置嵌入positional encoding(pe)对应的张量。有些是采用可以训练的参数表示的,有些干脆写死,就用序列位置和sin cos函数的组合直接给定(RoPE),也不需要参与优化(通过 self.register_buffer(‘pe’, pe_tensor) 存储,register_buffer() 自动设置 requires_grad=False)。

张量流过多个transformer层,最后输出结果。具体结构可见https://www.runoob.com/pytorch/transformer-model.html

以GPT为例,训练和评估时候的输出方式有所不同。训练时候,由于是对整个序列都计算注意力矩阵,可以一次性得到整个序列中任何位置的下一个位置上模型预测的词向量的概率分布,将这个和实际位置的词元对比,可以计算出loss。之后根据loss梯度求导,反向计算所有有关参数(所有nn.parameter默认requires_grad=True)的更新:loss = criterion(…)#计算loss, loss.backward() # 反向传播,optimizer.step() #更新所有可训练参数。

对于评估时候,或者说推理/生成时候,训练已结束,进入推理模式model.eval(),让其根据输入的prompt连续生成,直到输出EOS或者达到最大输出长度为止。模型需要反复地将自己上一步输出的词(输出词向量的概率分布之后,具体是输出哪个词,可以用不同的采样序列,例如Top-k/Top-p 采样)也作为输入,对这个新的词语也计算注意力矩阵。 在每一步生成时, 不需要重新计算 Prompt 部分 及其之前已生成部分的 $\mathbf{K}$ 和 $\mathbf{V}$ 矩阵。这些值可以被缓存(Key/Value 缓存,KV Caching),后续步骤只需计算 新增 Token 的 $\mathbf{Q}, \mathbf{K}, \mathbf{V}$,这大大提高了推理速度,但没有改变生成必须是自回归的本质。

线性注意力

线性注意力可以用于提升模型处理长文本的能力。

用N表示序列长度(Token 数量),用d表示每个 Head 的特征维度.

标准注意力 :Attn= Softmax(Q K^T) V, 这里Q K^T的维度是(N, d)×(d, N),得到一个N*N的巨大的注意力矩阵。下一步将这个注意力矩阵和V相乘,(N, N) × (N, d) ,复杂度O(N^2 * d)。

线性注意力,去掉 Softmax(或将其分解),使其可以先计算K^T* V, 其维度是(d,N)×(N, d),得到一个d*d的矩阵,生成的是一个特征相关性矩阵,与序列长短无关。之后再和Q相乘,复杂度O(d^2 * N)。

使用 Hugo 构建
主题 StackJimmy 设计