为何TransformerEncoder需实例化两个LayerNormalization层？

以下代码来自Francois Chollet所著《Python深度学习（第二版）》，是Transformer编码器自定义Layer子类的实现。我注意到代码里专门实例化了两个LayerNormalization层，从定义看二者完全一致，想了解作者这么实现的原因。

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class TransformerEncoder(layers.Layer):
   # 这是继承自keras.layers.Layer的自定义层，用于实现Transformer模型的编码器部分，采用多头注意力机制。

   def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
       super().__init__(**kwargs)
       self.embed_dim = embed_dim  # 输入嵌入的维度
       self.dense_dim = dense_dim  # 前馈网络中全连接层的维度
       self.num_heads = num_heads  # 多头注意力机制中的头数
       self.attention = layers.MultiHeadAttention(
           num_heads=num_heads, key_dim=embed_dim)  # 用于计算注意力权重的MultiHeadAttention实例
       self.dense_proj = keras.Sequential(
           [layers.Dense(dense_dim, activation="relu"),  # 输出维度为dense_dim，输入维度自动匹配
            layers.Dense(embed_dim),]
       )  # 包含两个全连接层的Sequential模型，用于处理注意力输出
       self.layernorm_1 = layers.LayerNormalization()  # 用于归一化输入数据的LayerNormalization层
       self.layernorm_2 = layers.LayerNormalization()  # 用于归一化前馈网络输出的另一个LayerNormalization层

   def call(self, inputs, mask=None):
       # 该方法接受两个参数：inputs（输入数据）和mask（可选的掩码，通常用于标记序列中需忽略的位置，比如填充序列）

       if mask is not None:
           mask = mask[:, tf.newaxis, :]  # 将掩码张量的形状从(batch_size, sequence_length)扩展为(batch_size, 1, sequence_length)

       attention_output = self.attention(inputs, inputs, attention_mask=mask)  # 使用MultiHeadAttention层计算注意力权重

       proj_input = self.layernorm_1(inputs + attention_output)  # 使用LayerNormalization层归一化输入与注意力输出的和

       proj_output = self.dense_proj(proj_input)  # 通过前馈网络处理归一化后的输入

       return self.layernorm_2(proj_input + proj_output)  # 使用另一个LayerNormalization层归一化proj_input与proj_output的和

   def get_config(self):
       # 该方法返回包含层配置的字典，包括层的类型、参数、权重等。当使用model.save()保存模型时，Keras会用此配置重建模型结构。

       config = super().get_config()  # 先调用父类的get_config方法，获取包含父类配置的字典

       config.update({
           "embed_dim": self.embed_dim,
           "num_heads": self.num_heads,
           "dense_dim": self.dense_dim,
       })
       return config  # 然后用自定义层的特定配置（embed_dim、num_heads、dense_dim）更新字典，确保保存模型时自定义层的配置被保存

核心原因解析： #

• LayerNormalization是带可学习参数的层：每个LayerNormalization实例都有独立的缩放参数（gamma）和偏移参数（beta），这些参数会在训练中根据所在位置的输入分布独立优化。两个LN层处理的输入数据分布完全不同：第一个处理的是「原始输入+注意力输出」的残差和，第二个处理的是「前馈输入+前馈输出」的残差和，独立的LN层能针对性适配不同分布，提升模型表达能力。
• 遵循Transformer原始设计：Transformer论文中，每个残差模块后的层归一化都是独立组件，各自维护参数。如果复用同一个LN层，相当于强制共享gamma和beta，会限制模型拟合能力——两个位置的输入统计特征（均值、方差）差异很大，共享参数无法适配这种差异。
• 避免训练时的参数冲突：Keras的Layer实例是有状态的，复用同一个LN层会导致训练时参数更新同时影响两个位置的归一化操作，引发梯度更新混乱，导致模型收敛变慢甚至不收敛。

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