文章目录
-
- 摘要
- [1 引言](#1 引言)
- [2 相关工作](#2 相关工作)
- [3 方法论](#3 方法论)
- [4 实验设置](#4 实验设置)
- [5 结果](#5 结果)
摘要
增加语言模型的参数数量是提升其性能的常见策略。然而，由于运行成本较低，小型语言模型依然具有重要价值。尽管如此，即使在使用相同的数据和计算资源条件下，小型模型的性能常常不及大型模型。具体来说，小型模型的性能往往在预训练的后期出现下降。这一现象通常被归因于其较低的表达能力，但造成性能下降的具体原因仍不清楚。我们使用Pythia模型组对该现象背后的训练动态进行了分析。在不同规模的模型中，我们研究了注意力机制（Attention）和多层感知机（MLP）激活值向最终状态的收敛情况，并考察了参数的有效秩（effective rank）对该过程的影响。结果发现，大型模型的几乎所有层都能在训练的前20%内早早稳定下来，而小型模型的层收敛速度较慢且不稳定，尤其当其参数的有效秩较低时表现更为明显。通过将层激活的收敛性与参数的有效秩联系起来，我们的分析可为未来解决小型模型学习动态中的低效问题提供指导。
1 引言
增加语言模型（LM）参数数量在多种任务中带来了显著的性能提升（Hendrycks 等，2021），并已成为模型设计进步的事实标准（例如，Chowdhery 等，2023）。然而，小型语言模型依然不可或缺，因为它们更为实用：较低的训练和推理成本降低了环境影响（Schwartz 等，2020）。小型语言模型使得个人能够使用较少的资源在专有数据上进行训练，从而增强数据隐私（Huang 等，2022）并促进语言建模技术的普及（Bender 等，2021）。不过，在相同的数据和计算预算下，小型语言模型（不足为奇地）表现不如大型模型（Biderman 等，2023），且其性能往往在预训练后期出现下降，这一现象被 Godey 等（2024）称为"饱和"（saturation）。饱和现象通常被归因于小型语言模型"有限的表示能力"，但除了这一经验性解释外，我们对其成因的理解仍然有限。
近期，Godey 等（2024）将饱和现象与语言模型输出嵌入的变异性降低联系起来，认为这是由于隐藏层维度与词汇表大小不匹配导致的（Yang 等，2018）。具体来说，语言模型最后一层会将随机词的隐藏表示映射到高余弦相似度的输出嵌入空间中。
本文基于 Pythia 模型组（Biderman 等，2023），提供了从训练动态角度的独立分析。首先，我们研究不同规模语言模型中注意力层（Attention）和多层感知机层（MLP）激活值向最终状态的收敛过程；接着，将不同模型规模中收敛行为差异与参数的有效秩（effective rank）联系起来：那些激活值收敛较晚的层，其激活空间涵盖的维度比例更小。
具体而言，我们首先采用中心核对齐（Centered Kernel Alignment，CKA；Kornblith 等，2019）指标测量各层激活在不同训练检查点间的相似性。观察发现，大型语言模型收敛更快、更平稳。如图1所示，在训练的前20%阶段，较大模型（2.8B参数）的几乎所有层激活已接近最终状态，而较小模型（160M参数）的大部分层在训练大部分时间内依然不同于最终状态。
随后，我们发现层激活的收敛模式与其参数和梯度的秩存在强烈相关性。为实现不同模型规模间的有效秩比较，我们引入了"比例有效秩"（proportional effective rank）概念（详见第3节）。我们的分析揭示了小型语言模型训练中的低效环节，为未来有针对性的改进提供了方向。
2 相关工作
此前已有研究分析了Pythia模型组的多种学习动态，包括记忆能力（Biderman 等，2023；Lesci 等，2024）、训练数据的影响（Liu 等，2024）和学习到的嵌入统计特性（Belrose 等，2024）。与本文相关的是，Godey 等（2024）研究了不同模型规模下反嵌入矩阵（即从隐藏表示映射到词元的矩阵）的秩差异，这被称为softmax瓶颈（Yang 等，2018）。但与他们不同，本文关注的是所有层的收敛动态。
相似性度量如中心核对齐（CKA）和奇异向量典型相关分析（SVCCA）被广泛用于分析语言模型特性。Nguyen 等（2021）发现模型的结构设计，如宽度和深度，会影响隐藏表示的相似性。Wu 等（2020）表明，同一类架构的模型具有相似的隐藏结构，这种相似性在微调模型中依然存在（Phang 等，2021）。此外，SVCCA曾用于研究多语言模型中的词元表示分布（Singh 等，2019）和单语言模型中的句法元素学习（Saphra 和 Lopez，2019）。与本文最相似的工作是Brown 等（2023），他们使用包括CKA在内的表示相似性度量来研究Pythia模型的泛化能力。但本文首次使用CKA度量来考察不同规模模型层激活的收敛动态。
3 方法论
我们首先介绍基于Transformer模型的残差流（residual stream）视角，并定义层的激活。随后介绍CKA和比例有效秩（proportional effective rank）指标。
残差流视角：Transformer架构（Vaswani 等，2017）中的残差连接为分析信息流动提供了框架（Elhage 等，2021）。残差连接集合称为残差流，每一层通过对残差流进行加法操作来更新信息。每层包含两个主要部分：Attention和MLP，依次更新残差流。形式上，输入序列 t = ⟨ t 1 , . . . , t T ⟩ t=\langle t_1, ..., t_T \rangle t=⟨t1,...,tT⟩首先通过嵌入层映射为矩阵 x 0 ∈ R T × D x_0 \in \mathbb{R}^{T \times D} x0∈RT×D，其中每列为一个维度为 D D D的词元表示。第 l l l层依次执行：
x ′ = x l − 1 + Attention ( x l − 1 ) x' = x_{l-1} + \text{Attention}(x_{l-1}) x′=xl−1+Attention(xl−1)
x l = x ′ + MLP ( x ′ ) x_l = x' + \text{MLP}(x') xl=x′+MLP(x′)
最终第 T T T列的 x L x_L xL用于预测第 T + 1 T+1 T+1个词元。
激活和参数：公式（1）中带下划线的更新即为该层的激活，维度与残差流一致，即 R T × D \mathbb{R}^{T \times D} RT×D。Attention和MLP均先将残差流投影到较低维的中间表示，再映射回残差流。本文研究的是将中间表示写回残差流的参数矩阵，分别记为 θ A T T \theta_{ATT} θATT和 θ M L P \theta_{MLP} θMLP，对应的激活分别为 a A T T a_{ATT} aATT和 a M L P a_{MLP} aMLP。
激活相似性：针对层 l l l在训练检查点 c c c的激活 a l , c a_{l,c} al,c，用中心核对齐（CKA）线性版本度量其与最终检查点 C C C激活 a l , C a_{l,C} al,C的相似度：
CKA ( a c , a C ) = ∥ a c ⊤ a C ∥ F 2 ∥ a c ⊤ a c ∥ F ⋅ ∥ a C ⊤ a C ∥ F \text{CKA}(a_c, a_C) = \frac{\| a_c^\top a_C \|_F^2}{\| a_c^\top a_c \|_F \cdot \| a_C^\top a_C \|_F} CKA(ac,aC)=∥ac⊤ac∥F⋅∥aC⊤aC∥F∥ac⊤aC∥F2
其中 a a a为中心化后的激活矩阵， ∥ ⋅ ∥ F \|\cdot\|_F ∥⋅∥F为Frobenius范数。该指标用于所有层和所有检查点的 a_{ATT} 和 和 和 a_{MLP} ，以分析激活的收敛动态。
参数的比例有效秩（PER）：设Attention或MLP中间表示维度为 H H H，第 l l l层对应参数矩阵为 θ l ∈ R D × H \theta_l \in \mathbb{R}^{D \times H} θl∈RD×H，代表从隐藏空间投影回残差流的权重。采用Roy 和 Vetterli (2007)提出的有效秩定义，通过参数矩阵奇异值的归一化熵计算：
E R ( θ l ) = exp ( − ∑ k = 1 K σ k ∥ σ ∥ 1 log σ k ∥ σ ∥ 1 ) ER(\theta_l) = \exp \left(-\sum_{k=1}^K \frac{\sigma_k}{\|\sigma\|_1} \log \frac{\sigma_k}{\|\sigma\|_1} \right) ER(θl)=exp(−k=1∑K∥σ∥1σklog∥σ∥1σk)
其中 σ = ⟨ σ 1 , . . . , σ K ⟩ \sigma = \langle \sigma_1, ..., \sigma_K \rangle σ=⟨σ1,...,σK⟩为奇异值向量， ∥ ⋅ ∥ 1 \|\cdot\|_1 ∥⋅∥1为 ℓ 1 \ell_1 ℓ1范数。为便于不同规模层的比较，定义比例有效秩为：
P E R ( θ l ) = E R ( θ l ) H PER(\theta_l) = \frac{ER(\theta_l)}{H} PER(θl)=HER(θl)
本文计算了 θ A T T \theta_{ATT} θATT和 θ M L P \theta_{MLP} θMLP及其梯度的PER，在所有层和训练检查点上进行分析。
4 实验设置
实验使用Pythia模型组（Biderman 等，2023），包含8个不同规模的Transformer模型，在去重后的Pile数据集（Gao 等，2020；Biderman 等，2022）上训练143k步。数据集存在一个标准版本和去重版本，本文采用去重版本。为控制计算成本，选取参数规模从7000万（70M）、1.6亿（160M）、4.1亿（410M）、14亿（1.4B）到28亿（2.8B）的模型。训练检查点选取了初始步数0、对数间隔步数（如1、2、4，直至512）、1000、3000步，以及之后每隔1万步至143k步。每个检查点使用训练集最后一个batch进行评估并收集激活。更多细节见附录B。
5 结果
我们的分析揭示了不同规模模型层学习动态上的定量差异。
结果1：较大模型的激活更快且更单调地收敛到最终状态。
如图2（第一列）所示，较大模型的Attention和MLP激活平均更早收敛。例如，在训练进行20%时，2.8B模型的MLP层CKA达0.8，Attention达0.7，而70M和160M模型约为0.5。这种快速收敛的趋势贯穿所有层（图1展示了相关分布）。
结果2：较早层的激活收敛更快，与模型规模无关。
不同规模模型中，靠前层的激活比靠后层更早地收敛到最终状态。如图3（附录C）所示，较大模型较多靠后层也较早收敛，而小模型的层多在训练末期才达到最终状态。
基于近期关于参数秩差异的工作（Godey 等，2024），接下来我们研究不同收敛行为是否与层参数和梯度的有效秩相关。
结果3：较大模型层的参数在比例上覆盖更多维度。
较大模型层参数覆盖的有效维度比例稍高（图2第二列），且较大模型的比例有效秩（PER）较早稳定，而小模型的PER在训练过程中持续下降。图4（附录D）进一步显示，小模型较后层的PER随训练下降，而大模型所有层的PER较早稳定。梯度的PER差异更为明显（图2第三列）。
结果4：较大模型层的梯度更新覆盖更多维度。
梯度的PER反映梯度学习信号相对于参数维度的覆盖比例。图2（第三列）显示，较大模型的梯度始终覆盖更大比例的参数维度，且比例随时间递减；小模型表现出更大波动。注意2.8B模型Attention层梯度PER平均值似乎偏小，但查看各层梯度PER（图5，附录E）发现，较大模型层梯度PER更早稳定，部分层梯度PER较低，而小模型层梯度尚未收敛。总体来看，较大模型层不仅收敛更快，也接收了比例更大的秩更新。
结果5：参数有效秩动态与激活收敛模式相关。
我们量化层激活收敛速度与参数和梯度有效秩的相关性。定义两个二元变量，分别表示（i）层是否早期收敛（激活CKA在训练10%前≥0.45），（ii）参数和梯度PER在训练中是否保持高且稳定。使用Matthew相关系数测量两者相关性，结果见表1。Attention层在所有模型规模中均表现出较强相关性。MLP层中，梯度PER与收敛相关性在1.4B及以下模型较强，参数PER相关性仅在70M模型明显。我们推测MLP层的PER整体较大，除了70M模型例外。
虽然这些结果为相关性分析，但为未来研究提供了基础，测试是否通过提升层参数和梯度的PER可以加速小模型层激活的收敛。
6 结论
本研究揭示了大小语言模型学习动态的差异。利用Pythia模型组，展示了大模型层激活更快且更单调地收敛至最终状态，并将此现象与大模型参数和梯度更高的比例有效秩相关联。我们的分析拓展了对小模型训练低效的理解，并为未来改进提供了思路，例如开发提升层参数PER的方法。
伦理影响
本工作是Green AI（Schwartz 等，2020）降低语言模型训练及使用环境影响的努力之一。我们同时认识到，小模型同样会继承训练数据中的偏见，例如Pile数据集中存在的性别和种族偏见（Gao 等，2020）。
局限性
实验仅限于Pythia模型组及Pile数据集，尚不明确结论能否迁移至其他模型或非英语数据集。因计算资源限制，无法深入研究更大规模模型，当前最大模型仅28亿参数，远小于百亿级别的开源大模型。此外，CKA相似度与PER的关系为相关性，未来工作计划基于当前发现开展针对性干预，检验是否存在因果关系，即提升有效秩是否能加速层收敛。
致谢
感谢匿名审稿人宝贵意见。实验使用剑桥大学数据驱动发现服务（CSD3）资源，相关资助来自EPSRC和STFC。Richard Diehl Martinez获盖茨剑桥信托基金资助（Bill & Melinda Gates Foundation）。