title: Layer Normalization Training vs Evaluation mode in TensorFlow/Keras
body: <p>I am trying to understand how Layer Normalization behaves in training vs evaluation mode. I thought it is similar to Batch Normalization (training updates the beta/gamma values and evaluation uses their global values). Hence I was expecting different results when <code>training=True/False</code> but the results are same. Am I missing something?</p>
  <p>Here is the working code:</p>
  <pre><code>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)
  
  </code></pre>
  <p>When I use <code>training=True vs False</code>, the results are same. But when I replace the layer with BatchNorm, the results are different as expected.</p>
  <pre><code>model(x_train[:1], training=False).numpy()
  </code></pre>
  <pre><code>model(x_train[:1], training=True).numpy()
  </code></pre>