LLM 模型架构
1. LLM 核心模型 Transformer
这是《大语言模型》这本书的学习笔记。详细内容可以参考书籍。
目前主流的大模型都是基于Transformer模型进行设计的。Transformer是由多层的多头注意力模块(Multi-Head Self-Attention) 堆叠而成的神经网络模型。
原始的Transformer是由Encoder 和Decoder两部分构成,而这两部分是可以独立使用的(如BERT 单独使用Encoder, GPT单独使用Decoder).
大语言模型与早期的预训练语言模型相比,主要是使用了更长的向量维度、更深的层数、更大的数据集、进而包含了更大规模的模型参数。LLM主要使用Decoder 架构,对Transformer本身结构和配置改变不大。
关于Transformer的详细介绍可以参考我之前的学习笔记:Attention
-
输入端:x = v + p (词向量 + 位置编码)
- 输入词元序列 首先经过一个Embedding Module转成词向量。
- 位置编码 (由于Transformer无法识别序列顺序所以 通过位置编码来表示序列中的位置)
- 位置编码也有绝对位置编码,相对位置编码。这个后续讨论。
-
多头注意力机制(MHA)
- 多头注意力机制是Transfomer的核心技术,能够直接建模任意距离的词元之间的相互关系。之前RNN迭代利用前一个时刻状态更新当前时刻状态,在处理长文本时会出现梯度爆炸或消失的问题。在卷积神经网络中只有同一个卷积核的窗口中的词元可以直接交互,通过堆叠层数来实现远距离词元信息交互。
- 多头注意力机制 = 多个自注意力模型组成。
- 每个模块将输入词元映射成Query,Key,Value 三个矩阵。然后对于每个query,与所有没有被mask的key计算点积,点积值除以 D \sqrt{D} D(D是key的向量维度),然后传入softmax中计算权重。权重与Value加权和得到最终输出。
- A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T D ) V Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{D}})V Attention(Q,K,V)=softmax(DQKT)V 其中 Q = X W Q Q=XW^Q Q=XWQ, K = X W K K=XW^K K=XWK, V = X W V V=XW^V V=XWV
- 多头注意力机制 使用了 H 组 结构相同但是参数不同的自注意模块。输入序列通过不同的权重映射为K Q V,每组KQV经过上面计算映射成一个Head得到自注意力输出,不同 Head 被拼接在一起,通过一个权重矩阵 W O W^O WO(H*H)进行映射,产生最终的输出。 M H A = C o n c a t ( h e a d 1 , . . h e a d N ) W O MHA=Concat(head_1,..head_N)W^O MHA=Concat(head1,..headN)WO 其中 h e a d n = A t t e n t i o n ( X W n Q , X W n K , X W n V ) head_n = Attention(XW^Q_n,XW^K_n,XW^V_n) headn=Attention(XWnQ,XWnK,XWnV).
- 自注意力机制能够直接建模序列中任意两个位置的关系,进而有效捕捉长程依赖关系,具有更强的序列建模能力。另一个优势是计算可以并行化,因此可以支持大规模参数的高效优化。
-
前馈网络层(FFN)
- FeedForwardNetwork.对每个位置的隐状态进行非线性变化和特征提取。
F
F
N
(
x
)
=
σ
(
X
W
U
+
b
1
)
W
D
+
b
2
FFN(x) = \sigma(XW^U+b_1)W^D+b_2
FFN(x)=σ(XWU+b1)WD+b2
- FeedForwardNetwork.对每个位置的隐状态进行非线性变化和特征提取。
F
F
N
(
x
)
=
σ
(
X
W
U
+
b
1
)
W
D
+
b
2
FFN(x) = \sigma(XW^U+b_1)W^D+b_2
FFN(x)=σ(XWU+b1)WD+b2
-
编码器 (Encoder = LayerNorm+MHA+FFN)
- 编码器=多个相同的层堆叠而成(每层包含MulitHeadAttention和FFN)
- 在MHA和FFN之后,使用LayerNormalization 和 Residual Connection (残差连接)来加强模型的训练稳定性。
- 残差连接将输入层与输出层相加,实现了信息在不同层的传递,缓解梯度爆炸和消失的问题。
- LayerNorm则对数据进行重新缩放,提升模型的训练稳定性。
- 编码器端,输入数据是完全可见的,自注意力模块采用的是双向注意力,每个词元表示能够有效融合上下问的语义关系。
- X l ′ = L a y e r N o r m ( M H A ( X l − 1 + X l − 1 X'_l = LayerNorm(MHA(X_{l-1}+X_{l-1} Xl′=LayerNorm(MHA(Xl−1+Xl−1
- X l = L a y e r N o r m ( F F N ( X l ′ + X l ′ ) X_l = LayerNorm(FFN(X'_l+X'_l) Xl=LayerNorm(FFN(Xl′+Xl′)
-
解码器(Decoder)
- Decoder基于Encoder的最后一层输出表示以及模型生成的词元序列,执行后续的生成任务。
- Decoder与Encoder的不同是 Decoder需要引入掩码自注意力模块(Masked Self-attention) 用来计算注意力分数时掩盖当前位置之后的词,以保证生成目标序列时不依赖未来的信息。
- 经过Decoder后,模型通过一个全连接层将输出映射到大小为V的目标词汇表的概率分布,并基于某种解码策略生成对应的词元。
2. 详细配置
2.1 归一化方法
大语言模型的预训练过程中经常会出现不稳定的问题。为了应对这一问题,深度学习方法通常会采用特定的归一化策略来加强神经网络训练过程的稳定性。原始的Transformer 模型主要使用了层归一化方法(Layer Normalization, LN)。随着研究工作的不断深入,基于层归一化的改进技术不断涌现,例如均方根层归一化(Root Mean Square Layer Normalization, RMSNorm) 和DeepNorm ,这些新技术已经在一些大语言模型中得到应用。
-
LayerNorm: 在早期的研究中,批次归一化(Batch Normalization, BN)是一种广泛采用的归一化方法。然而,该方法难以处理可变长度的序列数据和小批次数据。因此,相关研究提出了层归一化这一技术,针对数据进行逐层归一化.层归一化会计算每一层中所有激活值的均值𝝁 和方差𝝈,从而重新调整激活值的中心和缩放比例:
-
RMSNorm:为了提高层归一化的训练速度,RMSNorm 仅利用激活值总和的均方根RMS(𝒙) 对激活值进行重新缩放。使用RMSNorm 的Transformer 模型相比于之前LayerNorm 训练的模型在训练速度和性能上均具有一定优势。采用RMSNorm 的代表性模型包括Gopher 和Chinchilla.
-
DeepNorm: 由微软的研究人员提出,旨在稳定深层Transformer的训练。具体而言,DeepNorm 在LayerNorm 的基础上,在残差连接中对之前的激活值𝒙 按照一定比例𝛼 进行放缩。通过这一简单的操作,Transformer 的层数可以被成功地扩展至1,000 层,进而有效提升了模型性能与训练稳定性。DeepNorm(𝒙) = LayerNorm(𝛼 · 𝒙 + Sublayer(𝒙)),其中,Sublayer 表示Transformer 层中的前馈神经网络或自注意力模块,。GLM-130B采用了DeepNorm 作为归一化技术.
2.2 归一化模块位置
为了加强大语言模型训练过程的稳定性,除了归一化方法外,归一化模块的位置也具有重要的影响.,归一化模块的位置通常有三种选择,分别是层后归一化(Post-Layer Normalization, Post-Norm)、层前归一化(Pre-LayerNormalization, Pre-Norm)和夹心归一化(Sandwich-Layer Normalization, Sandwich-Norm)。
- Post-LayerNorm: Post-Norm 是在原始Transformer 模型中所使用的一种归一化技
术。其中,归一化模块被放置于残差计算之后。Post-Norm(𝒙) = Norm(𝒙 + Sublayer(𝒙)),。在原理上,后向归一化具有很多优势。首先,有助于加快神经网络的训练收敛速度,使模型可以更有效地传播梯度,从而减少训练时间。其次,后向归一化可以降低神经网络对于超参数(如学习率、初始化参数等)的敏感性,使得网络更容易调优,并减少了超参数调整的难度。然而,由于在输出层附近存在梯度较大的问题,采用Post-Norm 的Transformer 模型在训练过程中通常会出现不稳定的现象。因此,现有的大语言模型中,Post-Norm很少被单独使用,通常是与其他策略相结合应用。例如,GLM-130B 将Post-Norm与DeepNorm 结合使用. - Pre-Norm: 与Post-Norm 不同,Pre-Norm 将归一化模块应用在每个子层之前.Pre-Norm(𝒙) = 𝒙 + Sublayer(Norm(𝒙)),Pre-Norm 在最后一个Transformer层后还额外添加了一个LayerNorm。相较于Post-Norm,Pre-Norm 直接把每个子层加在了归一化模块之后,仅仅对输入的表示进行了归一化,从而可以防止模型的梯度爆炸或者梯度消失现象。虽然使用了Pre-Norm 的Transformer 模型在训练过程中更加稳定,但是性能却逊色于采用了Post-Norm 的模型.尽管对于性能有一定的影响,但由于其能够有效维持训练的稳定性,很多主流的大语言模型仍然采用Pre-Norm。
- Sandwich-Norm. 在Pre-Norm 的基础上在残差连接之前增加了额外的LayerNorm,旨在避免Transformer 层的输出出现数值爆炸的情况。Sandwich-Norm(𝒙) = 𝒙 + Norm(Sublayer(Norm(𝒙))).但是研究人员发现,Sandwich-Norm 有时仍然无法保证大语言模型的稳定训练,甚至会引发训练崩溃的问题.
问题:为什么pre-norm不如post-norm性能好?
请参考:https://spaces.ac.cn/archives/9009
Pre Norm的深度有“水分”!也就是说,一个L层的Pre Norm模型,其实际等效层数不如L层的Post Norm模型,而层数少了导致效果变差了。
为什么每Norm一次就削弱一次恒等分支的权重?
2.3 激活函数
激活函数主要是为神经网络引入非线性变化,从而提升神经网络的模型能力。
原始的Transformer 中采用了ReLU(Rectified Linear Unit)激活函数。该激活函数计算较为简单,仅仅是将对输入中每个神经元和“零值”进行比较,并将小于零的神经元的值设置为0。然而,ReLU 可能会产生神经元失效的问题,被置为0 的神经元将学习不到有用的信息。ReLU=max(x,0)
针对ReLU的不足,研究人员探索了ReLU的变种。Swish 激活函数将神经元和该神经元的sigmoid 激活的乘积作为新的激活函数。而GELU(Gaussian Error Linear Unit) 则利用标准高斯累积分布函数作为激活函数,被很多的Transformer 模型所采用。相比于原始的ReLU 函数,这些新的激活函数通常能够带来更好的性能并且收敛性更好,但是计算过程更为复杂。
S
w
i
s
h
(
x
)
=
x
∗
s
i
g
m
o
i
d
(
x
)
,
G
E
L
U
(
x
)
=
0.5
x
∗
[
1
+
e
r
f
(
x
2
)
]
,
e
r
f
(
x
)
=
2
π
∑
1
x
e
−
t
2
Swish(x)=x*sigmoid(x), GELU(x)=0.5x *[1+erf(x\sqrt{2})],erf(x)=\frac{2}{\sqrt{\pi}}\sum_1^x e^{-t^2}
Swish(x)=x∗sigmoid(x),GELU(x)=0.5x∗[1+erf(x2)],erf(x)=π21∑xe−t2
近来,大语言模型(例如PaLM 和LaMDA)也经常采用GLU(Gated LinearUnit)激活函数以及它的变种,特别是SwiGLU 和GeGLU。不同于其他激活函数,GLU 激活函数引入了两个不同的线性层。其中一个线性层的输出将被输入到一个激活函数(例如,GeGLU 采用GELU 激活函数)中,其结果将和另一个线性层的输出进行逐元素相乘作为最终的输出。相比于其他的激活函数,使用GLU激活函数变体通常能够带来更佳的性能表现.
2.4 位置编码
由于Transformer 模型中自注意力模块具有置换不变性,因此仅使用注意力机制无法捕捉序列中的顺序关系,从而退化为“词袋模型”。为了解决这一问题,需要引入位置编码(Position Embedding, PE)对于序列信息进行精确建模,从而将绝对或相对位置信息整合到模型中。
关于位置编码 苏神的介绍更详细:https://kexue.fm/archives/8130
2.4.1 绝对位置编码
在原始Transformer中,采样了绝对位置编码。根据输入词元在序列中的绝对位置生成唯一的位置嵌入,并与词元的嵌入表示进行相加来注入位置信息
x
t
=
v
t
+
p
t
x_t = v_t + p_t
xt=vt+pt。原始Transfomer采用了正余弦位置编码。
2.4.2 相对位置编码
相对位置编码是根据Key和query之间的偏移量计算得来的。计算得到的相对位置编码通常应用于注意力矩阵的计算中,而不是直接与词元本身的位置编码进行相加。相对位置编码的Transformer可以对更长序列的进行建模,具备一定的外推能力。
- Transformer-XL提出了一种相对位置编码,基于绝对位置编码分解得到的。
- T5提出了一种较为简化的相对位置编码。在注意力分数中引入了可学习的标量,这些标量是基于查询和键的位置之间的距离计算的。
A i j = x i W Q W K T x j T + r i − j A_{ij} = x_i W^Q W^{K^T} x_j^T + r_{i-j} Aij=xiWQWKTxjT+ri−j 其中 r i − j r_{i-j} ri−j表示基于query 和 key之间偏移的可学习标量。
2.4.3 旋转位置编码 RoPE
RoPE的出发点就是“通过绝对位置编码的方式实现相对位置编码”,通过公式推导得到 "向量旋转"的位置编码。 原作者的解读 https://kexue.fm/archives/8265
2.4.4 ALiBi位置编码
ALiBi 通过在键和查询之间的距离上施加相对距离相关的惩罚来调整注意力分数。
A
i
j
=
x
i
W
Q
W
K
T
x
j
T
−
m
(
i
−
j
)
A_{ij} = x_i W^Q W^{K^T} x_j^T - m(i-j)
Aij=xiWQWKTxjT−m(i−j)
i-j是query 和key之间的位置偏移量,m是每个注意力头独有的惩罚系数。ALiBi的惩罚分数是预先设定的,不需要引入任何可训练得参数。
2.5 注意力机制
注意力机制是Transformer 架构中的核心技术,它能够针对序列中的词元对构建交互关系,聚合来自于不同位置的语义信息。
2.5.1 完整自注意力机制
考虑序列中所有词元之间的相互关心。因此对于序列长度为T的序列需要O(T^2)的计算复杂度。
2.5.2 稀疏注意力机制
完整自注意力机制在处理长序列时O(T^2)的时间复杂度带来计算和存储开销较大。
为了降低注意力机制的计算复杂度,研究人员提出了多种高效的注意力变种。其中,滑动窗口注意力机制(Sliding Window Attention, SWA)是大语言模型中使用最多的一种稀疏注意力机制。不同于完整的注意力机制,滑动窗口注意力根据词元位置,仅仅将位置索引上距离该词元一定范围内的词元考虑到注意力的计算中。具体来说,滑动窗口注意力设置了一个大小为𝑤 的窗口,对每个词元𝑢𝑡,只对窗口内的词元[𝑢𝑡−𝑤+1, . . . , 𝑢𝑡 ] 进行注意力计算,从而将复杂度降低到𝑂(𝑤𝑇)。进一步,通过信息的逐层传递,模型实现了随着层数线性增长的感受野,从而获取远处词元的信息。
2.5.3 多查询/分组注意力机制
为了提升注意力机制的效率,多查询注意力(Multi-Query Attention, MQA)提出针对不同的头共享相同的键和值变换矩阵。这种方法减少了访存量,提高了计算强度,从而实现了更快的解码速度并且对于模型性能产生的影响也比较小。
分组查询注意力机制(Grouped-Query Attention, GQA)结合多查询注意力机制的效率与多头注意力机制的性能。GQA 将全部的头划分为若干组,并且针对同一组内的头共享相同的变换矩阵。这种注意力机制有效地平衡了效率和性能,被LLaMA-2 模型所使用。
2.5.4 硬件优化的注意力机制
利用硬件设施来优化注意力模块的速度和内存消耗,FlashAttention 通过矩阵分块计算以及减少内存读写次数的方式,提高注意力分数的计算效率;PagedAttention 则针对增量解码阶段,对于KV 缓存进行分块存储,并优化了计算方式,增大了并行计算度,从而提高了计算效率。
2.6 混合专家模型
来自于《Mixtral of Experts》,Mixtral 8x7B 是一种稀疏的混合专家(稀疏 MoE)模型,目前是性能最佳的大型语言模型(LLM)之一,同时也是最受人关注的一种公开可用的 LLM。具体解读可参考: Mixtral of Experts 机器之心解读
Mixtral 8x7B 的关键思想是用 8 个专家层替换 Transformer 架构中的每个前馈模块,8x 是指使用了 8 个专家子网络。7B 是指其组合了 Mistral 7B 模块。Mixtral 8x7B 总共有 47B 参数,明显少于 Llama 270B 等模型。此外,由于每个时间步骤仅有 2 个专家处于活动状态,因此对于每个输入 token,该模型仅使用 13B 参数。如此一来,它的效率就比常规的非 MoE 47B 参数模型高多了。
在MoE 中,每个混合专家层包含K个专家组件,每个组件Ei都是一个前馈神经网络,对于输入的每个词元xt,模型通过一个路由网络(门控函数)G来计算该词元对应于各个专家的权重。在路由函数中,首先通过线性层
W
G
W^G
WG映射为K各专家的得分,并基于此选择出概率最该的k个专家进行激活。随后这k个专家的得分被送入softmax函数计算他们的权重
G
(
x
t
)
=
[
G
(
x
t
)
1
,
.
.
G
(
x
t
)
k
]
G(x_t)=[G(x_t)_1,..G(x_t)_k]
G(xt)=[G(xt)1,..G(xt)k] 没有被选则的专家权重置为0.
G
(
x
t
)
=
s
o
f
t
m
a
x
(
t
o
p
k
(
x
t
∗
W
G
)
)
G(x_t)=softmax(topk(x_t * W^G))
G(xt)=softmax(topk(xt∗WG)),
之后,每个被选择的词元的输出的加权和将作为该混合专家网络层的最终输出𝒐𝑡:
o
t
=
M
o
E
l
a
y
e
r
(
x
t
)
=
∑
i
=
1
K
G
(
x
t
)
i
∗
E
i
(
x
i
)
o_t = MoElayer(x_t) = \sum_{i=1}^K G(x_t)_i * E_i(x_i)
ot=MoElayer(xt)=i=1∑KG(xt)i∗Ei(xi)
2.7 LLaMA实现和模型配置建议
关于模型详细配置的推荐建议:
- 为了增强稳定性:建议采用前置RMSNorm作为层归一化。
- 为了更好地模型性能,激活函数可以考虑SwiGLU或GeGLU
- 位置编码可以选择RoPE或ALiBi 这两种在建模长序列数据具有较好的性能。
对于一个LLaMA 模型,其首先将输入的词元序列通过词嵌入矩阵转化为词向量序列。之后,词向量序列作为隐状态因此通过𝐿 个解码器层,并在最后使用RMSNorm 进行归一化。归一化后的最后一层隐状态将作为输出。
3. 主流架构
在预训练时代,主要范式是 预训练+微调。Bert为代表的Encoder-only架构,以GPT为代表的Decoder-Only架构,以T5为代表的Encoder-Decoder架构。
随着GPT的发展,目前LLM生成式大语言模型上,主要以Decoder-only为主流的架构。Decoder架构还可以细分为两种变种:因果解码器(Causal Decoder)和前缀解码器(Prefix Decoder).一般提到解码器通常是因果解码器架构。
4. 长上下文模型
目前,增强大语言模型长文本建模能力的研究主要集中在两个方向:一个是扩展位置编码,一个调整上下文窗口。
4.1 扩展位置编码
目前比较主流的位置编码方法RoPE 在未经特殊修改的情况下并不具备良好的外推能力。具体来说,在处理更长的文本时,RoPE 在每个子空间上需要处理更大的旋转角度,而这些旋转角度可能会超过其训练中的角度分布范围。因此,很多研究工作在RoPE 的基础上进行了重要改进,旨在提升其在不经过训练或继续训练的情况下对于长文本的建模能力.