Transformer

Posted on Edited on In Views: Word count in article: 6k Reading time ≈ 5 mins.

超参数（hyperparameter）在机器学习的上下文中，超参数是在开始学习过程之前设置值的参数，而不是通过训练得到的参数数据

背景

先前的模型输入都是一个vector，现在考虑输入是一系列vector。
此时输出可能是：

  • 每个向量给一个label
  • 整个序列给一个label
  • 模型自己决定label的数量（seq2seq，如翻译，语音辨识）

下面只考虑第一种情况，例如，给句子I saw a saw的单词标词性。
简单的FC网络不会考虑词的位置信息，会将两个saw理解成同一个词性，即使考虑加入一个窗口考虑词的上下文，受窗口大小的限制，无法处理长度不定的序列。
FC 所以引入self-attention机制，考虑每一个向量与其他向量的相关性，以为例。

计算过程

self-attention
  1. 定义三个矩阵,矩阵的值由后期学习得到，这三个矩阵被所有向量共享。
  2. 相关性
  3. 对进行softmax归一化:
  4. 得到考虑上下文后的向量值：

考虑矩阵，


QKV

为了学习到不同的特征，类似CNN，考虑多个channel，即多头注意力,使用多组不同的矩阵参数，将结果拼起来
multi-head

架构

𝐿𝑁(⋅) Layer normalization 𝑀𝐻𝐴(⋅) Multi-head attention 𝐹𝑁(⋅) Feed-forward neural network 𝑥+𝑓(𝑥) Residual connection

除以根号d

原因： As d gets larger, the variance of the inner product increases, the softmax distribution gets very peaked，hence its gradient gets smaller.

推导：
var(x)=var(y)=1,E(x)=E(y)=0 => Var(x·y)=d

多头multi-head attention

在单注意力机制里，，这里矩阵的每一列为一个d维向量。

多头注意力就是先对QKV分别乘一个降维到再计算，h为头的数量。
公式为

Feed-Forward Neural Network

先扩展到4d维再扩展回去。

Layer Normalization

positional encoding

需要满足条件：

  1. 点乘结果只与i-j有关
  2. 位置相差越大，点积结果越小

改进：考虑旋转

FLOPS计算

Given the embedding dimension d,feedforward hidden dimension f, number of attention heads h, number of layers l, vocabulary size V, context length s, number of KV heads is g.

  1. Compute the compute FLOPs of matrix multiplication M @ N, M and N are in the shape of (m, k) and (k, n) respectively, (hint: add operation accounts 1 FLOP andmul operation acounts 1 FLOP).
  2. Compute per token FLOPs for each decoder layer (hint: ignore the norm layer, only attention layer and feedforward layer in consideration).
    1. Attention layer. softmax computation is not required. GQA is not required (g=h). (hint: W_k, W_q, W_v shape (d, d/h), W_o shape (d, d))
    2. What if GQA is required (g<h)?
    3. Feedforward layer. computation: (activation(XW1) * XW2)W3. W1, W2 shape (d, f), W3 shape (f, d).
  3. Compute per token FLOPs for decoder (with GQA).
  4. Compute per token FLOPs for the LLM head (hint: LLM head layer shape (d, V)).
  5. Compute total per token FLOPs

Q1

对于一个（m,k）的矩阵M和一个（k,n）的矩阵N，结果为一个（m，n）的矩阵P，那么P中的每一个元素都需要k次乘法和k-1次加法，所以总共需要2mnk次操作。

Q2

Q2.1

不考虑GQA，如图所示.
1. 淡黄色向量为句子的embedding，维度为（s，d），分别与三个维度为（d，d/h）的矩阵相乘，得到三个维度为（s，d/h）的矩阵，这一步计算量为。
2. 将QK两个维度为（s，d/h）的矩阵相乘，得到一个维度为（s，s）的矩阵，计算量为。
3. 将V与一个维度为（s，s）的矩阵相乘，得到一个维度为（s，d/h）的矩阵，计算量为。
4. 一个head的总计算量为。h个head的总计算量为。
5. 将concat得到的维度为（s，d）的矩阵与一个维度为（d，d）的矩阵相乘，得到一个维度为（s，d）的矩阵，计算量为。

所以单个注意力层的计算量为。

Q2.2

考虑GQA，省去的是KV的计算量，原来的计算量为，现在的计算量为，单层总计算量为

Q2.3

feedforward的计算为self.w2(F.silu(self.w1(x)) * self.w3(x)),

self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

对于单个token，计算量为，则总计算量为

Q3

一层decoder的计算量为
l层总计算量是

Q4

一个（s，d）的矩阵与一个（d，V）的矩阵相乘，计算量为，所以一个LLM head的计算量为。

Q5

前向过程的计算量为。

chatglm参数

d=4096
f=13696
h=32 g=2 l=28
V=65024
s=8192

计算量=128714721460224.0FLOPs

对比

Transformer与CNN

CNN可看做是简化版的Transformer。CNN中向量要与哪些向量寻找相关性是认为给定的超参数，即那个矩阵的大小，而Transformer是学出来的。
Transformer相较于CNN有更多参数，因此简单数据集上CNN效果好，数据多时Transformer效果好。

Transformer与RNN

Transformer相较于RNN的优势：

  • 可以注意到相隔很远的词，因为会对所有向量计算相关性。（也因此参数多，也有给Transformer设置窗口提升训练速度的操作）
  • 可以并行计算（本质是矩阵乘法）

DIT 参数

分析DiT (Diffusion Transformer) 模型的参数量与给定参数的关系：

主要影响DiT参数量的因素包括：

  • num_layers: 模型层数
  • num_attention_heads: 注意力头数
  • attention_head_dim: 每个注意力头的维度
  • joint_attention_dim: 联合注意力维度 (FFN隐藏层维度)
  • in_channels/out_channels: 输入/输出通道数
  • patch_size: 图像分块大小
  • sample_size: 样本大小
  • caption_projection_dim: 文本投影维度
  • pooled_projection_dim: 池化投影维度

DiT模型的参数量数学表达式：

首先计算隐藏维度： hidden_dim = num_attention_heads × attention_head_dim

总参数量主要由以下部分组成：

  • 图像patch嵌入参数：
    • patch_embed_params = in_channels × (patch_size²) × hidden_dim + hidden_dim
  • 位置嵌入参数（与位置嵌入最大尺寸相关）：
    • pos_embed_params = (sample_size/patch_size)² × hidden_dim
  • Transformer层参数（线性关系）：
    • 每层参数量 = 4 × hidden_dim² (多头注意力) + 2 × hidden_dim × joint_attention_dim + 3 × hidden_dim (FFN和层归一化)
    • transformer_params = num_layers × (4 × hidden_dim² + 2 × hidden_dim × joint_attention_dim + 3 × hidden_dim)
  • 输出投影参数：
    • output_params = hidden_dim × out_channels × patch_size²
  • 文本和池化投影参数：
    • projection_params = hidden_dim × caption_projection_dim + hidden_dim × pooled_projection_dim
  • 总参数量 = patch_embed_params + pos_embed_params + transformer_params + output_params + projection_params

参数量与各因素的关系：

  • 与num_layers成线性关系
  • 与num_attention_heads和attention_head_dim的乘积成二次关系
  • 与joint_attention_dim成线性关系
  • 与sample_size成二次关系（由于位置嵌入）
  • 与patch_size成平方关系（在patch嵌入部分）
  • 最主要的参数量贡献来自于Transformer层，特别是当num_layers、num_attention_heads和attention_head_dim很大时。