What is in Standard Transformers

  • 在 Input Embedding 上进行 Position Embedding 的嵌入
    • 为什么需要位置编码？什么样的位置编码是好的位置编码？
    • 对于序列中第 pospos 个位置（pospos 表示词语在序列中的索引，从 00 开始），其位置编码向量的第 ii 个维度（ii 表示向量的维度索引）的计算公式
PE(pos,2i)=sin⁡(pos100002idmodel)PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{\frac{2i}{d_{model}}}}\right)PE(pos,2i+1)=cos⁡(pos100002idmodel)PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{\frac{2i}{d_{model}}}}\right)
  • Nomalizations
    • Post-Norm（先做 forward 和残差，然后再归一化） && Layer-Norm (对 dmodeld_{\text{model}} 这个维度做归一化)
  • FFN
    • 使用 RELU 作为激活函数
    • FFN(x)=W2max⁡(0,W1x+b1)+b2\text{FFN}(x) = W_2 \max(0, W_1 x + b_1) + b_2

在 Attention 那一讲，我们更多探索的是现代大模型架构对 Attention 计算 (Attention Block) 的核心组件的改进和加速，在本文中，我们会思考模型架构中的一些其他组件，看看现代大模型架构相对于传统 Transformers 结构的改进。

在第三讲，我们将会从各家最前沿的技术报告出发，探索下一个时代的前卫的 LLM 架构。

Pre-Norm and Post-Norm

  • Pre-Norm: 先对输入做 Norm 然后再 forward && 做残差连接
  • Post-Norm: 先 forward && 做残差连接 再做 Norm

从梯度传递的角度，Pre-Norm 的训练稳定性更好 (xl+1=xl+F(N(xl))x_{l+1} = x_l + F(N(x_l))) 进行反求梯度会更好 -> 残差项在外面，因此会更稳定

从数学角度诠释：

  • PostNorm:
    • Forward: xl+1=N(xl+F(xl))x_{l+1} = N(x_l + F(x_l))
    • Backward: ∂xl+1∂xl=JN(I+JF)\frac{\partial x_{l+1}}{\partial x_l} = J_N(I + J_F)
  • PreNorm:
    • Forward: xl+1=xl+F(N(xl))x_{l+1} = x_l + F(N(x_l))
    • Backward: ∂xl+1∂xl=I⏟identity path+JFJN\frac{\partial x_{l+1}}{\partial x_l} = \underbrace{I}_{\text{identity path}} + J_FJ_N 从计算稳定性的角度，因为 Pre-Norm 对残差在外面，因此反向传播的时候会更加的稳定

https://arxiv.org/pdf/2002.04745 这篇论文给出了系统性的诠释：

  • 对于 PostNorm 的结构，学习率预热非常的重要（在训练初期，学习率线性上升到一个峰值，然后通过各种策略进行 decay）
  • Pre-Norm 从理论上和实验上分别证明了其训练稳定性往往高于 Post-Norm
    • Post-Norm 需要很谨慎的 learning rate warmup 的过程，稳定性比 Pre-Norm 差

llm/images/post-norm.png

除了标准的 Pre-Norm 和 Post-Norm，还有一些更加高级的操作：

  • 传统 Pre-Norm: xl+1=xl+F(Norm(xl))x_{l+1} = x_l + F(Norm(x_l))
  • Double Norm: xl+1=xl+Norm(F(Norm(xl)))x_{l+1} = x_l + Norm(F(Norm(x_l)))
    • 注意，仍然把残差连接放在了最外面，以保证稳定性
  • Non-residual Post-Norm: xl+1=xl+Norm(F(xl))x_{l+1} = x_l + Norm(F(x_l))

RMS Normalizations

我们主要讨论三种常见的 Norm 计算方式：

  • Batch Norm
  • Layer Norm
  • RMS Norm

Batch Norm

针对 batch-size 和 length 维度进行归一化，但是因为 LLM 需要处理可变长文本，因此不再采用 Batch Norm

μB=1m∑k=1mxk\mu_B = \frac{1}{m} \sum_{k=1}^{m} x_kσB2=1m∑k=1m(xk−μB)2\sigma_B^2 = \frac{1}{m} \sum_{k=1}^{m} (x_k - \mu_B)^2x^i=xi−μBσB2+ϵ\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}yi=γx^i+βy_i = \gamma \hat{x}_i + \beta

Layer Norm

μ=1d∑i=1dxi\mu = \frac{1}{d} \sum_{i=1}^{d} x_iσ2=1d∑i=1d(xi−μ)2\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2x^i=xi−μσ2+ϵ\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}yi=γi⊙x^i+βiy_i = \gamma_i \odot \hat{x}_i + \beta_i

在传统的 Attention 中，模型的 Norm 使用 LayerNorm 进行优化

RMSNorm

RMS(x)=1d∑i=1dxi2+ϵ\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon}x^i=xiRMS(x)\hat{x}_i = \frac{x_i}{\text{RMS}(x)}yi=γi⊙x^iy_i = \gamma_i \odot \hat{x}_i

RMS-Norm 仍然使用 dimension 这个维度进行归一化，同时，RMSNorm 省略了大量的中间变量（不需要计算均值和归一化的均值）

为什么 RMS Norm 的优势显著高于 LayerNorm ?

运算更简单，涉及更少中间变量的存储

激活函数

激活函数在 FFN 中承担的作用：引入非线性

常见的激活函数：

  • sigmoid
  • tanh
  • relu
  • leaky relu
  • gelu
  • silu/swish

Gated FFNs and SwiGLU

Gated FFN 在传统 FFN 层上引入了门控机制：

FFN(x)=W2max⁡(0,W1x+b1)+b2\text{FFN}(x) = W_2 \max(0, W_1 x + b_1) + b_2

h=σ(W1x) h = \sigma(W_1 x)

对于 Gated FFN:

h=σ(Wgx)⊙(Wux)h = \sigma (W_g x) \odot (W_u x)
  • Gate Branch 负责提供门控单元，需要过激活函数
  • Feature Branch 添加了一个可学习矩阵

SwiGLU:

SiLU(x)=zσ(z) \text{SiLU}(x) = z \sigma (z)

FFNSwiGLU(x)=Wd[SiLU(Wgx)⊙(Wux)] \text{FFN}_{\text{SwiGLU}}(x) = W_d [\text{SiLU}(W_g x) \odot (W_u x)]
  • 激活函数使用 SiLU
  • 添加 Gated FFN 同时在外面再套一层可学习矩阵

从参数角度分析，因为引入了一个 gated 参数 WgW_g，因此参数量会大 3/2 倍。

Position Information and RoPE

会考 Rope 的推导 && 模型长度的内推和外推

首先，考虑一个问题，为什么我们需要 Positiom Embedding？因为 transformer 本身的 Attention 不会考虑 position 关系，对 sequence 中的每一个 token 都做同等地位的注意力 attention score 计算。

然后，接下来，我们考虑第二个问题，对于一个 token 在训练 batch 中具备一个编码 ii，但是这个 ii 不可以按照绝对值的形式直接 append 到一个独立的维度，因为训练的 batch 在实际 inference 的时候，绝对位置可能会发生变化。因此，关注 token间相对位置 是使用 Position Enbedding 最关键的方法。

我们接下来介绍 RoPE 的基本思想，假设当前的 query token 是 qtq_t(Positon: t)，当前 token 需要和历史的 ss token 进行注意力计算。

  • Linear Projections: qt=xtWQq_t = x_t W_Q, ks=xsWKk_s = x_s W_K
  • Position-dependent rotations: qt′=Rtqt{q_t}^{'} = R_t q_t, ks′=Rsks{k_s}^{'} = R_s k_s

RoPE 需要解决的关键是，RsR_s, RtR_t 需要使用什么样矩阵表示？

Rotary Matrix

(x′y′)=(cos⁡θ−sin⁡θsin⁡θcos⁡θ)(xy)\begin{pmatrix} x' \\ y' \end{pmatrix} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x \\ y \end{pmatrix}

几何意义上，定义上述二维矩阵是 R(θ)R(\theta) 则表示为，在二维空间下按照正方向对向量进行旋转。

旋转矩阵具备如下的性质：

  • R(α)T=R(−α)R(\alpha)^T = R(- \alpha)
  • R(α)TR(β)=R(β−α)R(\alpha)^T R(\beta) = R(\beta - \alpha)

我们来看看如果我们使用旋转矩阵作为我们的 RR, 我们的 attention 计算如下:

qt′=R(t)qtq'_t = R(t) q_t, ks′=R(s)ksk'_s = R(s) k_s，展开 Attention 计算：

(qt′)⊤ks′=(R(t)qt)⊤(R(s)ks)=qt⊤R(t)⊤R(s)ks=qt⊤R(s−t)ks(q'_t)^\top k'_s = (R(t) q_t)^\top (R(s) k_s) = q_t^\top R(t)^\top R(s) k_s = q_t^\top R(s - t) k_s

如果我们考虑旋转“角速度”：

(qt′)⊤ks′=(R(t)qt)⊤(R(s)ks)=qt⊤R(t)⊤R(s)ks=qt⊤R(ω(s−t))ks (q'_t)^\top k'_s = (R(t) q_t)^\top (R(s) k_s) = q_t^\top R(t)^\top R(s) k_s = q_t^\top R(\omega(s - t)) k_s

这个式子给了我们一个很有意思的结论：加上 Rope 旋转矩阵后，attention 的计算分数只和两个 token 的相对位置有关，而和绝对位置无关。

对于简单的 Positional Embedding，向量的加法通过 sin 和 cos 实现不同频率的 position 位置的插入，也有类似的性质，但是没有 RoPE 性质那么好。

将二维的旋转矩阵 apply 到高维隐空间中，核心思想是：将高维向量拆分为多个独立的 2 维子空间，并在每个子空间上分别应用不同频率的二维旋转。

对于一个 dd 维的向量（假设 dd 是偶数），高维旋转矩阵 RtR_t 可以写成如下形式：

Rt=(R(θ1,t)0⋯00R(θ2,t)⋯0⋮⋮⋱⋮00⋯R(θd/2,t))R_t = \begin{pmatrix} R(\theta_{1, t}) & 0 & \cdots & 0 \\ 0 & R(\theta_{2, t}) & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & R(\theta_{d/2, t}) \end{pmatrix}

其中：

  • 整个矩阵大小为 d×dd \times d。
  • 对角线上的每一个小块 R(θi,t)R(\theta_{i, t}) 都是一个前面讨论过的 2×22 \times 2 二维旋转矩阵： R(θi,t)=(cos⁡θi,t−sin⁡θi,tsin⁡θi,tcos⁡θi,t)R(\theta_{i, t}) = \begin{pmatrix} \cos\theta_{i, t} & -\sin\theta_{i, t} \\ \sin\theta_{i, t} & \cos\theta_{i, t} \end{pmatrix}
  • 每一个 2 维小块的旋转角度 θi,t\theta_{i, t} 是位置 tt 与不同频率 ωi\omega_i 的乘积：θi,t=t⋅ωi\theta_{i, t} = t \cdot \omega_i。

其中 θi,t\theta_{i,t} 指的是在绝对位置 tt 处，针对第 ii 个二维子空间的旋转角度。ωi\omega_i 代表不同 dim 位置上的 rotation，具体的角速度分数和 positional embedding 几乎一致：

ωi=10000−2(i−1)d\omega_i = 10000^{-\frac{2(i-1)}{d}}

Positional Interpolation & Extrapolation

LLM 会设置一个最大上下文长度 max-content-length, 在 vLLM 等推理引擎进行模型推理时识别这个参数，在监测到长度超过上下文时，会拒绝这一次 LLM-Call 的请求。并且，面对超过预训练长度的请求，模型会因为缺乏对应的训练数据，导致PPL 陡增。

  • 外推：直接去让模型处理超过训练长度的 Positional Embedding，但是会导致效果 drop
  • 内推 (PI: Positional Interpolation)：假设模型训练时的最大上下文长度是 LL，当处理超过上下文长度 L′>LL' > L 的时候，需要将 L′L' 映射到一个在 LL 内部的值 L′′L''（可以是小数）

定义缩放因子 S=LL′′S = \frac{L}{L''}，则上述 Rope 被修正为: θi,t=tSωi\theta_{i,t} = \frac{t}{S} \omega_i，PI 本质上是一种线性内插的方法。

PI 通过固定的缩放因子，可以在推理过程中动态的拉长最大上下文到一个更大的值 （e.g. 拉长 2/4 倍）

NTK-Aware Methods

https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/

在 rope 中，不同位置的旋转矩阵具有不同的频率，低维高频，高维低频。

因此，在线性插值的过程中，会导致低维区域下的旋转变的异常拥挤，会导致效果 drop，可以看到，在低维高频区域，经过线性插值的 PI 和原来的 Rope 存在比较大的差异。

llm/images/ntk.png

在 PI 中，θ\theta 的格式如下:

(θpi)i,t=tS⋅b−2(i−1)d(\theta_{\text{pi}})_{i,t} = \frac{t}{S} \cdot b^{-\frac{2(i - 1)}{d}}

(θNTK)i,t=t⋅(b⋅Sdd−2)−2(i−1)d(\theta_{\text{NTK}})_{i,t} = t \cdot (b \cdot S^{\frac{d}{d - 2}})^{-\frac{2(i - 1)}{d}}

我们考虑 0-based 计数的 RoPE:

θi,t=t⋅ωi=t⋅b−2id\theta_{i, t} = t \cdot \omega_i = t \cdot b^{-\frac{2i}{d}}

这个旋转角可以通过 β\beta 进制的角度解释，define β=b2d\beta = b^{\frac{2}{d}}:

θi,t=t⋅ωi=t⋅b−2id=tβi\theta_{i, t} = t \cdot \omega_i = t \cdot b^{-\frac{2i}{d}} = \frac{t}{\beta^i}

因此，旋转角本质上就是求位置 tt 的 β\beta 进制，现在我们的目的是拉长推理时的最大上下文长度，在保持 dimension 不变的情况下，模型需要对应提升 β\beta ：

θi,t=t/Sβi=t(kβ)i \theta_{i,t} = \frac{t/S}{\beta^i} = \frac{t}{(k \beta)^i} k=S−1/i k = S^{-1/i}

考虑最大维度下的偏差 -> i=d/2−1i = d/2 - 1 -> k=S2/(d−2)k = S^{2/(d-2)}

更精彩的诠释可以查看苏老师的博客：https://normxu.github.io/Rethinking-Rotary-Position-Embedding/

同样的，在 NTK-Aware 方法之后，还出现了 NTK-by-parts 等等方法：

NTK-by-parts 吸收了 NTK 高频和低频的在压缩时不一致的问题，因此引入了一个混合加权的策略，在低维高频区域，不引入任何压缩，在高维低频区域，引入简单的线性插值，中间部分通过滑动加权进行平滑处理。具体的数学公式可以看经典论文 https://arxiv.org/pdf/2309.00071

此外，还有 Dynamic-NTK，只在推理上下文长度超过 LL 时引入对应的策略 (PI, NTK-Aware, NTK-Select, YaRN)，在正常的上下文长度时，保持为 LL 不变。

rope.png

YaRN

https://arxiv.org/pdf/2309.00071 是在 NTK-Aware 基础上提出的一种新的拉长模型上下文的方法。

在插值过程中，研究者发现模型因为上下文变长导致 Attention 计算后的 attention score 出现熵增的趋势（因为参与计算的 token 数量变多了）

Yarn 在 Attention 计算时引入了一种退火参数：

Attention(Q,K,V)=softmax(QKTdk⋅t)V\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k} \cdot t} \right)V

在论文中，作者也给出了 tt 的最佳取值

1t=0.1⋅ln⁡(s)+1\sqrt{\frac{1}{t}} = 0.1 \cdot \ln(s) + 1

当我们将一个向量乘以一个大于 1 的倍数 kk 之后再输入到 Softmax 函数中，输出的概率分布会变得更加“尖锐”（Sharp）——即最大的概率值会更接近 1，而其余较小的概率值会迅速被压低趋近于 0。 这个有点类似于大模型输出时候，对最后一层的词表做 temperature 的 调整，再经过最后的 softmax 操作。如果 temperature 很小的时候，会导致模型 softmax 输出概率变尖锐，导致最终输出文本更趋于确定性。