TensorFlow/Keras中LayerNormalization训练评估模式输出一致问题

问题现象 #

使用TensorFlow/Keras开发时，对Layer Normalization（LN）的运行逻辑存在初始认知偏差：

• 误以为LN机制与Batch Normalization（BN）一致：训练阶段更新beta/gamma参数，评估阶段使用全局统计值，预期传入training=True/training=False参数时会得到不同的推理结果
• 实际测试结果：模型训练完成后，两种模式下LN层的输出完全一致；将LN替换为BN后，两种模式输出存在明显差异，符合初始预期

可复现代码 #

模型训练代码 #

import tensorflow as tf

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

tf.random.set_seed(42)

class MLPModel(tf.keras.Model):
    def __init__(self):
        super(MLPModel, self).__init__()
        
        self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10)
        self.norm = tf.keras.layers.LayerNormalization()
        

    def call(self, inputs):
        x = self.flatten(inputs)
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.norm(x)
        return x

model = MLPModel()

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer='adam', loss=loss_fn, metrics=['accuracy'])

model.fit(x_train, y_train, epochs=2)

测试调用代码 #

训练完成后分别执行以下两行调用：

model(x_train[:1], training=False).numpy()

model(x_train[:1], training=True).numpy()

两次调用返回结果完全相同。

核心机制说明 #

出现这个现象的本质原因是LN和BN的归一化统计量计算逻辑完全不同，不存在认知里的「训练用批次统计、推理用全局统计」的分支：

• BN的归一化维度是批次维度：训练阶段使用当前输入批次的均值、方差做归一化，同时通过滑动平均累计全量数据的全局均值、方差；推理阶段直接使用训练阶段累计好的全局统计量做归一化，因此训练/推理模式下计算逻辑不同，输出存在差异。BN中的beta、gamma是可训练仿射参数，训练和推理阶段都使用最新训练得到的参数值，不存在模式切换带来的参数变化。
• LN的归一化维度是单个样本的特征维度：无论训练还是推理阶段，都是针对当前输入的单个样本，在层的特征维度上计算均值、方差完成归一化，不需要维护跨批次的全局滑动统计量，因此training参数不会改变LN层的任何计算逻辑，两种模式输出自然完全一致。LN中的beta、gamma同样是可训练仿射参数，训练完成后取值固定，不受模式切换影响。

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