Repository: OpenNMT/OpenNMT-tf

Issue:
Positional encoding with sinusoids
Implement a new `opennmt.layers.PositionEncoder` with sinusoids as described in "Attention Is All You Need" (https://arxiv.org/abs/1706.03762).

Discussion:

Patch:
diff --git a/opennmt/decoders/self_attention_decoder.py b/opennmt/decoders/self_attention_decoder.py
index fe1b3a6e..1e40d5d6 100644
--- a/opennmt/decoders/self_attention_decoder.py
+++ b/opennmt/decoders/self_attention_decoder.py
@@ -4,7 +4,7 @@ import tensorflow as tf
 import opennmt.layers.transformer as transformer
 
 from opennmt.decoders.decoder import Decoder, get_embedding_fn
-from opennmt.layers.position import PositionEmbedder
+from opennmt.layers.position import SinusoidalPositionEncoder
 from opennmt.utils.beam_search import beam_search
 
 
@@ -21,7 +21,7 @@ class SelfAttentionDecoder(Decoder):
                dropout=0.1,
                attention_dropout=0.1,
                relu_dropout=0.1,
-               position_encoder=PositionEmbedder()):
+               position_encoder=SinusoidalPositionEncoder()):
     """Initializes the parameters of the decoder.
 
     Args:
diff --git a/opennmt/encoders/self_attention_encoder.py b/opennmt/encoders/self_attention_encoder.py
index 69d299cc..3a942a76 100644
--- a/opennmt/encoders/self_attention_encoder.py
+++ b/opennmt/encoders/self_attention_encoder.py
@@ -4,7 +4,7 @@ import tensorflow as tf
 import opennmt.layers.transformer as transformer
 
 from opennmt.encoders.encoder import Encoder
-from opennmt.layers.position import PositionEmbedder
+from opennmt.layers.position import SinusoidalPositionEncoder
 
 
 class SelfAttentionEncoder(Encoder):
@@ -20,7 +20,7 @@ class SelfAttentionEncoder(Encoder):
                dropout=0.1,
                attention_dropout=0.1,
                relu_dropout=0.1,
-               position_encoder=PositionEmbedder()):
+               position_encoder=SinusoidalPositionEncoder()):
     """Initializes the parameters of the encoder.
 
     Args:
diff --git a/opennmt/layers/__init__.py b/opennmt/layers/__init__.py
index cdf13693..79a97413 100644
--- a/opennmt/layers/__init__.py
+++ b/opennmt/layers/__init__.py
@@ -10,3 +10,4 @@ from opennmt.layers.bridge import ZeroBridge
 from opennmt.layers.bridge import DenseBridge
 
 from opennmt.layers.position import PositionEmbedder
+from opennmt.layers.position import SinusoidalPositionEncoder
diff --git a/opennmt/layers/position.py b/opennmt/layers/position.py
index 358c5007..f344f225 100644
--- a/opennmt/layers/position.py
+++ b/opennmt/layers/position.py
@@ -1,5 +1,6 @@
 """Define position encoder classes."""
 
+import math
 import abc
 import six
 
@@ -96,10 +97,10 @@ class PositionEncoder(object):
     input_dim = inputs.get_shape().as_list()[-1]
 
     position = tf.tile([position], [batch_size])
+    position = tf.expand_dims(position, 1)
 
     with tf.variable_scope("position_encoding"):
       position_encoding = self.encode(position, input_dim, dtype=inputs.dtype)
-      position_encoding = tf.expand_dims(position_encoding, 1)
       return self.reducer.reduce([inputs, position_encoding])
 
   @abc.abstractmethod
@@ -157,3 +158,20 @@ class PositionEmbedder(PositionEncoder):
     embeddings = tf.get_variable(
         "w_embs", shape=[self.maximum_position + 1, depth], dtype=dtype)
     return tf.nn.embedding_lookup(embeddings, positions)
+
+
+class SinusoidalPositionEncoder(PositionEncoder):
+  """Encodes positions with sine waves as described in
+  https://arxiv.org/abs/1706.03762.
+  """
+
+  def encode(self, positions, depth, dtype=tf.float32):
+    batch_size = tf.shape(positions)[0]
+    positions = tf.cast(positions, dtype)
+
+    log_timescale_increment = math.log(10000) / (depth / 2 - 1)
+    inv_timescales = tf.exp(tf.range(depth / 2, dtype=dtype) * -log_timescale_increment)
+    inv_timescales = tf.reshape(tf.tile(inv_timescales, [batch_size]), [batch_size, -1])
+    scaled_time = tf.expand_dims(positions, -1) * tf.expand_dims(inv_timescales, 1)
+
+    return tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index f095998e..32bdc812 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -5,7 +5,7 @@ import tensorflow as tf
 from opennmt.models.sequence_to_sequence import SequenceToSequence
 from opennmt.encoders.self_attention_encoder import SelfAttentionEncoder
 from opennmt.decoders.self_attention_decoder import SelfAttentionDecoder
-from opennmt.layers.position import PositionEmbedder
+from opennmt.layers.position import SinusoidalPositionEncoder
 
 
 class Transformer(SequenceToSequence):
@@ -23,7 +23,7 @@ class Transformer(SequenceToSequence):
                dropout=0.1,
                attention_dropout=0.1,
                relu_dropout=0.1,
-               position_encoder=PositionEmbedder(),
+               position_encoder=SinusoidalPositionEncoder(),
                name="transformer"):
     """Initializes a Transformer model.
 


Test patch:
diff --git a/opennmt/tests/position_test.py b/opennmt/tests/position_test.py
index e71ec8c3..7cfb6e45 100644
--- a/opennmt/tests/position_test.py
+++ b/opennmt/tests/position_test.py
@@ -19,6 +19,15 @@ class PositionTest(tf.test.TestCase):
       positions = sess.run(positions)
       self.assertAllEqual([[1, 2, 3, 4, 0, 0, 0], [1, 2, 3, 4, 5, 6, 0]], positions)
 
+  def testSinusoidalPositionEncoder(self):
+    encoder = position.SinusoidalPositionEncoder()
+    positions = position.make_positions([4, 6])
+    depth = 10
+    encoding = encoder.encode(positions, depth)
+    with self.test_session() as sess:
+      encoding = sess.run(encoding)
+      self.assertAllEqual([2, 6, depth], encoding.shape)
+
 
 if __name__ == "__main__":
   tf.test.main()