理解 Transformer 神经网络中的注意力机制

图片

刚刚结束的 2022 年是人工智能领域取得诸多进步的不可思议的一年。人工智能领域最近取得的大多数里程碑式成就都是由一种称为 transformers 的特殊模型推动的,无论是 chatGPT 的惊人进步,还是稳定传播,它都为你的智能手机带来了科幻般的功能。即使是特斯拉的自动驾驶软件堆栈,也许是世界上部署最广泛的深度学习系统,也在底层使用了 transformer 模型(双关语)。“神经注意力机制”是 transformers 在各种任务和数据集上如此成功的秘诀。 

这是关于视觉转换器 (ViT) 的系列文章中的第一篇。在本文中,我们将了解注意力机制,并回顾导致该机制的思想演变。接下来,我们将直观地理解它。我们将用数学细节巩固直观理解,并最终通过在 PyTorch 框架中从头开始实现注意力机制将这种理解转化为代码。虽然我们将在本文结尾处专门讨论视觉转换器,但大部分讨论同样适用于大型语言模型 (LLM),例如 GPT-3 和最近发布的 chatGPT 模型。

  • 神经注意力简介
    1. 循环神经网络
    2. 长短期记忆网络
    3. 神经自注意力机制
  • 神经自注意力机制的直观解释
    1. 步骤1,查询,键值
    2. 第二步,计算注意力矩阵
    3. 步骤 3,规范化和注意力分数
    4. 第 4 步,参加
  • 多头自注意力
  • 自注意力机制的数学公式
  • PyTorch 实现 Self Attention
    1. 基本结果
  • 概括

神经注意力简介

为了对自注意力机制有最直观的理解,我们需要快速回顾一下过去十年自然语言处理 (NLP) 架构的发展情况。我们将保持讨论的独立性,不需要熟悉 NLP 领域。我们在这里主要关注的是思想的演变,而不是每个架构的精确数学细节。考虑到这一点,让我们回顾一下 NLP 中的思想流。

循环神经网络

图片
图 1. 循环块的示意图(来源:维基百科)

机器学习文献中已经充分认识到,自然语言中的句子是序列,因此能够处理序列的模型比“无状态”模型(例如普通的全连接深度神经网络)更适合 NLP 任务。因此,提出了循环块,其示意图如图 1 所示。基本思想是可以在深度神经网络中逐字处理句子,这样前一个单词的表示就会被注入到当前单词的隐藏层中。这使得 RNN 能够将数据的时间特性明确地纳入模型中。 

尽管 RNN 背后的理念非常深刻,但它们的数学公式使得它们无法对长序列进行建模。RNN 将序列建模为单词隐藏表示之间的乘法交互。因此,尝试对长序列进行建模会导致梯度爆炸和消失。需要修改 RNN 的设计以稳定训练。

长短期记忆网络

图片
图 2. LSTM 单元示意图。需要注意的关键区别是存在加法相互作用,而图 1 中则是乘法相互作用。

长短期记忆网络 (LSTM) 修改了原始 RNN 的架构,以稳定训练。Schmidhuber 教授与他的同事通过专门研究 RNN 的缺点提出了 LSTM。从概念上讲,LSTM 的关键创新是用加法交互取代 RNN 中的乘法交互。因此,梯度值在反向传播过程中不会衰减,因为它们是加法分布而不是乘法分布。 

因此,在 LSTM 中,梯度幅度不会爆发或消失,因此可以对长序列进行建模。作为计算机视觉专家,您可能已经注意到,He 等人在设计 ResNet(称为残差块)时使用了相同的加性相互作用思想。在这里,我们看到了使用相同思想改进 NLP 和计算机视觉模型的第一个证据。这一趋势在视觉转换器的设计中得到了明确的推动,我们将在本系列的下一篇文章中看到。

尽管 LSTM 比 RNN 有所改进,但建模超长序列的问题仍然存在。此外,LSTM 就像普通的 RNN 一样,按顺序(即一个接一个)处理单词。因此,对长序列的推理速度很慢。

神经自注意力机制

虽然 LSTM 旨在解决 RNN 的数学缺陷,但神经自注意力机制(或简称注意力机制)则从头开始,重新思考 RNN 背后的动机。虽然句子确实是单词序列,但人类不会按顺序、一个接一个地处理单词,而是以块为单位。此外,有些单词对于预测下一个单词高度相关,但大多数单词并非如此。 

LSTM 被迫按顺序处理所有单词,而不管它们的相关性如何,这限制了模型可以学习的内容。是否有一种机制可以让模型将特定单词的相关性纳入其自身? 

换句话说,在预测下一个单词时,我们应该关注哪些先前的单词?这是注意力机制起源背后的动机问题。事实证明,答案不仅与语言模型相关,而且还导致了 ViT 的发展。

神经自注意力机制的直观解释

注意机制不是将输入作为序列来处理,而是将其作为没有任何明确时间联系的块来处理。让我们在对句子中的下一个单词进行建模的背景下逐步了解它的工作原理。

图片
图 3. 注意力机制第一步

步骤 1,查询,键值:与 RNN 非常相似,句子中的每个单词都会“发出”其自身的隐藏表示。为了让该机制能够模拟相关性,我们让每个单词提出问题并得到答案。 

当一个单词提出的问题与另一个单词发送的答案匹配时,我们会将此解释为两个单词彼此具有高度相关性。一个单词使用“查询”向量向序列中的所有单词提出相同的问题。同样,它使用“关键”向量向所有单词提供相同的答案。 

在注意力机制最常用的变体中,还使用单独的“值”向量,以允许模型非线性地组合查询和键向量的输出。这增加了模型的表达能力。这三种类型的向量都是其单独密集层的输出,并且这三种类型的向量具有相同的大小,例如 d。总而言之,在注意力机制的第 1 步中,如图 3 所示,句子中的每个单词使用三个单独的密集层发出自身的三个独立表示,称为查询 Q、键 K 和值 V 向量。 

图片
图 4.第 2 步,根据查询和关键向量计算注意力矩阵。

第 2 步,计算注意力矩阵:如果我们用 N 个单词来建模一个句子,那么经过第 1 步之后,我们就有 N 个查询向量、N 个键向量和 N 个值向量。那么,一个单词究竟是如何提出问题并得到答案的呢?这是由注意力矩阵完成的,它是注意力机制的核心。我们将所有 N 个查询向量与所有 N 个键向量进行点积,如图 4 所示。由于两个向量的点积(无论多大)只是一个数字,因此这些 NxN 点积的结果是一个 NxN 矩阵图片,其每个元素图片是第 i 个查询向量与第 j 个关键字向量的点积。这里有三件重要的事情需要注意:

  1. 首先,由于在计算注意力矩阵时每个单词都会与其他单词交互,因此输入没有时间顺序。因此,原始注意力机制不考虑输入的顺序性。
  2. 其次,计算注意力矩阵的计算复杂度为图片。因此,如果我们想用两倍的单词数量来建模一个句子,就需要 4 倍的计算资源。当我们构建视觉变换器模型时,这将在本系列的下一部分中变得很重要。
  3. 第三,我们注意到自注意力机制中的“自我”指的是,我们还将给定单词的查询与其自己的关键向量进行点积运算。注意力机制的一些早期变体跳过了计算自相关性,导致输出包含 NxN – N 个元素(矩阵的所有对角线元素都缺失)。然而,自注意力机制的效果最好,并且在现代 GPU 上可高度并行化。这就是该机制被称为自注意力机制的原因。

步骤 3,规范化和注意力得分:如上所述,注意力矩阵将包含一些大数字和一些小数字。我们将大数字解释为代表高相关性,小相关性解释为低相关性。 

此时,我们可以通过多种方式使用该矩阵。我们可以使用最大池化操作挑选出每行中的最大元素,并将它们用于进一步处理。然而,这会丢弃大部分信息。此外,有时给定的单词可能需要来自多个先前单词的上下文才能正确地对数据进行建模。因此,我们应该尽可能保持表达能力。另一个重要要求是,无论我们选择使用注意矩阵,它都应该保持模型的可微性,因为模型需要通过反向传播进行训练。

softmax 层满足了所有这些要求。因此,我们对注意力矩阵进行逐行 softmax,得到一个 NxN 矩阵,其每行都是一个概率分布。这里的一个小实现细节是,在进行 softmax 之前,我们将注意力矩阵除以每个查询向量维度的平方根。这充当了一个合适的缩放因子,以“软化”得到的概率分布。 

缩放后,大概率值和小概率值之间的差异会减小,梯度可以相对均匀地流入多个单词位置。这类似于知识蒸馏算法如何将“温度”纳入教师模型的逻辑中。就像在知识蒸馏领域一样,通过温度因子缩放到 softmax 的输入可以改善梯度流并加快训练速度。我们将在下一篇文章中看到,这是视觉变换器成功的关键原因之一。

步骤 4,注意:我们将上一步中获得的归一化注意力分数与值向量进行点积。从数学上讲,这表示为

图片

由于使用密集层为每个单词获取值向量 V,因此注意力机制的输出具有 N 个向量,与输入单词的数量相同。最后,为了获得自注意力层的输出,我们使用密集层将 h 向量转换为与输入相同的大小。在上面的等式中,圆形符号表示点积,而不是矩阵乘法。就矩阵乘法而言,运算可能取决于查询和键向量的形状。我们将在下一节中更具体化符号。

多头自注意力

您可能对卷积网络非常熟悉。卷积运算是 CNN 的基本构建块。同样,上面介绍的自注意力机制是所有 Transformer 模型(包括大型语言模型和视觉 Transformer)的基本构建块。 

继续类比,我们永远不会设计一个只有一个卷积滤波器的卷积层,因为一个滤波器不足以模拟自然图像的复杂性。事实上,在一个层中看到超过 500 个卷积滤波器的卷积是很常见的。同样,如上所述,一个单一的注意力层不足以模拟自然语言的所有复杂性。 

因此,我们将多个注意力块(称为“头”)并行应用于同一序列。生成的层称为多头自注意力 (MHSA)。与卷积的区别在于,一个注意力头比单个卷积过滤器更具表现力。因此,我们不需要使用约 500 个注意力头,只需几十个即可。ViT 中注意力头数量的一些常见选择是 12、24 和 32。与卷积的另一个区别是,在 CNN 中,每个后续层中的过滤器数量都会增加,而最流行的 ViT 中的注意力头数量通常是恒定的,例如 12。

总而言之,自注意力机制超越了密集层、卷积层和 LSTM 单元,并提出了一种全新的通用计算机制来对数据中的关系进行建模。在 NLP 中,数据具有顺序性,但在计算机视觉中可能不是。然而,在本系列的下一篇中,我们将看到使用注意力机制构建的视觉转换器在计算机视觉应用中效果很好。

一些说明:我们在本系列文章中的目标是介绍视觉变换器。到目前为止,我们已经使用了自然语言处理中的例子来激发自我注意机制的引入。然而,从现在开始,我们将脱离 NLP,专门研究视觉变换器。MHSA在计算机视觉 (CV) 和 NLP 中的使用方式存在细微差别。特别是,现阶段有两个显着的差异:

  1. 视觉变换模型一般不使用掩蔽注意力(自监督学习除外)
  2. 视觉变换模型不使用交叉注意力(DETR 等物体检测模型除外)

随着本系列的进展,我们将在适当的时候介绍这些变体,但目前,请注意存在其他变体,本博文的其余部分将仅介绍视觉变换器中最常用的 MHSA 版本。

自注意力机制的数学公式

让我们用数学术语来表述单头注意力机制。

考虑一张以某种方式被分割成一系列块的图像图片,其中每个图片是维度为图片。这里将图像分割成块的具体方法并不重要。让我们构造矩阵图片第一维为长度 N,第二维为图片。如前所述,我们将使用密集层来获取查询、键和值表示。但是,这些密集层将在没有任何偏置项和激活函数的情况下使用,因此密集层只是简单的矩阵乘法。这三个层的参数分别为查询、键和值的 WQ 、WK 、WV ,使得 WQ、WK、WV图片. 另一个带权重的密集层图片和偏见图片也使用。让我们按照上一节直观解释的那样进行计算。

步骤 1:查询、键和值矩阵(这些是矩阵,因为我们在一个步骤中计算序列中所有单词的向量)计算如下:

图片
图片
图片

这里,图片是一个超参数,表示键、查询和值向量的维度。通常,图片远小于图片

第 2 步:注意力矩阵计算如下

图片

请注意,单个matmul操作会计算所有单词的注意力矩阵,包括同一单词的键与查询的点积。因此,它完全实现了自注意力机制。

步骤 3:缩放和 softmax 实现如下

图片

这里,dim=-1 表示 softmax 是在行而不是列中进行的(这没有意义)。

步骤 4:注意力的输出计算如下

图片

注意层的输出计算如下

图片

因此,输出的大小与输入的大小相同。

虽然我们在这里只制定了单头自注意力,但对多头版本的推广包含多个独立的查询、键和值权重矩阵,并且非常简单。

PyTorch 实现 Self Attention

现在我们对自注意力层有了直观和数学上的理解,让我们在 PyTorch 中实现它。在开始之前,有几点需要注意:

  1. 由于自注意力机制中使用了大量的密集层,因此存在过度拟合的危险。因此,在MHSA的实际实现中,大量使用了dropout层来避免过度拟合。
  2. PyTorch 有一个名为 nn.Transformer 的模块。它实现了完整的 Transformer 架构(编码器和解码器),而我们将仅实现模型的一部分(仅编码器)。

Einops 层: Einops 是一个出色的张量操作、重塑和调整大小库。我们可以使用 einops 中的 Rearrange 层,而不是使用 torch.reshape,它具有更出色的 API,并且可以与 PyTorch、TensorFlow、JAX 和 numpy 无缝集成。我们没有足够的篇幅在这里介绍 einops,但如果您以前从未见过 einops 代码,请不要担心,因为它非常易读,并且工作方式与您直观猜测的完全一样。

下载代码 为了轻松学习本教程,请点击下面的按钮下载代码。免费!

下载代码

考虑到这一点,让我们将 MHSA 层实现为 的子类nn.Module

123456789101112十三1415161718192021222324二十五二十六二十七二十八二十九三十31三十二33三十四三十五三十六三十七三十八三十九4041四十二43四十四四十五四十六四十七四十八4950515253545556from torch import nnfrom einops.layers.torch import Rearrange class MultiHeadedSelfAttention(nn.Module):  def __init__(self, indim, adim, nheads, drop):    '''    indim: (int) dimension of input vector    adim: (int) dimensionality of each attention head    nheads: (int) number of heads in MHA layer    drop: (float 0~1) probability of dropping a node         Implements QKV MSA layer    output = softmax(Q*K/sqrt(d))*V    scale= 1/sqrt(d), here, d = adim    '''    super(MultiHeadedSelfAttention, self).__init__()    hdim=adim*nheads    self.scale= hdim** -0.5 #scale in softmax(Q*K*scale)*V    self.key_lyr = self.get_qkv_layer(indim, hdim, nheads)    #nn.Linear(indim, hdim, bias=False)    #there should be nheads layers    self.query_lyr=self.get_qkv_layer(indim, hdim, nheads)    self.value_lyr=self.get_qkv_layer(indim, hdim, nheads)         self.attention_scores=nn.Softmax(dim=-1)    self.dropout=nn.Dropout(drop)         self.out_layer=nn.Sequential(Rearrange('bsize nheads indim hdim -> bsize indim (nheads hdim)'),    nn.Linear(hdim, indim),    nn.Dropout(drop))     def get_qkv_layer(self, indim, hdim, nheads):    '''    returns query, key, value layer (call this function thrice to get all of q, k & v layers)    '''    layer=nn.Sequential(nn.Linear(indim, hdim, bias=False),    Rearrange('bsize indim (nheads hdim) -> bsize nheads indim hdim', nheads=nheads))         return layer   def forward(self, x):    query=self.key_lyr(x)    key=self.query_lyr(x)    value=self.value_lyr(x)         dotp=torch.matmul(query, key.transpose(-1-2))*self.scale         scores=self.attention_scores(dotp)         scores=self.dropout(scores)         weighted=torch.matmul(scores, value)         out=self.out_layer(weighted)         return out

构造函数采用以下参数:

  • 每个输入的大小。 
  • 注意力层的大小
  • 注意力头的数量,以及
  • 辍学概率

步骤 1:该函数get_qkv_layer用于同时获取所有注意力头的查询、键和值向量的密集层。简单来说,这将在 for 循环中实现,每个注意力头进行一次迭代。但是,将多个矩阵乘法运算合并为一个可以减少延迟并提高 CUDA 占用率。einops 库中的 Rearrange 层重塑 QKV 矩阵,使头维度紧挨着批处理维度。

第 2 步:然后将注意力矩阵沿其最后两个维度转置,计算为matul查询和密钥的注意力矩阵,正如第 2 节和第 3 节中解释的那样。

步骤 3:将注意力矩阵按注意力维度的平方根缩放,然后采用 softmax 得到注意力得分矩阵。得分矩阵后面跟着一个 dropout 层,以防止过度拟合。

步骤4:利用得分矩阵计算值矩阵的加权和。最后,einops Rearrange 层将所有注意力头合并为一个,并使用线性层将注意力层的输出转换为与输入相同的形状。再次应用 Dropout 以防止过拟合。

4.1 基本结果

我们引入的 MHSA 层只是 Transformer 模型中使用的构建块。除了验证其自身的工作之外,我们无法用它做很多事情。我们将输入 x 传递到 MHSA 层,并验证输出是否具有与输入相同的形状。我们将在本系列的下一篇中使用此层来构建视觉 Transformer 模型。

图片
图5.验证MHSA层是否正常工作。

概括

这篇博文奠定了 transformer 模型的基础,尤其是视觉 transformer。机器学习的最新进展以 transformer 为核心,而 transformer 的核心是多头自注意力层。因此,直观和数学地理解注意力机制对于你未来的项目和职业生涯至关重要。 

我们首先回顾了循环网络和 LSTM 背后的核心概念。它们的成功和不足让我们了解了注意力机制的适用范围以及它与早期模型的不同之处。 

接下来,我们直观地了解了注意力层背后的机制。我们一步步地计算查询、键、值、注意力矩阵、分数,最后是注意力层的输出。在获得直观理解之后,我们将这些概念形式化为具体的数学形式。

最后,我们将数学和直观的理解转化为实现多头自注意力层的 PyTorch 代码。

现在我们已经打下了基础,我们可以在本系列的下一篇中轻松理解视觉转换器。我们希望在那里见到你。

源代码:

http://www.gitpp.com/datasets/learnopencv-cn/tree/master/Attention_Mechanism_Introduction

留下评论

您的电子邮箱地址不会被公开。 必填项已用 * 标注