PyTorch LayerNorm推理阶段均值和标准差为何不固定？如何提取？

我尝试从torch.LayerNorm的输出还原原始输入，原本以为LayerNorm在推理阶段的均值和标准差是固定的，打算提取这些参数完成输入还原。我成功提取出针对原输入的融合权重和偏置（和LayerNorm自身的weight、bias参数不同），用它能还原原输入，但将其应用到新输入张量时结果完全不符。我推测LayerNorm为新输入重新计算了均值和标准差，但疑惑为何推理阶段这些值不固定。

以下是我的代码及输出：

layer = layer().eval()
with torch.inference_mode():
    out = layer(input_data)

w = torch.zeros(len(out[0, :, 0]))
b = torch.zeros(len(out[0, :, 0]))

for i in range(len(out[0, :, 0])):
    w[i] = (input_data[0, i, 0] - input_data[0, i, 10]) / (out[0, i, 0] - out[0, i, 10])
    b[i] = (input_data[0, i, 0] * out[0, i, 10] - input_data[0, i, 10] * out[0, i, 0]) / (out[0, i, 10] - out[0, i, 0])

for i1 in range(len(input_remade[0, :, 0])):
    input_remade[0, i1, :] = out[0, i1, :] * w[i1] + b[i1]
print(torch.sum(input_remade - input_data))


input_data2 = torch.randn(1, 577, 768)
input_remade2 = torch.randn(1, 577, 768)
with torch.inference_mode():
    out2 = layer(input_data2)

for i1 in range(len(input_remade2[0, :, 0])):
    input_remade2[0, i1, :] = out2[0, i1, :] * w[i1] + b[i1]
print(torch.sum(input_remade2 - input_data2))

w1 = torch.zeros(len(out2[0, :, 0]))
b1 = torch.zeros(len(out2[0, :, 0]))

for i in range(len(out2[0, :, 0])):
    w1[i] = (input_data2[0, i, 0] - input_data2[0, i, 10]) / (out2[0, i, 0] - out2[0, i, 10])
    b1[i] = (input_data2[0, i, 0] * out2[0, i, 10] - input_data2[0, i, 10] * out2[0, i, 0]) / (out2[0, i, 10] - out2[0, i, 0])

for i1 in range(len(input_remade2[0, :, 0])):
    input_remade2[0, i1, :] = out2[0, i1, :] * w1[i1] + b1[i1]
print(torch.sum(input_remade2 - input_data2))

输出：

tensor(-0.0061)
tensor(1280.9966)
tensor(0.0014)

同时想请教：是否有方法从LayerNorm层提取固定的均值和标准差？

解答 #

核心结论：LayerNorm不存在固定的均值/标准差 #

LayerNorm和BatchNorm的核心区别在于：无论训练还是推理阶段，LayerNorm都会针对每个输入样本的指定归一化维度，实时计算均值和标准差，而非像BatchNorm那样在训练阶段统计全局滑动平均的均值/方差，推理时复用固定值。

这是LayerNorm的设计目标决定的——它是对单个样本内部的特征维度做归一化（比如NLP场景中对每个token的特征向量归一化），每个输入样本的特征分布不同，计算出的均值和标准差自然不同，不存在可以提前提取的固定全局参数。

你的代码为何只对原输入有效 #

你计算的w和b，本质是基于原输入的均值/方差，结合LayerNorm的weight（$\gamma$）和bias（$\beta$）推导的线性变换参数，仅适配原输入的分布。当输入换成input_data2时，其均值/方差已经改变，用原参数自然无法还原新输入；而重新计算w1和b1后能还原，正是因为这组参数匹配了新输入的实时均值/方差。

正确的输入还原方式 #

LayerNorm的正向计算逻辑为：
$$ out = \gamma * \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta $$
其中$\mu$是输入x在归一化维度上的均值，$\sigma^2$是方差，$\gamma$是LayerNorm的weight，$\beta$是bias，$\epsilon$是防止除零的小常数。

要从out还原x，必须使用对应输入的$\mu$和$\sigma^2$，公式为：
$$ x = \frac{(out - \beta)}{\gamma} * \sqrt{\sigma^2 + \epsilon} + \mu $$
这里的$\mu$和$\sigma^2$是当前输入的实时计算值，无法提前获取固定值。如果需要还原输入，你需要在计算LayerNorm输出的同时，同步记录该输入对应的均值和方差。

内容的提问来源于stack exchange，提问作者Hibama