Numerical Error Analysis of Large Language Models
Numerical Error Analysis of Large Language Models
作者/机构: Stanislav Budzinskiy (维也纳大学数学系), Wenyi Fang (华为技术有限公司), Longbin Zeng (华为技术有限公司), Philipp Petersen (维也纳大学数学系)
A1 主要贡献
本文旨在对大型语言模型(LLMs)中有限数值精度计算导致的舍入误差进行理论分析,以揭示其对训练不稳定性的潜在影响。
-
核心问题:基于 Transformer 架构的大型语言模型在训练过程中常出现“损失尖峰”(loss spikes)等不稳定性现象。尽管已有多种潜在原因被识别,但由海量浮点运算带来的有限精度舍入误差的影响尚未得到充分的理论分析。特别是在低精度计算成为必需的背景下,理解并量化这些误差至关重要。
-
研究目标:本文对 Transformer 架构的前向传播过程进行了一次前向误差分析(forward error analysis),旨在为舍入误差的影响提供严格的理论界限,并识别出影响数值稳定性的关键因素。
-
主要创新点与发现:
- 理论误差界限:本文为具有 L 个 Transformer 模块的深度神经网络的输出推导出了一个分量级相对舍入误差的上限(定理 4.30)。这是对 Transformer 架构前向传播中舍入误差的首次严格理论分析。
- 误差增长模式:
- 理论界限显示,误差随网络深度 L 指数级增长。图 1 的数值实验也证实了平均相对误差随层数增加呈现近乎指数的增长趋势。然而,实验同时表明中位数误差远低于平均误差,说明大的相对舍入误差较为罕见。
图1:应用深度 Transformer (7) 的相对舍入误差。我们通过将中间值四舍五入到不同位数来模拟低精度计算。我们设置 d, n, D = 20 和 L = 40,使用中心正态分布(单位方差)的条目生成自注意力矩阵 Wq、Wk、Wv,以及使用中心正态分布(方差为 1/d)的条目生成两层感知器的矩阵。为了增加 W⊺kWq 的条件数,我们在其两侧乘以对角矩阵,其条目为 [1/4, 4] 中的均匀随机数。输入向量使用中心正态分布(单位方差)的条目进行初始化。我们对 5000 次初始化计算统计数据。我们在上排图中绘制了每一层的平均相对误差:几乎线性的曲线表明接近指数增长。为了可视化误差的随机行为,我们在下排图中绘制了中位数误差,并用阴影区域表示第 5 和第 95 百分位数之间的区域。最后一层误差分布的直方图显示,大多数误差比平均误差(垂直黑线)低几个数量级。
- 关键影响因素识别:
- 自注意力矩阵:理论分析(定理 4.30)和数值实验(图 2)均表明,自注意力机制中的键矩阵和查询矩阵的乘积 $W_k^\top W_q$ 对舍入误差有显著影响。误差上界与该矩阵的条件数 $(\sigma_{max}/\sigma_{min})$ 和谱范数成正比。当所有 Transformer 模块相同时,误差界限的增长形式如下:
$$ \left(1+4\sqrt{d}\sigma_{\max}(|W_k|^\mathrm{T}|W_q|)\frac{\sigma_{\max}(W_k^\mathrm{T}W_q)}{\sigma_{\min}(W_k^\mathrm{T}W_q)}\right)^L $$ - 图 2 的实验通过缩放 $W_k^\top W_q$ 证实了误差随其谱范数的增长,尽管是以较慢的多项式速率。
- 归一化层:均方根层归一化(RMS Layer Normalization)比标准的层归一化(Layer Normalization)在数值上更稳定。后者的误差在其参数 $x$ 的分量接近样本均值时会发生爆炸。
- 自注意力矩阵:理论分析(定理 4.30)和数值实验(图 2)均表明,自注意力机制中的键矩阵和查询矩阵的乘积 $W_k^\top W_q$ 对舍入误差有显著影响。误差上界与该矩阵的条件数 $(\sigma_{max}/\sigma_{min})$ 和谱范数成正比。当所有 Transformer 模块相同时,误差界限的增长形式如下:
- 实践指导:基于分析结果,本文提出具体建议以减轻舍入误差,从而实现更鲁棒和稳定的推理:
- 在训练过程中应控制 $W_k^\top W_q$ 的谱范数和条件数。
- 建议在计算自注意力(特别是相似度分数和 softmax)时使用更高的精度。
A3 背景知识与关键定义
2.1 符号和辅助结果
向量与矩阵表示法。本文使用 $x \in R^m$ 表示向量,$Y \in R^{m \times n}$ 表示矩阵。$x$ 的第 $i$ 个元素记为 $x_i$,Y 的第 $j$ 列记为 $y_j$,Y 的第 $(i, j)$ 个元素记为 $y_{i,j}$。对于索引子集 $\Omega \subseteq \{1, . . . , m\}$, $x(\Omega) \in R^{|\Omega|}$ 和 $Y(\Omega, :) \in R^{|\Omega|\times n}$ 表示提取对应的子数组。$e_1, . . . , e_m \in R^m$ 表示单位矩阵 $I_m \in R^{m \times m}$ 的列向量,而 $e = \sum_{i=1}^m e_i$ 是全为1的向量。向量的绝对值 $|x|$ 或矩阵的 $|Y|$ 是逐元素计算的。$supp(x) = \{i : x_i \neq 0\}$ 是 $x$ 的支撑集,符号函数为 $sign(x)$,克罗内克δ为 $\delta_{i,j}$ 。
范数定义。对于 $1 \le p, q \le \infty$,$\ell_p$ 向量范数记为 $\|\cdot\|_p$,由 $\ell_p$ 和 $\ell_q$ 向量范数诱导的矩阵范数记为 $\|\cdot\|_{p,q}$。$\|Y\|_{2,2}$ 是 $Y$ 的最大奇异值。一个较少人知的表示法是【21, Problem 6.11】,如下所示:
$$\|\mathbf{Y}\|_{1, q}=\max _{1 \leq j \leq n}\left\|\mathbf{y}_{j}\right\|_{q},$$
这可以导出标准范例,如 $\|Y\|_{1,1} = \max_j \sum_i |y_{i,j}|$ 和 $\|Y\|_{1,\infty} = \max_{i,j} |y_{i,j}|$。作为 $\|Y\|_{\infty,\infty} = \|Y^\top\|_{1,1}$ 等式的推广,以下等式成立【27, Theorem 5.6.35】:
$$\left\|Y^{\mathrm{T}}\right\|_{p, q}=\|Y\|_{\frac{q}{q-1}, \frac{p}{p-1}}.$$
此外,本文定义了一个类范数的泛函 $\|\cdot\|_{-\infty}$:
$$\| \mathbf{x} \|_{-\infty} = \min_{1 \le i \le d} |x_i|.$$
克罗内克积。设 $Z \in R^{k \times l}$。克罗内克积 $Y \otimes Z \in R^{mk \times nl}$ 定义为:
$$Y \otimes Z = \begin{bmatrix} y_{1,1}Z & \cdots & y_{1,n}Z \\ \vdots & \ddots & \vdots \\ y_{m,1}Z & \cdots & y_{m,n}Z \end{bmatrix},$$
并且满足 $\|Y \otimes Z\|_{p,q} = \|Y\|_{p,q}\|Z\|_{p,q}$ 【32, Theorem 8】。对于可相乘的矩阵,以下混合乘积性质成立【26, Lemma 4.2.10】:
$$(Y_1 \otimes Z_1)(Y_2 \otimes Z_2) = Y_1 Y_2 \otimes Z_1 Z_2.$$
2.2 Transformer 架构
基本定义。基于 Transformer 架构的神经网络【44, Attention is all you need, 2017, NIPS】处理输入 token 序列,即任意长度的有序集合。本文中,token 是来自 $R^d$ 的向量。为了定义 Transformer,需要先引入一些符号和子模块。对于集合 G,令 $G^* = \cup_{n \in N} G^n$ 表示 G 值的集合。对于函数 $f : G \to H$,我们用 $f^* : G^* \to H^*$ 表示其对集合的扩展,该扩展逐元素作用,如下所示:
$f^*:(g_1,\dots,g_n) \mapsto (f(g_1),\dots,f(g_n)).$
本文中,这种扩展主要应用于两种情况:当 $G = H = R$ 且 $x \in R^d$ 时,可将 $f$ 逐元素应用于 $x$ 得到 $f^*(x) \in R^d$;当 $G=R^d$ 且 $H=R^D$ 时,通过同构关系 $R^{d \times n} \cong (R^d)^n$,可将 $f$ 按列应用于 $X \in R^{d \times n}$ 得到 $f^*(X) \in R^{D \times n}$。
层归一化 (Layer Normalization)。层归一化是一个函数 $N : R^d \to R^d$,定义为:
$$N(x)=\frac{x-\mu(x)e}{\sqrt{\nu(x)}}, \quad e = \begin{bmatrix} 1 & 1 & \cdots & 1 \end{bmatrix}^T \in \mathbb{R}^d.$$
其中样本均值 $\mu(x) = \frac{1}{d} \sum_{i=1}^d x_i$,样本方差 $\nu(x) = \frac{1}{d} \sum_{i=1}^d (x_i - \mu(x))^2$。
两层感知机 (Two-layer Perceptron)。两层感知机是一个函数 $M : R^d \to R^d$,是两个仿射函数 $L_1$ 和 $L_2$ 与 ReLU 激活函数 $\varrho(x) = \max\{0, x\}$ 的复合,即 $M = L_2 \circ \varrho^* \circ L_1$。其中仿射函数定义为:
$$ \begin{aligned} \mathcal{L}_1(x) &= A_1x + b_1, & A_1 &\in \mathbb{R}^{D \times d}, & b_1 &\in \mathbb{R}^D, \\ \mathcal{L}_2(y) &= A_2y + b_2, & A_2 &\in \mathbb{R}^{d \times D}, & b_2 &\in \mathbb{R}^d \end{aligned} $$
自注意力 (Self-attention)。自注意力是一个函数 $A : (R^d)^* \to (R^d)^*$,是 Transformer 中唯一一个对整个输入集合进行操作的部分,而不是对单个 token 操作。它由嵌入矩阵 $W_q, W_k, W_v \in R^{d \times d}$ 参数化。第一步是计算相似度分数。考虑一个函数 $S : (R^d)^* \to R^*$,它将 $X = [x_1 \cdots x_t] \in R^{d \times t}$ 映射为:
$S(X) = \frac{1}{\sqrt{d_k}}(W_k X)^T (W_q x_t) \in \mathbb{R}^t.$
第 $i$ 个相似度分数决定了第 $i$ 个键向量 $W_k x_i$ 是否关注第 $t$ 个查询向量 $W_q x_t$。然后,相似度分数向量通过 softmax 函数 $\psi : R^* \to R^*$ 转换为 $\{1, \dots, t\}$ 上的概率分布,该函数将 $s \in R^t$ 映射为:
$$\psi(s) = \left(\frac{\exp(s_1)}{\sum_{i=1}^t \exp(s_i)}, \dots, \frac{\exp(s_t)}{\sum_{i=1}^t \exp(s_i)}\right) \in (0,1)^t.$$
最后,这些概率被用来计算值向量的期望。对于输入 $X = [x_1 \cdots x_n] \in R^{d \times n}$ 及其子矩阵 $X_t = [x_1 \cdots x_t] \in R^{d \times t}$,自注意力的输出定义为:
$$
L(x) = \mathbb{E}_{y \sim P_y|x}(S(y)) + \mathbb{E}_{y \sim Q_y|x}(S(y)) - \mathbb{E}_{y \sim \mathbb{P}_y|x}(S(y)) \in \mathbb{R}^n
$$
Transformer 模块与架构。要构建一个 Transformer 模块,需要一个两层感知机 M 和一个自注意力 A。我们将它们的参数组合成一个称为 Transformer 配置的元组:
$\Theta = (D, A_1, b_1, A_2, b_2, W_q, W_k, W_v).$
它定义了一个 Transformer 模块 $T_\Theta : (R^d)^* \to (R^d)^*$,其构成为:
$\mathcal{T}_{\Theta} = (\text{Id} + \mathcal{M}^* \circ \mathcal{N}^*) \circ (\text{Id} + \mathcal{A} \circ \mathcal{N}^*).$
给定多个 Transformer 配置 $\Theta = (\Theta_l)_{l=1}^L$,我们将深度为 L 的 Transformer 定义为:
$\mathcal{T}_\Theta = \mathcal{T}_{\Theta_L} \circ \cdots \circ \mathcal{T}_{\Theta_1}$
这种特定的神经网络架构可以被描述为一个带有掩码(因果)单头自注意力的仅解码器(decoder-only)Transformer。虽然在原始文章【44, Vaswani et al., 2017, NIPS】和BERT【13, Devlin et al., 2019, NAACL-HLT】中归一化位置不同,但像公式(6)中的预归一化(pre-normalisation)【47, Xiong et al., 2020, ICML】已变得更加普遍(例如在GPT-3【7, Brown et al., 2020, NeurIPS】、LLaMA【43, Touvron et al., 2023】和PaLM【8, Chowdhery et al., 2023】中)。
均方根层归一化 (RMS Layer Normalization)。一些 Transformer 架构,如 LLaMA【43, Touvron et al., 2023】,在层归一化的方式上也有所不同,采用了均方根(RMS)层归一化【49, Zhang and Sennrich, 2019, NeurIPS】。它是一个函数 $R : R^d \to R^d$,定义为:
$$ \mathcal{R}(x) = \sqrt{d} \cdot \frac{x}{\|x\|_2}. $$
RMS 层归一化与标准层归一化(1)的区别在于,它在缩放前不从输入中减去均值。本文用以下符号表示使用 RMS 层归一化的 Transformer:
$$Q_\Theta = (\text{Id} + \mathcal{M}^* \circ \mathcal{R}^*) \circ (\text{Id} + A \circ \mathcal{R}^*),$$
$$Q_\Theta = Q_{\Theta_L} \circ \dots \circ Q_{\Theta_1}.$$
3. 条件数
本节计算 Transformer 模块(6)和(9)中各个组成部分的雅可比矩阵,并推导其范数和条件数的估计值。这些结果将用于第四节的前向舍入误差分析,并且本身也具有独立意义。例如,引理3.14表明,将归一化置于自注意力之前可以保证后者条件数的统一界限(见注3.15)。
3.1 预备知识
Fréchet 导数与雅可比矩阵。考虑定义在赋范线性空间 V 的开子集 U 上的函数 $f: V \to W$。若存在一个连续线性算子 $Df(u): V \to W$ 使得下式成立,则称 $f$ 在 $u \in U$ 处 Fréchet 可微:
$$\lim_{\|\delta\|_{\mathbb{V}} \to 0} \frac{\|f(u + \delta) - f(u) - D_f(u)\delta\|_{\mathbb{W}}}{\|\delta\|_{\mathbb{V}}} = 0.$$
当 $V = R^n$ 和 $W = R^m$ 时,Fréchet 导数是一个 $m \times n$ 的偏导数矩阵。本文中遇到的赋范空间是向量和矩阵空间,我们将使用标准正交基。当空间为矩阵空间时,为避免混淆,引入向量化算子:
$X = [x_1 \dots x_n] \in \mathbb{R}^{m \times n}$, $\text{vec}(X) = [x_1^T \dots x_n^T]^T \in \mathbb{R}^{mn}$
并定义函数 $f: R^{m \times n} \to R^{p \times q}$ 的向量化为:
$\bar{f}: \mathbb{R}^{pq} \rightarrow \mathbb{R}^{mn}, \quad \operatorname{vec}(f(X)) = \bar{f}(\operatorname{vec}(X)).$
此时,$f$ 的雅可比矩阵 $J_f(X)$ 是其向量化函数 $\tilde{f}$ 的 Fréchet 导数。例如,对于矩阵乘法 $f(X, Y) = XY$,其雅可比矩阵为:
$J_f(X,Y) = [Y^T \otimes \mathbf{I}_m, \mathbf{I}_n \otimes X] \in \mathbb{R}^{mn \times (mk+kn)}.$
相对距离与条件数。本文定义了范数相对距离和分量相对距离:
$$\rho_V(\hat{v}, v) = \frac{\|\hat{v} - v\|_V}{\|v\|_V}.$$
$$\rho_c(\hat{v}, v)=\max _{1 \leq i \leq \operatorname{dim} V}\left|\frac{\left(T_V \hat{v}\right)_i-\left(T_V v\right)_i}{\left(T_V v\right)_i}\right|,$$
基于这些距离,定义了三种条件数:范数条件数 $\kappa_{V,W}(f, u)$,混合条件数 $\kappa_{c,W}(f, u)$,以及分量条件数 $\kappa_{c,c}(f, u)$。当函数 Fréchet 可微时,这些条件数有明确的表达式(引理3.3)。例如,对于仿射函数 $L(x) = Ax+b$,其条件数可以明确计算(引理3.4)。
高阶项分析。条件数描述了函数对微小扰动的局部敏感性,即 $\rho_c(f(\hat{u}), f(u)) \lesssim \kappa_{c,c}(f, u)\rho_c(\hat{u}, u)$。一个更精确的不等式(引理3.5)表明,对于二次连续可微的函数,存在一个常数 $C(u, \beta)$ 使得:
$n(f(\bar{u}), f(u)) \leq \kappa_c(f, u) \rho_c(\bar{u}, u) + C(\bar{u}, 3) \mathcal{O}_q(\bar{u}, u)^2$.
此引理在第四节的舍入误差分析中至关重要。
3.2 层归一化 (Layer Normalization)
RMS 层归一化分析。首先分析 RMS 层归一化 $R(x)$。引理 3.6 推导了其雅可比矩阵 $J_R(x)$、特征分解、范数界限以及各种条件数。关键结论是其分量条件数有上界2,即 $\kappa_{c,c}(R, x) \le 2$,这表明 RMS 归一化本身是数值稳定的。
$\kappa_{c,c}(\mathcal{R}, \mathbf{x}) = 2\left(1 - \frac{\|\mathbf{x}\|_{- \infty}^{2}}{\|\mathbf{x}\|_{2}^{2}}\right) \leq 2.$
标准层归一化分析。标准层归一化 $N(x)$ 可以看作是中心化函数 $\mathring{c}(x) = x - \mu(x)e$ 和 RMS 归一化 $R(x)$ 的复合,即 $N(x) = R(\mathring{c}(x))$。引理 3.7 分析了中心化函数 $\mathring{c}$,引理 3.8 分析了复合后的 $N(x)$。与 RMS 归一化不同,标准层归一化的分量条件数依赖于 $\frac{\|x\|_\infty}{\|\mathring{c}(x)\|_{-\infty}}$,当输入向量 $x$ 的分量值彼此接近时(即 $\|\mathring{c}(x)\|_{-\infty}$ 很小),其条件数会变得非常大,导致数值不稳定。
$$\kappa_{cc}(\mathcal{N}, x) \le \frac{|\hat{v}(x)|^T|x|}{\|\hat{v}(x)\|^2_2} + \max_{1 \le i \le d} \frac{\|x\|_1 + (d-2)|x_i|}{\|\hat{v}(x)\|_1} < \frac{3\|x\|_\infty}{\|\hat{v}(x)\|_\infty}$$
3.3 两层感知机 (Two-layer Perceptron)
可微性分析。由于 ReLU 函数在零点不可微,两层感知机 $M$ 通常不是处处 Fréchet 可微的。引理 3.10 指出,如果 $L_1(x)$ 的所有分量均非零(即 $L_1(x)$ 是通用的),那么 $M$ 在 $x$ 的一个开邻域内是仿射函数,因此是可微的,其雅可比矩阵为 $J_M(x) = A_2(:,\Omega)A_1(\Omega,:)$,其中 $\Omega = \text{supp}_+(L_1(x))$。
3.4 自注意力 (Self-Attention)
相似度分数分析。引理 3.12 分析了相似度分数函数 $S(X)$,推导了其雅可比矩阵 $J_S(X)$ 及其范数和条件数的界限。其分量条件数依赖于 $\frac{|x_i|^\top|B|x_n| + |x_i^\top B||x_n|}{|x_i^\top B x_n|}$,其中 $B = \frac{1}{\sqrt{d_k}} W_k^\top W_q$。
Softmax 分析。引理 3.13 分析了 softmax 函数 $\psi(s)$,推导了其雅可比矩阵 $J_\psi(s)$、特征分解、范数界限和条件数。其分量条件数上界为 $2\|s\|_\infty$,表明当输入 $s$ 的范数较大时,softmax 可能变得敏感。
自注意力整体分析。引理 3.14 结合上述结果分析了整个自注意力函数 $A(X)$。其雅可比矩阵 $J_A(X)$ 是一个块下三角矩阵。其分量条件数的上界可以简化为:
$$ e_{\text{out}}(A,X) \leq C(A,X) + 40 \| B \|_{1,2} \max_{1 \leq i \leq k_B} \| x_i \| \quad C(A,X) = \max_{1 \leq i \leq d, i \notin \mathcal{A}} \frac{\| W_i \cdot X_i(S(X_i)) \|}{\| W_i \cdot X_i(S(X_i)) \|} \geq 1 $$
其中 $\xi(A, X)$ 是一个与输出和中间值相关的因子,而 $B = \frac{1}{\sqrt{d_k}}W_k^\top W_q$。
归一化的重要性。注 3.15 强调了归一化在 Transformer 中的重要性。上述自注意力的条件数界限表明,如果输入矩阵 $X$ 的列范数很大,条件数可能会很高。然而,当在自注意力之前使用(RMS)层归一化时,输入到自注意力的向量的范数得到控制($\|N^*(X)_t\|_2^2=d$ 或 $\|R^*(X)_t\|_2^2=d$),从而可以得到一个几乎统一的条件数上界,这大大增强了数值稳定性。
A2 方法细节
本节对深度 Transformer (7) 和 (10) 的前向舍入误差进行确定性分析。所有陈述均在浮点运算模型 4.1 下进行。
4.1 预备知识
浮点运算模型。本文采用标准的浮点运算模型(模型 4.1),其中假定基本算术运算(+,-,×,÷)以及指数和平方根函数,其计算结果 fl(·) 与真实结果之间满足关系 $fl(x \text{ op } y) = (x \text{ op } y)(1+\delta)$,其中 $|\delta| \le u$,$u$ 为单位舍入误差。
$$
\begin{array}{rcll}
\mathbf{fl}(x \text{ op } y) & = & (x \text{ op } y)(1+\delta_1), & |\delta_1| \leq \mathbf{u}, \\
\mathbf{fl}(\exp(x)) & = & \exp(x)(1+\delta_2), & |\delta_2| \leq \mathbf{u}, \\
\mathbf{fl}(\sqrt{y}) & = & \sqrt{y}(1+\delta_3), & |\delta_3| \leq \mathbf{u}, \quad y \geq 0.
\end{array}
$$
误差累积。对于 $n$ 次连续运算产生的误差,乘积项 $\prod_{k=1}^n (1+\delta_k)$ 可以通过引理 4.2 进行简化:如果 $nu < 1$,则 $\prod_{k=1}^n(1+\delta_k)^{\rho_k} = 1+\theta_n$,其中 $|\theta_n| \le \gamma_n = \frac{nu}{1-nu}$。
算法稳定性。本文关注前向舍入误差,即寻求对相对分量误差 $\rho_c(fl(f(x)), f(x))$ 的界定。分析将基于标准数值分析中对求和、矩阵-向量乘法等基本操作的舍入误差界限(引理 4.3),这些操作具有后向稳定性。
4.2 层归一化
RMS 层归一化的舍入误差。引理 4.6 分析了 RMS 层归一化 $R(x)$ 的前向稳定性。证明过程通过追踪计算过程中的每一步浮点运算,包括求范数、除法和乘法,最终得出其相对分量误差界限为 $\rho_c(fl(R(x)), R(x)) \le (\frac{d}{2}+3)u + O(u^2)$。这个界限不依赖于输入 $x$ 的值,表明其无条件的数值稳定性。
$\rho_c(\text{fl}(\mathcal{R}(x)), \mathcal{R}(x)) \leq \left(\frac{d}{2}+3\right) \mathbf{u}+\mathcal{O}\left(\mathbf{u}^2\right)$
受扰动输入的误差传播。当函数 $f$ 的输入 $\hat{x}$ 本身就是先前计算的结果(即 $\hat{x}=fl(g(z))$),总误差不仅包含 $f$ 本身的计算误差,还包括由输入扰动 $\rho_c(\hat{x}, x)$ 传播的误差。引理 4.8 结合了函数的条件数和本身的计算误差,给出了总误差的界限:
$\mu_N(R(x), f(x)) \leq \mu_N(R(x), f(x)) + \omega_\mathcal{L}(f, x)\sigma(x, x) + O(u^2 + \sigma(x, x^2))$
推论 4.9 将此应用于 RMS 归一化,得到其在输入受扰动时的误差界限。由于其条件数有界,传播的误差项也被良好控制。
标准层归一化的舍入误差。引理 4.10 首先分析了中心化函数 $\mathring{c}(x)$ 的前向稳定性,其误差界限依赖于 $\frac{\|x\|_1}{\|\mathring{c}(x)\|_{-\infty}}$。这表明当输入 $x$ 的各分量值非常接近时(即分母很小),中心化操作的相对误差可能很大。
标准层归一化的条件稳定性。引理 4.11 进一步分析了整个标准层归一化 $N(x) = R(\mathring{c}(x))$ 的误差。由于中心化步骤的潜在不稳定性,只有当 $\|\mathring{c}(x)\|_{-\infty}$ 足够大时(即输入向量的各分量不那么集中),$N(x)$ 才是前向稳定的。其误差界限同样依赖于 $\frac{\|x\|_1}{\|\mathring{c}(x)\|_{-\infty}}$。推论 4.12 扩展了这一分析到受扰动输入的情况,结论类似。
4.3 两层感知机
MLP的舍入误差。引理 4.14 分析了两层感知机 $M(x)$ 的后向稳定性。由于 ReLU 的存在,分析要求 $L_1(x)$ 的输出“足够通用”,即 $|L_1(x)|$ 的每个分量都必须大于计算 $L_1(x)$ 时的绝对前向误差界限。这个条件保证了在浮点数计算中,fl(L1(x)) 的符号与真实值 L1(x) 的符号一致,从而使得 ReLU 的作用区域不变。在此条件下,可以得到一个后向误差界限,最终的前向误差界限为:
$p_e(\mathbf{n}, \mathcal{M}(x), \mathcal{M}(x)) \le \xi(\mathcal{M}, x)(d + |\Omega| + 2\mathbf{u} + \mathcal{O}\mathbf{u}^2),$
其中 $\xi(M,x)$ 是一个与 $M(x)$ 的分量大小相关的因子。推论 4.16 将此分析扩展到输入受扰动的情况。
4.4 自注意力
相似度分数和 Softmax 的误差。引理 4.17 分析了相似度分数 $S(X)$ 的混合前向-后向稳定性,其误差界限为 $\rho_c(fl(S(X)), S(X)) \le \xi(S,X)(3d+1)u + O(u^2)$,其中 $\xi(S,X)$ 因子在 $W_k^\top W_q$ 非奇异时,与该矩阵的条件数成正比(注4.18)。引理 4.19 引用了【4,Blanchard et al., 2021, IMA J Numer Anal】的结果,给出了 softmax 函数 $\psi(s)$ 的前向稳定误差界限。推论 4.20 结合条件数分析了当 softmax 输入受扰动时的误差传播。
自注意力整体误差分析。引理 4.21 综合了以上结果,对整个自注意力函数 $A(X)$ 进行了混合前向-后向稳定性分析。其误差界限受到多个因素的共同影响,包括输入 $X$ 的范数、相似度分数 $S(X_t)$ 的范数、以及 $\xi(S, X_t)$ 因子。注 4.22 简化了此界限,明确指出了它与 $W_k^\top W_q$ 的条件数以及输入 $X$ 的列范数 $\max_t \|x_t\|_2^2$ 的关系。
$$ \rho_b[\mathcal{A}(X), \mathcal{A}(X)] \leq \xi(\mathcal{A}, X)(2\alpha + 2\delta) \left(1 + \frac{2}{\sqrt{d}} \sigma_{\max}(W_b^T W_b) \frac{\sigma_{\max}(W_b^T W_b)}{\sigma_{\min}(W_b^T W_b)} \max_i \{s_i(\frac{1}{2})\} \right) u + O(\alpha^2). $$
推论 4.23 将此分析扩展到受扰动输入的情况。
归一化对自注意力误差的影响。注 4.24 进一步阐述了归一化的作用。自注意力的舍入误差界限表明,当输入列的范数较大时,计算会变得不准确。图 3 的数值实验验证了这一点,显示了当输入范数增大时,舍入误差呈现最坏情况下的二次增长。图 4 则通过实验对比了预归一化(本文模型)和后归一化(将归一化层放在自注意力之后)的误差累积情况,结果表明预归一化能有效控制误差增长,而没有预归一化的模型误差会随层数增加而急剧放大。
4.5 Transformer
Transformer 模块误差分析。本节将各子模块的分析组合起来,推导整个 Transformer 模块的误差界限。
- 自注意力子模块:引理 4.25 (RMS) 和 4.26 (标准LN) 分析了包含残差连接的自注意力子模块 $X + A(\text{Norm}(X))$ 的误差。误差界限是先前推导的归一化误差和自注意力误差的组合。
- 前馈网络子模块:引理 4.27 (RMS) 和 4.28 (标准LN) 以同样方式分析了前馈网络子模块 $x + M(\text{Norm}(x))$ 的误差。
- 完整 Transformer 模块:引理 4.29 将上述两个子模块的误差界限串联起来,给出了单个 Transformer 模块(RMS 版本的 $Q_\Theta$ 和标准LN版本的 $T_\Theta$)在前向传播中的舍入误差上界。这个界限明确地展示了误差如何受到 $W_k^\top W_q$ 的条件数和范数、以及标准LN中与中心化相关的项的影响。
深度 Transformer 的总误差。定理 4.30 是本文的核心结果,它通过递归应用单层模块的误差分析,推导了经过 L 个 Transformer 模块后的总累积舍入误差。
1. 对于使用 RMS 归一化的 Transformer ($Q_\Theta$),其输出的相对分量误差界限为:
$$ \rho_k(m(X_k^{(t)}), Q_k(X_k^{(t)})) \le \prod_{i=1}^k \left[ u \left( 1 + \sqrt{d_{max}(W_k^{(t)} \vert W_k^{(t)}) } \frac{\sigma_{max}(W_k^{(t)} \vert W_k^{(t)})}{\sigma_{min}(W_k^{(t)} \vert W_k^{(t)})} \right) \right] \cdot \frac{e^u}{u^k} (16m + 23d + 17u + O(u^t)) $$
其中 $\chi_l$ 是一个与第 $l$ 层的中间输出相关的因子。
2. 对于使用标准层归一化的 Transformer ($T_\Theta$),其误差界限形式类似,但包含了额外的与中心化向量范数相关的项,反映了其潜在的不稳定性:
$$
\rho(\mathbf{n}, \mathcal{T}_{\theta}(\mathbf{X}^{(0)}), \mathcal{T}_{\theta}(\mathbf{X}^{(0)})) \leq \prod_{k=1}^{L}\left[\kappa\left(1+4 \sqrt{d_{\max}(\mathbf{W}_{k}^{(0)^{T}} \mathbf{W}_{k}^{(0)})} \frac{\sigma_{\max}((\mathbf{W}_{k}^{(0)^{T}} \mathbf{W}_{k}^{(0)}))}{\sigma_{\min}((\mathbf{W}_{k}^{(0)^{T}} \mathbf{W}_{k}^{(0)}))}\right)\right. \\
\left.\left(1+\max_{1 \leq i \leq n} \frac{|\mathcal{E}_{k}^{(i)}-U_{k}^{(i)}|}{\beta^{(i)}(\mathbf{X}_{k}^{(i)})|_{\infty}}\right)\left(1+\max_{1 \leq i \leq n} \frac{|\mathcal{E}_{k}^{(i)}-U_{k}^{(i)}|}{|\mathbf{Y}_{k}^{(i)}|_{\infty}}\right)\right] \frac{16^{L}}{15}[8n+3d+D]u+\mathcal{O}(u^2).
$$
注 4.31 指出,这个理论界限的指数增长特性以及对 $W_k^\top W_q$ 属性的依赖性,与图 1 和图 2 的数值实验结果一致,同时也反映了该分析的最坏情况性质。
A4 实验
本文没有独立的实验章节,所有数值实验均穿插在引言和方法章节中,用以验证和说明理论分析的结果。
实验环境
- 通用设置:所有实验均为数值模拟,通过将计算的中间结果舍入到指定位数(digits)来模拟低精度浮点运算,而非在真实硬件上训练。
- 硬件/软件:未具体说明,为通用数值计算环境。
- 实验1 (图1):
- 模型与参数:深度 Transformer,层数 $L=40$,维度 $d, n, D = 20$。
- 权重初始化:自注意力矩阵 $W_q, W_k, W_v$ 从标准正态分布 $N(0,1)$ 采样;MLP 矩阵从 $N(0, 1/d)$ 采样。为了放大效应,通过对角矩阵乘法增大了 $W_k^\top W_q$ 的条件数。
- 输入数据:输入向量从标准正态分布 $N(0,1)$ 采样。
- 统计:结果基于 5000 次不同的随机初始化。
- 实验2 (图2):
- 模型与参数:与实验1类似,但层数 $L$ 分别为 10, 15, 20。
- 变量控制:通过一个缩放因子 $\lambda$ 控制 $W_k^\top W_q$ 的谱范数,即 $\|W_k^\top W_q\|_{2,2} \sim \lambda$,以研究误差与其谱范数的关系。
- 实验3 (图3):
- 模型与参数:单个自注意力层,维度 $d, n = 10$。$W_q, W_k, W_v$ 均为单位矩阵。
- 输入数据:输入 $X$ 从均值为 1、方差为 0.01 的正态分布采样。通过控制输入的列范数来研究其对误差的影响。
- 统计:结果基于 1000 次初始化。
- 实验4 (图4):
- 模型与参数:深度 Transformer,层数最高到 100,维度 $d, n, D = 10$。所有权重矩阵从 $N(0, 0.1)$ 采样。
- 对比项:对比了两种架构:(a) 预归一化(本文分析的模型);(b) 后归一化(将LN置于自注意力之后)。
- 输入数据:输入向量从标准正态分布 $N(0,1)$ 采样。
- 统计:结果基于 1000 次初始化。
实验结果
-
误差随深度的增长 (图1):实验结果表明,平均相对舍入误差随着 Transformer 层数 $L$ 的增加呈近指数增长,这与理论分析的最坏情况界限一致。然而,中位数误差远低于平均值,且误差分布直方图显示大误差是稀疏事件,说明理论界限是悲观的,但在趋势上是正确的。
-
误差与 $W_k^\top W_q$ 谱范数的关系 (图2):实验证实,舍入误差随着 $\|W_k^\top W_q\|_{2,2}$ 的增大而增大。对于不同的层数 $L$,误差增长率不同($L=10$ 时近似线性,$L=20$ 时近似二次),这有力地支持了理论界限中关于 $W_k^\top W_q$ 的谱范数和条件数是误差关键影响因素的结论。
-
误差与输入范数的关系 (图3):对于单层自注意力,舍入误差随输入最大列范数的增加而增加,其增长趋势接近二次方,这验证了理论分析中指出的输入范数对数值稳定性的影响(注 4.24)。
-
归一化位置的影响 (图4):实验清晰地展示了预归一化(pre-normalization)的优越性。采用预归一化的模型 (a) 的误差增长缓慢且稳定,而采用后归一化(post-normalization)的模型 (b) 的误差随层数增加急剧累积,很快达到不可接受的水平。这为在 Transformer 架构中采用预归一化提供了强有力的经验证据。
A5 结论
本文对 Transformer 架构的前向传播过程进行了严谨的确定性舍入误差分析。通过分析 Transformer 的各个组成部分(层归一化、多层感知机、自注意力)的条件数和数值稳定性,最终推导出了一个适用于深度 Transformer 网络的端到端相对误差上界。
核心结论:
1. Transformer 前向传播的舍入误差会随着网络深度 L 呈指数级增长。
2. 自注意力模块中的键-查询矩阵乘积 $W_k^\top W_q$ 的谱范数和条件数是控制误差增长的关键因素。较大的谱范数和条件数会显著放大舍入误差。
3. 均方根层归一化(RMS Layer Normalization) 在数值上比标准层归一化更稳定,后者在特定输入下可能导致误差爆炸。
4. 将归一化层置于自注意力和前馈网络之前(预归一化),对于维持整个网络的数值稳定性至关重要。
实践意义:
本文的理论和实验结果为设计和训练更稳定的 LLMs 提供了具体指导。为了减轻舍入误差的影响,建议:
- 在训练过程中对 $W_k^\top W_q$ 的谱范数和条件数进行正则化或约束。
- 在执行自注意力计算(尤其是相似度分数和 softmax)时,考虑使用更高的计算精度。
未来工作展望:
本文的工作可从以下几个方面扩展:
- 将误差分析扩展到反向传播和整个训练过程,以更全面地理解训练不稳定性。
- 采用随机舍入(stochastic rounding)框架来优化误差界限,使其更贴近实际平均情况。
- 将分析扩展到支持混合精度计算的场景,这在现代 LLM 训练中非常普遍。
- 分析其他 Transformer 变体,例如多头注意力(multi-head attention)。
💬 评论讨论
欢迎在这里分享您的想法和见解!