Repository: keras-team/keras-nlp

Issue:
Add a sinusoidal embedding layer
Along with a learned positional embedding (https://github.com/keras-team/keras-nlp/issues/23), we should add a fixed sin/cos embedding as described in [Attention is all you need](https://arxiv.org/pdf/1706.03762.pdf).


Discussion:
I am familiar with transformer and TensorFlow, interested in taking up this issue.
@amantayal44 that would be great!

@hertschuh is currently working on #23 which should be up for review shortly. It would probably be good to wait for us to at least discuss the overall design on that embedding layer before starting on this, so we can align on naming and the like. Then the code for that layer could be a good launching off point for this.
OK, I think we have a clear enough path forward here. Let's name this `SinusoidalPositionalEmbedding` and include a `max_length` argument (though we may still tweak these). Mimic the inputs and outputs of https://github.com/keras-team/keras-nlp/pull/47, you should be able to swap one for the other.

A possible starting point for code for this is [here](https://github.com/tensorflow/models/blob/e0ed507504b4efa6780e244677ded563b36b0ee7/official/nlp/modeling/layers/position_embedding.py#L98).

@amantayal44 would you still like to work on this?
> OK, I think we have a clear enough path forward here. Let's name this `SinusoidalPositionalEmbedding` and include a `max_length` argument (though we may still tweak these). Mimic the inputs and outputs of #47, you should be able to swap one for the other.
> 
> A possible starting point for code for this is [here](https://github.com/tensorflow/models/blob/e0ed507504b4efa6780e244677ded563b36b0ee7/official/nlp/modeling/layers/position_embedding.py#L98).
> 
> @amantayal44 would you still like to work on this?

yes sure.

I have few doubts

1. In TensorFlow implementation, there are using min_timescale and max_timescale, I am not clear about sinusoidal embedding using these two variables, In an original paper (https://arxiv.org/abs/1706.03762), there is only one variable, min_freq = 1e-4.

2.  In sinusoidal encoding we don't need to know max_length beforehand, or should I pre-calculate the positional embedding matrix and avoid calculation at each call?
 



Thank you!

re 1) good catch, let's stick to the original paper details when possible.

re 2) also good point, that's my bad. We can drop the max_length argument.

A couple other points from https://github.com/keras-team/keras-nlp/issues/23 
 - The learned position embedding layer will be assume the `-2` axis is the sequence axis. We can do the same here for now (and possibly open this as a config option in the future).
 -  Also, the naming discussion continues over on [this thread](https://github.com/keras-team/keras-nlp/pull/47/files/3d2bf5c5ccd0791ef726e23acbeba8d51e5df3cb#r828478237). Tentative new name `SinePositionEncoding`, we will see if that sticks 😁 

Resolving patch:
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -13,5 +13,6 @@
 # limitations under the License.
 
 from keras_nlp.layers.fnet_encoder import FNetEncoder
+from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
 from keras_nlp.layers.transformer_decoder import TransformerDecoder
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
diff --git a/keras_nlp/layers/sine_position_encoding.py b/keras_nlp/layers/sine_position_encoding.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/sine_position_encoding.py
@@ -0,0 +1,94 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sinusoidal position embedding layer."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class SinePositionEncoding(keras.layers.Layer):
+    """Sinusoidal positional encoding layer.
+
+    This layer calculates the position encoding as a mix of sine and cosine
+    functions with geometrically increasing wavelengths. Defined and formulized
+    in [Attention is All You Need](https://arxiv.org/abs/1706.03762).
+
+    Takes as input an embedded token tensor. The input must have shape
+    [batch_size, sequence_length, feature_size]. This layer will return a
+    positional encoding the same size as the embedded token tensor, which
+    can be added directly to the embedded token tensor.
+
+    Args:
+        max_wavelength: The maximum angular wavelength of the sine/cosine
+            curves, as described in Attention is All You Need. Defaults to
+            10000.
+
+    Example:
+    ```python
+    # create a simple embedding layer with sinusoidal positional encoding
+    seq_len = 100
+    vocab_size = 1000
+    embedding_dim = 32
+    inputs = keras.Input((seq_len,), dtype=tf.float32)
+    embedding = keras.layers.Embedding(
+        input_dim=vocab_size, output_dim=embedding_dim
+    )(inputs)
+    positional_encoding = keras_nlp.layers.SinePositionEncoding()(embedding)
+    outputs = embedding + positional_encoding
+    ```
+
+    References:
+      [Attention is All You Need](https://arxiv.org/abs/1706.03762)
+    """
+
+    def __init__(
+        self,
+        max_wavelength=10000,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.max_wavelength = max_wavelength
+
+    def call(self, inputs):
+        input_shape = tf.shape(inputs)
+        # length of sequence is the second last dimension of the inputs
+        seq_length = input_shape[-2]
+        hidden_size = input_shape[-1]
+        position = tf.cast(tf.range(seq_length), self.compute_dtype)
+        min_freq = tf.cast(1 / self.max_wavelength, dtype=self.compute_dtype)
+        timescales = tf.pow(
+            min_freq,
+            tf.cast(2 * (tf.range(hidden_size) // 2), self.compute_dtype)
+            / tf.cast(hidden_size, self.compute_dtype),
+        )
+        angles = tf.expand_dims(position, 1) * tf.expand_dims(timescales, 0)
+        # even indices are sine, odd are cosine
+        cos_mask = tf.cast(tf.range(hidden_size) % 2, self.compute_dtype)
+        sin_mask = 1 - cos_mask
+        # embedding shape is [seq_length, hidden_size]
+        positional_encodings = (
+            tf.sin(angles) * sin_mask + tf.cos(angles) * cos_mask
+        )
+
+        return tf.broadcast_to(positional_encodings, input_shape)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "max_wavelength": self.max_wavelength,
+            }
+        )
+        return config


Test patch:
diff --git a/keras_nlp/layers/sine_position_encoding_test.py b/keras_nlp/layers/sine_position_encoding_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/sine_position_encoding_test.py
@@ -0,0 +1,122 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for Sinusoidal Positional encoding."""
+
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import sine_position_encoding
+
+
+class SinePositionEncodingTest(tf.test.TestCase):
+    def test_valid_call(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        model = keras.Sequential(
+            [
+                keras.Input(shape=(4, 6)),
+                pos_encoding,
+            ]
+        )
+        input = tf.random.uniform(shape=[2, 4, 6])
+        model(input)
+
+    def test_static_layer_output_shape(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        seq_length = 100
+        hidden_size = 32
+        inputs = keras.Input(shape=(seq_length, hidden_size))
+        outputs = pos_encoding(inputs)
+
+        # When using static positional encoding shapes, the output is expected
+        # to be the same as the input shape in all dimensions.
+        expected_output_shape = [None, seq_length, hidden_size]
+        self.assertEqual(expected_output_shape, outputs.shape.as_list())
+
+    def test_dynamic_layer_output_shape(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        hidden_size = 32
+        inputs = keras.Input(shape=(None, hidden_size))
+        outputs = pos_encoding(inputs)
+
+        # When using dynamic positional encoding shapes, the output is expected
+        # to be the same as the input shape in all dimensions but may be None.
+        expected_output_shape = [None, None, hidden_size]
+        self.assertEqual(expected_output_shape, outputs.shape.as_list())
+
+    # do multi dimension before sequence length
+    def test_multi_dimension_layer_output_shape(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        seq_length = 100
+        hidden_size = 32
+        inputs = keras.Input(shape=(None, seq_length, hidden_size))
+        outputs = pos_encoding(inputs)
+
+        # When using muliple dimensions before sequence length, the output is
+        # expected to be the same as the input shape in all dimensions.
+        expected_output_shape = [None, None, seq_length, hidden_size]
+        self.assertEqual(expected_output_shape, outputs.shape.as_list())
+
+    def test_output_correct_values(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        model = keras.Sequential(
+            [
+                keras.Input(shape=(4, 6)),
+                pos_encoding,
+            ]
+        )
+        input = tf.random.uniform(shape=[1, 4, 6])
+        output = model(input)
+
+        # comapre position encoding values for position 0 and 3
+        expected_encoding_position_0 = [0.0, 1.0, 0.0, 1.0, 0.0, 1.0]
+        expected_encoding_position_3 = [
+            0.14112,
+            -0.9899925,
+            0.1387981,
+            0.9903207,
+            0.00646326,
+            0.99997914,
+        ]
+        self.assertAllClose(output[0, 0, :], expected_encoding_position_0)
+        self.assertAllClose(output[0, 3, :], expected_encoding_position_3)
+
+    def test_get_config_and_from_config(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding(
+            max_wavelength=1000,
+        )
+        config = pos_encoding.get_config()
+        expected_config_subset = {
+            "max_wavelength": 1000,
+        }
+        self.assertEqual(config, {**config, **expected_config_subset})
+        restored_pos_encoding = (
+            sine_position_encoding.SinePositionEncoding.from_config(config)
+        )
+        self.assertEqual(
+            restored_pos_encoding.get_config(),
+            {**config, **expected_config_subset},
+        )
+
+    def test_float16_dtype(self):
+        pos_encoding = sine_position_encoding.SinePositionEncoding(
+            dtype="float16"
+        )
+        seq_length = 100
+        hidden_size = 32
+        inputs = keras.Input(shape=(seq_length, hidden_size))
+        outputs = pos_encoding(inputs)
+
+        # output dtype for this layer should be tf.float16.
+        self.assertEqual(outputs.dtype, tf.float16)