Fix TransformerMemory model

Fixes a bug with training and enables relative attention. Absolute attention works poorly when the timing signal is reset at the start of each chunk.

PiperOrigin-RevId: 239686861

Author: T2T Team

tensor2tensor/layers/common_attention.py

@@ -1667,6 +1667,20 @@ def dot_product_attention_relative(q,
     return _relative_attention_inner(weights, v, relations_values, False)
 
 
+def dot_product_attention_relative_memory(q, k, v, bias, *args, **kwargs):
+  """Wrapper of dot_product_attention_relative to use with recurrent memory."""
+
+  q_len = tf.shape(q)[2]
+  k_len = tf.shape(k)[2]
+  num_memory_items = k_len - q_len
+
+  q = tf.pad(q, [[0, 0], [0, 0], [num_memory_items, 0], [0, 0]])
+  bias = tf.pad(bias, [[0, 0], [0, 0], [num_memory_items, 0], [0, 0]])
+  output = dot_product_attention_relative(q, k, v, bias, *args, **kwargs)
+
+  return output[:, :, num_memory_items:, :]
+
+
 def _relative_position_to_absolute_position_masked(x):
   """Helper to dot_product_self_attention_relative_v2.
 
@@ -4152,6 +4166,18 @@ def multihead_attention(query_antecedent,
           save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
           cache=cache is not None)
+    elif attention_type == "dot_product_relative_memory":
+      x = dot_product_attention_relative_memory(
+          q,
+          k,
+          v,
+          bias,
+          max_relative_position,
+          dropout_rate,
+          image_shapes,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          cache=cache is not None)
     elif attention_type == "dot_product_unmasked_relative_v2":
       x = dot_product_unmasked_self_attention_relative_v2(
           q,

tensor2tensor/layers/transformer_memory.py

@@ -90,17 +90,21 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
         [tf.stop_gradient(previous_vals), query_antecedent], 1)
     new_bias = tf.concat([previous_bias, bias], -1)
 
-    cancel_update = tf.equal(self.previous_segment, segment[0])
     remember_segment = segment[0]
-    remember_vals = tf.cond(
-        cancel_update,
-        lambda: self.previous_vals,
-        lambda: tf.pad(query_antecedent, [[0, amount_to_pad], [0, 0], [0, 0]]))
-    remember_bias = tf.cond(
-        cancel_update,
-        lambda: self.previous_bias,
-        lambda: tf.zeros_like(bias) + tf.reduce_max(bias, -1, keep_dims=True))
-
+    # TODO(kitaev): The code assumes that we always either increment the chunk
+    # number or reset it to zero, which is checked by the assertion. This
+    # assumption will not hold if we re-run the model for each token, e.g. for
+    # autoregressive greedy/beam/sampling decode.
+    with tf.control_dependencies(
+        [tf.Assert(tf.math.logical_or(
+            tf.equal(remember_segment, 0),
+            tf.equal(remember_segment, self.previous_segment + 1)),
+                   [self.previous_segment, remember_segment])]):
+      remember_segment = tf.identity(remember_segment)
+    remember_vals = tf.pad(query_antecedent,
+                           [[0, amount_to_pad], [0, 0], [0, 0]])
+    remember_bias = tf.zeros_like(bias) + tf.reduce_max(
+        bias, -1, keep_dims=True)
     token = (remember_segment, remember_vals, remember_bias)
 
     return token, query_antecedent, new_memory_antecedent, new_bias

tensor2tensor/models/transformer.py

@@ -2624,17 +2624,21 @@ def transformer_wikitext103_l4k_v0():
 
 
 @registry.register_hparams
-def transformer_wikitext103_l4k_memory():
+def transformer_wikitext103_l4k_memory_v0():
   """HParams for training languagemodel_wikitext103_l4k with memory."""
   hparams = transformer_wikitext103_l4k_v0()
 
-  hparams.split_targets_chunk_length = 8
-  hparams.split_targets_max_chunks = 512
+  hparams.split_targets_chunk_length = 64
+  hparams.split_targets_max_chunks = 64
 
   # The hparams specify batch size *before* chunking, but we want to have a
   # consistent 4K batch size *after* chunking to fully utilize the hardware.
   target_tokens_per_batch = 4096
-  hparams.batch_size = target_tokens_per_batch * (
-      hparams.max_length / hparams.split_targets_chunk_length)  # 2097152
+  hparams.batch_size = int(target_tokens_per_batch * (
+      hparams.max_length / hparams.split_targets_chunk_length))  # 262144
+
+  hparams.pos = None
+  hparams.self_attention_type = "dot_product_relative_memory"
+  hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
   return hparams