No public description

PiperOrigin-RevId: 675335745

Author: tensorflower-gardener

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -114,6 +114,7 @@ def __init__(self,
                tgt_block_size=None,
                use_sigmoid_attn=False,
                sigmoid_attn_bias=None,
+               linformer_dim=None,
                **kwargs):
     """Initializes `TransformerEncoderBlock`.
 
@@ -191,6 +192,8 @@ def __init__(self,
         `block_sparse_attention.MultiHeadAttention`
       sigmoid_attn_bias: This param is only used in
         `block_sparse_attention.MultiHeadAttention`
+      linformer_dim: Applies low-rank factorization on keys/values as in
+        https://arxiv.org/pdf/2006.04768.
       **kwargs: keyword arguments.
     """
     util.filter_kwargs(kwargs)
@@ -230,6 +233,7 @@ def __init__(self,
     self._tgt_block_size = tgt_block_size
     self._use_sigmoid_attn = use_sigmoid_attn
     self._sigmoid_attn_bias = sigmoid_attn_bias
+    self._linformer_dim = linformer_dim
     if self._num_kv_heads is not None and self._src_block_size is not None:
       raise ValueError(
           "Block sparse attention does not support Multi-query attention."
@@ -366,16 +370,31 @@ def build(self, input_shape):
         name="output",
         kernel_initializer=tf_utils.clone_initializer(self._kernel_initializer),
         bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
-        **common_kwargs)
+        **common_kwargs,
+    )
     self._output_dropout = tf_keras.layers.Dropout(
-        rate=self._output_dropout_rate)
+        rate=self._output_dropout_rate
+    )
     # Use float32 in layernorm for numeric stability.
     self._output_layer_norm = tf_keras.layers.LayerNormalization(
         name="output_layer_norm",
         axis=-1,
         epsilon=self._norm_epsilon,
-        dtype=tf.float32)
-
+        dtype=tf.float32,
+    )
+    if self._linformer_dim is not None:
+      # Current implementation uses the same weights for keys and values.
+      # TODO(akandoor): Explore using different weights for keys and values.
+      self._lowrank_kv_projection = tf_keras.layers.EinsumDense(
+          "...bc,cd->...bd",
+          output_shape=(None, self._linformer_dim),
+          kernel_initializer=tf_utils.clone_initializer(
+              self._kernel_initializer
+          ),
+          bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
+          name="lowrank_kv_projection",
+          **common_kwargs,
+      )
     super().build(input_shape)
 
   def get_config(self):
@@ -480,6 +499,19 @@ def call(self, inputs: Any, output_range: Optional[tf.Tensor] = None) -> Any:
     if key_value is None:
       key_value = input_tensor
 
+    if self._linformer_dim is not None:
+      if attention_mask is not None:
+        # Applying mask before the low rank factorization so that padding is
+        # accounted for.
+        query_mask = tf.cast(attention_mask[:, :, 0], dtype=target_tensor.dtype)
+        target_tensor = target_tensor * tf.expand_dims(query_mask, axis=-1)
+        key_mask = tf.cast(attention_mask[:, 0, :], dtype=target_tensor.dtype)
+        key_value = key_value * tf.expand_dims(key_mask, axis=-1)
+        attention_mask = None
+      key_value = tf.transpose(key_value, [0, 2, 1])
+      key_value = self._lowrank_kv_projection(key_value)
+      key_value = tf.transpose(key_value, [0, 2, 1])
+
     if self._return_attention_scores:
       attention_output, attention_scores = self._attention_layer(
           query=target_tensor,

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -800,6 +800,42 @@ def test_block_sparse_attention(self, use_sigmoid_attn):
         output_tensor[1].shape.as_list(), expected_attention_scores_shape
     )
 
+  def test_low_rank_attention(self):
+    num_attention_heads = 8
+    sequence_length = 21
+    linformer_dim = 7
+    width = 80
+
+    test_layer = TransformerEncoderBlock(
+        num_attention_heads=num_attention_heads,
+        inner_dim=2048,
+        inner_activation='relu',
+        return_attention_scores=True,
+        linformer_dim=linformer_dim,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    output_tensor = test_layer(data_tensor)
+
+    expected_layer_output_shape = [None, sequence_length, width]
+    expected_attention_scores_shape = [
+        None,
+        num_attention_heads,
+        sequence_length,
+        linformer_dim,
+    ]
+
+    self.assertIsInstance(output_tensor, tuple)
+    self.assertLen(output_tensor, 2)
+    # First is the standard output.
+    self.assertEqual(
+        output_tensor[0].shape.as_list(), expected_layer_output_shape
+    )
+    # Second is the attention scores.
+    self.assertEqual(
+        output_tensor[1].shape.as_list(), expected_attention_scores_shape
+    )
+
 
 if __name__ == '__main__':
   tf.test.main()