4 – Transformer块

前文提到,原始数据经过嵌入层处理后下一步会作为输入给到transformer 块。本篇就详细讲解一下transformer 的原理和实际应用场景。

2017 年,Google 的研究人员发表了一篇题为《Attention is All You Need》的论文,提出了 Transformer 架构。这一架构彻底改变了 NLP 领域的发展和研究,其核心思想是基于注意力机制(Attention Mechanism)来取代传统的序列处理方法。

论文中展示的transformer 架构如图4-1:

图4-1 transformer 架构

其核心结构包括:

  • 自注意力机制(Self-Attention)

Transformer 使用自注意力机制来计算输入序列中每个位置信息与其他所有位置信息的关系。

这种机制使得模型能够并行处理序列中的所有位置,从而加速训练过程。

  • 编码器解码器结构(Encoder-Decoder Structure)

Transformer 由多个编码器(Encoder)和解码器(Decoder)层组成,图4-1中左右虚线框内即分别为编码器和解码器。

编码器层负责处理输入序列,而解码器层负责生成输出序列。

  • 多头注意力(Multi-head Attention):

为了增强模型的表达能力,Transformer 在每个注意力层中使用多个注意力头(Head),每个头关注输入的不同部分。多头注意力使得模型可以从不同的角度捕捉输入序列中的依赖关系,同时能够直接建模任意距离的词元之间的交互关系。具体对比,参见文末引用说明[1]

  • 位置编码(Positional Encoding):

由于 Transformer 没有显式的序列位置信息,因此引入了位置编码来提供位置信息。位置编码通常使用正弦函数来表示位置信息。这于位置嵌入稍有不同,位置编码是有固定函数生成,而位置嵌入是将位置信息视为可学习的参数,通过嵌入层(Embedding Layer)生成。

Transformer论文中的位置编码公式
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

# 使用PyTorch实现位置嵌入(可学习的嵌入层)
self.position_embeddings = nn.Embedding(max_seq_len, hidden_size)
  • 前馈神经网络Feed-Forward Networks):

对每个位置的输出进行变换,通常包含两个线性层和一个 ReLU 激活函数。

  • 残差连接与层归一化Residual Connections and Layer Normalization)

Transformer 使用残差连接来缓解梯度消失问题,并通过层归一化来稳定训练过程。

Transformers架构最初是为序列-序列任务(如机器翻译)设计的,但编码器和解码器模块很快就被改编为独立的模型。虽然有数百种不同的Transformers模型,但大多数属于下列三种类型之一:

  • Only-encoder

  这类架构将输入的文本序列转换为丰富的数字表示,适合于文本分类或命名实体识别等任务。BERT及其变种,如RoBERTa和DistilBERT,属于这一类架构。在这类架构中,使用双向编码器来学习词的上下文表示,这意味着它同时考虑了词的左侧和右侧的上下文信息。

  • Only-decoder

  这类架构主要用于生成任务,例如文本生成、补全文本等,它仅考虑当前词左侧的上下文信息,通常又被称为因果或者自回归注意,GPT就属于这类架构。

  • Encoder-decoder

  这类架构用于建模从一个文本序列到另一个序列的复杂映射,它们适用于机器翻译和总结任务等,典型的T5(Text-to-Text Transfer Transformer)属于这类架构。

注意力机制

注意力是一种人类不可或缺的复杂认知功能,指人可以在关注一些信息的同时忽略另一些信息的选择能力.在日常生活中,我们通过视觉、听觉、触觉等方式接收大量的感觉输入。但是人脑还能在这些外界的信息轰炸中有条不紊地工作,是因为人脑可以有意或无意地从这些大量输入信息中选择小部分的有用信息来重点处理,并忽略其他信息.这种能力就叫作注意力(Attention)。

注意力一般分为两种:

(1) 自上而下的有意识的注意力,称为聚焦式注意力(Focus Attention。聚焦式注意力也常称为选择性注意力。聚焦式注意力是指有预定目的、依赖任务的,主动有意识地聚焦于某一对象的注意力。

(2) 自下而上的无意识的注意力,称为基于显著性的注意力。基于显著性的注意力是由外界刺激驱动的注意,不需要主动干预,也和任务无关。

一个和注意力有关的例子是鸡尾酒会效应:

当一个人在吵闹的鸡尾酒会上和朋友聊天时,尽管周围噪音干扰很多,他还是可以听到朋友的谈话内容,而忽

略其他人的声音(聚焦式注意力)。同时,如果背景声中有重要的词(比如他的名字),他会马上注意到(显著性注意力)。

在机器学习任务中,注意力机制最初是为了改进序列到序列(Seq2Seq)模型中的编码器-解码器架构而提出的。传统的 Seq2Seq 模型使用一个编码器将输入序列编码为一个固定长度的向量,然后使用一个解码器根据这个向量生成输出序列。这种方法在处理长序列时效果不佳,因为很难通过一个固定长度的向量捕捉输入序列中的所有信息。

注意力机制通过允许解码器在生成输出时动态地关注输入序列的不同部分来解决这个问题。具体来说,注意力机制允许模型在处理序列数据时,动态地选择最相关的信息,并忽略不相关的信息。这种机制使得模型可以更好地处理长序列数据,并且提高模型的表达能力和泛化能力,其本质是信息筛选和聚焦。

  • 信息筛选:

注意力机制通过计算每个输入位置的重要性得分(Attention Scores),选择出最重要的部分进行处理。

这些得分反映了输入序列中各部分对当前任务的重要性,模型可以据此做出更有针对性的决策。

  • 信息加权:

注意力机制通过对输入信息进行加权平均(Weighted Sum),突出那些更重要的部分。

这种加权操作使得模型可以更好地捕捉输入序列中的依赖关系,并且能够处理长序列数据。

  • 动态分配:

注意力机制允许模型根据当前任务的需求动态地分配注意力,而不是固定地关注某些位置。

这种动态性使得模型可以更好地适应不同的输入数据,并且提高模型的泛化能力。

注意力机制通常包含以下几个步骤:

  1. 将每个标记嵌入投射到三个向量中,称为查询(Query)、键(Key)和值(Value)。
  2. 计算注意力得分(Attention Scores):

对于输入序列中的每个位置,计算其与当前解码器状态的相关性得分。

  如何计算注意力得分呢?

常用的注意力得分计算函数包括:

  • 加性模型

\(s(x,q) =v^Ttanh(Wx +Uq)\)

  • 点积模型

$$s(x,q)=x^Tq$$

  • 缩放点积模型

$$s(x,q)=x^Tq/( \sqrt D)$$,当输入项量维度D比较大时,点积模型的值通常有比较大的方差,从而导致Softmax函数的梯度会比较小.因此,缩放点积模型可以较好地解决这个问题。

  • 双线性模型

$$s(x,q)=x^TWq$$,相比点积模型,双线性模型在计算相似度时引入了非对称性。

以上函数中,$$W,U,v$$为可学习的参数,$$D$$为输入向量的维度。

当查询和键是不同长度的矢量时,可以使用可加性注意力评分函数。当它们的长度相同时,使用缩放的“点-积”注意力评分函数的计算效率更高。

  1. 应用 softmax 函数:

对注意力得分应用 softmax 函数,将其转换为概率分布。这些概率值反映了输入序列中各部分的相对重要性。

  1. 加权平均(Weighted Sum):

根据注意力得分对输入序列进行加权平均,得到一个上下文向量(Context Vector)。这个上下文向量包含了输入序列中最相关信息的综合表示。

下面是一个注意力机制的简单示例(代码段4-1):

import torch
import torch.nn as nn
import torch.nn.functional as F

class Attention(nn.Module):
    def __init__(self, embed_size):
        super().__init__()
        self.embed_size = embed_size
        # 1. 线性投影层(修正拼写)
        self.query_projection = nn.Linear(embed_size, embed_size)
        self.key_projection = nn.Linear(embed_size, embed_size)
        self.value_projection = nn.Linear(embed_size, embed_size)
        
    def forward(self, queries, keys, values, mask=None):
        # 2. 投影到查询/键/值空间
        Q = self.query_projection(queries)  # [B, T, D]
        K = self.key_projection(keys)       # [B, T, D]
        V = self.value_projection(values)   # [B, T, D]
        
        # 3. 计算缩放点积注意力
        d_k = K.size(-1)
        energy = torch.bmm(Q, K.transpose(1, 2)) / (d_k ** 0.5)  # 添加缩放因子
        
        # 4. 计算注意力权重
        attention_scores = F.softmax(energy, dim=-1)
        
        # 5. 加权求和
        context_vector = torch.bmm(attention_scores, V)
        return context_vector, attention_scores

# 示例用法
embed_size = 512  
attention = Attention(embed_size)
queries = torch.randn(32, 10, embed_size)  # [B, T, D]
keys = torch.randn(32, 10, embed_size)
values = torch.randn(32, 10, embed_size)

# 可选:创建掩码(如处理填充符)
mask = torch.ones(32, 10, 10)  # 示例中全1掩码(无屏蔽)

context_vector, attention_scores = attention(queries, keys, values, mask)
print(context_vector.shape)    # [32, 10, 512]
print(attention_scores.shape)  # [32, 10, 10]

自注意力机制

自注意力机制是注意力机制的一种特殊形式,它允许序列中的每个位置通过计算其与其他所有位置的注意力得分来生成自己的表示。这种机制使得模型能够并行处理序列中的所有位置。

在自注意力机制(self-attention)中,“self”指的是该机制能够通过关联单一输入序列内的不同位置来计算注意力权重。这与传统的注意力机制形成了对比,在传统的注意力机制中,焦点在于两个不同序列之间的元素关系,例如在序列到序列(seq2seq)模型中,注意力可能存在于输入序列和输出序列之间。

此处我们使用缩放点积模型来计算注意力得分。在机器学习中,点积运算可以用于衡量两个向量在数值上的相似程度,在多维张量(Tensor)的情况下,点积运算可以扩展为计算两个张量之间的相似度,代码段4-2。

import torch

# 定义两个向量
vector_a = torch.tensor([1, 2, 3])
vector_b = torch.tensor([4, 5, 6])

# 计算点积
dot_product = torch.dot(vector_a, vector_b)
print(f"Dot product: {dot_product}")

# 定义两个矩阵
matrix_a = torch.tensor([[1, 2], [3, 4]])
matrix_b = torch.tensor([[5, 6], [7, 8]])

# 计算矩阵乘积
matrix_product = torch.matmul(matrix_a, matrix_b)
print(f"Matrix product:\n{matrix_product}")

而在自注意力机制中,注意力得分是用来衡量输入序列中不同位置之间相关性的指标。具体来说,注意力得分反映了序列中一个位置的“查询”(Query)与另一个位置的“键”(Key)之间的相似度,那我们就可以通过Q和K的点积,获得对应的注意力得分,相似的查询和键将有一个大的点积。这正如代码段4-1第22行展示的那样:

22 energy = torch.bmm(Q, K.transpose(1, 2)) / (d_k ** 0.5)  # 添加缩放因子

torch.bmm 表示批量矩阵乘法(Batch Matrix Multiplication)

其实,自注意力机制也被称为缩放点积注意力(scaled dot-product attention)。

值得一提的是,“键”(Key)、“查询”(Query)和“值”(Value)这几个术语借鉴了信息检索和数据库领域的概念,

  • 查询(Query)

查询类似于数据库中的搜索查询。

  • 键(Key)

键类似于数据库中的键,用于索引和搜索。

  • 值(Value)

值类似于数据库中的键值对中的值。一旦模型确定了哪些键(以及输入序列中的哪些部分)与查询(即当前的关注项)最为相关,它就会检索出相应的值。这些值随后会被用来形成最终的输出表示。

那原始输入数据是如何映射成Q、K和V的呢, 我们通过代码来展示这一过程,代码段4-3。

import torch
inputs = torch.tensor(
  [[0.13, 0.15, 0.89],
   [0.25, 0.87, 0.66], 
   [0.57, 0.85, 0.64],
   [0.22, 0.58, 0.33], 
   [0.77, 0.25, 0.10],
   [0.05, 0.80, 0.55]] 
)

x_2 = inputs[1]            #
input_dim = inputs.shape[1] # shape [6,3]
out_dim = 3


torch.manual_seed(168)

#初始化 query 、key、value 为可训练参数
query = torch.nn.Parameter(torch.rand(input_dim , out_dim ), requires_grad=True)
key   = torch.nn.Parameter(torch.rand(input_dim , out_dim ), requires_grad=True)
value = torch.nn.Parameter(torch.rand(input_dim , out_dim ), requires_grad=True)


query_2 = x_2 @ query 
key_2 = x_2 @  key 
value_2 = x_2 @ value

#生成带可训练权重系数的 x_2
print(query_2)
>>> tensor([0.9990, 0.8667, 0.5560], grad_fn=<SqueezeBackward4>)

有了上面的基础,以下代码实现一个带有可训练权重的自注意力框架,代码段4-4。

import torch
import torch.nn as nn

class Attention(nn.Module):
    def __init__(self, input_dim, out_dim):
        super().__init__()
        self.W_query = nn.Parameter(torch.empty(input_dim, out_dim))
        self.W_key = nn.Parameter(torch.empty(input_dim, out_dim))
        self.W_value = nn.Parameter(torch.empty(input_dim, out_dim))

        # 使用Xavier均匀分布初始化权重矩阵
        nn.init.xavier_uniform_(self.W_query)
        nn.init.xavier_uniform_(self.W_key)
        nn.init.xavier_uniform_(self.W_value)

    def forward(self, inputs):
        keys = inputs @ self.W_key
        queries = inputs @ self.W_query
        values = inputs @ self.W_value
        
        attn_scores = queries @ keys.transpose(-2, -1)  # 计算注意力得分
        scaling_factor = keys.size(-1) ** 0.5
        attn_weights = torch.softmax(attn_scores / scaling_factor, dim=-1)  # 计算注意力权重
        context_vec = attn_weights @ values
        return context_vec
        

一般来说,在计算注意力得分的部分,点积可能产生任意大的数字,这会破坏训练过程的稳定性。为了处理这个问题,首先将注意力分数乘以一个比例因子(上述代码为 0.05),使其方差正常化,然后用softmax进行归一 化,以确保所有的列值相加为1。

因果注意力机制

也称为遮蔽注意力(Masked Attention),主要应用于自回归(autoregressive)模型中,在这种情况下,模型在生成序列中的每一个词时,只能参考它前面的词,而不能看到当前位置及其之后的任何信息。这样做的目的是模拟自然语言生成的过程,即在说话或写作时,我们只能基于已有的信息来决定下一个词是什么。

在Transformer模型中,因果注意力常用于解码器部分,以确保在生成序列时不会泄露未来的信息。实现因果注意力的方法是在计算注意力权重之前,给注意力得分矩阵应用一个掩码(mask),这样模型就不会对那些不应该看到的位置分配注意力。

代码段4-5中,在计算注意力得分之前,我们创建了一个下三角矩阵作为掩码。这个掩码确保在计算注意力得分时,对于每个位置i ,只有位置$$ j<i$$的得分会被保留,其余位置的得分会被设置为 -inf。

import torch
import torch.nn as nn

class CausalSelfAttention(nn.Module):
    def __init__(self, input_dim):
        super(CausalSelfAttention, self).__init__()

        # 创建线性层用于查询、键和值的变换
        self.W_query = nn.Linear(input_dim, input_dim, bias=False)
        self.W_key = nn.Linear(input_dim, input_dim, bias=False)
        self.W_value = nn.Linear(input_dim, input_dim, bias=False)

    def forward(self, inputs):
        batch_size, seq_len, input_dim = inputs.size()

        # 对输入进行线性变换得到查询、键和值
        queries = self.W_query(inputs)
        keys = self.W_key(inputs)
        values = self.W_value(inputs)

        # 计算注意力得分
        attn_scores = torch.matmul(queries, keys.transpose(-2, -1)) / (input_dim ** 0.5)

        # 创建一个三角形的mask来实现因果注意力
        mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, seq_len, seq_len).to(inputs.device)
        attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        # 应用softmax函数获得注意力权重
        attn_weights = torch.softmax(attn_scores, dim=-1)

        # 加权求和得到上下文向量
        context_vec = torch.matmul(attn_weights, values)

        return context_vec

多头注意力机制

到目前为止,我们只是 “按原样 “使用嵌入来计算注意力分数和权重,每个嵌入进行了三个独立的线性转换,以生成Q、K和V向量。 这些变换对嵌入进行投射,每个投射都有自己的一套可学习的参数,这使得自我注意层能够关注序列的不同语义方面。 事实证明,拥有多组线性投射也是有益的,每组投射代表一个所谓的注意头,由此产生了多头注意力。

多头注意力允许模型从不同的表示子空间中并行地关注输入的不同部分。通过将输入分割成多个头(head),每个头都可以学习输入的不同表示,从而增强模型捕捉复杂依赖关系的能力。

与单头注意力相比,多头注意力机制的主要区别在于它使用了多组结构相同但映射参数不同的自注意力模块。输入序列首先通过不同的权重矩阵被映射为一 组查询、键和值。每组查询、键和值的映射构成一个“头”,并独立地计算自注意 力的输出。最后,不同头的输出被拼接在一起,并通过一个权重矩阵进行映射,产生最终的输出。

代码段4-6展示了多头注意力的代码示例。

import torch
import torch.nn as nn

class MultiHeadCausalSelfAttention(nn.Module):
    def __init__(self, input_dim, num_heads=1):
        super(CausalSelfAttention, self).__init__()
        assert input_dim % num_heads == 0

        self.input_dim = input_dim
        self.num_heads = num_heads
        self.head_dim = input_dim // num_heads

        self.W_query = nn.Linear(input_dim, input_dim, bias=False)
        self.W_key = nn.Linear(input_dim, input_dim, bias=False)
        self.W_value = nn.Linear(input_dim, input_dim, bias=False)

    def forward(self, inputs):
        batch_size, seq_len, _ = inputs.size()

        # 对输入进行线性变换得到查询、键和值
        queries = self.W_query(inputs)
        keys = self.W_key(inputs)
        values = self.W_value(inputs)

        # 将输入分为多个头
        queries = queries.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = keys.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        values = values.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力得分
        attn_scores = torch.matmul(queries, keys.transpose(-2, -1)) / (self.head_dim ** 0.5)

        # 创建一个三角形的mask来实现因果注意力
        mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len).to(inputs.device)
        attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        # 应用softmax函数获得注意力权重
        attn_weights = torch.softmax(attn_scores, dim=-1)

        # 加权求和得到上下文向量
        context_vec = torch.matmul(attn_weights, values)

        # 拼接不同的头
        context_vec = context_vec.transpose(1, 2).contiguous().view(batch_size, seq_len, self.input_dim)

        return context_vec

具体的, 通过视图重排和转置操作,将原始张量划分为具有多头的张量。

queries = queries.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

原始的张量形状为 [batch_size, seq_len, num_heads, input_dim],我们希望将其转换为 [batch_size, num_heads, seq_len, head_dim],以便每个头可以独立地执行注意力计算。

假设我们有一个小批量的输入张量 inputs,其形状为 [batch_size, seq_len, input_dim]。例如,假设 batch_size = 2(有两个样本),seq_len = 4(每个样本有四个时间步),input_dim = 8(输入维度为8)。我们还假设 num_heads = 2(有两个注意力头),这意味着每个头的维度 head_diminput_dim // num_heads = 4

  1. 输入张量
inputs: shape (2, 4, 8)
  1. 线性变换

首先,我们需要对输入进行线性变换,以得到查询、键和值张量。假设经过线性变换后的形状仍然是 (2, 4, 8)

  1. 视图重排

如果我们有一个形状为 (2, 4, 8) 的张量,我们可以将其重排为 (2, 4, 2, 4)。这意味着每个样本有4个时间步,每个时间步被分成两个头,每个头的维度为4。

queries: shape (2, 4, 8) -> (2, 4, 2, 4)
  1. 转置操作

现在我们需要将张量的形状从 [batch_size, seq_len, num_heads, head_dim] 转换为 [batch_size, num_heads, seq_len, head_dim]。为此,我们需要调换第二维和第三维的位置。

queries: shape (2, 4, 2, 4) -> (2, 2, 4, 4)

以上,就是多头注意力机制的实现。

前馈神经网络

前馈神经网络通常由两个全连接(线性)层组成,中间夹着一个非线性激活函数。这个结构的主要目的是为模型提供非线性的映射能力,使得模型能够捕捉更复杂的特征。

在Transformer中,前馈神经网络位于编码器和解码器的每个子层中,并且通常与自注意力机制并行工作,然后通过残差连接和层归一化(Layer Normalization)连接起来。

代码段4-7是一个典型的前馈神经网络的实现。

import torch
import torch.nn as nn

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(FeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        # 第一层线性变换
        x = self.linear1(x)
        # 非线性激活函数
        x = self.relu(x)
        # Dropout层
        x = self.dropout(x)
        # 第二层线性变换
        x = self.linear2(x)
        return x

# 示例
d_model = 512  # 输入/输出维度
d_ff = 2048    # 中间维度
dropout = 0.1

ffn = FeedForward(d_model, d_ff, dropout)

# 创建一个随机输入
batch_size = 2
seq_len = 10
input_tensor = torch.randn(batch_size, seq_len, d_model)

# 前向传播
output = ffn(input_tensor)
print("Output shape:", output.shape)

>>>Output shape: torch.Size([2, 10, 512])

残差连接与层归一化

残差连接(Residual Connections)

残差连接是一种用于深度神经网络的技术,通过将输入直接加到层的输出上来缓解梯度消失问题,并有助于加速收敛。

核心思想是允许输入信号直接传递到后面的层,与经过非线性变换后的输出相加,使得网络只需学习输入与输出之间的残差(即变化部分),而非完整的映射。

数学表达:$$输出=F(x)+x$$ 其中,x是输入,F(x)是网络层(如卷积层、全连接层)对输入的非线性变换。

在Transformer中,残差连接主要用于自注意力层和前馈神经网络层。

层归一化(Layer Normalization)

大语言模型的预训练过程中经常会出现不稳定的问题,通常会采用特定的归一化策略来加强和改善训练过程的稳定性。

常用的归一化技术包括批归一化、层归一化、均方根归一化、深层归一化。

  • 批次归一化(Batch Normalization,简称 BatchNorm 或 BN),其主要思想是在每个 mini-batch 的训练过程中对激活值进行归一化,以使它们具有零均值和单位方差。这样做的目的是为了减少内部协变量偏移(internal covariate shift),即在网络训练过程中,由于前面层的参数变化导致后面层的输入分布发生变化,从而影响后面的层的学习效果。
  • 层归一化(Layer Normalization),它通过归一化神经网络中每个样本的特征值来减少内部协变量偏移(internal covariate shift),从而加速训练过程并提高模型的稳定性。与批次归一化(Batch Normalization)不同的是,层归一化是在每个样本的特征维度上进行的,而不是在 mini-batch 维度上进行的,这样可以避免由于 mini-batch 间统计特性的差异导致的不稳定问题。

层归一化的结构如代码段4-1

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super(LayerNorm, self).__init__()
        # 设置一个小的正数作为方根分母的平滑因子,防止除以零
        self.eps = 1e-5
        # 可学习的缩放参数
        self.scale = nn.Parameter(torch.ones(emb_dim))
        # 可学习的偏移参数
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        # 计算沿着最后一个维度(特征维度)的均值
        mean = x.mean(dim=-1, keepdim=True)
        # 计算沿着最后一个维度的方差
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        # 对输入进行归一化
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        # 返回经过缩放和平移变换后的归一化结果
        return self.scale * norm_x + self.shift

# 示例使用:
ln = LayerNorm(emb_dim=10)  # 假设嵌入维度为10
input_tensor = torch.randn(32, 10)  # 批次大小为32,特征维度为10
output_tensor = ln(input_tensor)
print(output_tensor.shape)  # 应该打印: torch.Size([32, 10])
  • 均方根归一化(RMSNorm),该方法仅利用激活值总和的均方根 RMS(𝒙) 对激活值进行重新缩放。
  • 深层归一化(DeepNorm),该方法在 LayerNorm 的基础上,在残差连接中对之前的激活值 𝒙 按照一定比例 𝛼 进行放缩。通过这一简单的操作,Transformer 的层数可以被成功地扩展至 1,000 层。

归一化模块的位置

归一化模块的位置通常有三种选择,分 别是层后归一化(Post-Layer Normalization, Post-Norm)、层前归一化(Pre-Layer Normalization, Pre-Norm)和夹心归一化(Sandwich-Layer Normalization, SandwichNorm)。具体的,

归一化类型归一化位置操作顺序优点缺点
层后归一化(Post-Norm)激活函数之后x = self.linear(x) x = self.activation(x) x = self.norm(x)1. 有助于加快神经网络的训练收敛速度,使模型可以更有效地传播梯度,从而减少训练时间。2. 后向归一化可以降低神经网络对于超参数(如学习率、初始化参数等)的敏感性,使得网络更容易调优,并减少了超参数调整的难度。训练稳定性:在某些情况下,激活后的数据可能具有较大的动态范围,导致训练不稳定。
层前归一化(Pre-Norm)激活函数之前x = self.norm(x) x = self.linear(x) x = self.activation(x)训练稳定性:归一化在激活之前进行,可以更好地控制输入数据的分布,缓解模型的
梯度爆炸或者梯度消失现象,提高训练稳定性。
激活函数的特性可能会受到归一化的影响,导致某些激活函数的特性不能完全发挥。训练性能逊色于层后归一化。
夹心归一化(Sandwich-Norm)激活函数前后x = self.norm1(x) x = self.linear(x) x = self.norm2(x) x = self.activation(x)双重控制:通过两次归一化操作,可以更好地控制输入和输出数据的分布,提高训练的稳定性和收敛速度。灵活性:可以同时利用激活前和激活后的归一化优势。计算开销:由于需要进行两次归一化操作,计算开销较大。复杂性:实现和调试可能会更加复杂。

综合以上各部分,我们得到了完整的transformer 架构示例,如代码段4-8

class Transformer(nn.Module):
    def __init__(self, d_model, d_ff, num_heads, dropout=0.1):
        super(TransformerSublayer, self).__init__()
        self.self_attn = MultiHeadCausalSelfAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        # 自注意力机制
        src2 = self.self_attn(src)
        # 残差连接 + 层归一化
        src = src + self.dropout(src2)
        src = self.norm1(src)

        # 前馈神经网络
        src2 = self.ffn(src)
        # 残差连接 + 层归一化
        src = src + self.dropout(src2)
        src = self.norm2(src)
        
        return src

附录

torch.view 函数

在 PyTorch 中,view 函数用于改变张量的形状而不改变其数据。view 是一种非常有用的操作,因为它允许你在不改变底层数据的情况下重新组织张量的维度。这对于将张量重塑成所需的形状以便进行下一步处理非常有用。

假设我们有一个形状为 (2, 4, 8) 的张量 x,我们想将其重塑为 (2, 4, 2, 4) 的形状:

import torch

# 创建一个形状为 (2, 4, 8) 的随机张量
x = torch.randn(2, 4, 8)
print("Original tensor shape:", x.shape)

# 使用 view 函数重塑张量
reshaped_x = x.view(2, 4, 2, 4)
print("Reshaped tensor shape:", reshaped_x.shape)

有时候,我们可能不知道某个维度的确切大小,但知道其他维度的大小。在这种情况下,可以使用 -1 来自动推断未知的尺寸:

# 使用 -1 自动推断未知尺寸
reshaped_x = x.view(2, 4, -1, 4)
print("Reshaped tensor shape with -1:", reshaped_x.shape)

>>> Reshaped tensor shape with -1: torch.Size([2, 4, 2, 4])

扩展

其他注意力机制形式

  • 稀疏注意力,滑动窗口注意力机制(Sliding Window Attention, SWA)是大语言模型中使用 最多的一种稀疏注意力机制。不同于完整的注意力机制,滑动窗口注意力根据词元位置,仅仅将位置索引上距离该词元一定范围内的词元考虑到注意力的计算中。 具体来说,滑动窗口注意力设置了一个大小为 𝑤 的窗口,对每个词元 𝑢𝑡,只对窗 口内的词元 [𝑢𝑡−𝑤+1, . . . , 𝑢𝑡] 进行注意力计算,从而将复杂度降低到 𝑂(𝑤𝑇)。进一 步,通过信息的逐层传递,模型实现了随着层数线性增长的感受野,从而获取远处词元的信息。
  • 多查询注意力(MultiQuery Attention, MQA), 为了提升注意力机制的效率,多查询注意力(MultiQuery Attention, MQA)提出针对不同的头共享相同的键和值变换矩阵 。这种方法减少了访存量,提高了计算强度,从而实现了更快的解码速度,并且对于模型性能产生的影响也比较小。
  • 分组查询注意力(Grouped-Query Attention, GQA),将全部的头划分为若干组,并且针对同一组内的头共享相同的变换矩阵。这种注意力机制有效地平衡了效率和性能,被 LLaMA-2 模型所使用。
  • 硬件优化的注意力,除了在算法层面上提升注意力机制的计算效率,还 可以进一步利用硬件设施来优化注意力模块的速度和内存消耗。其中,两个具有 代表性的工作是 FlashAttention 与 PagedAttention。相比于传统的注意力实现方式,FlashAttention 通过矩阵分块计算以及减少内存读写次数的方式,提高注意力分数的计算效率;PagedAttention 则针对增量解码阶段,对于 KV 缓存进行分块存储,并优化了计算方式,增大了并行计算度,从而提高了计算效率。

归一化示例

批次归一化(BN)计算

  1. 计算每个特征通道的均值和方差
  • 特征1:样本1=1,样本2=4
    $$(\mu_{\text{feature1}} = \frac{1+4}{2} = 2.5)$$
    $$(\sigma^2_{\text{feature1}} = \frac{(1-2.5)^2 + (4-2.5)^2}{2} = \frac{2.25 + 2.25}{2} = 2.25)$$
  • 特征2:样本1=2,样本2=5
    $$(\mu_{\text{feature2}} = \frac{2+5}{2} = 3.5)$$
    $$(\sigma^2_{\text{feature2}} = \frac{(2-3.5)^2 + (5-3.5)^2}{2} = \frac{2.25 + 2.25}{2} = 2.25)$$
  • 特征3:样本1=3,样本2=6
    $$(\mu_{\text{feature3}} = \frac{3+6}{2} = 4.5)$$
    $$(\sigma^2_{\text{feature3}} = \frac{(3-4.5)^2 + (6-4.5)^2}{2} = \frac{2.25 + 2.25}{2} = 2.25)$$
  1. 归一化每个特征,假设 $$(\gamma=1), (\beta=0), (\epsilon=0)$$
  • 样本1的归一化
    $$[
    \text{BN}(X_1) = \left[\frac{1-2.5}{\sqrt{2.25}}, \frac{2-3.5}{\sqrt{2.25}}, \frac{3-4.5}{\sqrt{2.25}}\right] = [-1, -1, -1]
    ]$$
  • 样本2的归一化
    $$[
    \text{BN}(X_2) = \left[\frac{4-2.5}{\sqrt{2.25}}, \frac{5-3.5}{\sqrt{2.25}}, \frac{6-4.5}{\sqrt{2.25}}\right] = [1, 1, 1]
    ]$$

最终结果

BN(X) = [[-1, -1, -1],
         [ 1,  1,  1]]

层归一化(LN)计算

  1. 对每个样本计算所有特征的均值和方差
  • 样本1:特征值=[1, 2, 3]
    $$(\mu_{\text{layer1}} = \frac{1+2+3}{3} = 2)$$
    $$(\sigma^2_{\text{layer1}} = \frac{(1-2)^2 + (2-2)^2 + (3-2)^2}{3} = \frac{1 + 0 + 1}{3} ≈ 0.6667)$$
  • 样本2:特征值=[4, 5, 6]
    $$(\mu_{\text{layer2}} = \frac{4+5+6}{3} = 5)$$
    $$(\sigma^2_{\text{layer2}} = \frac{(4-5)^2 + (5-5)^2 + (6-5)^2}{3} = \frac{1 + 0 + 1}{3} ≈ 0.6667)$$
  1. 归一化每个样本,假设 $$(\gamma=1), (\beta=0), (\epsilon=0)$$。
  • 样本1的归一化
    $$[
    \text{LN}(X_1) = \left[\frac{1-2}{\sqrt{0.6667}}, \frac{2-2}{\sqrt{0.6667}}, \frac{3-2}{\sqrt{0.6667}}\right] ≈ [-1.2247, 0, 1.2247]
    ]$$
  • 样本2的归一化
    $$[
    \text{LN}(X_2) = \left[\frac{4-5}{\sqrt{0.6667}}, \frac{5-5}{\sqrt{0.6667}}, \frac{6-5}{\sqrt{0.6667}}\right] ≈ [-1.2247, 0, 1.2247]
    ]$$

最终结果

LN(X) ≈ [[-1.2247,  0.0,  1.2247],
         [-1.2247,  0.0,  1.2247]]

引用说明

[1]与多头注意力机制相比,循环神经网络迭代地利用前一个时刻的状态更新当前时刻的状态, 因此在处理较长序列的时候,常常会出现梯度爆炸或者梯度消失的问题。而在卷积神经网络中,只有位于同一个卷积核的窗口中的词元可以直接进行交互,通过堆叠层数来实现远距离词元间信息的交换。

具体来说

1. 循环神经网络(RNN)的长期依赖问题

核心机制:
  • RNN 按时间步迭代处理序列:每个时刻的状态 ( h_t ) 依赖前一个时刻的状态 ( h_{t-1} ) 和当前输入 ( x_t ),即:

    $$h_t = f(h_{t-1}, x_t) $$
    其中 ( f ) 是非线性激活函数(如 tanh)。
梯度问题根源
  • 反向传播通过时间(BPTT):训练时需沿时间轴反向传播梯度,梯度需通过链式法则逐时间步相乘。
    梯度消失:若每个时间步的梯度 $$( \frac{\partial h_t}{\partial h_{t-1}} < 1 )$$,梯度会指数级衰减,较早时间步的参数无法有效更新。 • 梯度爆炸:若 $$( \frac{\partial h_t}{\partial h_{t-1}} > 1 )$$,梯度会指数级增长,导致数值不稳定。
长序列的挑战
  • 信息衰减:随着时间步增加,早期输入对后续状态的影响逐渐消失(例如,第100个词几乎无法影响第1000个词的状态)。
  • 示例:在句子“The cat, which ate the fish, was happy”中,RNN可能难以关联“cat”和“was”(相距较远的主谓一致)。

2. 卷积神经网络(CNN)的局部交互与层级传递

核心机制

  • 局部感受野:每个卷积核仅覆盖输入序列的一个局部窗口(如3个词),窗口内的词元通过权重共享进行交互。
  • 堆叠卷积层:通过多层卷积,逐步扩大感受野(如3层3-gram卷积的感受野为7个词),从而捕捉远距离依赖。
  • 层级抽象:浅层捕捉局部特征(如词组合),深层融合更大范围的上下文(如句子结构)。
  • 并行计算:不同位置的卷积窗口独立计算,无需依赖前一时刻的状态,避免梯度传播路径过长。
特性RNNCNN
信息传递方式时间步迭代依赖(序列式)局部窗口并行计算(层次式)
长距离依赖捕捉困难(梯度消失/爆炸)通过多层堆叠实现
训练效率低(无法并行处理时间步)高(窗口内并行计算)
典型改进方案LSTM/GRU(门控机制缓解梯度问题)空洞卷积(扩大感受野)、残差连接
适用场景短序列建模(如文本分类)长序列建模(如机器翻译、语音识别