想要更好地理解大模型架构？从计算参数量快速入手

2023-06-12
湖北
本文字数：6853 字
阅读完需：约 22 分钟

编者按：要理解一种新的机器学习架构（以及其他任何新技术），最有效的方法就是从头开始实现它。然而，还有一种更简单的方法——计算参数数量。
通过计算参数数量，读者可以更好地理解模型架构，并检查其解决方案中是否存在未被发现的错误。
该文章提供了精确的 Transformers 模型的参数量计算公式和不太准确的简略公式版本，使读者能够快速估算基于 Transformer 的任何模型中参数的数量。
以下是译文，Enjoy!

作者 | Dmytro Nikolaiev (Dimid)

编译 | 岳扬

要理解一种新的机器学习架构（以及其他任何新技术），最有效的方法就是从头开始实现它。 虽然这可能会非常复杂、耗时，并且有时几乎不可能实现，但这是帮助我们理解每个技术细节的最佳方法。例如，如果没有类似的计算资源或数据，我们将无法确保我们的解决方案中没有未被发现的错误。

然而，还有一种更简单的方法——计算参数数量。 这比仅仅阅读论文要困难得多，但可以让我们深入挖掘并检查是否完全理解了新架构的构件（在本文的例子是 Transformer 的编码器（Encoder）和解码器（Decoder）构件）。

我们可以通过下面这幅图表来思考这个问题，这张图表展示了三种理解新 ML 架构的方法——圆圈的大小表示对该架构的理解程度。

本文主要研究著名的 Transformer 架构，并考虑如何计算 PyTorch TransformerEncoderLayer[1]和 TransformerDecoderLayer[2]类中的参数数量。因此，我们需要确保对于该架构由哪些部分组成不再充满神秘感。

TLDR（总结）

（该文篇幅比较长，如果不想深入探讨或时间有限，可以直接看总结部分）

您可以阅读“结论 Conclusions”部分，所有参数量计算公式都总结在“结论 Conclusions”部分。

本文不仅提供精确的参数量计算公式，还能够提供不太准确的公式近似版本，将使您能够快速估算基于 Transformer 的任何模型中参数的数量。

01 Transformer 架构

著名的 Transformer 架构于 2017 年在《Attention Is All You Need[3]》这篇论文中提出，并因其具有能够有效捕捉长距离的依赖关系（long-range dependencies）的能力而成为自然语言处理和计算机视觉任务中的标准架构。

早在 2023 年初，扩散模型（Diffusion）[4]由于文转图生成模型[5]的大火而变得极其流行。也许，很快扩散模型将成为各种任务的最先进技术，就像 Transformer 与 LSTM 和 CNN 一样。但我们先来看看 Transformer……

本文并不试图去解释 Transformer 架构，因为已经有很多足够好的文章做到了这一点。这篇文章只是让我们能够从不同的角度去看待它，或者讲解一些细节问题。所以如果你正在寻找更多有关此架构的学习资源，我可以向你推荐一些；否则，您可以继续阅读下去。

1.1 了解更多 Transformer 的资源

如果你正在寻找更加详细的 Transformer 架构概述，可以阅读以下材料（请注意，互联网上有很多技术内容，我只是个人喜欢这些）：

首先，可以阅读官方论文[3]。第一次接触 Transformer 就阅读论文可能不是最佳方式，但这并不像看起来那么复杂。可以尝试使用 Explainpaper 来帮助您阅读此论文[6]或其他论文（这是一种基于 AI 的工具，可以解释用鼠标标记的文本）。
Jay Alammar 的“Great Illustrated Transformer[7]”。如果您不喜欢阅读文章，可以观看同一作者的 YouTube 视频[8]。
Lukasz Kaiser 在 Google Brain 的 “Awesome Tensor2Tensor” 讲座[9]。
如果想直接进行实操并使用各种 Transformer 模型构建应用程序，请查看 Hugging Face 课程[10]。

1.2 Original Transformer

首先，让我们回顾一下 Transformer 的基础知识。

Transformer 的架构由两个组件组成：编码器（在左边）和解码器（在右边）。编码器接受输入 token 序列并生成隐藏状态序列（sequence of hidden states），而解码器则接受这个隐藏状态序列并生成输出 token 序列。

Transformer 架构图，来自https://arxiv.org/pdf/1706.03762.pdf

编码器和解码器都由一堆相同的层组成。对于编码器，该层包括多头注意力（multi-head attention）（1——此处及下文中的数字指的是下面的图片中标序号的部分）和一个带有一些层归一化（3）和跳跃连接（skip connections）的前馈神经网络（feed-forward neural network）（2）。

解码器也类似于编码器，但除了第一个多头注意力（4）（在机器翻译任务中被屏蔽，所以解码器不会通过查看未来的 tokens 进行舞弊）和一个前缀网络（5）之外，它还具有第二个多头注意力机制（6）。它允许解码器在生成输出时使用编码器提供的上下文（context）。与编码器一样，解码器也有一些层归一化（layer normalization）（7）和跳跃连接组件。

带有序号标记组件的 Transformer 架构图

来自https://arxiv.org/pdf/1706.03762.pdf

我不会将输入嵌入层（带有位置编码）和最终输出层（linear+softmax）视为 Transformer 组件，而只关注编码器和解码器块。这样做是因为这些组件是适用于某些特定任务和嵌入方法的，而编码器和解码器栈是其他体系结构的基础。

这种架构的例子包括用于编码器的基于 BERT 的模型（BERT、RoBERTa、ALBERT、DeBERTa 等），用于解码器的基于 GPT 的模型（GPT、GPT-2、GPT-3、ChatGPT），以及构建在完整的编码器-解码器框架上的模型（T5、BART 等）。

尽管我们在该架构中标记了七个组件，但我们可以看到，其中仅有三个独特的组件：

多头注意力（Multi-head attention）；
前馈网络（Feed-forward network）；
层的归一化（Layer normalization）。

Transformer 构件来自论文https://arxiv.org/pdf/1706.03762.pdf

02 Transformer 构件块

让我们考虑一下每个模块的内部结构以及它需要多少参数。在本节中，我们还将开始使用 PyTorch[11]来验证我们的计算结果。

为了检查某个模型块的参数数量，我将使用以下这行函数[12]：

import torch
# https://discuss.pytorch.org/t/how-do-i-check-the-number-of-parameters-of-a-model/4325/9def count_parameters(model: torch.nn.Module) -> int: """ Returns the number of learnable parameters for a PyTorch model """ return sum(p.numel() for p in model.parameters() if p.requires_grad)

复制代码

在我们开始之前，请注意一个事实，即所有构件块都是标准化的，并且使用跳跃连接。这意味着所有输入和输出的 shape（更确切地说，是其最后一个数字 ， 因为 batch size 和 tokens 数量可能会有所不同）必须相同 。 对于原论文，这个数字（d_model）为 512。

2.1 多头注意力

著名的注意力机制是 Transformer 架构的关键。但是，无论设计动机和技术细节如何，它只涉及几个矩阵乘法。

Transformer 多头注意力架构图

来自论文https://arxiv.org/pdf/1706.03762.pdf

计算了每个 head 的注意力后，我们将所有 head 连接起来，并通过一个线性层（W_O 矩阵）进行传递。反过来，每个 head 都是用三个独立的矩阵乘以 query、key 和 value（分别为 W_Q、W_K 和 W_V 矩阵）的 Scaled dot-product attention(缩放点积注意力)。这三个矩阵对每个 head 都是不同的，这就是下标 i 出现的原因。

最终线性层（final linear layer）（W_O）的 shape 为 d_model 到 d_model。其余三个矩阵（W_Q、W_K 和 W_V）的 shape 相同：d_model 到 d_qkv。

请注意，在上面的图像中，d_qkv 被表示为原论文中的 d_k 或 d_v。我认为这个名称更直观，因为尽管这些矩阵可能具有不同的 shape，但几乎总是相同的。

此外，请注意，d_qkv = d_model / num_heads (文中的 h)。这就是为什么 d_model 必须能够被 num_heads 整除的原因：以确保后面的连接正确。

可以通过检查上图中的所有中间阶段的 shape（正确的 shape 在右下角标出）来自行测试。

因此，我们需要每个 head 有三个较小的矩阵和一个大的最终矩阵。那么我们需要多少参数（不要忽略偏差）？

用于计算 Transformer 注意力模块中参数数量的公式。图片由作者提供

我希望这个公式不会太繁琐——我试图让推导的结果尽可能的清晰。不要担心! 未来的公式会更加简短。

参数的大致数量是这样的，因为与 4 * d_model 相比，我们可以忽略 4 * d_model^2。让我们现在用 PyTorch 进行测试。

from torch import nn
d_model = 512n_heads = 8 # must be a divisor of `d_model`
multi_head_attention = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads)print(count_parameters(multi_head_attention)) # 1050624print(4 * (d_model * d_model + d_model)) # 1050624

复制代码

数字匹配，这意味着我们做得很好！

2.2 前馈网络

Transformer 中的前馈网络由两个全连接层（fully connected layers）组成，其中间有一个 ReLU 激活函数。该网络的内部部分比输入和输出（input and output）更具表现力（输入和输出必须相同）。

在一般情况下，它是 MLP(d_model, d_ff) -> ReLU -> MLP(d_ff, d_model)，对于原始论文，d_ff = 2048。

前馈神经网络描述图来自论文https://arxiv.org/pdf/1706.03762.pdf

稍微进行一下可视化不会有坏处。

Transformer 中的前馈网络。作者提供的图像。

参数的计算相当容易，主要的还是不要被弄混。

用于计算 Transformer 前馈网络中参数数量的公式。图像由作者提供。

我们可以使用以下代码描述这样一个简单的网络并检查其参数的数量（请注意，官方的 PyTorch 实现也使用了 dropout，我们将在后面的编码器/解码器代码中看到。但是正如我们所知，dropout 层没有可训练的参数，因此为了简单起见，我在这里省略它）：

from torch import nn
class TransformerFeedForward(nn.Module): def __init__(self, d_model, d_ff): super(TransformerFeedForward, self).__init__()        self.d_model = d_model        self.d_ff = d_ff
        self.linear1 = nn.Linear(self.d_model, self.d_ff)        self.relu = nn.ReLU()        self.linear2 = nn.Linear(self.d_ff, self.d_model)
 def forward(self, x):        x = self.linear1(x)        x = self.relu(x)        x = self.linear2(x) return x
d_model = 512d_ff = 2048
feed_forward = TransformerFeedForward(d_model, d_ff)print(count_parameters(feed_forward)) # 2099712print(2 * d_model * d_ff + d_model + d_ff) # 2099712

复制代码

再次看看图中的数字，仅剩下一个组件没有介绍啦。

2.3 层归一化

Transformer 架构的最后一个构件块是层归一化。简单地说，只是一种智能的（即可学习的）归一化方式，具有缩放功能，可以提高训练过程的稳定性。

Transformer 的层归一化，图片由作者提供

这里的可训练参数是两个向量 gamma 和 beta，每个向量的维度都是 d_model。

用于计算 Transformer 层归一化模块中参数数量的公式。作者提供的图像。

让我们使用代码来检验我们的假设。

from torch import nn
d_model = 512
layer_normalization = nn.LayerNorm(d_model)print(count_parameters(layer_normalization)) # 1024print(d_model * 2) # 1024

复制代码

很好! 在近似计算中，这个数字可以忽略不计，因为层归一化的参数大大少于前馈网络或多头注意力块（尽管这个模块出现了几次）。

03 推导出完整的公式

现在我们有了一切，可以计算整个编码器/解码器模块的参数了！

3.1 用 PyTorch 实现的编码器和解码器

请让我们记住，编码器是由一个注意力块、前馈网络和两个层归一化组成。

Transformer 编码器。来源于论文https://arxiv.org/pdf/1706.03762.pdf

我们可以查看 PyTorch 代码中的细节来验证所有组件是否都已就位。其中多头注意力机制用红色标注（左侧），前馈网络用蓝色标注，层归一化用绿色标注（在 PyCharm 中的 Python 控制台截图）。

PyTorch TransformerEncoderLayer。图片由作者提供

3.2 最终公式

确认好之后，我们可以编写以下函数来计算参数数量。实际上，这只是三行代码，甚至可以合并为一行。函数的其余部分是文档字符串以作说明。

def transformer_count_params(d_model=512, d_ff=2048, encoder=True, approx=False): """    Calculate the number of parameters in Transformer Encoder/Decoder.    Formulas are the following:        multi-head attention: 4*(d_model^2 + d_model)            if approx=False, 4*d_model^2 otherwise        feed-forward: 2*d_model*d_ff + d_model + d_ff             if approx=False, 2*d_model*d_ff otherwise        layer normalization: 2*d_model if approx=False, 0 otherwise    Encoder block consists of:         1 multi-head attention block,         1 feed-forward net, and         2 layer normalizations.    Decoder block consists of:         2 multi-head attention blocks,         1 feed-forward net, and         3 layer normalizations.    :param d_model: (int) model dimensionality    :param d_ff: (int) internal dimensionality of a feed-forward neural network    :param encoder: (bool) if True, return the number of parameters of the Encoder,         otherwise the Decoder    :param approx: (bool) if True, result is approximate (see formulas)    :return: (int) number of learnable parameters in Transformer Encoder/Decoder    """
    attention = 4 * (d_model ** 2 + d_model) if not approx else 4 * d_model ** 2    feed_forward = 2 * d_model * d_ff + d_model + d_ff if not approx else 2 * d_model * d_ff    layer_norm = 2 * d_model if not approx else 0
 return attention + feed_forward + 2 * layer_norm \ if encoder else 2 * attention + feed_forward + 3 * layer_norm

复制代码

现在是测试它的时候了。

from torch import nn
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)print(count_parameters(encoder_layer))  # 3152384print(transformer_count_params(d_model=512, d_ff=2048, encoder=True, approx=False))  # 3152384print(transformer_count_params(d_model=512, d_ff=2048, encoder=True, approx=True))   # 3145728# ~0.21% difference
decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)print(count_parameters(decoder_layer))  # 4204032print(transformer_count_params(d_model=512, d_ff=2048, encoder=False, approx=False))  # 4204032print(transformer_count_params(d_model=512, d_ff=2048, encoder=False, approx=True))   # 4194304# ~0.23% difference