attention与FFN复杂度&模型内存占用估计

按照定义一个a×b的矩阵乘以一个b×c的矩阵要做abc次乘法,所以abc就是两个矩阵相乘的复杂度了,这是我们估算Transformer复杂度的依据

SA复杂度

设n为序列长度,d为head_size(base版是64),h为head的数目(base版是12),那么hd就是我们通常说的“hidden_size”(base版是768)。
对于SA来说:
Q,K,V的投影变换,即n×hd的矩阵乘以hd×hd的矩阵做3次,因此计算量是3n(hd)2
h个Attention头的运算,每个头先是n×d的Q与d×n的KT相乘得到n×n的Attention矩阵(softmax和归一化的计算量暂且忽略),然后n×n的矩阵与n×d的V相乘得到n×d的矩阵,这两步的计算量都是n2d,所以总计算量是h(n2d+n2d);
输出投影变换,也是n×hd的矩阵乘以hd×hd的矩阵,计算量是n(hd)2
所以,SA的总计算量是
3n(hd)2+h(n2d+n2d)+n(hd)2=4nh2d2+2n2hd

FFN复杂度

FFN就是两个全连接层,也就是两个矩阵变换(激活函数的计算量也忽略不计),一般的参数设置是:第一层是n×hd的矩阵乘以hd×4hd的矩阵,第二层就是n×4hd的矩阵乘以4hd×hd的矩阵。所以总计算量是
n×hd×4hd+n×4hd×hd=8nh2d2

如果SA的计算量比FFN大

4nh2d2+2n2hd > 8nh2d2 ==> n>2hd
对于base版来说,这意味着n>1536!也就是说,只有当序列长度超过1536时,SA的计算量才大于FFN,在这之前,都是线性复杂度的FFN占主导

总复杂度

4nh2d2+2n2hd + 8nh2d2 = 12nh2d2+2n2hd
它是关于n的一次项和二次项的求和,当n足够大时,复杂度自然是𝒪(n2),然而二次项占主导的条件是
2n2hd>12nh2d2==> n>6hd
对于base版来说,这意味着n>4608!也就是说,当序列长度接近5000时,Transformer的复杂度才真正体现出二次性!

对于base版来说,当序列长度不超过1536时,Transformer的复杂度都是近乎线性的;
当序列长度超过1536时,Transformer的计算量逐渐以Attention为主,复杂度慢慢趋于二次方,直到长度超过4608,才真正以二次项为主

这些改进工作所关心的序列长度主要都是以千为单位的,有明显计算效率提升的序列长度基本上都要好几千;当然,我们前面的讨论主要针对的还是时间复杂度,对于空间复杂度,也就是显存占用量,降低的幅度一般要比时间复杂度提升的幅度的要大,但总体而言都是长序列才有价值。

怎样计算大模型的内存占用

GLM-130B模型有1300亿个参数。每个参数的大小:通常模型参数存储为32位浮点数
模型大小 = (参数个数)(每个参数大小,单位为字节)
一个32位浮点数占用4个字节的存储空间,因此对于具有1300亿个参数的
GLM-130B模型:模型大小=(130 * 10^9参数)
(4字节/参数)模型大小= 520 * 10^9字节
我们带进去计算也就是520 * 10^9字节/ (1024 * 1024 * 1024)≈ 484.5 GB
因此,具有1300亿个参数的GLM-130B模型将占用大约484.5 GB的硬盘空间。

### Transformer 模型FFN 前馈神经网络详解 #### FFN 的定义功能 位置前馈网络是一个全连接前馈网络,每个位置的词都单独经过这个完全相同的前馈神经网络[^3]。此设计使得模型能够逐位置独立地处理序列中的每一个元素,而不会引入额外的位置依赖关系。 #### 架构细节 FFN 由两个线性变换组成,即两个全连接层构成。具体来说: - **第一层**:输入先被映射到一个更高的维度空间(通常称为中间维度),这一过程通过第一个线性变换完成,并紧接着应用 ReLU 激活函数来增加非线性特性。 - **第二层**:接着,输出再次经由另一个线性变换返回原始特征空间大小,最终形成该位置上的新表示形式[^4]。 这种架构允许 FFN 对每一时刻的数据点执行局部且固定的转换操作,有助于提升整体系统的灵活性和表现力。 #### 数学描述 设 \( \mathbf{X} \in \mathbb{R}^{d_{\text{model}}} \) 表示输入向量,则 FFN 可以写作如下公式: \[ \text{FFN}(\mathbf{X}) = \max(0, \mathbf{XW}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2 \] 这里, - \( \mathbf{W}_1 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}} \),\( \mathbf{W}_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}} \) 分别代表两层间的权重矩阵; - \( \mathbf{b}_1,\mathbf{b}_2 \) 则是对应的偏置项; - \( d_{\text{model}} \) 和 \( d_{\text{ff}} \) 分别指代模型尺寸和内部扩展后的宽度。 #### Python 实现代码片段 下面给出一段简单的 PyTorch 版本实现 FFN 层的例子: ```python import torch.nn as nn class PositionwiseFeedForward(nn.Module): "Implements FFN equation." def __init__(self, d_model, d_ff, dropout=0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 = nn.Linear(d_model, d_ff) self.w_2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.w_2(self.dropout(torch.relu(self.w_1(x)))) ```
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值