从矩阵乘法探秘Transformer

前言

学习连博的另外一篇文章，从矩阵乘法的角度来理解 transformer，仅供自己参考

refer1：从矩阵乘法探秘transformer+代码讲解

refer2：深入理解transformer

以下内容来自于连博的博客：深入理解transformer，强烈建议阅读原文🤗

1. transformer背景

1.1 回顾线性代数的知识

我们先来回顾下线性代数的一些知识，这是因为《Attention Is All Your Need》这篇文章中 attention 的公式全都是一些矩阵相乘比较晦涩难懂，我们把矩阵剖解从行向量来看可能更容易理解

原文公式如下：

$$\begin{array}{r}\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{Q{K}^T}{\sqrt{d_k}})V\\ \mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_h)W^O\\ {\mathrm{head}}_i=\mathrm{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)\end{array}$$

1.1.1 矩阵和行向量

我们先来看矩阵和行向量

定义矩阵 $X\in R^{N\times F}$，其中 $X=\left(\begin{array}{c}{X}_{11},X_{12},\dots ,X_{1F}\\ X_{21},X_{22},\dots ,X_{2F}\\ ⋮\\ X_{N1},X_{N2},\dots ,X_{NF}\end{array}\right)$

矩阵 $X$ 每一行定义为行向量 $X_i$，其中 $X_i=\left(\begin{array}{c}{X}_{i1},X_{i2},\dots ,X_{iF}\end{array}\right),X_i\in R^{1\times F}$

矩阵 $X$ 可以看作是 $N$ 个行向量堆叠的结果，也就是说 $X=\left(\begin{array}{c}{X}_1\\ X_2\\ ⋮\\ X_N\end{array}\right)$

比如 pytorch 中的 nn.Embedding 它其实就是按照行向量来组织数据的

import torch
import torch.nn as nn
N = 3
F = 8
embed = nn.Embedding(N, F)
idx = torch.tensor([0, 1, 2])
X = embed(idx)
print(X.shape)  # torch.Size([3, 8])

我们举个简单的例子，假设有 $N$ 个 token，$F$ 是 embedding 的维度，每行对应于一个 token 的 embedding 行向量，那么对应的矩阵如下所示：

$$\begin{gathered} tokens=\left(\begin{array}{c}\text{hello}\\ \text{world}\\ \text{pad}\\ \text{pad}\\ \text{pad}\end{array}\right) \\ X=\left(\begin{array}{c}[0.59,0.20,0.04,0.96]\\ [0.96,0.30,0.16,0.63]\\ [0.02,0.19,0.34,0.25]\\ [0.02,0.19,0.34,0.25]\\ [0.02,0.19,0.34,0.25]\end{array}\right) \end{gathered}$$

1.1.2 矩阵相乘和算子作用

我们接着看矩阵相乘和算子作用的一些知识

• 定义线性算子 $\mathcal{A}$
  • 它可以作用到行向量上 $\mathcal{A}(X_i)=X_{i}A$
  • 也可以作用到矩阵上 $\mathcal{A}(X)=XA$
• 右乘矩阵等于对每个行向量逐个施加行变换

$$XA=\left(\begin{array}{c}{X}_1\\ X_2\\ ⋮\\ X_N\end{array}\right)A=\left(\begin{array}{c}{X}_{1}A\\ X_{2}A\\ ⋮\\ X_{N}A\end{array}\right)=\left(\begin{array}{c}\mathcal{A}(X_1)\\ \mathcal{A}(X_2)\\ ⋮\\ \mathcal{A}(X_N)\end{array}\right)=\mathcal{A}(X)$$

代码对应于 nn.Linear

import torch
import torch.nn as nn

F = 6
linear = nn.Linear(in_features=F, out_features=F)
X_i = torch.rand(1, 6)
X = torch.rand(3, 6)
print(linear(X_i).shape)    # torch.Size([1, 6])
print(linear(X).shape)      # torch.Size([3, 6])

Note：pytorch/tensorflow 的代码都是按照作用于行向量来组织的

1.1.3 从分块矩阵的乘法来看$Q{K}^{T}V$

我们从分块矩阵乘法看看 $Q{K}^{T}V$ 具体做了什么事情

首先 $S=Q{K}^T$ 是行向量两两计算点积相似性

$$\left(\begin{array}{c}{Q}_1\\ Q_2\\ ⋮\\ Q_N\end{array}\right)\left(\begin{array}{c}{K}_1^T,K_2^T,\dots ,K_N^T\end{array}\right)=(Q_i{K}_j^T{)}_{ij}=S$$

接着 $SV$ 是对 $V$ 行向量做加权求和

$$\left(\begin{array}{c}{S}_{11},S_{12},\dots ,S_{1N}\\ S_{21},S_{22},\dots ,S_{2N}\\ ⋮\\ S_{N1},S_{N2},\dots ,S_{NN}\end{array}\right)\left(\begin{array}{c}{V}_1\\ V_2\\ ⋮\\ V_N\end{array}\right)=\left(\begin{array}{c}\sum _j{S}_{1j}{V}_j\\ \sum _j{S}_{2j}{V}_j\\ ⋮\\ \sum _j{S}_{Nj}{V}_j\end{array}\right)$$

因此我们可以认为 attention 的计算首先是基于 $Q,K$ 计算相似性，然后基于 $V$ 来加权求和。其中 $Q{K}^{T}V$ 的每个行向量都是 $V$ 行向量的一个加权求和

值得注意的是：

• 论文：一般会有行/列向量两种表示方式
  • 列向量表现为左乘以一个矩阵
  • 左乘以一个矩阵相当于对每个列向量来施加变化
• 代码：基本都是行向量来作为数据组织的标准
• 本文：
  • 向量都按照行向量的形式来组织
  • 按照作用于单个行向量的方式来讲解 transformer

1.2 encoder-decoder

接着来看下 encoder-decoder

大部分 seq2seq 的任务建模为 encoder-decoder 的结构，如机器翻译、语音识别、文本摘要、问答系统等等，原论文 《Attention Is All Your Need》 中的 Transformer 结构就是 encoder-decoder 的结构，如下图所示：

Transformer 中的 encoder 用于把离散的 token 序列 $x_1,x_2,\dots ,x_N$ 转化为语义向量序列 $Y_1,Y_2,\dots ,Y_N$，一般组织为多层的网络的形式：

• 第一层：基础语义向量序列 $x_1,x_2,\dots ,x_N\to (X_1,X_2,\dots ,X_N)$
• 其它层：从低阶语义向量转化为高阶语义向量序列 $(X_1,X_2,\dots ,X_N)\to (Y_1,Y_2,\dots ,Y_N)$

而 decoder 则基于 $Y_1,Y_2,\dots ,Y_N$ 自回归式的逐个 token 解码

那像翻译这类的任务通常涉及输入（源语言）和输出（目标之间的映射）之间的映射，因此需要基于 encoder-decoder 这样的架构，encoder 用于处理输入数据（源数据），decoder 则生成输出数据（目标语言）。而像 GPT、DeepSeek 这样的语言模型，它们的任务主要是生成文本（比如对话生成、文本补全等），这些任务并不需要明确的输入和输出对，而只需要基于一个上下文来生成接下来的文本，因此，单独的 decoder 就足够处理这些任务了。

这些模型也被称为自回归模型（Autoregressive Models），因为它们在生成每个词时，依赖于之前生成的词。换句话说，它们是一步一步地生成文本，每生成一个词，就把它作为上下文输入到模型中预测下一个词。例如，GPT、DeepSeek 都是通过给定一段文本（输入），然后依次预测每一个后续词语。这种逐步生成的过程使得它们成为了自回归的模型。

Note：本文主要聚焦到 encoder 部分来理解 transformer

1.3 低阶到高阶语义向量的转换

encoder 的主要工作是寻找算子 $\mathcal{T}$ 将低阶的语义向量序列变换为高阶的语义向量序列即

$$\mathcal{T}\left(\begin{array}{c}{X}_1\\ X_2\\ ⋮\\ X_N\end{array}\right)\to \left(\begin{array}{c}{Y}_1\\ Y_2\\ ⋮\\ Y_N\end{array}\right)$$

• 输入：$X$ 低阶语义向量序列
• 输出：$Y$ 高阶语义向量序列
• 意义
  • $Y_i=f(X_1,X_2,\dots ,X_N)$
  • 对低阶语义向量做加工组合处理和抽象，变换为一个高阶的语义向量序列
  • 高阶语义向量考虑了上下文的语义向量表达
• 用算子作用来表达
  • $Y=\mathcal{T}(X)$
  • $X\in R^{N\times F},Y\in R^{N\times F}:R^{N\times F}\to R^{N\times F}$
  • 这个算子天然可以复合嵌套，形成多层的网络结构 $Y={\mathcal{T}}_L\circ {\mathcal{T}}_{L-1}\circ \dots \circ {\mathcal{T}}_1(X)$

1.4 核心的问题

我们现在的核心问题是如何设计 $Y_i=f(X_1,X_2,\dots ,X_N)$，满足：

• $Y_1,\dots ,Y_N$ 能够并行得到
• $Y_i$ 能够高效的建立起对周围 token 的远程依赖

我们可以先看下 RNN，看它是如何做的：

RNN 的特性如下：

• 递归语义序列 $Y_0\to Y_1\to \dots \to Y_N$
• $Y_i=\tanh (X_{i}W+Y_{i-1}U)$
• 串行
• 单方向的依赖关系，例如 $Y_3$ 直接依赖于 $Y_2,X_3$，间接依赖于 $X_1$

接着再看下 CNN：

CNN 的特性如下：

• $Y_i=(X_{i-1},X_i,X_{i+1})W$
• 并行
• 假设窗口宽度是 3，即 kernel_size = 3
• 它不能长距离依赖，一层卷积只能依赖于当前窗口内，不能对窗口外的形成依赖
• 例如 $Y_3$ 依赖于 $X_2,X_3,X_4$，但它没有办法和 $X_1$ 建议起依赖关系

transformer 要解决的问题就是设计 $Y_i=f(X_1,X_2,\dots ,X_N)$ 使得：

• $Y_1,\dots ,Y_N$ 可以做并行计算
• 同时解决长距离依赖的问题

如上图所示，我们在计算 $Y_2$ 时就希望对所有 token 的低阶语义向量序列都能够建议起依赖关系来

整体思路的话就是做两次矩阵的变换即 $Y^{\prime }=\mathcal{F}(Y)=\mathcal{F}\circ \mathcal{A}(X)$

• $Y=\mathcal{A}(X)$
  • 第一次矩阵变换
  • MultiHead Attention 多头注意力机制
  • 高阶的语义等于对全部的低阶语义向量基于相似性（Attention）做加权平均
  • $\begin{array}{rl}\mathcal{A}(X_i)& =\frac{{\sum }_{j=1}^{N}sim(X_i,X_j)X_j}{{\sum }_{j=1}^{N}sim(X_i,X_j)}\end{array}$
  • attention = 相似性
• $Y^{\prime }=\mathcal{F}(Y)$
  • 第二次矩阵变换
  • Position-wise Feedforward 前馈神经网络层
  • 再施加若干线性变换

2. transformer网络结构

下面我们就来看看 transformer 的网络结构

2.1 基于KV查询的相似性计算

首先看 transformer 第一部分相似性（attention）的计算即

$$\begin{array}{rl}\mathcal{A}(X_i)& =\frac{{\sum }_{j=1}^{N}sim(X_i,X_j)X_j}{{\sum }_{j=1}^{N}sim(X_i,X_j)}\end{array}$$

前面我们说了 transformer 的 motivation 就是把 $X_i$ 这样一个低阶语义向量和周围所有的低阶语义向量去做一个相似性，然后再做一个加权平均

那我们怎么来算 $\mathcal{A}(X_i)$ 这个相似性呢？如果直接计算相似性会发现参数太少，模型复杂度低无法有效学习。那一种自然而然的想法就是我们投影到别的空间来计算相似度即 $X_i\to X_{i}W$，而不是直接来计算

因此我们可以在原有公式基础上都乘以相应的矩阵 $W$，投影到更高维的空间，那此时相似性公式如下：

$$\begin{array}{rl}\mathcal{A}(X_i)& =\frac{{\sum }_{j=1}^{N}sim(X_i{W}_1,X_j{W}_2)X_j{W}_3}{{\sum }_{j=1}^{N}sim(X_i{W}_1,X_j{W}_2)}\end{array}$$

如果我们记 $Q_i=X_i{W}_1,K_i=X_i{W}_2,V_i=X_i{W}_3$ 则有：

$$\begin{array}{rl}\mathcal{A}(X_i)& =\frac{{\sum }_{j=1}^{N}sim(Q_i,K_j)V_j}{{\sum }_{j=1}^{N}sim(Q_i,K_j)}\end{array}$$

那这个公式和原文中的类似，那我们怎么去理解 KV 查询呢？

首先我们把 $X_i$ 投影出三个向量 $Q_i,K_i,V_i$，其中 $K,V$ 是大家熟悉的 key-value 存储，$K_j\to V_j$ 相互对应，而 $Q$ 是查询使用的 query 向量 $Q_i$

$Q,K,V$ 的查询方法是 query 查询多个 key，获取多个 value，最后把这些 value 加权平均，即

$$Q_i\Rightarrow \left(\begin{array}{c}{K}_1\to V_1\\ K_2\to V_2\\ ⋮\\ K_N\to V_N\end{array}\right)\Rightarrow \left(\begin{array}{c}sim(Q_i,K_1)V_1\\ sim(Q_i,K_2)V_2\\ ⋮\\ sim(Q_i,K_N)V_N\end{array}\right)\Rightarrow \sum _{j=1}^{N}sim(Q_i,K_j)V_j$$

那我们怎么理解呢？举个简单的例子，假设我们现在有 3 个 token 对应 3 个低阶语义向量 $X_1,X_2,X_3$，接着我们会把 $X_1$ 投影出三个向量来分别是 $Q_1,K_1,V_1$，同理 $X_2$ 投影出三个向量分别是 $Q_2,K_2,V_2$，$X_3$ 投影出三个向量分别是 $Q_3,K_3,V_3$。投影完成之后我们可以把 $K_1\to Y_1,K_2\to Y_2,K_3\to Y_3$ 当成一个 $K\to V$ 查询体系

假设我们要计算低阶语义向量 $X_2$ 对应的高阶语义向量 $Y_2$，那么我们先要用 $Q_2$ 查询 $K\to V$ 体系中的 $K$ 即 $K_1,K_2,K_3$，然后分别计算它们的 $sim$，最后把对应的 $V$ 做一个加权平均，也就是 $\begin{array}{rl}{Y}_2& =\frac{{\sum }_{j=1}^{3}sim(Q_2,K_j)V_j}{{\sum }_{j=1}^{3}sim(Q_2,K_j)}\end{array}$

也就是对应到前面的公式：

$$\begin{array}{rl}\mathcal{A}(X_i)& =\frac{{\sum }_{j=1}^{N}sim(Q_i,K_j)V_j}{{\sum }_{j=1}^{N}sim(Q_i,K_j)}\end{array}$$

那做完这些后我们会发现已经加了一些有效的参数出来了，也就是对应于 $Q,K,V$ 产生的三个投影矩阵 $W_Q,W_K,W_V$

2.2 在一个低维空间做attention

下面看一下它的一个实现，单个行向量做 attention 的流程如下：

step 1. 把 $X_i$ 从 $F$ 维空间投影到 $D$ 维空间

• $Q_i=X_i{W}_Q,W_Q\in R^{F\times D}$
• $K_i=X_i{W}_K,W_K\in R^{F\times D}$
• $V_i=X_i{W}_V,W_V\in R^{F\times M}$

step 2. $Q_i$ 和所有的 $K_j$ 做基于点积的相似度计算

• $Q_i{K}^T=Q_i(K_1^T,\dots ,K_N^T)=(Q_i{K}_1^T,\dots ,Q_i{K}_N^T)$
• Note：简单起见，我们这里省略了 scaling 缩放因子 $\frac{1}{\sqrt{D}}$

step 3. 对相似度的分布做 softmax

• $S=\mathrm{soft}(Q_i{K}_1^T,\dots ,Q_i{K}_N^T)=(s_{i1},\dots ,s_{iN})$
• $s_{i,j}=\frac{exp(Q_i{K}_j^T)}{{\sum }_{j=1}^{N}exp(Q_i{K}_j^T)}$

step 4. 加权平均

• $\mathcal{A}(X_i)={\sum }_{j=1}^N{s}_j{V}_j=(s_{i1},\dots ,s_{iN})\left(\begin{array}{c}{V}_1\\ V_2\\ ⋮\\ V_N\end{array}\right)$
• $\mathcal{A}(X_i)=\mathrm{soft}(Q_i{K}^T)V=\mathrm{soft}(X_i{W}_Q{W}_K^T{X}^T)X{W}_V$

扩展到多个行向量即对应的矩阵表达式如下：

$$Y=\mathcal{A}(X)=\left(\begin{array}{c}\mathcal{A}(X_1)\\ \mathcal{A}(X_2)\\ ⋮\\ \mathcal{A}(X_N)\end{array}\right)=\left(\begin{array}{c}\mathrm{soft}(Q_1{K}^T)V\\ \mathrm{soft}(Q_2{K}^T)V\\ ⋮\\ \mathrm{soft}(Q_N{K}^T)V\end{array}\right)=\mathrm{soft}(Q{K}^T)V$$

对应的代码实现如下：

import math
import torch.nn as nn
from torch.nn import functional as F

class SingleHeadAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.F = config["hidden_dim"]   # F
        self.D = config["subspace_dim"] # D
        self.q_proj = nn.Linear(self.F, self.D)
        self.k_proj = nn.Linear(self.F, self.D)
        self.v_proj = nn.Linear(self.F, self.D)

    def forward(self, x):
        # x->[B, N, F]
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = F.softmax(att, dim=-1)
        y = att @ v
        return y

Note：当 $D\ne F$ 时，$\mathcal{A}(X)$ 不可用

2.3 在多个低维空间做attention

我们思考下为什么要在多个低维空间中去做 attention 呢？在单个低维空间做不就行了吗？🤔

原文中的描述是：

Multi-Head atttention allows the model to jointly attend to information from different representation subspaces at different positions.

也就是说多头注意力可以让模型从不同低维空间表达中去学习不同的语义。我们知道单个词往往有多个含义，把 $F$ 维的语义向量投影到 $H$ 个不同的子空间中去计算相似加权组合可能会得到完整的语义

具体的做法如下：

• 每个头做独立的 attention 变换 ${\mathcal{A}}^h(X)$
  • 假设有 $H$ 个头，每个头作用的低维空间维度是 $D$
  • $D\times H=F$
• 对 $H$ 个 $D$ 维行向量拼接，之后再做一次矩阵变换
  • $\mathcal{A}(X)=\mathrm{concat}({\mathcal{A}}^1(X),{\mathcal{A}}^2(X),\dots ,{\mathcal{A}}^H(X))W_O$
  • $W_O\in R^{F\times F}$
• 对前面的符号简化
  • 在第 $j$ 个子空间做单头注意力 $Y^j=sim(Q^j,K^j)V^j$
  • 合并 $Y=(Y^1,\dots ,Y^H)W_o$

代码实现如下：

import math
import torch.nn as nn
from torch.nn import functional as F

class SelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.H = config["n_head"]       # H
        self.F = config["hidden_dim"]   # F
        self.D = self.F // self.H       # D
        # 一次把 qkv 全部映射完成, 对应 $W_Q$, $W_K$, $W_V$
        self.qkv_proj = nn.Linear(self.F, 3 * self.F)
        # 最后的投影, 对应于 $W_O$
        self.out_proj = nn.Linear(self.F, self.F)

    def forward(self, x):
        # x->[B, N, F]
        B, N, _ = x.size()
        q, k, v = self.qkv_proj(x).split(self.F, dim=-1)
        # matmul 只能在最后两个维度相乘, 需要对 NxD 的矩阵相乘, 做 1,2 维度的交换
        # [B, H, N, D]
        q = q.view(B, N, self.H, self.D).transpose(1, 2)
        k = k.view(B, N, self.H, self.D).transpose(1, 2)
        v = v.view(B, N, self.H, self.D).transpose(1, 2)
        # 一次把多个头的映射全部完成, 对任意的 (batch, head)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = F.softmax(att, dim=-1)
        # [B, H, N, D]
        y = att @ v
        # [B, N, H, D]
        y = y.transpose(1, 2)
        # 最后两个维度合并
        y = y.contiguous().view(B, N, F)
        y = self.out_proj(y)
        return y

Note：代码参考自：https://github.com/karpathy/minGPT/tree/master/mingpt

代码示意图如下所示：

输入 $X$ 是 (B,N,F) 的三维矩阵，经过 linear(F,3F) 算子后维度变成了 (B,N,3F)，通过 split 拿到我们独立的 $Q,K,V$ 矩阵。接着通过 view 和 transpose 之后将 $Q,K,V$ 的最后两维变成我们关心的子矩阵，此时维度是 (B,H,N,D)

然后对不同子空间的 $Q^i{K}^i{V}^i$ 做 attention，做完之后再经过 transpose 和 view 得到我们的输出 $Y$，这样我们的输入 $X$ 和输出 $Y$ 是完全能够 match 上的

2.4 位置无关的全连接

前面我们已经讲完了 transformer 中的 attention 变换，下面我们来看 transformer 中另外一个变换 Feedforward 即位置无关的全连接

公式表达如下：

$$\mathcal{F}(X_i)=(g(X_i{W}_1)+b_1)W_2+b_2$$

作用到每个行向量 $X_i$ 上时是先右乘一个矩阵 $W_1$，然后做个激活 $g$ 再加个偏置 $b_1$，接着再乘以一个 $W_2$ 加上 $b_2$，其实也就是两层的全连接

代码实现如下：

import torch.nn as nn

class PWiseFeedForward(nn.Module):
    
    def __init__(self, config):
        super().__init__()
        self.F = config["F"]
        self.proj_wide   = nn.Linear(self.F, 4 * self.F)
        self.proj_narrow = nn.Linear(4 * self.F, self.F)
        self.act = nn.ReLU()
    
    def forward(self, x):
        x = self.proj_wide(x)
        x = self.act(x)
        x = self.proj_narrow(x)
        return x

2.5 归一化+残差网络

前面我们已经把 Transformer 的核心架构部分讲完了，也就是 $\mathcal{T}(X)=\mathcal{F}\circ \mathcal{A}(X)$，它包含两部分先做 $\mathcal{A}$ 变换再做一个 $\mathcal{F}$ 变换

那其实在 Transformer 网络结构的中间部分还加入了一些归一化和残差网络，下面我们简单说明下

Transformer 中的 Normalization 层一般都是采用 LayerNorm 来对 Tensor 进行归一化，LayerNorm 的公式如下：

$$\begin{array}{rl}{A}^{\prime }(X)& =\mathcal{N}\circ \mathcal{A}(X)\\ LayerNorm:y& =\frac{x-\mu }{\sqrt{\sigma }}\gamma +\beta \\ \mu & =\frac{1}{d}\sum _{i=1}^d{x}_i\\ \sigma & =\sqrt{\frac{1}{d}\sum _{i=1}^d(x_i-\mu {)}^2}\end{array}$$

LayerNorm 和 BatchNorm 比较像，区别是一个是在行上面做归一化一个是在列上面做归一化，而 LayerNorm 可以看作是作用在行向量上的算子。在 NLP 的序列建模里面一般使用 LayerNorm，而在 CV 里面一般使用 BatchNorm

这主要是因为 padding 的影响，以下面的输入矩阵为例，不同 batch 中 <pad> 个数不同，沿着 token 方向做归一化并没有意义，而每个位置做独立的归一化更有意义

$$\left(\begin{array}{c}\text{hello}\\ \text{world}\\ \text{pad}\\ \text{pad}\\ \text{pad}\end{array}\right)\to X=\left(\begin{array}{c}[0.59,0.20,0.04,0.96]\\ [0.96,0.30,0.16,0.63]\\ [0.02,0.19,0.34,0.25]\\ [0.02,0.19,0.34,0.25]\\ [0.02,0.19,0.34,0.25]\end{array}\right)$$

其他的可能选择 RMSNorm 归一化方法，例如 LLaMA 中使用的就是 RMSNorm，RMSNorm 是 LayerNorm 的变体，RMSNorm 省去了求均值的过程，也没有了偏置 $\beta$，公式如下：

$$\begin{array}{rl}RMSNorm:y& =\frac{x}{\sqrt{Mean(x^2)+ϵ}}\ast \gamma \\ Mean(x^2)& =\frac{1}{N}\sum _{i=1}^N{x}_i^2\end{array}$$

其中 $\gamma$ 为可学习的参数

大家感兴趣的可以看看：RMSNorm算子的CUDA实现

2.6 整体的变换

最后我们看下 transformer 的整体变换 $Y=\mathcal{T}(X)$，它主要分为以下几个部分：

• Attention $Z=N\circ (X+\mathcal{A}(X))$
• 位置无关的全连接 $Y=\mathcal{N}\circ (X+\mathcal{F}(Z))$
• 残差网络
  • ${\mathcal{A}}^{\prime }(X)=\mathcal{N}\circ (X+\mathcal{A}(X))$
  • ${\mathcal{F}}^{\prime }=\mathcal{N}\circ (X+\mathcal{F}(X))$

前面我们主要是解释了 transformer 中一层的网络结构，实际上我们是多层，可以任意的去嵌套，对于一个 $L$ 层的 transformer 模型表达如下：

$$\begin{array}{c}\begin{array}{rl}\mathcal{T}(X)& ={\mathcal{T}}_L\circ \dots {\mathcal{T}}_2\circ {\mathcal{T}}_1(X)\end{array}\end{array}$$

代码实现如下：

import torch.nn as nn

class Block(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.attn   = SelfAttention(config)
        self.norm_1 = nn.LayerNorm(config["hidden_dim"])
        self.mlp    = PWiseFeedForward(config)
        self.norm_2 = nn.LayerNorm(config["hidden_dim"])

    def forward(self, x):
        x = self.norm_1(x + self.attn(x))
        x = self.norm_2(x + self.mlp(x))
        return x

3. transformer参数和计算量

最后我们来回顾下 transformer 的参数量

3.1 关于参数量

一般模型增加复杂度的方式包括：

• 增加深度，增加宽度
• 增加 embedding 的维度
• 增加词典的大小

各种 DNN 主要参数的位置：

• CNN：$Y_i=(X_{i-1},X_i,X_{i+1})W$
• RNN：$Y_i=\tanh (X_{i}W+Y_{i-1}U)$

3.2 参数的分布

我们来看下 transformer 它的参数分布是什么样子的：

1. 多头注意力（Multi-Head Attention）

• 每个头有
  • 3 个投影矩阵 $W_Q,W_K,W_V$
  • 1 个投影 concat 结果的矩阵 $W_O$
• 假设投影到的子空间维度是 $D$，有 $H$ 个子空间且 $D\times H=F$
• 参数量：
  • $F\times D\times 3\times H=3F^2$
  • $F^2$

2. 前馈网络层（Feedforward）

• 两个矩阵，先从 $F$ 变宽到 $4F$，再收窄回到 $F$
• 参数量：$F\times 4F+4F\times F=8F^2$

3. word embedding

• $E$ 是 token 字典的大小
• $E\times F$

总共：

• $L(12F^2)+EF$
• $L$ 表示模型的层数

例如：

3.3 linear transformer

transformer 中两个算子的计算量分别是：

• $\mathcal{A}(X)$ 计算量 $O(N^2)$
• $\mathcal{F}(X)$ 计算量 $O(N)$

softmax 的存在导致 $\mathcal{A}(X)$ 计算量是 $O(N^2)$，我们知道 attention 核心的计算量在 $Q{K}^{T}V$ 三个矩阵的相乘上，而乘法的计算量密切依赖于矩阵组合的方式

有 softmax 的存在的话只能先计算 $H=Q{K}^T$，对 $H$ 做 softmax 变换后再计算 $HV$ 乘法，这个计算量是 $N^{2}D+N^{2}M$，整体的复杂度是 $O(N^2)$

$$Q{K}^{T}V=(Q{K}^T)V=\left(\begin{array}{c}{H}_{11},H_{12},\dots ,H_{1N}\\ ⋮\\ H_{N1},H_{N2},\dots ,H_{NN}\end{array}\right)V$$

如果没有 softmax 的话，可以先计算后两个矩阵相乘 $H=K^{T}V$，再计算 $QH$ 乘法，这时计算量是 $NDM+DMN=2NDM$，而当 $N\gg D$ 的时候，计算量可以是 $O(N)$，因为 $K^{T}V$ 可以提前算出来缓存，大致如下面这个表达式所示：

$$Q(K^{T}V)=\left(\begin{array}{c}{Q}_1\\ Q_2\\ ⋮\\ Q_N\end{array}\right)(K^{T}V)$$

接着我们看下 kernel 的表达形式，前面我们提到过很多次 attention 可以表示成下面这种加权平均的形式

$$\mathcal{A}(X_i)=\frac{{\sum }_{j=1}^{N}sim(Q_i,K_j)V_j}{{\sum }_{j=1}^{N}sim(Q_i,K_j)}$$

这里的 $sim$ 其实是可以用非负的 kernel 来替换掉，对于 kernel 函数可以映射到其他空间 $k(x,y)=<\varphi (x),\varphi (y)>$，从而将 $sim$ 变成内积的形式 $k(x,y)=(x\cdot z{)}^2,\varphi (x)=(x_1^2,x_2^2,\sqrt{2}{x}_1{x}_2)$

当前的 sim 函数 $sim(x,y)=\exp (x{y}^T/\sqrt{D})$

Note：kernel 对应一个 feature map

linear transformer 其实就是用 kernel 来替换掉 sim，公式如下：

$$\begin{array}{rl}\mathcal{A}(X_i)& =\frac{{\sum }_{j=1}^{N}sim(Q_i,K_j)V_j}{{\sum }_{j=1}^{N}sim(Q_i,K_j)}\\ & =\frac{{\sum }_{j=1}^N\varphi (Q_i)\varphi (K_j{)}^T{V}_j}{{\sum }_{j=1}^N\varphi (Q_i)\varphi (K_j{)}^T}\\ & =\frac{\varphi (Q_i){\sum }_{j=1}^N\varphi (K_j{)}^T{V}_j}{\varphi (Q_i){\sum }_{j=1}^N\varphi (K_j{)}^T}\end{array}$$

• ${\sum }_{j=1}^N\varphi (K_j{)}^{T}V,{\sum }_{j=1}^N\varphi (K_j{)}^T$ 可以提前算好
• $O(N)$ 复杂度，Linear Transformer
• $\varphi (x)=\mathrm{elu}(x)+1$

更多细节大家可以参考原始论文：《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》

4. 补充—线性Attention的探索

在前面的 3.3 小节我们介绍了 linear transformer，用 kernel 核函数代替原有 $sim$ 中的 softmax，博主一头雾水，怎么就突然提到了 kernel 核函数呢？🤔

在苏神的 线性Attention的探索：Attention必须有个Softmax吗？ 文章中就有详细介绍，我们来简单了解下

原始的 Attention 机制是 Scaled-Dot Attention，形式为：

$$Attention(\mathbf{Q},\mathbf{K},\mathbf{V})=softmax(\mathbf{Q}{\mathbf{K}}^{\top })\mathbf{V}$$

其中 $\mathbf{Q}\in {\mathbb{R}}^{n\times d_k},\mathbf{K}\in {\mathbb{R}}^{m\times d_k},\mathbf{V}\in {\mathbb{R}}^{m\times d_v}$，简单起见这里我们省略了缩放因子 scaling

这里我们只关注 self-attention 场景，为了介绍方便统一设 $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{n\times d}$，一般场景下都有 $n>d$ 甚至 $n\gg d$

前面我们提到过制约 attention 性能的关键因素是定义里面边的 softmax，${\mathbf{QK}}^{\top }$ 这一步我们得到一个 $n\times n$ 的矩阵，就是这一步决定了 attention 的复杂度是 $O(n^2)$；如果没有 softmax，那么就是三个矩阵连乘 $(\mathbf{Q}{\mathbf{K}}^{\top })\mathbf{V}$，而矩阵乘法是满足结合律的，所以我们可以先算 ${\mathbf{K}}^{\top }\mathbf{V}$ 得到一个 $d\times d$ 的矩阵，然后再用 $\mathbf{Q}$ 左乘它，由于 $d\ll n$，所以这样算的复杂度只有 $O(n)$

也就是说，去掉 softmax 的 attention 的复杂度可以降到最理想的线性级别 $O(n)$！这显然就是我们的终极追求：Linear Attention，复杂度为线性级别的 attention

问题是，直接去掉 softmax 还能算是 attention 吗？它还能有标准 attention 的效果吗？为了回答这个问题，我们先将 Scaled-Dot Attention 的定义等价地改写为：

$$Attention{(\mathbf{Q},\mathbf{K},\mathbf{V})}_i=\frac{\sum _{j=1}^n{e}^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}{\mathbf{v}}_j}{\sum _{j=1}^n{e}^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}}$$

Note：苏神文章中提到的向量都是列向量

所以 Scaled-Dot Attention 其实就是以 $e^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}$ 为权重对 ${\mathbf{v}}_j$ 做加权平均，因此我们可以提出一个 Attention 的一般化定义：

$$Attention{(\mathbf{Q},\mathbf{K},\mathbf{V})}_i=\frac{\sum _{j=1}^{n}sim({\mathbf{q}}_i,{\mathbf{k}}_j){\mathbf{v}}_j}{\sum _{j=1}^{n}sim({\mathbf{q}}_i,{\mathbf{k}}_j)}$$

也就是把 $e^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}$ 换成 ${\mathbf{q}}_i,{\mathbf{k}}_j$ 的一般函数 $sim({\mathbf{q}}_i,{\mathbf{k}}_j)$，为了保留 attention 相似的分布特性，我们要求 $sim({\mathbf{q}}_i,{\mathbf{k}}_j)\ge 0$ 恒成立。也就是说，如果我们要定义新式的 attention，那么要保留上述公式的形式，并且满足 $sim({\mathbf{q}}_i,{\mathbf{k}}_j)\ge 0$

这种一般形式的 attention 在 CV 中也被称为 Non-Local 网络，出自论文《Non-local Neural Networks》

如果直接去掉 softmax，那么就是 $sim({\mathbf{q}}_i,{\mathbf{k}}_j)={\mathbf{q}}_i^{\top }{\mathbf{k}}_j$，问题是内积无法保证非负性，所以这还不是一个合理的选择，下面我们简单介绍几种可取的方案

一个自然的想法是：如果 ${\mathbf{q}}_i,{\mathbf{k}}_j$ 的每个元素都是非负的，那么内积自然也就是非负的。为了完成这点，我们可以给 ${\mathbf{q}}_i,{\mathbf{k}}_j$ 各自加个激活函数 $\varphi ,\phi$，即

$$sim({\mathbf{q}}_i,{\mathbf{k}}_j)=\varphi ({\mathbf{q}}_i{)}^{\top }\phi ({\mathbf{k}}_j)$$

其中 $\varphi (\cdot ),\phi (\cdot )$ 是值域非负的激活函数，在论文《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》中选择的是 $\varphi (x)=\phi (x)=\text{elu}(x)+1$

另一篇更早的文章《Efficient Attention: Attention with Linear Complexities》则给出了一个更有意思的选择。它留意到在 ${\mathbf{QK}}^{\top }$ 中，$\mathbf{Q},\mathbf{K},\in {\mathbb{R}}^{n\times d}$，如果 $\mathbf{Q}$ 在 $d$ 那一维是归一化的、并且 $\mathbf{K}$ 在 $n$ 那一维是归一化的，那么 ${\mathbf{QK}}^{\top }$ 就是自动满足归一化了，所以它给出的选择是：

$$Attention(\mathbf{Q},\mathbf{K},\mathbf{V})=softma{x}_2(\mathbf{Q})softma{x}_1(\mathbf{K}{)}^{\top }\mathbf{V}$$

其中 $softma{x}_1$、$softma{x}_2$ 分别指在第一个（$n$）、第二个维度（$d$）进行 softmax 运算。也就是说，这时候我们是各自给 $\mathbf{Q},\mathbf{K}$ 加 softmax，而不是 ${\mathbf{QK}}^{\top }$ 算完之后才加 softmax

如果直接取 $\varphi ({\mathbf{q}}_i)=softmax({\mathbf{q}}_i),\phi ({\mathbf{k}}_j)=softmax({\mathbf{k}}_j)$，那么很显然这个形式也是前面我们说的核函数形式的一个特例。

最后，苏神给出了他自己的一个构思，这个构思的出发点源于对原始 attention 公式的近似，由泰勒展开我们有：

$$e^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}\approx 1+{\mathbf{q}}_i^{\top }{\mathbf{k}}_j$$

如果 ${\mathbf{q}}_i^{\top }{\mathbf{k}}_j\ge -1$，那么就可以保证右端的非负性，从而可以让 $sim({\mathbf{q}}_i,{\mathbf{k}}_j)=1+{\mathbf{q}}_i^{\top }{\mathbf{k}}_j$。想要保证 ${\mathbf{q}}_i^{\top }{\mathbf{k}}_j\ge -1$，只需要分别对 ${\mathbf{q}}_i,{\mathbf{k}}_j$ 做 $l_2$ 归一化，所以苏神最终提出的方案就是：

$$sim\left({\mathbf{q}}_i,{\mathbf{k}}_j\right)=1+{\left(\frac{{\mathbf{q}}_i}{\Vert {\mathbf{q}}_i\Vert }\right)}^{\top }\left(\frac{{\mathbf{k}}_j}{\Vert {\mathbf{k}}_j\Vert }\right)$$

这不同于核函数形式，但理论上它更加接近原始的 Scaled-Dot Attention

结语

本篇文章从矩阵乘法的角度来探究 transformer，首先从 encoder 的角度去观察 attention，其本质是将低阶语义向量转换为高阶语义向量的形式，其动机是要对低阶语义向量 $X_i$ 周围的每一个低阶语义向量做 similarity（相似性），然后再把它们的低阶语义向量基于相似性做一个加权平均，如果直接做的话是没有参数的，也无法进行学习，因此我们给它投影到多个子空间去做 attention 最后做一个拼接

transformer 的核心变换是两次，先做一次 self-attention，最后做 Feedforward。self-attention 核心的计算量在 $Q{K}^{T}V$ 三个矩阵的相乘上，先计算 $Q{K}^T$ 计算量是 $O(N^2)$，先计算 $K^{T}V$ 计算量是 $O(N)$

最后我们介绍了下 linear transformer，利用 kernel 表达可以将原来的 $Q,K$ 映射到新空间 $\varphi (Q_i),\varphi (K_j)$ 上，把 $\varphi (Q_i)$ 提取出来从而使得 attention 的计算量只有 $O(N)$

大家可以多看看连博的讲解，非常的不错🤗

参考

• 从矩阵乘法探秘transformer+代码讲解
• 深入理解transformer
• 《Attention Is All Your Need》
• https://github.com/karpathy/minGPT/tree/master/mingpt
• 《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》
• 线性Attention的探索：Attention必须有个Softmax吗？