LayerNorm 前向实现：数值稳定性陷阱

回答与解析

NumPy 实现

import numpy as np

class LayerNorm:
    def __init__(self, normalized_shape, eps=1e-5):
        self.normalized_shape = tuple(
            normalized_shape if isinstance(normalized_shape, (tuple, list))
            else (normalized_shape,)
        )
        self.eps = eps
        self.gamma = np.ones(self.normalized_shape, dtype=np.float32)
        self.beta = np.zeros(self.normalized_shape, dtype=np.float32)

    def forward(self, x):
        if tuple(x.shape[-len(self.normalized_shape):]) != self.normalized_shape:
            raise ValueError("input tail shape does not match normalized_shape")
        axes = tuple(range(x.ndim - len(self.normalized_shape), x.ndim))

        input_dtype = x.dtype
        input_is_bfloat16 = input_dtype.name == "bfloat16"
        input_is_float = input_is_bfloat16 or np.issubdtype(
            input_dtype, np.floating
        )
        low_precision = input_dtype == np.float16 or input_is_bfloat16
        if low_precision or not input_is_float:
            work = x.astype(np.float32)
        else:
            work = x

        mean = np.mean(work, axis=axes, keepdims=True)
        var = np.mean((work - mean) ** 2, axis=axes, keepdims=True)
        x_hat = (work - mean) / np.sqrt(var + self.eps)
        gamma = self.gamma.astype(work.dtype, copy=False)
        beta = self.beta.astype(work.dtype, copy=False)
        out = x_hat * gamma + beta
        return out.astype(input_dtype, copy=False) if input_is_float else out

关键点

PyTorch LayerNorm 对最后 len(normalized_shape) 个维度共同求均值和方差，不只是最后一维；gamma/beta 的 shape 与 normalized_shape 相同，并通过广播作用于前导 batch 维。使用 mean((x-mean)^2) 比 E[x^2]-E[x]^2 更不易遭受消减误差。epsilon 加在方差后、开根前，用来避免极小方差造成除零或过大缩放；不是因为 sqrt(var) 自己会产生负数。低精度浮点输入可用 FP32 累积统计量后转回输入 dtype；float64 输入不应无条件降为 float32。

口语版讲法(约3分钟)

• 确认 PyTorch 的多维归一化语义
• 计算均值、方差和仿射变换
• 解释 epsilon 与数值稳定性
• 说明 shape 广播和 dtype 边界

手写 LayerNorm 时，我先确认 normalized shape 的语义。它可以是一个整数，也可以是多维 tuple。PyTorch 的定义是在输入最后 len(normalized shape) 个维度上共同计算均值和方差。例如输入是 batch、channel、height、width，normalized shape 如果是 channel、height、width，就要在后三个维度一起归一化；只写 axis=-1 会与框架语义不一致。

实现里我先把 normalized shape 统一成 tuple，再检查输入尾部 shape 是否完全匹配。归一化 axes 就是从 x.ndim-len(shape) 到最后一维。gamma 初始化为一，beta 初始化为零，二者 shape 与 normalized shape 相同，前面的 batch 维会自动广播。

数值计算时，我会按输入 dtype 选择累积精度。FP16 或环境支持的 BF16 先转成 FP32 统计，float32 保持 float32，float64 则保持 float64，不能把所有输入一律降成 FP32。先求均值，再计算 mean((x-mean)^2)。这种两步中心化形式通常比直接用 E[x^2]-E[x]^2 更不容易出现大数相减的消减误差。完成归一化和仿射变换后，浮点输入再按约定转回原 dtype。

epsilon 的作用也要说准确：当方差非常小或为零时，它防止除零和过大的缩放。不是说 sqrt(var) 会因为浮点精度自己变成负数。当前的方差由非负平方项求均值，对有限输入不会产生负方差；只有采用 E[x^2]-E[x]^2 等存在消减误差的公式时，才需要讨论微小负值与 clamp。

这道题只要求 NumPy 前向，所以重点是 axis、shape、广播、累积精度和输出 dtype；若要完全模拟训练框架，还需实现反向传播和参数 dtype 管理。LayerNorm 不依赖 batch 统计，适合序列和小 batch，但这不代表它和 BatchNorm 在图像 Transformer 中混用就天然更好，是否组合要由具体架构和实验决定。

测试代码时，我会准备几类输入：普通二维张量、多维 normalized shape、前置 batch 维可变的张量、所有元素相同的张量，以及 float16 和 float64。结果要和可信框架在相同 epsilon、gamma、beta 下比较，同时检查输出 shape、dtype 和广播是否一致。常量输入尤其能验证 epsilon 与 beta 的行为。

还要防止一个常见实现错误：先把所有轴都求均值，再用错误 shape 的 gamma 做广播，表面上能运行，语义却已经变成全局归一化。正确做法是只归一化尾部指定维，并保留这些维用于广播。若面试官继续问反向，我会从归一化输出对均值和方差的依赖推导，但不会把 NumPy 前向冒充完整训练层。

在 NumPy 实现里还应拒绝整数输入或先明确转换策略，因为归一化结果本质是浮点数。normalized shape 为空、尾部尺寸不匹配或 gamma、beta shape 错误时也应尽早报错。把错误静默广播掉，会让测试通过表面 shape，却产生难排查的数值偏差。 这些边界测试应纳入自动回归。

关键一句:多维 normalized_shape、统计量累积精度和输出 dtype 决定手写实现能否与框架语义对齐。

核验来源

1. Layer Normalization
2. PyTorch torch.nn.LayerNorm documentation

面试官还可能这样问

1. 问法 1 · 场景切入

   假设输入张量形状为 [batch, seq, hidden]，请不用深度学习框架 API，用 NumPy 写一个支持 normalized_shape、eps 和可学习仿射参数的 LayerNorm forward。

2. 问法 2 · 层层追问

   LayerNorm 沿哪些轴计算统计量？……方差怎样稳定计算？……eps 加在方差还是标准差上？……FP16 输入是否需要 FP32 累积？……gamma 与 beta 怎样广播？

3. 问法 3 · 直球技术

   手写 Python+NumPy 的 LayerNorm 前向实现，并解释 normalized_shape、归一化轴、仿射参数、dtype 与数值稳定性处理。