Repository: keras-team/keras-nlp

Issue:
Investigate recreating a GPT-2 forward pass with KerasNLP
With the addition of a roughly compatible tokenizer with gpt-2 in https://github.com/keras-team/keras-nlp/pull/303, it should be possible to start experimenting with recreating an entire gpt forward pass with KerasNLP components.

This should be more an investigation and validation for now, as we continue to answer some core design questions with our BERT modeling code.

Rough flow:

 - Load a reference version of the gpt-2 tokenizer and network, either from the [released github](https://github.com/openai/gpt-2) or huggingface.
 - Recreate the gpt-2 tokenizer using https://github.com/keras-team/keras-nlp/pull/303
 - Recreate the decoder models itself using KerasNLP components, and convert weights.
 - Verify that our output can match with a high degree of accuracy.


Discussion:
Tried copying the weights present in the official GPT-2 repo here: https://colab.research.google.com/drive/1zjmKuv--nwTbxHM1jeL0hPL-brHJXwCd?usp=sharing. There is a big difference in outputs though.

On further inspection, there seems to be a difference between KerasNLP's `TransformerDecoder` and Hugging Face's GPT-2 decoder.


KerasNLP TransformerDecoder:

https://github.com/keras-team/keras-nlp/blob/b88ff1ba859846c35d9701f5acaf90fecbb7f908/keras_nlp/layers/transformer_decoder.py#L269-L279

1. Apply Attn.
2. Add original input and attention output.
3. Apply LayerNorm to the sum.


Hugging Face:


https://github.com/huggingface/transformers/blob/65fb71bc762c46bb067306c1fd083b1cba87a095/src/transformers/models/gpt2/modeling_gpt2.py#L405-L417

1. Apply LayerNorm.
2. Apply Attn.
3. Add original input and attention output.

We will have to implement a custom GPT-2 decoder layer.

Discussed offline. Tweaked the `TransformerDecoder` layer. Able to reproduce output: https://colab.research.google.com/drive/1zjmKuv--nwTbxHM1jeL0hPL-brHJXwCd?usp=sharing
I think adding a `normalize_first=False` argument to our `TransformerEncoder` and `TransformerDecoder` blocks makes sense. And this is the only change we need for GPT-2 with `TransformerDecoder`, correct?

We want to be careful about making these basic transformer layers "god objects", and generally not venture too far from the *Attn is all you need* formulation here. But this seems like a common and simple enough change to just add a parameter. It seems somewhat silly to have a `keras_nlp.layers.TransformerDecoderWhichNormalizesFirst` class.

@jbischof sound good to you re adding a parameter?
I think this makes sense @mattdangerw! The reasoning is that `TransformerDecoder` is a low-level and heavy-lifting module that would be very hard for users to recreate. It's also enabling important new use cases for a large number of users. The reservations I expressed in #343 were about high level convenience wrappers offering many usage modalities that might confuse new users without offering much to advanced users.

@abheesht17 it would be helpful if you could import your proposed changes from a git branch rather than re-creating all the code in the notebook because it's hard to figure out what's going on at a glance otherwise. Thanks!
Actually, since I think we have alignment here... @abheesht17 want to just open a PR to add a parameter `normalize_first=False` to both `TranformerEncoder` and `TransformerDecoder`?

We can discuss any naming concerns on the PR.
@mattdangerw, @jbischof - sure, will open a PR for `normalize_first` first! Do we also want to add this parameter: https://github.com/tensorflow/models/blob/master/official/nlp/modeling/layers/transformer_encoder_block.py#L124-L126, or should we save that for later?
`diff_q_and_kv_attention_layer_norm`? I would very much hope we do not ever add that to the main TransformerEncoder/TransformerDecoder, we don't even support a separate `key_value` input like the block you are linking. That's the level of customization I would want to expose as a separate layer (either model specific, or generic if we really think it's a general purpose one).

Is there a reason we need that? Is that needed for gpt? If so, I might have an incomplete understanding of this problem.
Nah, we don't need it for GPT-2, just asked because it is related to `norm_first`. They define a separate LayerNorm layer for the **cross-attention layer**, where they have different QKVs. I guess we can skip it, lel.
Got it thanks! Yeah a cross attention encoder (e.g. cross attention without the causal mask), sounds to me like something we could absolutely split out into its own layer. And in fact it might be clearer to do so. `CrossAttentionEncoder` or something something like that. For now, we can ignore for sure.

Resolving patch:
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -63,6 +63,10 @@ class TransformerDecoder(keras.layers.Layer):
         bias_initializer: string or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for
             the dense and multiheaded attention layers.
+        normalize_first: bool. Defaults to False. If True, the inputs to the
+            attention layer(s) and the intermediate dense layer are normalized
+            (similar to GPT-2). If set to False, outputs of attention layer and
+            intermediate dense layer are normalized (similar to BERT).
         name: string, defaults to None. The name of the layer.
         **kwargs: other keyword arguments.
 
@@ -100,6 +104,7 @@ def __init__(
         layer_norm_epsilon=1e-05,
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
+        normalize_first=False,
         name=None,
         **kwargs,
     ):
@@ -116,6 +121,7 @@ def __init__(
         self.layer_norm_epsilon = layer_norm_epsilon
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
+        self.normalize_first = normalize_first
         self._built = False
         self.supports_masking = True
 
@@ -191,10 +197,7 @@ def _build(self, input_shape, has_cross_attention):
         )
         self._output_dropout = keras.layers.Dropout(rate=self.dropout)
 
-    def _add_and_norm(self, input1, input2, norm_layer):
-        return norm_layer(input1 + input2)
-
-    def _feed_forward(self, input):
+    def _feedforward(self, input):
         x = self._intermediate_dense(input)
         x = self._output_dense(x)
         return self._output_dropout(x)
@@ -266,6 +269,11 @@ def call(
         else:
             decoder_mask = tf.minimum(decoder_mask, causal_mask)
 
+        residual_decoder_sequence = decoder_sequence
+        if self.normalize_first:
+            decoder_sequence = self._decoder_attention_layernorm(
+                decoder_sequence
+            )
         # Decoder input self-attention.
         self_attended = self._self_attention_layer(
             decoder_sequence,
@@ -274,14 +282,21 @@ def call(
             attention_mask=decoder_mask,
         )
         self_attended = self._self_attention_dropout(self_attended)
-        attention_output = self._add_and_norm(
-            self_attended, decoder_sequence, self._decoder_attention_layernorm
-        )
+        attention_output = residual_decoder_sequence + self_attended
+        if not self.normalize_first:
+            attention_output = self._decoder_attention_layernorm(
+                attention_output
+            )
 
         if self._cross_attention_layer is not None:
             encoder_mask = merge_padding_and_attention_mask(
                 encoder_sequence, encoder_padding_mask, encoder_attention_mask
             )
+            residual_attention_output = attention_output
+            if self.normalize_first:
+                attention_output = self._cross_attention_layernorm(
+                    attention_output
+                )
             # Cross attention.
             cross_attended = self._cross_attention_layer(
                 query=attention_output,
@@ -292,19 +307,21 @@ def call(
             cross_attended = self._cross_attention_dropout(
                 cross_attended,
             )
-            attention_output = self._add_and_norm(
-                cross_attended,
-                attention_output,
-                self._cross_attention_layernorm,
-            )
-
+            attention_output = residual_attention_output + cross_attended
+            if not self.normalize_first:
+                attention_output = self._cross_attention_layernorm(
+                    attention_output
+                )
+
+        residual_attention_output = attention_output
+        if self.normalize_first:
+            attention_output = self._feedforward_layernorm(attention_output)
         # Feedforward.
-        feed_forward_output = self._feed_forward(attention_output)
-        return self._add_and_norm(
-            attention_output,
-            feed_forward_output,
-            self._feedforward_layernorm,
-        )
+        feedforward_output = self._feedforward(attention_output)
+        feedforward_output = residual_attention_output + feedforward_output
+        if not self.normalize_first:
+            feedforward_output = self._feedforward_layernorm(feedforward_output)
+        return feedforward_output
 
     def get_config(self):
         config = super().get_config()
@@ -321,6 +338,7 @@ def get_config(self):
                 "bias_initializer": keras.initializers.serialize(
                     self.bias_initializer
                 ),
+                "normalize_first": self.normalize_first,
                 "build_input_shape": self._input_shape,
                 "has_cross_attention": self._has_cross_attention,
             }
diff --git a/keras_nlp/layers/transformer_encoder.py b/keras_nlp/layers/transformer_encoder.py
--- a/keras_nlp/layers/transformer_encoder.py
+++ b/keras_nlp/layers/transformer_encoder.py
@@ -51,6 +51,10 @@ class TransformerEncoder(keras.layers.Layer):
         bias_initializer: string or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for
             the dense and multiheaded attention layers.
+        normalize_first: bool. Defaults to False. If True, the inputs to the
+            attention layer and the intermediate dense layer  are normalized
+            (similar to GPT-2). If set to False, outputs of attention layer and
+            intermediate dense layer are normalized (similar to BERT).
         name: string, defaults to None. The name of the layer.
         **kwargs: other keyword arguments.
 
@@ -84,6 +88,7 @@ def __init__(
         layer_norm_epsilon=1e-05,
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
+        normalize_first=False,
         name=None,
         **kwargs
     ):
@@ -98,6 +103,7 @@ def __init__(
         self.layer_norm_epsilon = layer_norm_epsilon
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
+        self.normalize_first = normalize_first
         self._built = False
         self.supports_masking = True
 
@@ -144,10 +150,7 @@ def _build(self, input_shape):
         )
         self._output_dropout = keras.layers.Dropout(rate=self.dropout)
 
-    def _add_and_norm(self, input1, input2, norm_layer):
-        return norm_layer(input1 + input2)
-
-    def _feed_forward(self, input):
+    def _feedforward(self, input):
         x = self._intermediate_dense(input)
         x = self._output_dense(x)
         return self._output_dropout(x)
@@ -179,21 +182,27 @@ def call(self, inputs, padding_mask=None, attention_mask=None):
             attention_mask,
         )
 
+        residual_inputs = inputs
+        if self.normalize_first:
+            inputs = self._attention_layernorm(inputs)
         # Self attention.
         attended = self._multi_head_attention_layer(
             inputs, inputs, inputs, attention_mask=mask
         )
         attended = self._attention_dropout(attended)
-        attended = self._add_and_norm(
-            inputs,
-            attended,
-            self._attention_layernorm,
-        )
+        attended = residual_inputs + attended
+        if not self.normalize_first:
+            attended = self._attention_layernorm(attended)
+
+        residual_attended = attended
+        if self.normalize_first:
+            attended = self._feedforward_layernorm(attended)
         # Feedforward.
-        feed_forward_output = self._feed_forward(attended)
-        return self._add_and_norm(
-            attended, feed_forward_output, self._feedforward_layernorm
-        )
+        feedforward_output = self._feedforward(attended)
+        feedforward_output = residual_attended + feedforward_output
+        if not self.normalize_first:
+            feedforward_output = self._feedforward_layernorm(feedforward_output)
+        return feedforward_output
 
     def get_config(self):
         config = super().get_config()
@@ -210,6 +219,7 @@ def get_config(self):
                 "bias_initializer": keras.initializers.serialize(
                     self.bias_initializer
                 ),
+                "normalize_first": self.normalize_first,
                 "build_input_shape": self._input_shape,
             }
         )


Test patch:
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -23,12 +23,17 @@
 
 
 class TransformerDecoderTest(tf.test.TestCase, parameterized.TestCase):
-    def test_valid_call(self):
+    @parameterized.named_parameters(
+        ("without_norm_first", False),
+        ("with_norm_first", True),
+    )
+    def test_valid_call(self, normalize_first):
         encoder_input = keras.Input(shape=[4, 6])
         decoder_input = keras.Input(shape=[4, 6])
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
+            normalize_first=normalize_first,
         )
         output = decoder(decoder_input, encoder_input)
         model = keras.Model(
@@ -39,11 +44,16 @@ def test_valid_call(self):
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model([decoder_sequence, encoder_sequence])
 
-    def test_valid_call_without_cross_attention(self):
+    @parameterized.named_parameters(
+        ("without_norm_first", False),
+        ("with_norm_first", True),
+    )
+    def test_valid_call_without_cross_attention(self, normalize_first):
         decoder_input = keras.Input(shape=[4, 6])
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
+            normalize_first=normalize_first,
         )
         output = decoder(decoder_input)
         model = keras.Model(
@@ -83,6 +93,7 @@ def test_get_config_and_from_config(self):
             num_heads=2,
             kernel_initializer="HeNormal",
             bias_initializer="Zeros",
+            normalize_first=True,
         )
 
         config = decoder.get_config()
@@ -98,6 +109,7 @@ def test_get_config_and_from_config(self):
             "bias_initializer": keras.initializers.serialize(
                 keras.initializers.Zeros()
             ),
+            "normalize_first": True,
         }
         self.assertEqual(config, {**config, **expected_config_subset})
         self.assertEqual(config, {**config, **expected_config_subset})
@@ -264,6 +276,7 @@ def test_save_model(self, format):
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
+            normalize_first=True,
         )
         output = decoder(encoder_input, decoder_input)
         model = keras.Model(
@@ -287,6 +300,7 @@ def test_save_model_without_cross_attention(self, format):
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
+            normalize_first=True,
         )
         output = decoder(decoder_input)
         model = keras.Model(
diff --git a/keras_nlp/layers/transformer_encoder_test.py b/keras_nlp/layers/transformer_encoder_test.py
--- a/keras_nlp/layers/transformer_encoder_test.py
+++ b/keras_nlp/layers/transformer_encoder_test.py
@@ -23,10 +23,15 @@
 
 
 class TransformerEncoderTest(tf.test.TestCase, parameterized.TestCase):
-    def test_valid_call(self):
+    @parameterized.named_parameters(
+        ("without_norm_first", False),
+        ("with_norm_first", True),
+    )
+    def test_valid_call(self, normalize_first):
         encoder = transformer_encoder.TransformerEncoder(
             intermediate_dim=4,
             num_heads=2,
+            normalize_first=normalize_first,
         )
         model = keras.Sequential(
             [
@@ -53,6 +58,7 @@ def test_get_config_and_from_config(self):
             num_heads=2,
             kernel_initializer="HeNormal",
             bias_initializer="Zeros",
+            normalize_first=True,
         )
 
         config = encoder.get_config()
@@ -69,6 +75,7 @@ def test_get_config_and_from_config(self):
             "bias_initializer": keras.initializers.serialize(
                 keras.initializers.Zeros()
             ),
+            "normalize_first": True,
         }
 
         self.assertEqual(config, {**config, **expected_config_subset})
@@ -161,6 +168,7 @@ def test_save_model(self, format):
                 transformer_encoder.TransformerEncoder(
                     intermediate_dim=4,
                     num_heads=2,
+                    normalize_first=True,
                 ),
             ]
         )