Repository: OpenNMT/OpenNMT-tf

Issue:
Relative Positional Encoding
Hi,

Instead of Sinusoidal Positional Encoding, Relative Positional Encoding is becoming dominant in Transformer models. I would like to ask for adding this feature to the project.



Discussion:

Patch:
diff --git a/opennmt/decoders/self_attention_decoder.py b/opennmt/decoders/self_attention_decoder.py
index 1fa052ee..897ca2ff 100644
--- a/opennmt/decoders/self_attention_decoder.py
+++ b/opennmt/decoders/self_attention_decoder.py
@@ -23,6 +23,7 @@ class SelfAttentionDecoder(decoder.Decoder):
                ffn_activation=tf.nn.relu,
                position_encoder_class=SinusoidalPositionEncoder,
                num_sources=1,
+               maximum_relative_position=None,
                **kwargs):
     """Initializes the parameters of the decoder.
 
@@ -42,6 +43,8 @@ class SelfAttentionDecoder(decoder.Decoder):
         class to use for position encoding (or a callable that returns an
         instance).
       num_sources: The number of source contexts expected by this decoder.
+      maximum_relative_position: Maximum relative position representation
+        (from https://arxiv.org/abs/1803.02155).
       **kwargs: Additional layer arguments.
     """
     super(SelfAttentionDecoder, self).__init__(num_sources=num_sources, **kwargs)
@@ -61,7 +64,8 @@ class SelfAttentionDecoder(decoder.Decoder):
             dropout=dropout,
             attention_dropout=attention_dropout,
             ffn_dropout=ffn_dropout,
-            ffn_activation=ffn_activation)
+            ffn_activation=ffn_activation,
+            maximum_relative_position=maximum_relative_position)
         for i in range(num_layers)]
 
   @property
diff --git a/opennmt/encoders/self_attention_encoder.py b/opennmt/encoders/self_attention_encoder.py
index 9f930d39..6235d5ea 100644
--- a/opennmt/encoders/self_attention_encoder.py
+++ b/opennmt/encoders/self_attention_encoder.py
@@ -24,6 +24,7 @@ class SelfAttentionEncoder(Encoder):
                ffn_dropout=0.1,
                ffn_activation=tf.nn.relu,
                position_encoder_class=SinusoidalPositionEncoder,
+               maximum_relative_position=None,
                **kwargs):
     """Initializes the parameters of the encoder.
 
@@ -42,6 +43,9 @@ class SelfAttentionEncoder(Encoder):
       position_encoder_class: The :class:`opennmt.layers.PositionEncoder`
         class to use for position encoding (or a callable that returns an
         instance).
+      maximum_relative_position: Maximum relative position representation
+        (from https://arxiv.org/abs/1803.02155).
+      **kwargs: Additional layer arguments.
     """
     super(SelfAttentionEncoder, self).__init__(**kwargs)
     self.num_units = num_units
@@ -58,7 +62,8 @@ class SelfAttentionEncoder(Encoder):
             dropout=dropout,
             attention_dropout=attention_dropout,
             ffn_dropout=ffn_dropout,
-            ffn_activation=ffn_activation)
+            ffn_activation=ffn_activation,
+            maximum_relative_position=maximum_relative_position)
         for i in range(num_layers)]
 
   def call(self, inputs, sequence_length=None, training=None):
diff --git a/opennmt/layers/transformer.py b/opennmt/layers/transformer.py
index 6d085713..cc3e23b4 100644
--- a/opennmt/layers/transformer.py
+++ b/opennmt/layers/transformer.py
@@ -65,6 +65,44 @@ def combine_heads(inputs):
   outputs = tf.reshape(outputs, [shape[0], shape[2], shape[1] * shape[3]])
   return outputs
 
+def relative_positions(length, maximum_position, with_cache=False):
+  """Builds the relative positions.
+
+  Args:
+    length: The maximum length of the sequence.
+    maximum_position: The maximum relative position to represent.
+    with_cache: Set to ``True`` if this function is called from a multi-head
+      attention layer with cache states.
+
+  Returns:
+    Positive relative positions with shape :math:`[T or 1, T]`.
+  """
+  if with_cache:
+    distance = tf.expand_dims(tf.range(-length + 1, 1, delta=1), 0)
+  else:
+    arange = tf.range(length)
+    distance = tf.expand_dims(arange, 0) - tf.expand_dims(arange, 1)  # Distance to the diagonal.
+  distance = tf.clip_by_value(distance, -maximum_position, maximum_position)
+  return distance + maximum_position  # Return positive indices.
+
+def matmul_with_relative_representations(a, b, transpose_b=False):  # pylint: disable=invalid-name
+  """Multiplies :obj:`a` with the relative representations :obj:`b`.
+
+  Args:
+    a: Tensor with shape :math:`[B, H, T, _]`.
+    b: Tensor with shape :math:`[T, T, _]`.
+
+  Returns:
+    Tensor with shape :math:`[B, H, T, T]`.
+  """
+  batch, head, time, _ = misc.shape_list(a)
+  a = tf.transpose(a, perm=[2, 0, 1, 3])
+  a = tf.reshape(a, [time, batch * head, -1])
+  c = tf.matmul(a, b, transpose_b=transpose_b)
+  c = tf.reshape(c, [time, batch, head, -1])
+  c = tf.transpose(c, perm=[1, 2, 0, 3])
+  return c
+
 
 class FeedForwardNetwork(tf.keras.layers.Layer):
   """Implements the Transformer's "Feed Forward" layer.
@@ -120,6 +158,7 @@ class MultiHeadAttention(tf.keras.layers.Layer):
                num_units,
                dropout=0.1,
                return_attention=False,
+               maximum_relative_position=None,
                **kwargs):
     """Initializes this layer.
 
@@ -129,6 +168,8 @@ class MultiHeadAttention(tf.keras.layers.Layer):
       dropout: The probability to drop units from the inputs.
       return_attention: If ``True``, also return the attention weights of the
         first head.
+      maximum_relative_position: Maximum relative position representation
+        (from https://arxiv.org/abs/1803.02155).
       kwargs: Additional layer arguments.
     """
     super(MultiHeadAttention, self).__init__(**kwargs)
@@ -136,13 +177,14 @@ class MultiHeadAttention(tf.keras.layers.Layer):
       raise ValueError("Multi head attention requires that num_units is a"
                        " multiple of %s" % num_heads)
     self.num_heads = num_heads
-    self.num_units = num_units
+    self.num_units_per_head = num_units // num_heads
     self.linear_queries = common.Dense(num_units)
     self.linear_keys = common.Dense(num_units)
     self.linear_values = common.Dense(num_units)
     self.linear_output = common.Dense(num_units)
     self.dropout = dropout
     self.return_attention = return_attention
+    self.maximum_relative_position = maximum_relative_position
 
   def map_v1_weights(self, weights):
     # V1 used conv1d layers that have a leading dimensions.
@@ -167,6 +209,16 @@ class MultiHeadAttention(tf.keras.layers.Layer):
       m += self.linear_output.map_v1_weights(weights["conv1d_2"])
     return m
 
+  def build(self, input_shape):
+    if self.maximum_relative_position is not None:
+      relative_window_size = self.maximum_relative_position * 2 + 1
+      relative_repr_shape = [relative_window_size, self.num_units_per_head]
+      self.relative_position_keys = self.add_weight(
+          name="relative_position_keys", shape=relative_repr_shape)
+      self.relative_position_values = self.add_weight(
+          name="relative_position_values", shape=relative_repr_shape)
+    super(MultiHeadAttention, self).build(input_shape)
+
   def call(self, inputs, memory=None, mask=None, cache=None, training=None):  # pylint: disable=arguments-differ
     """Runs the layer.
 
@@ -195,7 +247,7 @@ class MultiHeadAttention(tf.keras.layers.Layer):
     # Compute queries.
     queries = self.linear_queries(inputs)
     queries = split_heads(queries, self.num_heads)
-    queries *= (self.num_units // self.num_heads)**-0.5
+    queries *= self.num_units_per_head**-0.5
 
     # Compute keys and values.
     if memory is None:
@@ -212,10 +264,26 @@ class MultiHeadAttention(tf.keras.layers.Layer):
       else:
         keys, values = _compute_kv(memory)
 
+    if self.maximum_relative_position is not None:
+      if memory is not None:
+        raise ValueError("Relative position representations only supports self-attention")
+      keys_length = tf.shape(keys)[2]
+      relative_pos = relative_positions(
+          keys_length,
+          self.maximum_relative_position,
+          with_cache=bool(cache))
+      relative_repr_keys = tf.gather(self.relative_position_keys, relative_pos)
+      relative_repr_values = tf.gather(self.relative_position_values, relative_pos)
+    else:
+      relative_repr_keys = None
+      relative_repr_values = None
+
     cache = (keys, values)
 
     # Dot product attention.
     dot = tf.matmul(queries, keys, transpose_b=True)
+    if relative_repr_keys is not None:
+      dot += matmul_with_relative_representations(queries, relative_repr_keys, transpose_b=True)
     if mask is not None:
       mask = tf.cast(mask, tf.float32)
       if mask.shape.rank == 2:
@@ -225,6 +293,8 @@ class MultiHeadAttention(tf.keras.layers.Layer):
     attn = tf.cast(tf.nn.softmax(tf.cast(dot, tf.float32)), dot.dtype)
     drop_attn = common.dropout(attn, self.dropout, training=training)
     heads = tf.matmul(drop_attn, values)
+    if relative_repr_values is not None:
+      heads += matmul_with_relative_representations(drop_attn, relative_repr_values)
 
     # Concatenate all heads output.
     combined = combine_heads(heads)
@@ -277,6 +347,7 @@ class SelfAttentionEncoderLayer(tf.keras.layers.Layer):
                attention_dropout=0.1,
                ffn_dropout=0.1,
                ffn_activation=tf.nn.relu,
+               maximum_relative_position=None,
                **kwargs):
     """Initializes the layer.
 
@@ -291,11 +362,16 @@ class SelfAttentionEncoderLayer(tf.keras.layers.Layer):
         the feed forward layer.
       ffn_activation: The activation function to apply between the two linear
         transformations of the feed forward layer.
+      maximum_relative_position: Maximum relative position representation
+        (from https://arxiv.org/abs/1803.02155).
       kwargs: Additional layer arguments.
     """
     super(SelfAttentionEncoderLayer, self).__init__(**kwargs)
     self.self_attention = MultiHeadAttention(
-        num_heads, num_units, dropout=attention_dropout)
+        num_heads,
+        num_units,
+        dropout=attention_dropout,
+        maximum_relative_position=maximum_relative_position)
     self.self_attention = TransformerLayerWrapper(
         self.self_attention, dropout)
     self.ffn = FeedForwardNetwork(
@@ -331,6 +407,7 @@ class SelfAttentionDecoderLayer(tf.keras.layers.Layer):
                attention_dropout=0.1,
                ffn_dropout=0.1,
                ffn_activation=tf.nn.relu,
+               maximum_relative_position=None,
                **kwargs):
     """Initializes the layer.
 
@@ -346,13 +423,16 @@ class SelfAttentionDecoderLayer(tf.keras.layers.Layer):
         the feed forward layer.
       ffn_activation: The activation function to apply between the two linear
         transformations of the feed forward layer.
+      maximum_relative_position: Maximum relative position representation
+        (from https://arxiv.org/abs/1803.02155).
       **kwargs: Additional layer arguments.
     """
     super(SelfAttentionDecoderLayer, self).__init__(**kwargs)
     self.self_attention = MultiHeadAttention(
         num_heads,
         num_units,
-        dropout=attention_dropout)
+        dropout=attention_dropout,
+        maximum_relative_position=maximum_relative_position)
     self.self_attention = TransformerLayerWrapper(
         self.self_attention, dropout)
     self.attention = []
diff --git a/opennmt/models/catalog.py b/opennmt/models/catalog.py
index 69fc1989..2b816ce1 100644
--- a/opennmt/models/catalog.py
+++ b/opennmt/models/catalog.py
@@ -219,37 +219,56 @@ class LstmCnnCrfTagger(sequence_tagger.SequenceTagger):
         }
     })
 
-class Transformer(transformer.Transformer):
-  """Defines a Transformer model as decribed in https://arxiv.org/abs/1706.03762."""
-  def __init__(self):
-    super(Transformer, self).__init__(
-        source_inputter=inputters.WordEmbedder(
-            embedding_size=512),
-        target_inputter=inputters.WordEmbedder(
-            embedding_size=512),
+class _DefaultTransformer(transformer.Transformer):
+  def __init__(self, big=False, relative=False):
+    if big:
+      num_units = 1024
+      num_heads = 16
+      ffn_inner_dim = 4096
+    else:
+      num_units = 512
+      num_heads = 8
+      ffn_inner_dim = 2048
+    if relative:
+      position_encoder_class = None
+      maximum_relative_position = 20
+    else:
+      position_encoder_class = layers.SinusoidalPositionEncoder
+      maximum_relative_position = None
+    super(_DefaultTransformer, self).__init__(
+        source_inputter=inputters.WordEmbedder(embedding_size=num_units),
+        target_inputter=inputters.WordEmbedder(embedding_size=num_units),
         num_layers=6,
-        num_units=512,
-        num_heads=8,
-        ffn_inner_dim=2048,
+        num_units=num_units,
+        num_heads=num_heads,
+        ffn_inner_dim=ffn_inner_dim,
         dropout=0.1,
         attention_dropout=0.1,
-        ffn_dropout=0.1)
+        ffn_dropout=0.1,
+        position_encoder_class=position_encoder_class,
+        maximum_relative_position=maximum_relative_position)
+
+class Transformer(_DefaultTransformer):
+  """Defines a Transformer model as decribed in https://arxiv.org/abs/1706.03762."""
+
+class TransformerRelative(_DefaultTransformer):
+  """Defines a Transformer model using relative position representations as
+  described in https://arxiv.org/abs/1803.02155.
+  """
+  def __init__(self):
+    super(TransformerRelative, self).__init__(relative=True)
 
-class TransformerBig(transformer.Transformer):
+class TransformerBig(_DefaultTransformer):
   """Defines a large Transformer model as decribed in https://arxiv.org/abs/1706.03762."""
   def __init__(self):
-    super(TransformerBig, self).__init__(
-        source_inputter=inputters.WordEmbedder(
-            embedding_size=1024),
-        target_inputter=inputters.WordEmbedder(
-            embedding_size=1024),
-        num_layers=6,
-        num_units=1024,
-        num_heads=16,
-        ffn_inner_dim=4096,
-        dropout=0.3,
-        attention_dropout=0.1,
-        ffn_dropout=0.1)
+    super(TransformerBig, self).__init__(big=True)
+
+class TransformerBigRelative(_DefaultTransformer):
+  """Defines a large Transformer model using relative position representations as
+  described in https://arxiv.org/abs/1803.02155.
+  """
+  def __init__(self):
+    super(TransformerBigRelative, self).__init__(big=True, relative=True)
 
 class GPT2Small(language_model.LanguageModel):
   """GPT-2 language model (small version) as described in:
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index fc2e74a0..6376ba5e 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -28,7 +28,8 @@ class Transformer(SequenceToSequence):
                ffn_activation=tf.nn.relu,
                position_encoder_class=SinusoidalPositionEncoder,
                share_embeddings=EmbeddingsSharingLevel.NONE,
-               share_encoders=False):
+               share_encoders=False,
+               maximum_relative_position=None):
     """Initializes a Transformer model.
 
     Args:
@@ -55,6 +56,8 @@ class Transformer(SequenceToSequence):
         :class:`opennmt.models.EmbeddingsSharingLevel` for possible values.
       share_encoders: In case of multi source architecture, whether to share the
         separate encoders parameters or not.
+      maximum_relative_position: Maximum relative position representation
+        (from https://arxiv.org/abs/1803.02155).
     """
     encoders = [
         SelfAttentionEncoder(
@@ -66,7 +69,8 @@ class Transformer(SequenceToSequence):
             attention_dropout=attention_dropout,
             ffn_dropout=ffn_dropout,
             ffn_activation=ffn_activation,
-            position_encoder_class=position_encoder_class)
+            position_encoder_class=position_encoder_class,
+            maximum_relative_position=maximum_relative_position)
         for _ in range(source_inputter.num_outputs)]
     if len(encoders) > 1:
       encoder = ParallelEncoder(
@@ -85,7 +89,8 @@ class Transformer(SequenceToSequence):
         ffn_dropout=ffn_dropout,
         ffn_activation=ffn_activation,
         position_encoder_class=position_encoder_class,
-        num_sources=source_inputter.num_outputs)
+        num_sources=source_inputter.num_outputs,
+        maximum_relative_position=maximum_relative_position)
 
     self._num_units = num_units
     super(Transformer, self).__init__(


Test patch:
diff --git a/opennmt/tests/transformer_test.py b/opennmt/tests/transformer_test.py
index 464ab928..03aaddae 100644
--- a/opennmt/tests/transformer_test.py
+++ b/opennmt/tests/transformer_test.py
@@ -100,6 +100,12 @@ class TransformerTest(tf.test.TestCase):
     inputs, combined = self.evaluate([inputs, combined])
     self.assertAllEqual(inputs, combined)
 
+  def testRelativePositions(self):
+    positions = transformer.relative_positions(4, 2)
+    self.assertAllEqual(
+        self.evaluate(positions),
+        [[2, 3, 4, 4], [1, 2, 3, 4], [0, 1, 2, 3], [0, 0, 1, 2]])
+
   def testFeedForwardNetwork(self):
     ffn = transformer.FeedForwardNetwork(20, 10)
     x = tf.random.uniform([4, 5, 10])
@@ -124,6 +130,18 @@ class TransformerTest(tf.test.TestCase):
     self.assertEqual(cache[0].shape[2], 2)
     self.assertEqual(cache[1].shape[2], 2)
 
+  def testMultiHeadSelfAttentionRelativePositions(self):
+    attention = transformer.MultiHeadAttention(4, 20, maximum_relative_position=6)
+    x = tf.random.uniform([2, 9, 10])
+    mask = tf.sequence_mask([9, 7])
+    y = attention(x, mask=mask)
+
+  def testMultiHeadSelfAttentionRelativePositionsWithCache(self):
+    attention = transformer.MultiHeadAttention(4, 20, maximum_relative_position=6)
+    x = tf.random.uniform([4, 1, 10])
+    cache = (tf.zeros([4, 4, 0, 5]), tf.zeros([4, 4, 0, 5]))
+    _, cache = attention(x, cache=cache)
+
   def testMultiHeadAttention(self):
     attention = transformer.MultiHeadAttention(4, 20)
     queries = tf.random.uniform([4, 5, 10])