斯坦福大学 | CS336 | 从零开始构建语言模型 | Spring 2025 | 笔记 | Lecture 3: Architectures,Hyperparameters
目录
前言
学习斯坦福 CS336 课程,本篇文章记录课程第三讲:语言模型架构和超参数设置,记录下个人学习笔记,仅供自己参考😄
website:https://stanford-cs336.github.io/spring2025
video:https://www.youtube.com/playlist?list=PLoROMvodv4rOY23Y0BoGoBGgQ1zmU_MT_
materials:https://github.com/stanford-cs336/spring2025-lectures
course material:https://github.com/stanford-cs336/spring2025-lectures/tree/main/nonexecutable/3-architecture.pdf
1. Outline and goals
本次讲座为《关于语言模型架构和训练的那些冷知识》,因为我们要深入探讨大多数课程都会略过的技术细节,比如 “我的超参数该如何设置?”,诸如此类的问题
我们接下来要讲的内容如下:
- Quick recap of the ‘standard’ transformer (what you implement)
- What do most of the large LMs have in common?
- What are common variations to the architecture / training process?
首先快速回顾一下 Transformer 架构,我们将介绍标准 Transformer 的两种变体,其中一种是原版的 Transformer,也就是 CS224n 课程中讲的那个版本,然后说说实际实现的版本,再讲讲目前业界公认的 Transformer 改进版
接着我们会从数据角度深入分析 Transformer 架构的本质,我们要探讨的核心问题是:目前已经训练出了大量的大语言模型,通过研读这些论文可以发现哪些要素发生了变化,哪些设计始终保持不变,通过这种近乎演化分析的方法我们试图揭示哪些核心要素对 Transformer 的有效运行真正至关重要
今天我们要探讨的课程主题是:实践出真知的核心要义在于—既然我们无法实际训练所有这些 Transformer 模型,就要学会借鉴他人的实践经验
2. Starting point: the original transformer
我们的起点正是最初的 Transformer 架构,这相当于是个知识回顾,希望大家还记得 CS224n 或其他 NLP 课程中学过的这部分内容

模型底层采用了简单的位置编码,架构中包含了多头注意力机制,随后接入了层归一化模块,残差连接以垂直方向贯穿整个架构,最后经过 MLP 处理并以 softmax 函数作为最终输出层。接下来我们将探讨这些核心组件的各种变体,直至呈现最前沿的 Transformer 架构演变形态
作业中要求实现的并非原论文的标准 Transformer 架构,我们做了一些改进,如下图所示:

我们将层归一化(LayerNorm)置于残差连接中每个模块的前端,并要求你实现旋转位置编码(RoPE),前馈网络采用了 SwiGLU 激活函数,并且线性层不会输出偏置项。你可能会疑惑为何会要求实现这个特殊的 Transformer 变体而非经典论文 “Attention Is All You Need” 中的标准版本呢?接下来我们就来探讨这些问题
当前模型数量呈现爆发式增长,去年一年就发布了约 19 个新的密集 Transformer 模型,其中大多数都做了架构微调。一方面,研读这些论文搞清楚每篇论文的创新点确实有点烦人,但另一方面,这其实是宝贵的信息财富,因为这些研究各有所长,并非千篇一律
这里整理了一份简明表格,梳理了这些模型的演变历程:

从 2017 年最初的 Transformer 架构开始一直到 2025 年最新模型的突破性进展,我们会在后续讨论中逐步展开这些内容,但从上表你可以看出研究者们正在探索多种架构变革的方向
例如上表展示的位置编码(Position Embedding)技术的演进,研究人员尝试过各种方法:从绝对位置编码、相对位置编码到旋转位置编码,部分团队还曾短暂采用过 ALiBi(注意力线性偏置位置编码)方案,但自 2023 年起,旋转位置编码已成为行业标准方案,接下来我们会详细探讨这些不同的技术方案
本次讲座包含三大核心板块:
- 通用架构变体
- 激活函数(Activation)、前馈网络(FFN)
- 注意力机制(Attention)
- 位置编码(Position embeddings)
- 超参数
- 稳定性技巧
首先我们将探讨架构设计的各种变体,这是本次讲座的重点所在,包括激活函数、前馈网络、注意力机制等各个模块的变体以及位置编码等关键技术
在确定架构之后,我们接下来需要做什么呢?当然是确定模型中各项超参数的配置,例如隐藏层维度该如何设定?MLP 内部投影层的的维度数量该如何设置?词表大小该如何确定?这些都是在训练语言模型时必须确定的关键配置,你不能随意拍脑袋决定,而应该采取科学合理的方式进行选择
最后我们将探讨模型训练时的一些稳定性技巧,若时间允许的话,我们还将补充讲解各种注意力机制的变体实现
3 Architecture variations
我们先来探讨下架构设计的各种变体,这里先重点强调两个方面,后续讲解中我们会再次深入探讨:
- Low consensus(except pre-norm)
- Trends toward
LLaMA-likearchitectures
首先,在众多技术方案的选择上大家意见并不统一,最近几年出现了一种趋同演化的现象,过去几年在最基础的架构层面形成了所谓的类 llama 架构。但研究者们的做法五花八门,有的采用 LayerNorm,有的则使用 RMSNorm,有的使用串行层结构,有的选择并行层结构,从初代 GPT 模型开始只有一项技术选择获得了普遍认可,这一点我们稍后会详细讲解,但其他方面存在诸多变体方案,这些都值得我们借鉴研究
3.1 Pre-vs-post norm
其中最关键的一点在 CS224n 课程中已经讲解过,目前有一点是业界普遍认同的,而且这个共识从最初就形成了,那就是 采用前归一化(Pre-norm)而非后归一化(Post-norm)

原始 Transformer 论文采用了后一种方案(Post-norm),如上图左侧部分所示,图中灰色部分表示残差流,除了残差连接之外,在每个子组件后都接有层归一化操作。具体流程是:先进行多头注意力计算,再将结果叠加到残差连接上,然后进行层归一化,在全连接层重复相同操作,然后执行层归一化,如下所示:

在模型架构探索的初期,研究者们 [Xiong+ 2020] 就发现:若将层归一化前置,置于非残差部分的前端(即上图右侧部分所示)时在多项指标上表现更优,因此几乎所有现代语言模型都采用的是 pre-norm 这种预归一化方法
关于 pre-norm 和 post-norm 之争,追溯其设计初衷就会发现:若采用 post-norm 方案,其训练稳定性会大幅降低,因此必须精心调整学习率,采用预热等策略才能确保训练过程稳定
翻阅早期提倡 pre-norm 的论文 [Nguyen+ 2019][Xiong+ 2020] 可以发现:若采用 pre-norm 配合其他稳定训练的技巧就能省去预热环节,其表现完全不逊色于需要精心设计预热方案的 post-norm 方法,甚至往往更胜一筹,这个现象在机器翻译场景中尤为明显:

这种现象在其他任务中同样存在,尤其是采用 post-norm 训练的 BERT 模型上体现得更为突出:

关于这种架构优势的讨论可谓众说纷纭,有学者 [Xiong+ 2020] 从跨层梯度衰减的角度展开分析,认为采用 pre-norm 时,梯度幅度能够保持相对恒定,而采用无预热的 post-norm 时梯度便会如下图橙色曲线所示出现爆炸性增长

这个观点确实有其合理之处,但我认为更贴切现代认知的解释或许是:pre-norm 本质上构建了更稳定的训练架构。[Nguyen+ 2019] 的早期研究就曾发现采用 post-norm 训练时会出现大量损失值突增现象(如下图紫色点所示),导致整个训练过程极不稳定,梯度范数不仅频繁出现尖峰,整体数值也也普遍高于 pre-norm

因此如今我们能看到 pre-norm 及其他层归一化技巧已成为训练大型神经网络时确保稳定性的重要手段。这就引出了近期一项颇具创新性的技术突破—double norm,记得去年在讲授这门课程时,这项技术尚未问世,这个变体至今尚未有统一的命名,不过目前我们暂且称之为 double-norm

我们都知道在残差流中直接插入层归一化效果不好(图左侧部分所示),因此我们考虑采用 pre-norm 将层归一化放在模块前面,但有人可能会困惑,为什么采用 pre-norm 时一定要把层归一化放在模块前面呢,可以将其置于模块后面吗?
当然,这个方案完全可行,不仅如此,近期研究者们还尝试直接在模块末端追加层归一化操作(图右侧部分所示),Grok 和 Gamma 2 都采用了这种前后双重层归一化的架构设计,而 Olmo 2 则采用了不同的策略,他们仅在前馈网络和多头注意力机制之后加入层归一化,这实际上是个相当有趣的变化
pre-norm 曾长期占据主导地位,一度成为行业标准配置,但如今架构设计已悄然生变,一种新型变体(double-norm)应运而生。针对这种创新方法,学界也展开多项评估研究,研究者们发现,这种架构在训练超大规模模型时展现出更优的稳定性,整体训练过程也更为顺畅
3.2 LayerNorm vs RMSNorm
原始 Transformer 架构中层归一化采用的是 LayerNorm 技术,其计算公式如下:
y = x − E [ x ] Var [ x ] + ϵ ∗ γ + β {y}=\frac{{x}-\text{E}[{x}]}{\sqrt{\text{Var}[{x}]+{ \epsilon}}}*{\gamma}+\beta y=Var[x]+ϵx−E[x]∗γ+β
具体操作流程是:输入激活值 x x x 后,先减去其均值 E [ x ] \text{E}[{x}] E[x] 即对 x x x 向量各维度取算术平均值,接着除以标准差(即方差加上微小调节系数 ϵ \epsilon ϵ 后的平方根),这样就能实现对激活值 x x x 的标准化处理,随后通过可训练参数 γ \gamma γ 进行缩放,再通过 β \beta β 参数进行平移调整
这个逻辑很清晰,首先对激活值进行归一化处理,再根据需要调整到目标位置,许多模型都采用了这种层归一化技术,实际表现相当出色。不过现在基本上所有主流模型都已转向 RMSNorm,具体要怎么操作呢,直接舍弃 LayerNorm 中的均值调整环节,无需进行减均值操作也不再添加偏置项,其计算公式如下:
y = x ∣ ∣ x ∣ ∣ 2 2 + ε ∗ γ y=\frac{x}{\sqrt{\big||x|\big|_{2}^{2}+\varepsilon}}*\gamma y= ∣x∣ 22+εx∗γ
这一做法已被众多知名模型采用,包括 LLaMA 系列、PaLM、Chinchilla 以及 T5 等模型都已转向使用 RMSNorm。这种转变的原因何在呢?原因之一在于实际效果差异不大。实验表明,使用 RMSNorm 归一化训练模型取得的成效与 LayerNorm 归一化相当,因此从简化模型的角度考虑使用 RMSNorm 更好
RMSNorm 归一化不仅速度更快,效果也同样出色,那么具体快在哪里呢?速度优势体现在省略均值计算同时省去了偏置项 β \beta β 的添加步骤,这样能减少从内存加载到计算单元的参数量
那大家可能会有所困惑,上次讲座不是说过从运行时角度来看只有矩阵乘法最关键吗?但这并不是矩阵乘法操作呀,似乎无需关注,这种观点确实有其合理性。
我们来分析下 Transformer 模型中各类运算所占浮点计算量 flop 占比:

Note:表格来自于 [Ivanov+ 2020],该论文对 Transformer 各组件进行了全面的性能剖析
表格数据显示在 Transformer 模型中张量收缩(tensor contraction)即矩阵乘法占据了 Transformer 中 99.8% 的浮点运算量,因此,节省 0.17% 的浮点运算量似乎收效甚微,但当前架构设计的关键在于不能仅以浮点运算量为考量标准,虽然 flops 至关重要,但并非唯一考量,内存移动效率同样至关重要。
尽管张量收缩(例如矩阵乘法)运算占比高达 99.8%,而像 softmax 以及层归一化等操作在 Transformer 中仅占 0.17% 的 flop,但它们实际上占据了 25% 的运行时长,这背后的根本原因在于这些归一化操作仍然会带来巨大的内存搬运开销,因此对这些底层操作进行优化确实很有必要。因为性能评估不仅要看浮点运算能力(flop),还要考虑内存访问效率。在后续系统架构的讲座中,我们会重点强调这个关键问题,特别是在讲解 GPU 架构时,我们必须高度重视内存优化,其重要性远超单纯的浮点运算能力考量,这正是 RMSNorm 如今受到青睐的核心原因之一

关于 RMSNorm 研究文献中 [Narang+ 2021] 的消融实验很有代表性,图中对比了标准 Transformer 与 RMSNorm 改进版的性能差异,实验结果验证了我们先前的观点
标准 Transformer 为每秒 3.5 次迭代,而 RMSNorm 版本达到 3.68 次,虽然提升幅度不大,但它无需任何成本。更重要的是最终损失函数值低于标准 Transformer,这非常好,可以说我们实现了运行效率的提升
最后要强调的一个点与 RMSNorm 的设计理念一脉相承,那就是 大多数现代 Transformer 架构都移除了偏置项。如果观察原始 Transformer 的 FFN 前馈网络,其结构大致是下面这样的:
FFN ( x ) = max ( 0 , x W 1 + b 1 ) W 2 + b 2 \text{FFN} \left( x \right)= \max \left( 0,xW_{1}+b_{1} \right) W_{2}+b_{2} FFN(x)=max(0,xW1+b1)W2+b2
输入 x x x 首先通过带有偏置项的线性层,然后进行 ReLU 激活计算,最后再通过第二个线性层输出结果。不过目前大多数实现方案(非门控单元结构)采用了下面这种结构:
FFN ( x ) = σ ( x W 1 ) W 2 \text{FFN}(x) = \sigma(xW_1)W_2 FFN(x)=σ(xW1)W2
这些实现直接移除了偏置项,其性能表现同样出色,事实证明仅靠矩阵乘法就足够了。另一个精妙之处在于优化稳定性,研究人员发现移除这些偏置项后往往能提升大型神经网络的训练稳定性,因此当前多数实现方案都完全舍弃了偏置项,仅保留纯粹的矩阵乘法运算进行训练
以上就是层归一化的相关内容,总结来说需要重点关注两个方面:
第一,基本上所有模型都会采用 pre-norm,无论无何也会在残差流之外进行层归一化,这可以说是铁律了。这样做能获得更优的梯度传播效果,模型训练过程也会更加稳定
第二,绝大多数模型都采用 RMSNorm 归一化方案,实际应用中和 LayerNorm 效果几乎不相上下,需要调整的参数也更少,省略偏置项的设计理念也具有普适性,当前多数模型在绝大多数层结构中均已摒弃偏置项
3.3 Activations
接下来要讨论的是激活函数,激活函数类型五花八门,例如 ReLU、GeLU、Swish、ELU、GLU、GeGLU、ReGLU、SeLU、SwiGLU、LiGLU 等等,但 SwiGLU 这类 GLU 变体总能稳定发挥出色效果,下面我们来解释其背后的原因
想必 ReU 和 GeLU 这些激活函数大家已经耳熟能详了,ReLU 这种基础激活函数,在深度学习入门课程中就会提到,简单来说就是取 max ( 0 , x ) \max (0, x) max(0,x):
F F ( x ) = max ( 0 , x W 1 ) W 2 FF(x)=\max(0,xW_1)W_2 FF(x)=max(0,xW1)W2

我们以 MLP 应用为例,先计算 x W 1 xW_1 xW1,接着经过 ReLU 激活函数,最后再乘以 W 2 W_2 W2,非常简单,注意这里我们省略了偏置项
GeLU(高斯误差线性单元)是一种基于高斯误差函数的激活函数:
F F ( x ) = G E L U ( x W 1 ) W 2 G E L U ( x ) : = x Φ ( x ) FF(x)=GELU(xW_1)W_2 \\ GELU(x) := x\Phi(x) FF(x)=GELU(xW1)W2GELU(x):=xΦ(x)

GeLU 激活函数会将线性输出与高斯分布的累积分布函数相乘,所以它本质上类似于 ReLU,但会在拐点处形成一个更平滑的过渡区域,如上图所示,这种设计略微提升了函数的可微性,至于实际效果可能因情况而异
而 GPT 系列模型包括 GPT1/2/3 以及 GPT-J 等均采用了 GeLU 激活函数,而原始 Transformer 模型及部分早期模型则采用的是 ReLU 激活函数。事实上,几乎所有现代模型都已转向使用门控线性单元(Gated Linear Unit,GLU)例如 SwiGLU、GeGLU 等变体
Google 团队在这方面起到了关键推动作用,比如 PaLM、T5 等模型都率先采用了这类架构,经过实践验证后,2023 年后推出的模型几乎都采用了门控线性单元结构
原来标准的全连接层使用的是 ReLU 激活函数:
F F ( x ) = max ( 0 , x W 1 ) W 2 FF(x) = {\color{red}\max (0,xW_1)} W_2 FF(x)=max(0,xW1)W2
现在我们不再简单地使用线性变换加 ReLU 激活,而是要 通过逐元素线性运算对输出进行门控处理:
max ( 0 , x W 1 ) → max ( 0 , x W 1 ) ⊗ ( x V ) {\color{red}\max (0,xW_1)} \rightarrow \max (0, xW_1) \otimes (xV) max(0,xW1)→max(0,xW1)⊗(xV)
我们的做法是: x x x 点乘 V V V 将生成一个向量 x V xV xV,然后将 x V xV xV 这个向量与 MLP 内部的原始项 max ( 0 , x W 1 ) \max (0,xW_1) max(0,xW1) 进行逐元素相乘,最后再将整个结果与 W 2 W_2 W2 相乘:
F F R e G L U ( x ) = ( max ( 0 , x W 1 ) ⊗ x V ) W 2 FF_{ReGLU}(x) = (\max (0,xW_1) \otimes xV)W_2 FFReGLU(x)=(max(0,xW1)⊗xV)W2
我们可以这样来理解:我们对 MLP 的隐藏层部分进行了门控处理,原始激活函数将输入映射到隐藏空间,然后用 x x x 点乘 V V V 的结果对这个映射进行门控,最后通过门控 W 2 W_2 W2 将这个门控结果重新投影回隐藏维度空间,这就是整个机制的核心原理,这里我们还额外引入了一个专门用于门控运算的参数 V V V,该架构被广泛应用于 Google 的多个模型,比如 T5v1.1、Gemma2、Gemma3 等等。
另一个广受欢迎的变体是 SwiGLU 架构:
F F N S w i G L U ( x , W , V , W 2 ) = ( S w i s h ( x W ) ⊗ x V ) W 2 FFN_{{SwiGLU}}(x,W,V,W_2)=(Swish(xW)\otimes xV)W_2 FFNSwiGLU(x,W,V,W2)=(Swish(xW)⊗xV)W2
Note: S w i s h ( x ) = x ∗ sigmoid ( x ) Swish(x) = x * \text{sigmoid}(x) Swish(x)=x∗sigmoid(x),Swish 函数形态与高斯误差单元(GeLU)高度相似
SwiGLU 其实就是对 Swish 函数施加了门控机制,后续采用相同的处理流程。如今大多数模型采用的都是 SwiGLU 架构,比如 LLaMA 系列、PaLM 和 OLMo 等模型,你会发现 SwiGLU 的应用确实非常普遍。
这里有个点需要注意,那就是新增的这个 V V V 参数,需要考虑如何设置这个参数的维度。通常的做法是将权重矩阵 W W W 的输出维度略微缩小至原来的 2/3,这样就能确保整个门控矩阵的总参数量与非门控版本保持一致。
这是目前业界普遍采用的做法,如果你对这个概念还不太理解,稍后我们会再详细讲解一遍。简单来说,使用门控线性单元时只需将各维度适当缩小就能确保整体参数量与传统结构保持一致
那么,门控线性单元到底管不管用呢?

GLU 原始论文 [Shazeer 2020] 系统评估了所有门控线性单元的变体,如上表所示,这些研究其实已经有些年头了,你现在看到的是 GoLA 和 SST-2 的性能表现,能明显看出各类门控线性单元变体的性能表现始终更胜一筹

[Narang+ 2021] 等人发表的论文也颇具价值,该研究基于类 T5 模型架构,系统考察了各类门控线性单元变体。数据再次表明门控线性单元变体始终能取得比传统结构更低的损失值
在门控机制与激活函数领域,不同模型衍生出诸多变体,但门控线性单元已然成为主流选择,这种优势地位的形成自然有其深层原因。当然,优秀的模型并非必须使用 GLU,区分这两个概念至关重要,不能仅仅因为它的表现略胜一筹或者大家都在使用,就认为它不可或缺
事实上,确实存在不使用 GLU 却表现卓越的模型案例,比如 GPT3 就是一个典型案例。此外,较新的 Nematon 340B 模型采用了 Squared ReLU,而 Falcon 2 11b 模型则采用了 ReLU,这些都是性能优异的模型,由此可见 GLU 这种做法并非必需
现有证据确实表明 SwiGLU 和 GeGLU 能持续提升模型性能,因此作业 1 中才要求必须准确实现该变体
3.4 Serial vs Parallel layers
接下来我们来讨论架构设计的最后一个重点,传统 Transformer 模块采用串行结构:

即每个模块的输入都来自前一级的输出,先进行注意力计算再将计算结果传递给下一层,接着执行 MLP 运算,最后将运算结果传递至下一模块,这种设计本质上是串行架构,先进行注意力计算再执行 MLP 运算
但这种串行架构可能存在特定的并行化限制,若要在海量 GPU 集群上实现并行计算,串行架构会大幅增加并行化难度,系统层面的挑战也会更加严峻,GPU 的利用率也会显著降低
因此,部分模型(例如 GPT-J、PaLM、GPT-NeoX)采用了所谓的 并行层(Parallel layers) 架构,不再串行执行注意力机制和 MLP 计算,而是并行处理这两个计算任务。你可以从上一层获取输入 x x x 然后并行计算 MLP 和注意力机制,最后将两者相加合并到残差流,这样就得到了最终输出
这一做法最早由 GPT-J 开源项目提出,Googla PaLM 团队更是大胆地将其应用于超大规模模型,此后众多研究者纷纷效仿,如果实现得当可以共享大量计算资源,例如层归一化和矩阵乘法可以融合处理,从而显著提升系统效率。不过最近一年来,这种架构已不再流行,目前主流模型多采用串行层结构而非并行结构

现在我们可以回过头来再看这张复杂难懂的大图表,回顾我们最初强调的内容,pre-norm 这一栏中对勾标记表示使用的是 pre-norm,早期只有原始 Transformer 和 GPT 两个模型采用了 post-norm,而其他几乎所有模型采用的都是 pre-norm,注意表格中其余未打勾的选项都是未公开技术细节的私有模型,
最左侧显示的是 RMSNorm 和 LayerNorm 的对比,灰色方框代表的是 LayerNorm,蓝色的那些是 RMSNorm,基本上大多数人都已转向使用 RMSNorm。Norm 旁边的这栏标注的是串行层与并行层的对比,主流方案普遍采用串行结构,但你也能看到其他变体结构
3.5 Position embeddings
接下来我们要讲的是位置编码(position embeddings)的各种变体。这个议题特别有意思,因为在语言模型发展的头几年,研究者们尝试了各种不同的方法,正弦/余弦位置编码(Sine embeddings)最早出现在原始 Transformer 模型中:
E m b e d ( x , i ) = v x + P E p o s Embed(x, i) = v_x + PE_{pos} Embed(x,i)=vx+PEpos
其中:
P E ( p o s , 2 i ) = s i n ( p o s / 1000 0 2 i / d model ) P E ( p o s , 2 i + 1 ) = c o s ( p o s / 1000 0 2 i / d model ) \begin{split} PE_{(pos,2i)}&=sin(pos/10000^{2i/d_{ \text{{model}}}})\\ PE_{(pos,2i+1)}&=cos(pos/10000^{2i/d_{\text{{model}}} })\end{split} PE(pos,2i)PE(pos,2i+1)=sin(pos/100002i/dmodel)=cos(pos/100002i/dmodel)
许多模型(如 GPT 系列和 OPT)采用了绝对位置编码:
E m b e d ( x , i ) = v x + u i Embed(x, i) = v_x + u_i Embed(x,i)=vx+ui
本质上都是在词嵌入基础上叠加一个可学习的位置向量。
另一些模型例如 T5、Gopher 等则采用了各类相对位置编码方法,通过向注意力计算中引入位置向量来实现:
e i j = x i W Q ( x j W K + a i j K ) T d z e_{ij}=\frac{x_{i}W^{Q}(x_{j}W^{K}+a_{ij}^{K})^{T}}{\sqrt{d_{z}}} eij=dzxiWQ(xjWK+aijK)T
现在看来,大多数模型最终都采用了 RoPE 旋转位置编码方法,这项技术实际上最早出现在 GPT-J 模型中,后来迅速被绝大多数模型所采用。RoPE 方法背后的核心思想在于:真正重要的是这些向量之间的相对位置关系。
因此,若用 f ( x , i ) f(x,i) f(x,i) 表示待嵌入词 x x x 在位置 i i i 的嵌入向量,其数学表达式应满足如下形式:
< f ( x , i ) , f ( y , j ) > = g ( x , y , i − j ) \left<f \left( x,i \right),f \left( y,j \right) \right>=g \left( x,y,i-j \right) ⟨f(x,i),f(y,j)⟩=g(x,y,i−j)
必然存在一个映射函数使得 f ( x , i ) f(x,i) f(x,i) 和 f ( y , j ) f(y,j) f(y,j) 之间满足特定关系,当我们计算这些嵌入向量的内积时,就可以将其表示为另一个函数 g g g,该函数同时接收两个 token 及其位置差作为输入参数,这个定义本质上强制实现了位置不变性或者说绝对位置不变性的数学约束,这意味着模型仅关注两个 token 之间的相对距离
旋转位置编码(RoPE)的精妙之处在于:它基于一个我们熟知的,对绝对位置保持不变的特性-旋转不变性,我们将充分利用这个旋转特性来设计我们的位置嵌入方法

众所周知,内积具有旋转不变性,这正是我们要利用的关键特性,左侧图示就是我们的出发点,假设单词 we 的嵌入向量对应图示中的这个箭头,而单词 know 的嵌入向量则对应图示中的另一个箭头。现在需要嵌入 we know that 这个序列,根据已知条件,我们只关注 we 和 know 这两个词,具体该怎么操作呢?
这里我们默认位置索引从 0 开始计算,首先将 we 放在位置 0,因此这个向量不需要旋转,know 位于位置 1,需要将其旋转 1 个单位向量,这样我们就得到了 we know 的嵌入表示。
假设接下来我们要处理新序列 of course we know,显然 we 和 know 之间的相对位置关系仍保持不变,让我们具体看看会发生什么变化
整个序列的位置偏移了两个单位,旋转 we 向量:从初始垂直位置开始,先旋转一次,再旋转第二次,共执行两次旋转,,接着将 know 向量旋转三个单位量,因为它的位置索引是第 3 位,现在观察这两个向量箭头,可以看到它们的相对夹角保持一致,因此它们的内积保持不变
这就是 RoPE 的精妙之处,只需根据词位置决定旋转角度来转动向量,内积运算不受相对旋转影响,内积值仅与位置间距相关

在二维空间中理解这个概念非常直观,因为旋转操作在二维中显而易见,向量的旋转方式具有唯一性,但在高维空间旋转方式的唯一性就变得复杂了,因此 RoPE 研究者创新性地提出了既简洁又高效的实现方式,具体做法是:将高维向量中的元素两两配对,这里用 d 表示维度,将向量分割为若干个二维子空间,每个二维子空间都会按特定角度 θ \theta θ 进行旋转,因此需要设定相应的旋转频率
通过成对旋转各维度分量,现在每对维度都在编码所有这些相对位置信息,就像正弦余弦编码中那样,选取一组 θ \theta θ 值使得一些嵌入快速旋转,而另一些则缓慢旋转,这样既能捕捉到高频的邻近位置信息,也能编码低频的远距离位置信息

具体到旋转位置编码(RoPE)的数学实现,其实就是用不同的正弦余弦旋转矩阵进行乘法运算,相信大家在线性代数和三角函数的课程里都接触过这个原理,这个过程可以理解为将嵌入向量与这些 2x2 的分块旋转矩阵进行乘法运算,整个过程不会产生任何加法项或交叉项,这种位置编码完全基于相对关系构建
与绝对位置编码或传统正弦余弦编码不同之处在于,旋转位置编码直接在注意力层进行运算,并非在模型底层进行叠加,每当执行注意力计算时都会在该层直接施加位置编码操作,这种机制就能为模型提供所需的位置信息

上面这段代码摘录自 LLaMA 对 RoPE 的具体实现,最上方是标准的注意力计算模块,包括查询向量 Query、键向量 Key 和值向量 Value 的常规处理,这些是常规的线性投影变换,接着会计算出对应的余弦和正弦角度,这些旋转角度决定了查询向量和键向量各分块需要旋转的幅度
接着使用计算出的余弦和正弦角度分别对查询向量和键向量进行旋转操作,这样就得到了旋转后的查询向量和键向量,这些旋转后的向量将作为输入参与后续注意力机制的计算。这个操作不是在网络底层进行的,而是在生成查询向量和键向量时执行的,这一机制对于确保仅保留相对位置信息至关重要
这里要特别指出,旋转位置编码实际上是业界已形成普遍共识的计算方案,因为旋转位置编码具备支持上下文窗口扩展的多种算法实现,这使其成为现代工业化语言模型的关键组件,实际研究也表明即便在较小模型和较短文本场景下其效果也非常显著,因此在位置编码方案竞争中脱颖而出
4. Hyperparameters
接下来我们要讨论的是超参数设置问题
当你突然接手一个新语言模型的训练任务时,关于模型超参数的设置你可能会有很多疑问,毕竟超参数的数量确实不少。但实际上,各种成功模型间真正需要调整的超参数并不多,业界其实已经形成了一些明确的经验法则以及大家普遍遵循的清晰准则
常见的超参数设置问题包括:
- FFN 前馈层的维度应该扩展多少
- 注意力头的数量设置
- 词表大小的合理取值
接下来我们会逐一探讨这些问题,尝试缩小超参数的选择范围,帮助大家理清超参数设置
4.1 Model-dim hyperparam
我们先从最基础的开始,来看一个简单的前馈神经网络层:
FFN ( x ) = max ( 0 , x W 1 + b 1 ) W 2 + b 2 \text{FFN}(x) = \max (0, xW_1+b_1)W_2 + b_2 FFN(x)=max(0,xW1+b1)W2+b2
假设选用的激活函数是 ReLU,这里涉及两个关键超参数,一个是 d m o d e l d_{model} dmodel,表示输入 x x x 的维度,也就是 MLP 的输入维度,另一个是 d f f d_{ff} dff,代表前馈网络的维度即 MLP 输出层的隐藏维度,它最终会投影回 d m o d e l d_{model} dmodel 维度
那么 d f f d_{ff} dff 通常应该如何设置呢,实际上这已经形成普遍共识了,绝大多数采用 ReLU MLP 架构的研究者,都会将 d f f d_{ff} dff 维度设为 d m o d e l d_{model} dmodel 的 4 倍,也就是:
d f f = 4 d m o d e l d_{ff} = 4 \ d_{model} dff=4 dmodel
待会儿我们会用实验数据证明这个数值设定的合理性,不过这个规则也存在特例:需要注意的是各类 GLU 变体会将这个数值按 2/3 比例进行缩放,若采用 2/3 的比例缩放,最终参数量将大致保持相当,通过简单计算就能验证,若将 GLU 变体的维度按 2/3 比例缩放,可推导出设置前馈网络维度 d f f = 8 3 d m o d e l d_{ff} = \frac{8}{3} \ d_{model} dff=38 dmodel 的结论

事实上,许多主流模型确实遵循着这一经验法则,像 PaLM、Mistral 和 LLaMA-2 等模型的维度配置会略高于这个基准值,这些虽然都是采用门控线性单元(GLU)架构的模型,但并未遵循
8
3
\frac{8}{3}
38 倍的比例规则。但如果你看 LLaMA、Qwen、DeepSeek、Yi 和 T5 等模型时,可以发现它们基本上都遵循着这个约为 2.6 倍的参数比例规则
。绝大多数模型都落在这个约 2.6 倍的参数区间内,这正是门控线性单元(GLU)的标准参数化形式
我们再讲一个例外情况,如果你看拥有 110 亿参数的 T5 模型时,会发现有一个相当不可思议的配置,它们的隐藏层维度设定为 d m o d e l = 1024 d_{model}=1024 dmodel=1024,但其前馈网络的上投影维度竟然高达 d f f = 65536 d_{ff}=65536 dff=65536,这使得前馈网络维度与模型维度的比例达到了 64 倍的放大系数,与前面我们提到的 2.6 倍的比例系数相比这是一个非常大的差异
近期还有其他研究采用了更大的放大系数,Gemma 2 沿袭了这一思路,采用了 8 倍的放大系数,稍后我们会专门讨论这个特例。毋庸置疑,T5 模型本身是相当优秀的,由此可见,以如此悬殊的比例训练模型是完全可行的。
我们可以观察到 4 倍的乘数效应,这种做法真的合理吗?是否有人通过更严谨的量化实验来验证过这个方案的可行性呢?关于缩放定律的论文中 [Kaplan+ 2020] 有个关键图表,虽然这篇论文以缩放定律理论闻名,但其超参数研究部分同样极具价值
论文中恰好验证了我们的观点,关于 d f f d_{ff} dff 与 d m o d e l d_{model} dmodel 的比例关系,他们通过图表展示了损失函数的变化趋势

数据显示存在一个最优区间,这个比例在 1 到 4 之间波动,最高可延伸至 10 左右。由此可见,前馈网络比例的选择范围相当宽泛—从 1 到 10 之间的任意数值都能取得不错的效果且基本都能达到近似最优的效果。而 4 这个数值与最优选择区间也相当接近,因此这个选择非常合理
从这些超参数研究中我们能得到什么启示?大量研究数据表明,在不使用门控线性单元的情况下,可以采用一个通用默认值,将维度扩大 4 倍。若采用门控线性单元,则建议使用约 2.66 倍的缩放系数
这些参数设置对绝大多数现代大语言模型都相当有效,T5 模型再次证明这些规则并非金科玉律,你完全可以打破常规,随心所欲地调整参数。超参数的选择并非一成不变,采用不同的超参数组合同样能训练出优秀的大语言模型
不过这个故事最有趣的后续发展是 T5 后来推出了改进版 T5 v1.1 模型,这个改进版采用了更标准的 GeGLU 2.5 倍系数配置。由此可见,我们不难推测,研发团队可能在评估原始 T5 模型后认为当初设定的 64 倍系数或许应该调整,改用更标准的参数设置,最终他们确实开发出了更优的模型,这确实很有意思
4.2 Head-dim and num-heads
另一个共识性的超参数是模型维度(头维度乘以头数)之间的比例:

上面这段内容摘自 CS224n 课程中的讲义,业界普遍采用的标准做法是:设定隐藏维度 d 后,若采用多头注意力机制,只需将该维度均匀分配给每个注意力头即可。因此,在增加注意力头数量时,各头的维度将保持恒定不变。
当然,这并不是必须遵循的规则,随着注意力头数量的增加,开发者完全可以选择保持每个头的维度不变,这样注意力机制部分的参数量就会持续增长,你可以这么做,但大多数模型仍然遵循上述设计准则

我们看到 GPT3、T5、LaMDA、PaLM 以及 LLaMA2 这些主流模型它们的头维度与总维度比值都严格保持或接近 1:1 的比例关系,T5 模型是一个例外,它尝试了高达 16 的极端比例,但除此之外,其他模型基本都严格遵循这一业界共识

确实有研究论文 [Bhojanapalli+ 2020] 对 1:1 的比例提出了质疑,他们经过实验发现当注意力头数量不断增加时,这些注意力头的秩会呈现递减趋势,而当每个注意力头的维度被压缩得过低时,就会开始削弱注意力机制的表达能力
但在实际应用中,我们似乎并未真正遇到太多显著的低秩瓶颈问题,而大多数比例为 1:1 的模型似乎表现都还不错,这确实是一个在现有主流模型中普遍保持恒定的参数
4.3 Aspect ratios
在超参数中宽深比是一个非常重要的参数,我们可以考虑深度网络,我们可以不断增加网络层数,也可以扩展网络的宽度。一般来说,若要通过单一参数控制网络宽度,残差流(residual stream)的隐藏维度就是最佳选择,这个参数能同时控制几乎所有运算操作的宽度,因此这个参数显然是需要重点调优的关键参数。
人们可能会认为更深的网络更智能、表现力更强,而更宽的网络计算效率更高,业界通常会发现一个最佳的比例区间,当然也存在例外情况:

早期部分模型采用了明显更小的比例系数,这意味着这些模型的宽度远大于深度,而另一些模型则走向了另一个极端—网络宽度极大,其模型维度 d m o d e l d_{model} dmodel 远超层数 n l a y e r n_{layer} nlayer
业内普遍发现一个黄金比例, d m o d e l / n l a y e r = 128 d_{model}/n_{layer}=128 dmodel/nlayer=128 时最为理想,这一标准已被 GPT3 及 LLaMA 系列模型广泛采用,关于这一标准的实证依据,我们稍后详细阐述
模型宽深比的考量至关重要,这些参数直接决定了模型可实现的并行计算能力,当采用流水线并行(pipeline parallel)技术时,标准操作流程是对神经网络各层进行切割划分,然后分配到不同计算设备或设备计算集群上,这时因为每层内部同样可以实现并行计算,因此在模型设计时必然存在特定的架构约束条件
此外,对于宽度较大的模型,可以采用张量并行技术—通过矩阵切片的方式将计算任务分配到多个 GPU 上执行。在后续讲座中我们将深入探讨这些不同的并行计算范式各自存在独特的约束条件。张量并行技术需要极高的网络带宽,流水线并行对网络延迟的容忍度相对较高,在网络延迟较高的情况下仍可保持运行,因此网络性能参数将直接决定模型在宽度与深度之间的取舍
抛开这些限制不谈,我们可以从理论层面探讨模型的宽深如何影响其性能表现?[Kaplan+ 2020] 的研究再次提供了直观的图示说明,清晰展现了宽深比对模型性能的影响:

图示展示了三种不同规模的模型,分别是 5000 万参数、2.74 亿参数和 15 亿参数,横轴表示模型的宽深比,纵轴表示损失变化的百分比差值。如图所示,当宽深比约为 100 时(这个数值正如我们之前所说,业界普遍采用的超参数范围),不同规模模型的性能损失均达到最低值。这一结论得到了 Kaplan 团队发布的大规模超参数实验数据的支持,与行业经验判断高度吻合
这里有个非常有意思的发现,在跨越多个数量级的模型规模中,宽深比的最优区间似乎保持相对稳定,如果这一规律能进一步得到验证,那将是个好消息,这意味着我们可以始终采用固定的宽深比进行模型训练
谷歌研究院的 Yi Tay 团队近期发表了一篇极具启发性的论文 [Tay+ 2021],系统研究了模型深度与宽度对上下游任务的影响:

他们的研究发现:若仅以损失函数为评估标准,模型深度与宽度的配比影响其实微乎其微,真正起决定性作用的只有参数量这个单一变量,增加模型深度并不能带来实质性的性能提升。但若着眼于下游任务准确率,情况就变得复杂微妙起来,当时研究团队重点分析了模型在 SuperGlue 基准上的微调准确率表现,他们认为在相同 FLOPs 的前提下,增加模型深度可能更具优势
4.4 Vocabulary sizes
在这个深入探讨超参数的环节,最后要讨论的是:如何选择合适的词汇表规模。
总体而言,词汇表规模呈现出持续扩大的趋势,很大程度上是因为大语言模型正在实际应用场景中广泛部署,这些模型正逐渐发展为更具实用价值的服务。当模型真正投入使用时,它们需要处理各种预料之外的情况:不同语言的用户交流、表情符号的使用以及各种近乎模态化或特殊形式的语言表达
因此,早期模型(尤其是单语种模型)的词汇量通常在 3 万到 5 万 token 之间,这一点在早期的 GPT 和 LLaMA 系列模型中就能看出来:

但如果你看多语言模型,他们的词汇量看起来非常大,像 GPT4 及其众多效仿者,它们的词汇量基本都在 10 万 token 这个量级:

因此 10 万至 20 万 token 已成为业内普遍采用的标准。现有研究表明,随着模型规模扩大,这些模型的处理能力会相应提升,对词汇元素的利用率也会提高,因此可以预见,随着模型规模扩展及训练数据量持续增长,token 数量呈现上升趋势将成为常态。
4.5 Dropuout and other regularization
最后要说明的是,虽然这不属于具体超参数范畴,但在启动模型运行前还需完成两项关键配置:dropout 及其他正则化策略的设置。
考虑到预训练的特性,预训练看似最不需要正则化,因为预训练通常仅需单轮迭代,数据量庞大到无法完整遍历,因此只能进行单次训练,这种情况下基本不会出现过拟合,因此很容易得出预训练无需正则化的结论,直接让优化器自由发挥即可,核心目标就是最小化损失函数
这些观点确实有力论证了为何无需进行正则化处理,但观察实际应用情况时,结果却呈现出矛盾的现象:

早期阶段,研究人员普遍采用了大量 dropout 技术,同时权重衰减(weight decay)的应用也相当普遍。如今看来,dropout 已经有点过时了,但权重衰减仍是众多研究者持续采用的主流方法,这时为何呢?这种做法着实令人费解,请大家仔细思考,当使用 SGD 对海量数据进行单次大规模神经网络训练时,为何仍要采用权重衰减?
原因在于:权重衰减的作用并非控制过拟合,具体而言,当我们观察权重衰减时,不同的衰减强度似乎并不会显著改变训练损失与验证损失之间的比例关系,因此,你可以尝试不同强度的权重衰减进行训练,只要训练时间足够长,或是恰当地调整超参数,最终得到的训练损失与验证损失之间的差距将会趋于一致
由此可见,即便完全不使用权重衰减,过拟合现象也并未发生实质性变化,但有趣的是,权重衰减似乎与优化器的学习率调度机制产生了某种微妙的异常交互作用

实际观察 [Andriushchenko+ 2023] 发现,当模型采用恒定学习率进行训练时会出现以下现象:当你突然将学习率骤降至接近 0 时,可以看到随着学习率下降,模型性能出现了断崖式下跌。接下来我们观察不同类型的权重衰减策略效果,实验结果表明:当启用权重衰减时,模型在较高学习率下的训练效果明显欠佳,而当学习率降低时,模型性能会急速衰退。
采用余弦退火学习率策略时,高权重衰减模型的初始训练进度会显著迟缓,但随着学习率逐渐冷却下降,这些模型会进入快速优化阶段,由此可见,优化器与权重衰减之间存在复杂的交互作用,特别是在训练接近尾声时,系统会自发产生某种隐式加速效应,这种机制最终能训练出更优质的模型
因此,我们最初提出的问题的答案其实是:使用权重衰减的根本目的并非如设计初衷那样仅用于模型正则化,真正目的其实是为了获得更优的训练损失值,这种现象源于训练末期的特殊学习动态—当学习率逐渐衰减至 0 时,整个优化过程就会自然呈现出这种特性
从某些角度来看,这种做法在语言模型领域确实暗藏隐忧,但现在大家就能理解,为何查阅诸多技术报告时,总会发现权重衰减的普遍应用,这正是该现象背后的深层成因
4.6 Summary: hyperparameters
综上所述,有些优化策略的选择其实是显而易见的,因此在为模型选择各类超参数时,这些参数选择其实无需过度纠结,它们早已通过实践验证,并已成为行业通用做法,这类参数包括:
- Feedforward
- Head dim
- Aspect ratio
- Regularization
MLP 的隐藏层维度、多头注意力机制中每个注意力头的维度、模型的宽深比以及通过权重衰减实现的正则化策略选择,这些超参数的选择标准,业界已形成相当成熟的共识方案,这些默认设计基本与我们作业中建议的配置方案相吻合,因此你只需参照执行,就能获得与此类似的配置效果
5. Stability tricks
近期有哪些新的架构和相关进展呢?实际上,核心架构并没有太大变化,不过在许多发布内容中最引入注目的,当属那些被反复强调的稳定性优化技巧。这些优化技巧的核心目标,就是让模型训练过程变得更加稳定可控,随着模型规模不断扩大,训练周期持续延长,这类稳定性问题就会愈发频繁地呈现出来

上图摘自 [OLMo 2] 论文,看上图的蓝色曲线,观察这个梯度向量的 L2 范数图,这张图表看起来相当触目惊心,损失曲线看起来还算正常,但时不时会出现一些异常尖峰。当你展开梯度范数图时,会看到一幅糟糕透顶的图表—到处都是尖峰,范数值完全失控。用这种方式训练模型的话,想要获得合理的收敛效果会非常困难,因为会出现梯度范数突然爆炸的问题,而你束手无策,训练只能被迫终止。
所以业界投入了大量精力,试图将这条蓝色曲线优化成更接近橙色曲线的形态。当然这个损失值会更高,不过这个差异可以忽略(猜测可能是他们在这两次训练之间更换了数据集导致的),但这条橙色曲线全程保持着理想的低梯度范数,这才是我们真正希望看到的理想曲线形态。
你可能会问:Transformer 模型中的稳定性问题究竟出在哪里呢?事实上,稳定性问题可能出现在模型的任何环节,但观察业界常用的干预措施时,确实有个环节特别突出,堪称问题高发区,那就是令人头疼的 softmax 函数。这确实可能引发问题,因为 softmax 中涉及指数运算,这些运算在数值上往往表现不佳,这就可能导致除数为零的情况,正因如此,softmax 这个环节往往会导致各种问题发生
那么 Transformer 中的 softmax 分布在哪些位置呢?

首先在模型末端有一个,因此我们需要特别注意输出层的 softmax 处理,此外,自注意力机制中也存在 softmax 运算,因此我们需要重点关注这两个 softmax 运算。针对每一处 softmax 我们都会介绍相应的稳定性优化方法,这些方法在实践中普遍表现良好
5.1 Output softmax stability - the ‘z-loss’
首先介绍的是 z-loss 归一化方法,该方法是在一篇机器翻译论文 [Devlin+ 2014] 中提出的,其核心目标是确保下面这个归一化因子 Z 值趋近于 1:
log ( P ( x ) ) = log ( e U r ( x ) Z ( x ) ) = U r ( x ) − log ( Z ( x ) ) Z ( x ) = Σ r ′ = 1 ∣ V ∣ e U r ′ ( x ) \begin{array}{rcl}\log(P(x))&=&\log\left(\dfrac{e^{U_{r}(x)}}{Z(x)}\right) \\ &=&U_{r}(x)-\log(Z(x))\\ Z(x)&=&\Sigma_{r^{\prime}=1}^{|V|}e^{U_{r^{\prime}}(x)}\end{array} log(P(x))Z(x)===log(Z(x)eUr(x))Ur(x)−log(Z(x))Σr′=1∣V∣eUr′(x)
因此,如果你观察概率分布 P(x),它实际上就是输出端的 softmax 结果,输出层的 softmax 由两项组成,首先对 logits 进行指数运算然后除以归一化因子 Z,而 Z 这个归一化因子,就是对整个词表所有数值进行求和的结果
因此,若要让 Z(x) 趋近于 1,就需要通过训练使网络的归一化因子 Z(x) 稳定在这个值附近,此时可以重构损失函数,在原有公式中加入一个附加项,通过约束 log ( Z ( x i ) ) \log(Z(x_i)) log(Z(xi)) 逼近 0 来实现优化目标,最终就会得到一个辅助损失项,也就是下面所示的 α log 2 ( Z ( x i ) ) \alpha \log^2(Z(x_i)) αlog2(Z(xi)) 这一项:
L = ∑ i [ log ( P ( x i ) ) − α ( log ( Z ( x i ) ) − 0 ) 2 ] = ∑ i [ log ( P ( x i ) ) − α log 2 ( Z ( x i ) ) ] \begin{split} L&\ \ =\ \ \sum_{i}\left[\log(P(x_{i}))- \alpha(\log(Z(x_{i}))-0)^{2}\right]\\ &\ \ =\ \ \sum_{i}\left[\log(P(x_{i}))-\alpha\log^{2}(Z(x_{i})) \right]\end{split} L = i∑[log(P(xi))−α(log(Z(xi))−0)2] = i∑[log(P(xi))−αlog2(Z(xi))]
上面展示了完整的推导过程,这其实就是业界常说的 z-loss,Jacob Devlin 等人在机器翻译领域引入这个方法时,其初衷与现今的应用场景完全不同。不过在语言模型领域首次应用这个方法的当属 PaLM 模型:

他们将其称为 z-loss 辅助损失项,采用 z _ l o s s = 1 0 − 4 ⋅ log 2 Z z\_{loss} = 10^{-4} \cdot \log ^2{Z} z_loss=10−4⋅log2Z 的形式来约束 softmax 归一化因子的数值稳定性。
通过分析不难理解这个正则项的作用机理,当正则项生效时,它会强制使 log Z ( x ) \log Z(x) logZ(x) 趋近于 0,此时对数与指数运算相抵消,最终得到的其实就是 U r ′ ( x ) U_{r'}(x) Ur′(x),这种状态正是我们期望达到的理想效果,这种处理方式在数值计算上非常稳定,如此一来,所有存在数值隐患的运算环节都得到了妥善解决。因此可以认为,当 Z(x) 接近 1 或者说当 log Z ( x ) \log Z(x) logZ(x) 趋近于 0 时,softmax 函数就处于数值稳定的理想状态,
从某种意义上说,PaLM 堪称是这一领域的开拓者,因为他们率先采用了这种 z 值损失技巧,而其他很多研究团队在很长一段时间里都未曾使用。但随后出现了一系列采用该技术的论文 Baichun 2 (2023)、DCLM (2024)、OLMo 2(2025),这项技术堪称提升模型稳定性的绝佳方案,既优雅又实用
5.2 Attention softmax stability - the ‘QK norm’
接下来我们还需要处理另一个 softmax 运算,它来自注意力机制环节,这项发现源自 NVIDIA 团队发表的研究论文 [Rybakov+ 2024]

上图展示的是注意力机制的工作原理框图,开头先进性层归一化处理,接着生成 Q、K、V 向量,暂时先忽略这部分,接下来会将查询向量 Q 与键向量 K 进行矩阵乘法运算,然后通过 softmax 函数进行归一化处理,再将结果与值向量 V 相乘后进行投影变换,最终输出将经过全连接层处理并生成最终结果
如果忽略图示中的 LN 这个小模块的话,整个结构看起来就和标准的多头注意力机制完全一样,那么这里的区别究竟在哪里呢?多位研究者提出了名为 QK norm 的方法:在计算 softmax 所需的点积之前,先让查询向量 Q 和键向量 K 通过层归一化 LayerNorm 的处理
这种方法为调控 softmax 的行为提供了一种全新的思路,此处并非通过控制归一化因子 z 来实现调控,而是通过约束 softmax 输入值的范围来实现调控,这种方法能自然而然地抑制 softmax 的不良表现
正如我们之前提到的,这项创新最初源自于视觉与多模态模型研究领域,2023 年 Dehghani 团队发表的论文 [Dehghani+ 2023],正是关于训练超大规模视觉 Transformer 的研究。随后,HuggingFace 推出的 Chameleon 和 IDEFICS 模型在其多模态训练模块中采用了这些技术方案。此后,包括 Gemma 2、DCLM 和 OLMo 2 在内的多个模型都采用了类似技术方案来提升训练稳定性
5.3 Logit soft-capping
最后要提到的技巧是对输入 softmax 函数的逻辑值 logits 进行软性截断,这个方法目前尚未被广泛采用,但值得关注。
从某种意义上说,QK norm 是一种相当强力的干预手段,因为它会对整个向量进行全局操作,但还有一种处理方式:在完成自注意力机制的内积运算后,可以通过软最大值(softmax maximum)操作对结果进行非线性转换

因此可以将原始逻辑值 logits 作为输入,将其除以软性阈值上限 soft_cap,这样做会实现什么效果呢?当 logits 值远超软性阈值上限时,双曲正切函数 tanh 会将其截断至 1,因此最终输出值不会超过此处设定的软性阈值上限,这实际上实现了对 logits 值得柔性截断控制。
Gemma 2 模型采用了这种机制,OLMo 2 也同样运用了这一技术,不过在其他模型中,这种处理方法似乎并未得到广泛应用。先前提到的 NVIDIA 团队其实尝试过多种提升稳定性干预措施:

他们的研究发现,基准模型的困惑度为 11.19,采用软性截断反而导致模型性能下降,QK norm 确实能提升效果—它允许采用更激进的学习率,从而将优化器的性能推向更高水平
OK,以上就是关于稳定性优化措施的全部内容
6. Attention heads
最后要讨论的是注意力头的几种变体,关于注意力头的研究还远未饱和,但也有一些重要的改进,了解这些改进对理解当前模型架构很重要
首先要讨论的是分组查询注意力(GQA)和多头查询注意力(MQA)这两种机制,虽然这些机制对模型训练阶段的影响不大,但对于理解模型的推理成本和推理性能却至关重要
另外一个值得关注的新进展是:以最新发布的 LLaMA 4 为例,据称其上下文窗口已突破千万级 token 规模,它是如何实现这一突破的呢?其核心在于对注意力机制进行了结构化改造,稍后也会重点讲解这个技术突破
6.1 GQA/MQA
在分析 LLaMA 等大型模型时,你应该接触过 GQA 和 MQA 这些专业术语,接下来我们将详细解析这两种注意力机制的核心原理
要理解其中的计算原理,我们首先需要剖析注意力机制的核心运算过程,以下是 CS224n 课程的讲义内容:

首先提取查询向量 XQ 和键向量 XK,接着构建大型二次注意力矩阵,我们可以逐步解析这些矩阵乘法运算,通过逐步推导可以确认,总算术运算次数为 b n d 2 bnd^2 bnd2,其中 b 代表 batch 维度,n 代表序列长度,而 d^2 对应隐藏层维度的平方。
我们还可以进一步计算总内存访问量,其计算结果为 b n d + b h n 2 + d 2 bnd+bhn^2+d^2 bnd+bhn2+d2,其中 b n d bnd bnd 这里仅涉及访问该矩阵的操作,而 softmax 的内存访问量将达到 b h n 2 bhn^2 bhn2,通过分析 softmax 矩阵的维度就能理解这一点—其尺寸为 batch 大小乘以注意力头数再乘以所有 softmax 激活值的数量,因此其计算复杂度为 n 2 n^2 n2 量级,最后阶段还需要执行 d 2 d^2 d2 投影运算
因此,我们可以计算总内存访问量与算术操作量的比值这个概念即 算术强度,这个概念将在后续几节课中变得至关重要,因为我们需要确保算术强度足够高。这意味着,我们希望每次内存访问都能触发大量计算操作,这是因为在 GPU 上,内存访问的相对开销非常高昂,而计算操作则相对廉价
这里的算术强度计算式为 ( 1 k + 1 b n ) − 1 (\frac{1}{k}+\frac{1}{bn})^{-1} (k1+bn1)−1,这意味着我们能让 GPU 持续保持高效运转,因为当注意力头数量、批处理规模和序列长度都足够大时,这些参数都将成为显著提升性能的关键因子。当然,这只是训练阶段的情况,问题在于推理阶段—我们无法获得这些庞大的矩阵进行批量乘法运算,这将从根本上改变我们算法的运行特性。
请注意,当我们生成文本时,必须逐个生成每个 token,随后,Transformer 模型必须读取该 token 进行处理,此时才能获得下一个 token 的概率分布,整个过程以自回归方式逐 token 推进。正因如此,文本生成过程无法实现并行化,每个 token 都必须逐步生成,在此过程中,我们需要增量式计算注意力机制,这一思路被业界称为 KV Cache 技术

那么该如何处理呢?上面是一个演示 KV Cache 工作原理的生动动画,整个过程始于一个查询 token,这里的查询 token 代表新生成的 token,模型以此为条件进行运算,此时系统需要决策应该从该查询 token 之前的上下文中检索哪些信息
随着新 token 逐个生成,查询 token 的位置会从 1 逐步推移到 n,此时正在构建的 key cache 实质上是在累积所有历史 token 的 key 值,历史 token 的 key 值始终保持不变,因为它们仅依赖于已生成的上下文。随着 token 的逐个生成,系统会持续累积这些历史 key 值,每次迭代都能计算出新的 QK 矩阵元素,因此最终形成的注意力矩阵将是一个下三角矩阵,每次迭代只计算矩阵的一行,当前计算出的这一行数据,正是生成下一个 token 所需的全部信息
这个 KV Cache 的核心思路是:在逐步生成每个 token 的过程中,我们会增量式地计算并存储当前的 K 和 V 向量,并且仅计算当前运算真正必须的 QK 矩阵部分
因此,我们可以再次详细计算各个算术环节,总共需要执行多少次浮点运算?内存访问总次数是多少?当我们考虑 KV Cache 时,只会计算绝对必要的键值对相乘操作,由于我们保持了所有中间计算结果,因此不会浪费任何矩阵或向量乘法运算。
算术运算总量 b n d 2 bnd^2 bnd2 始终保持不变,但内存访问模式已发生变化,这是为什么呢?因为采用 KV Cache 之后,需要反复将各类参数调入调出内存,每次与键矩阵相乘时,都需要先将其载入内存,再执行与 K 的乘法运算,随后又需将其移出内存,转而计算某些激活值。
因此需要循环加载不同的矩阵,这会导致总内存访问量大幅增加,达到 b n 2 d + n d 2 bn^2d+nd^2 bn2d+nd2,如此一来,计算这个比率时,算术强度就变得不太理想了,最终得到的算术强度将是 ( n d + 1 b ) − 1 (\frac{n}{d}+\frac{1}{b})^{-1} (dn+b1)−1。因此,如果我们仔细分析—若想获得较高的算术强度,就需要让 ( n d + 1 b ) (\frac{n}{d}+\frac{1}{b}) (dn+b1) 这个值尽可能小,因此必须使用非常大的批次规模,同时需要确保 n d \frac{n}{d} dn 的比值足够小
这意味着什么呢?意味着要么采用极短的序列长度,要么大幅增加模型维度,这本质上就是当前推理成本面临的核心取舍问题,这种极其低效的内存访问模式中, n d \frac{n}{d} dn 这个致命参数正在严重拖垮整个系统的吞吐性能。
6.2 MQA
这便推动了多查询注意力 MQA 的诞生,关键问题其实出在键值对这部分,这才是真正的性能瓶颈所在,它们会不断构建 KV Cache,同时伴随着内存数据的频繁进出

MQA 的具体做法是:查询可以设置多头结构,但键和值仅保留单一维度或单头结构,这极大地简化了整体架构。采用这种设计后,键和值需要传输的数据量就大幅减少了,因此键和值实现了共享,但查询仍保持多头结构,这样既保留了多头注意力机制同时仅需维度单组键值对,正因如此,这种机制才被称为多查询注意力
此时进行相同运算时,由于键和值实现了共享,所需的内存访问次数也随之减少,这使得算术强度特性得到了显著改善,其具体计算强度是 ( 1 d + n d h + 1 b ) − 1 ({\color{red}{\frac{1}{d}}} + \frac{n}{d{\color{red}{h}}} + \frac{1}{b})^{-1} (d1+dhn+b1)−1,第一项的数值降低了 n 倍,如此一来,更长的序列长度现在变得可行了,而第二项现在被注意力头的数量所均摊,因此这一项的影响也变得可控了
至此,所有关键参数都得到了有效调控,而多查询注意力机制能带来更优越的性能表现
6.3 GQA
组查询注意力(GQA)本质上是对该机制进行了微调优化:

不同于传统的单键值对查询模式,该方案通过倍数规约机制减少了键向量的数量,这种设计让你能在推理速度与模型表达能力之间实现灵活取舍
因为直接从多头注意力切换到多查询注意力可能过于激进,多项研究 [Shazeer+ 2019][Ainslie+ 2023] 表明 GQA 不会降低性能,而多头注意力反而会造成负面影响,这部分内容就不展开讨论了
6.4 Sparse / sliding window attention
最后想提一个近几个月非常有趣的新进展,2019 年 OpenAI 曾发表过一篇开创性论文 [Child+ 2019],核心是探讨如何构建更长的注意力模型。他们的核心观点是:实现这一目标的可行性方案是设计 稀疏注意力模式,不再对完整序列施加注意力,而是改为对每个数据块采用局部窗口注意力机制,同时配合对角线式注意力模式实现跨区块的信息传递

通过这种结构化稀疏注意力机制,可以在模型表达能力与计算效率之间实现动态取舍,GPT-3 在最初发布时就采用了这类技巧,从而实现了更大规模的注意力机制窗口化处理

滑动窗口注意力是该理念的另一种实现方式—在每一层网络仅关注当前位置周边的有限区域,这种机制还能有效控制所需资源总量,实现更长上下文处理时所需的资源总量。因此,有效感受野范围等于局部窗口大小乘以网络层数
6.4 Current standard trick
以上都属于早期技术方案,但现代实现方式已迭代升级,近期 LLaMA 4、Gemma 和 Cohere Command A 等论文提出精妙设计:将四个 Transformer 块组合成模块化单元,最底层的模块采用完全自注意力机制,且不添加位置编码,既没有旋转位置编码也不含任何其他位置信息,该模块完全不具备位置感知能力,但它仍采用完整的自注意力机制

这种设计每四个模块才会出现一次,而上方三个模块则采用带旋转位置编码的滑动窗口注意力机制。这种设计堪称精妙:通过让全局注意力机制仅间歇性触发(每四个模块出现一次),既有效控制了系统资源开销又通过交替使用滑动窗口注意力(带 RoPE)维持了位置感知能力,这种张弛有度的架构设计,在计算效率和模型性能间取得了绝佳平衡
此外,还需考虑长度外推问题,因为旋转位置编码仅能处理局部上下文窗口,而那些需要处理超长距离依赖关系的部分则完全不使用位置编码,这使得模型能够进行及其激进的长度外推,因为你无需使用旋转位置编码那样进行位置外推计算,这确实是过去几个月里涌现的绝妙突破
OK,以上就是本次讲座的全部内容了
结语
第三讲我们主要讲解了语言模型的架构变化、超参数设置、训练时的稳定性策略以及注意力头的各种变体。
在模型架构变体小节中,我们首先探讨了 pre-norm 和 post-norm 的选择,现有共识是采用 pre-norm 而非 post-norm,因为 pre-norm 本质上构建了更稳定的训练架构。接着讨论了层归一化方法的选择,包括 LayerNorm 和 RMSNorm,激活函数方面 SwiGLU 这类门控线性单元变体总能稳定发挥出色效果。最后我们探讨了位置编码,并说明了 RoPE 旋转位置编码设计的精妙之处在于利用旋转不变性。
超参数小节我们重点讨论了模型隐藏层维度的设置、注意力头维度和注意力头数量的权衡、网络宽深比、词汇表大小以及 dropout 等正则化方法的选择,这些参数大部分在业界都有一些明确的经验法则。
稳定性技巧小节我们主要探讨了 Transformer 模型训练时的稳定性问题,主要出在 softmax 函数上,为此基于两处 softmax 函数业界分别提出了对应的解决方案。针对末端的 softmax 函数采用 z-loss 方法来约束 softmax 归一化因子的数值稳定性,而针对注意力机制中的 softmax 函数则采用 QK norm 的方式,在计算点积之前让 Q、K 向量先通过层归一化的处理。最后该小节还探讨了对 logits 值进行软性截断方法。
注意力头变体小节我们首先对原始注意力机制进行了分析,在推理阶段因为 KV Cache 的存在导致算术强度并不高,由此推动了 MQA 和 GQA 的诞生。MQA 核心在于 key 和 value 仅保留单头而 query 仍保留多头,这样键和值传输的数据量大幅减少了,有利于提高算术强度。而 GQA 则取了一个均衡,通过倍数规约机制来减少键向量的数量。最后我们还谈到了稀疏和窗口注意力,用来构建更长的注意力模型,不过现有方案采用的是四个 Transformer 模块组合来间歇触发全局注意力。
整个讲解非常通俗易懂,大家感兴趣的可以看看
下节课我们将深入探讨混合专家模型 MoE(Mixture of Experts),敬请期待🤗
参考
更多推荐



所有评论(0)