实践中layernormal的计算方式

实践中layernormal的计算方式

LayerNorm（Layer Normalization，层归一化）需要对一组元素计算均值和方差，再完成归一化。数学公式很简单，但在 GPU kernel 中，如何计算统计量会影响运算速度、显存读取次数和数值精度。

实践中常见的三种方法是：Welford 算法、平方和法，以及先求均值再计算中心化平方和的两阶段方法。本文说明它们的计算过程、公式证明和实际取舍。

1. LayerNorm 要计算什么

假设一组需要归一化的数据为 x1,x2,...,xNx_1,x_2,\ldots,x_Nx1,x2,...,xN。实际张量可能沿最后一个或多个维度归一化，这里统一将其视为包含 NNN 个元素的一行。

均值为：

μ=1N∑i=1Nxi \mu=\frac{1}{N}\sum_{i=1}^{N}x_i μ=N1i=1∑Nxi

方差为：

σ2=1N∑i=1N(xi−μ)2 \sigma^2=\frac{1}{N}\sum_{i=1}^{N}(x_i-\mu)^2 σ2=N1i=1∑N(xi−μ)2

最终输出为：

yi=(xi−μ)1σ2+ϵγi+βi y_i=(x_i-\mu)\frac{1}{\sqrt{\sigma^2+\epsilon}}\gamma_i+\beta_i yi=(xi−μ)σ2+ϵ 1γi+βi

这里使用总体方差，分母是 NNN，不是样本方差中的 N−1N-1N−1。ϵ\epsilonϵ 用于避免分母为零，γi\gamma_iγi 和 βi\beta_iβi 是可学习的缩放与偏移参数。

即使输入是 FP16 或 BF16，统计量也通常使用 FP32 累加。

2. 两阶段中心化法：先求均值，再求方差

这是一种直接按定义计算的方法。

第一阶段计算均值：

μ=1N∑ixi \mu=\frac{1}{N}\sum_i x_i μ=N1i∑xi

第二阶段计算相对于均值的平方偏差：

σ2=1N∑i(xi−μ)2 \sigma^2=\frac{1}{N}\sum_i(x_i-\mu)^2 σ2=N1i∑(xi−μ)2

示意代码如下：

float sum = 0.0f;
for (int i = 0; i < N; ++i) {
    sum += float(input[i]);
}
float mean = sum / N;

float m2 = 0.0f;
for (int i = 0; i < N; ++i) {
    float delta = float(input[i]) - mean;
    m2 += delta * delta;
}
float variance = m2 / N;

它的优点是直观，并且避免了两个很大的平方量相减。均值和平方偏差和都可以通过并行归约计算。

它需要两阶段处理数据，但这不等于必须从显存读取输入两次。如果整行输入能够留在寄存器或 shared memory 中，kernel 可以先算均值，再复用缓存的输入计算方差和输出。若无法缓存，则可能需要重新读取输入。

因此，这种方法是否高效，取决于行长度、缓存方式、寄存器占用和归约成本。

3. 平方和法：同时累加 sum 和 sum_sq

平方和法在一次遍历中维护两个量：

S=∑ixi,Q=∑ixi2 S=\sum_i x_i,\qquad Q=\sum_i x_i^2 S=i∑xi,Q=i∑xi2

然后计算：

μ=SN,σ2=QN−μ2 \mu=\frac{S}{N},\qquad \sigma^2=\frac{Q}{N}-\mu^2 μ=NS,σ2=NQ−μ2

3.1 平方和公式的证明

从方差定义展开：

σ2=1N∑i(xi−μ)2=1N∑i(xi2−2μxi+μ2)=1N∑ixi2−2μN∑ixi+1N∑iμ2. \begin{aligned} \sigma^2 &=\frac{1}{N}\sum_i(x_i-\mu)^2\\ &=\frac{1}{N}\sum_i(x_i^2-2\mu x_i+\mu^2)\\ &=\frac{1}{N}\sum_i x_i^2 -\frac{2\mu}{N}\sum_i x_i +\frac{1}{N}\sum_i\mu^2. \end{aligned} σ2=N1i∑(xi−μ)2=N1i∑(xi2−2μxi+μ2)=N1i∑xi2−N2μi∑xi+N1i∑μ2.

因为 ∑ixi=Nμ\sum_i x_i=N\mu∑ixi=Nμ，而 μ\muμ 对所有元素都是同一个常数，所以 ∑iμ2=Nμ2\sum_i\mu^2=N\mu^2∑iμ2=Nμ2。代入得到：

σ2=1N∑ixi2−μ2 \boxed{\sigma^2=\frac{1}{N}\sum_i x_i^2-\mu^2} σ2=N1i∑xi2−μ2

这个等式在精确算术下完全成立。

3.2 kernel 中如何计算

float sum = 0.0f;
float sum_sq = 0.0f;

for (int i = 0; i < N; ++i) {
    float x = float(input[i]);
    sum += x;
    sum_sq += x * x;
}

float mean = sum / N;
float variance = sum_sq / N - mean * mean;

GPU 上可以让各线程处理一部分元素，得到局部的 sum 和 sum_sq，然后分别求和归约。它的统计量简单，合并成本较低，适合并行计算。

3.3 为什么可能算出负方差

数学上的等价，不意味着浮点运算的结果也完全等价。

记：

A=QN,B=μ2. A=\frac{Q}{N},\qquad B=\mu^2. A=NQ,B=μ2.

浮点计算得到的两个量带有误差：

A^=A+eA,B^=B+eB. \widehat A=A+e_A,\qquad \widehat B=B+e_B. A =A+eA,B =B+eB.

暂时忽略最后一次减法本身的舍入，算出的方差为：

σ2^=σ2+eA−eB. \widehat{\sigma^2}=\sigma^2+e_A-e_B. σ2 =σ2+eA−eB.

当真实方差很小，而 eB−eA>σ2e_B-e_A>\sigma^2eB−eA>σ2 时，结果就可能为负数。浮点计算并不保证分别算出的 A^\widehat AA 和 B^\widehat BB 仍满足 A^≥B^\widehat A\ge\widehat BA ≥B 。

例如：

x1=10000,x2=10002. x_1=10000,\qquad x_2=10002. x1=10000,x2=10002.

精确计算得到：

μ=10001,σ2=1. \mu=10001,\qquad \sigma^2=1. μ=10001,σ2=1.

但平方和公式需要计算：

QN=100020002,μ2=100020001. \frac{Q}{N}=100020002,\qquad \mu^2=100020001. NQ=100020002,μ2=100020001.

两个约为 10810^8108 的数相减，才能得到 111。FP32 在这个数量级上相邻可表示数的间隔为 888，无法保留这样的个位差异。按上述普通 FP32 计算顺序，这个例子可以得到 000；其他输入和累加顺序还可能产生负值。

关键在于：平方、累加、除法等步骤已经引入误差。最后一次减法即使本身精确，也无法恢复此前丢失的信息。两个接近的大数相减后，共同的大部分被抵消，留下的误差相对于真实方差可能非常大。这就是灾难性消减。

若将输入写成 xi=C+ϵix_i=C+\epsilon_ixi=C+ϵi，其中 CCC 很大而 ϵi\epsilon_iϵi 的波动很小，则真实方差只与 ϵi\epsilon_iϵi 有关。平方和法却先计算包含 C2C^2C2 的大数，再通过相减得到小方差，因此尤其容易丢失精度。

使用 FP32 累加、树形归约或融合乘加可以改善部分误差，但不能根除这个问题。将结果截断为 max(variance, 0) 可以避免负方差引发后续异常，却无法恢复丢失的方差。加上 ϵ\epsilonϵ 同样不能修复统计量的计算误差。

4. Welford：一次遍历更新均值和平方偏差和

Welford 在处理数据时维护三个量：

• 已处理元素数量 nnn。
• 当前均值 μn\mu_nμn。
• 当前平方偏差和 M2,n=∑i=1n(xi−μn)2M_{2,n}=\sum_{i=1}^{n}(x_i-\mu_n)^2M2,n=∑i=1n(xi−μn)2。

加入新元素 xxx 时，令 δ=x−μn\delta=x-\mu_nδ=x−μn，更新公式为：

μn+1=μn+δn+1 \boxed{\mu_{n+1}=\mu_n+\frac{\delta}{n+1}} μn+1=μn+n+1δ

M2,n+1=M2,n+δ(x−μn+1) \boxed{M_{2,n+1}=M_{2,n}+\delta(x-\mu_{n+1})} M2,n+1=M2,n+δ(x−μn+1)

处理完所有元素后：

σ2=M2,NN. \sigma^2=\frac{M_{2,N}}{N}. σ2=NM2,N.

4.1 均值更新的证明

根据均值定义：

μn+1=nμn+xn+1=μn+x−μnn+1=μn+δn+1. \begin{aligned} \mu_{n+1} &=\frac{n\mu_n+x}{n+1}\\ &=\mu_n+\frac{x-\mu_n}{n+1}\\ &=\mu_n+\frac{\delta}{n+1}. \end{aligned} μn+1=n+1nμn+x=μn+n+1x−μn=μn+n+1δ.

记均值移动量为：

a=μn+1−μn=δn+1. a=\mu_{n+1}-\mu_n=\frac{\delta}{n+1}. a=μn+1−μn=n+1δ.

4.2 旧元素的平方偏差和如何变化

均值改变后，每个旧元素相对于新均值的偏差变成：

xi−μn+1=(xi−μn)−a. x_i-\mu_{n+1}=(x_i-\mu_n)-a. xi−μn+1=(xi−μn)−a.

展开平方并求和：

∑i=1n(xi−μn+1)2=∑i=1n(xi−μn)−a2=M2,n−2a∑i=1n(xi−μn)+na2. \begin{aligned} \sum_{i=1}^{n}(x_i-\mu_{n+1})^2 &=\sum_{i=1}^{n}(x_i-\\mu_n)-a^2\\ &=M_{2,n}-2a\sum_{i=1}^{n}(x_i-\mu_n)+na^2. \end{aligned} i=1∑n(xi−μn+1)2=i=1∑n(xi−μn)−a2=M2,n−2ai=1∑n(xi−μn)+na2.

旧元素相对于旧均值的偏差之和为零：

∑i=1n(xi−μn)=∑i=1nxi−nμn=0. \sum_{i=1}^{n}(x_i-\mu_n)=\sum_{i=1}^{n}x_i-n\mu_n=0. i=1∑n(xi−μn)=i=1∑nxi−nμn=0.

因此，旧元素相对于新均值的平方偏差和为：

M2,n+na2. M_{2,n}+na^2. M2,n+na2.

这解释了为什么不必重新遍历旧元素：均值移动对它们造成的总变化可以直接用 na2na^2na2 表示。

4.3 加入新元素的贡献

新元素相对于新均值的偏差为：

x−μn+1=δ−a=nδn+1. x-\mu_{n+1}=\delta-a=\frac{n\delta}{n+1}. x−μn+1=δ−a=n+1nδ.

所以新的平方偏差和为：

M2,n+1=M2,n+na2+(x−μn+1)2=M2,n+nδ2(n+1)2+n2δ2(n+1)2=M2,n+n(n+1)δ2(n+1)2=M2,n+nn+1δ2. \begin{aligned} M_{2,n+1} &=M_{2,n}+na^2+(x-\mu_{n+1})^2\\ &=M_{2,n}+\frac{n\delta^2}{(n+1)^2} +\frac{n^2\delta^2}{(n+1)^2}\\ &=M_{2,n}+\frac{n(n+1)\delta^2}{(n+1)^2}\\ &=M_{2,n}+\frac{n}{n+1}\delta^2. \end{aligned} M2,n+1=M2,n+na2+(x−μn+1)2=M2,n+(n+1)2nδ2+(n+1)2n2δ2=M2,n+(n+1)2n(n+1)δ2=M2,n+n+1nδ2.

由于 x−μn+1=nδ/(n+1)x-\mu_{n+1}=n\delta/(n+1)x−μn+1=nδ/(n+1)，最终得到：

M2,n+1=M2,n+δ(x−μn+1). \boxed{M_{2,n+1}=M_{2,n}+\delta(x-\mu_{n+1})}. M2,n+1=M2,n+δ(x−μn+1).

从第一个元素的 μ1=x1\mu_1=x_1μ1=x1、M2,1=0M_{2,1}=0M2,1=0 开始，这个递推在精确算术下始终保持 M2M_2M2 等于已处理元素的平方偏差和。

4.4 实现示例

int count = 0;
float mean = 0.0f;
float m2 = 0.0f;

for (int i = 0; i < N; ++i) {
    float x = float(input[i]);
    ++count;
    float delta = x - mean;
    mean += delta / count;
    float delta2 = x - mean;
    m2 += delta * delta2;
}

float variance = m2 / N;

这里的 delta 相对于旧均值，delta2 相对于新均值。二者的乘积同时计入了旧元素因均值移动产生的变化，以及新元素的贡献。

Welford 避免了最后用两个很大的平方量相减，因此通常比平方和法稳定。它仍然存在舍入误差，也无法恢复输入量化时已经丢失的差异。

5. Welford 如何在 GPU 上并行

上面的代码是串行示例。实际 GPU kernel 可以让各线程处理一部分元素，再合并局部统计量。

假设两组数据的统计量分别为 (nA,μA,M2,A)(n_A,\mu_A,M_{2,A})(nA,μA,M2,A) 和 (nB,μB,M2,B)(n_B,\mu_B,M_{2,B})(nB,μB,M2,B)。令：

n=nA+nB,Δ=μB−μA. n=n_A+n_B,\qquad \Delta=\mu_B-\mu_A. n=nA+nB,Δ=μB−μA.

合并公式为：

μ=μA+ΔnBn \mu=\mu_A+\Delta\frac{n_B}{n} μ=μA+ΔnnB

M2=M2,A+M2,B+Δ2nAnBn. \boxed{M_2=M_{2,A}+M_{2,B}+\Delta^2\frac{n_A n_B}{n}}. M2=M2,A+M2,B+Δ2nnAnB.

这个公式也可以用"均值移动"证明。A 组均值移动量为 μ−μA=ΔnB/n\mu-\mu_A=\Delta n_B/nμ−μA=ΔnB/n，B 组为 μ−μB=−ΔnA/n\mu-\mu_B=-\Delta n_A/nμ−μB=−ΔnA/n。各组内部偏差之和为零，因此合并后的平方偏差和是：

M2=M2,A+nA(μA−μ)2+M2,B+nB(μB−μ)2=M2,A+M2,B+Δ2nAnB2+nBnA2n2=M2,A+M2,B+Δ2nAnBn. \begin{aligned} M_2 &=M_{2,A}+n_A(\mu_A-\mu)^2 +M_{2,B}+n_B(\mu_B-\mu)^2\\ &=M_{2,A}+M_{2,B} +\Delta^2\frac{n_A n_B^2+n_B n_A^2}{n^2}\\ &=M_{2,A}+M_{2,B}+\Delta^2\frac{n_A n_B}{n}. \end{aligned} M2=M2,A+nA(μA−μ)2+M2,B+nB(μB−μ)2=M2,A+M2,B+Δ2n2nAnB2+nBnA2=M2,A+M2,B+Δ2nnAnB.

令 B 组只有新元素 xxx，即 nB=1n_B=1nB=1、μB=x\mu_B=xμB=x、M2,B=0M_{2,B}=0M2,B=0，就得到前面的逐元素更新公式。

并行归约的过程可以理解为：

各线程处理自己的元素
        ↓
得到局部 (count, mean, M2)
        ↓
在 warp 或 block 内树形合并
        ↓
得到整行均值和方差
        ↓
计算归一化输出

浮点合并顺序会影响末尾精度，因此不同并行布局的结果可能存在微小差异。

6. 一次遍历统计，不等于一次读取完成 LayerNorm

Welford 和平方和法都能在一次遍历中计算统计量，但每个元素的最终输出依赖整行的均值和方差。统计完成前，通常无法直接确定输出。

如果输入保存在寄存器或 shared memory 中，统计完成后可以复用；如果没有缓存，就需要再次读取输入来计算输出。两阶段中心化法也可以通过缓存实现一次显存读取。

因此，比较 kernel 时，应分别看统计阶段的处理次数、输入的显存读取次数，以及缓存带来的寄存器或 shared memory 成本。

7. 实践中的选择与总结

方法	统计过程	数值稳定性	主要成本与限制
两阶段中心化法	先求均值，再求平方偏差和	通常较好	两阶段归约；能否复用输入取决于缓存
平方和法	同时累加和与平方和	大均值、小方差时较差	运算与合并简单，但有消减风险
Welford	更新并合并均值与平方偏差和	通常较好	状态和合并运算更复杂，逐元素更新有依赖

实践中并非所有 LayerNorm 都使用 Welford。同一个库也可能根据行长度、数据类型和目标硬件选择不同 kernel。

精度优先、需要处理大均值和小方差的输入时，Welford 是常见选择。如果整行输入可以高效保存在寄存器中，两阶段中心化法也值得考虑。平方和法结构简单、可能具有性能优势，但应针对实际输入分布验证误差，不能仅凭代数等价判断它可靠。

测试时，除了普通随机输入，还应覆盖常量输入、大均值与小波动、不同长度和不同数据类型，并与更高精度的参考结果比较。

三种方法的核心区别是：两阶段中心化法直接按定义计算；平方和法通过大数相减得到方差；Welford 通过递推与合并维护平方偏差和。选择时需要同时衡量速度、数值误差和内存成本，没有对所有场景都最优的方法。